/*
 * Copyright © 2013 Marek Olšák <maraeo@gmail.com>
 *
 * Permission is hereby granted, free of charge, to any person obtaining a
 * copy of this software and associated documentation files (the "Software"),
 * to deal in the Software without restriction, including without limitation
 * the rights to use, copy, modify, merge, publish, distribute, sublicense,
 * and/or sell copies of the Software, and to permit persons to whom the
 * Software is furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice (including the next
 * paragraph) shall be included in all copies or substantial portions of the
 * Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
 * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
 * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
 * DEALINGS IN THE SOFTWARE.
 */

/**
 * \file opt_dead_builtin_varyings.cpp
 *
 * This eliminates the built-in shader outputs which are either not written
 * at all or not used by the next stage. It also eliminates unused elements
 * of gl_TexCoord inputs, which reduces the overall varying usage.
 * The varyings handled here are the primary and secondary color, the fog,
 * and the texture coordinates (gl_TexCoord).
 *
 * This pass is necessary, because the Mesa GLSL linker cannot eliminate
 * built-in varyings like it eliminates user-defined varyings, because
 * the built-in varyings have pre-assigned locations. Also, the elimination
 * of unused gl_TexCoord elements requires its own lowering pass anyway.
 *
 * It's implemented by replacing all occurences of dead varyings with
 * temporary variables, which creates dead code. It is recommended to run
 * a dead-code elimination pass after this.
 *
 * If any texture coordinate slots can be eliminated, the gl_TexCoord array is
 * broken down into separate vec4 variables with locations equal to
 * VARYING_SLOT_TEX0 + i.
 */

#include "main/imports.h" /* for snprintf */
#include "ir.h"
#include "ir_rvalue_visitor.h"
#include "ir_optimization.h"
#include "ir_print_visitor.h"
#include "glsl_types.h"
#include "link_varyings.h"


/**
 * This obtains detailed information about built-in varyings from shader code.
 */
class varying_info_visitor : public ir_hierarchical_visitor {
public:
   /* "mode" can be either ir_var_shader_in or ir_var_shader_out */
   varying_info_visitor(ir_variable_mode mode)
      : lower_texcoord_array(true),
        texcoord_array(NULL),
        texcoord_usage(0),
        color_usage(0),
        tfeedback_color_usage(0),
        fog(NULL),
        has_fog(false),
        tfeedback_has_fog(false),
        mode(mode)
   {
      memset(color, 0, sizeof(color));
      memset(backcolor, 0, sizeof(backcolor));
   }

   virtual ir_visitor_status visit_enter(ir_dereference_array *ir)
   {
      ir_variable *var = ir->variable_referenced();

      if (var && var->mode == this->mode &&
          var->location == VARYING_SLOT_TEX0) {
         this->texcoord_array = var;

         ir_constant *index = ir->array_index->as_constant();
         if (index == NULL) {
            /* There is variable indexing, we can't lower the texcoord array.
             */
            this->texcoord_usage |= (1 << var->type->array_size()) - 1;
            this->lower_texcoord_array = false;
         }
         else {
            this->texcoord_usage |= 1 << index->get_uint_component(0);
         }

         /* Don't visit the leaves of ir_dereference_array. */
         return visit_continue_with_parent;
      }

      return visit_continue;
   }

   virtual ir_visitor_status visit(ir_dereference_variable *ir)
   {
      ir_variable *var = ir->variable_referenced();

      if (var->mode == this->mode && var->type->is_array() &&
          var->location == VARYING_SLOT_TEX0) {
         /* This is a whole array dereference like "gl_TexCoord = x;",
          * there's probably no point in lowering that.
          */
         this->texcoord_usage |= (1 << var->type->array_size()) - 1;
         this->lower_texcoord_array = false;
      }
      return visit_continue;
   }

   virtual ir_visitor_status visit(ir_variable *var)
   {
      if (var->mode != this->mode)
         return visit_continue;

      /* Handle colors and fog. */
      switch (var->location) {
      case VARYING_SLOT_COL0:
         this->color[0] = var;
         this->color_usage |= 1;
         break;
      case VARYING_SLOT_COL1:
         this->color[1] = var;
         this->color_usage |= 2;
         break;
      case VARYING_SLOT_BFC0:
         this->backcolor[0] = var;
         this->color_usage |= 1;
         break;
      case VARYING_SLOT_BFC1:
         this->backcolor[1] = var;
         this->color_usage |= 2;
         break;
      case VARYING_SLOT_FOGC:
         this->fog = var;
         this->has_fog = true;
         break;
      }

      return visit_continue;
   }

   void get(exec_list *ir,
            unsigned num_tfeedback_decls,
            tfeedback_decl *tfeedback_decls)
   {
      /* Handle the transform feedback varyings. */
      for (unsigned i = 0; i < num_tfeedback_decls; i++) {
         if (!tfeedback_decls[i].is_varying())
            continue;

         unsigned location = tfeedback_decls[i].get_location();

         switch (location) {
         case VARYING_SLOT_COL0:
         case VARYING_SLOT_BFC0:
            this->tfeedback_color_usage |= 1;
            break;
         case VARYING_SLOT_COL1:
         case VARYING_SLOT_BFC1:
            this->tfeedback_color_usage |= 2;
            break;
         case VARYING_SLOT_FOGC:
            this->tfeedback_has_fog = true;
            break;
         default:
            if (location >= VARYING_SLOT_TEX0 &&
                location <= VARYING_SLOT_TEX7) {
               this->lower_texcoord_array = false;
            }
         }
      }

      /* Process the shader. */
      visit_list_elements(this, ir);

      if (!this->texcoord_array) {
         this->lower_texcoord_array = false;
      }
   }

   bool lower_texcoord_array;
   ir_variable *texcoord_array;
   unsigned texcoord_usage; /* bitmask */

   ir_variable *color[2];
   ir_variable *backcolor[2];
   unsigned color_usage; /* bitmask */
   unsigned tfeedback_color_usage; /* bitmask */

   ir_variable *fog;
   bool has_fog;
   bool tfeedback_has_fog;

   ir_variable_mode mode;
};


/**
 * This replaces unused varyings with temporary variables.
 *
 * If "ir" is the producer, the "external" usage should come from
 * the consumer. It also works the other way around. If either one is
 * missing, set the "external" usage to a full mask.
 */
class replace_varyings_visitor : public ir_rvalue_visitor {
public:
   replace_varyings_visitor(exec_list *ir,
                            const varying_info_visitor *info,
                            unsigned external_texcoord_usage,
                            unsigned external_color_usage,
                            bool external_has_fog)
      : info(info), new_fog(NULL)
   {
      void *const ctx = ir;

      memset(this->new_texcoord, 0, sizeof(this->new_texcoord));
      memset(this->new_color, 0, sizeof(this->new_color));
      memset(this->new_backcolor, 0, sizeof(this->new_backcolor));

      const char *mode_str =
         info->mode == ir_var_shader_in ? "in" : "out";

      /* Handle texcoord outputs.
       *
       * We're going to break down the gl_TexCoord array into separate
       * variables. First, add declarations of the new variables all
       * occurences of gl_TexCoord will be replaced with.
       */
      if (info->lower_texcoord_array) {
         for (int i = MAX_TEXTURE_COORD_UNITS-1; i >= 0; i--) {
            if (info->texcoord_usage & (1 << i)) {
               char name[32];

               if (!(external_texcoord_usage & (1 << i))) {
                  /* This varying is unused in the next stage. Declare
                   * a temporary instead of an output. */
                  snprintf(name, 32, "gl_%s_TexCoord%i_dummy", mode_str, i);
                  this->new_texcoord[i] =
                     new (ctx) ir_variable(glsl_type::vec4_type, name,
                                           ir_var_temporary);
               }
               else {
                  snprintf(name, 32, "gl_%s_TexCoord%i", mode_str, i);
                  this->new_texcoord[i] =
                     new(ctx) ir_variable(glsl_type::vec4_type, name,
                                          info->mode);
                  this->new_texcoord[i]->location = VARYING_SLOT_TEX0 + i;
                  this->new_texcoord[i]->explicit_location = true;
                  this->new_texcoord[i]->explicit_index = 0;
               }

               ir->head->insert_before(new_texcoord[i]);
            }
         }
      }

      /* Create dummy variables which will replace set-but-unused color and
       * fog outputs.
       */
      external_color_usage |= info->tfeedback_color_usage;

      for (int i = 0; i < 2; i++) {
         char name[32];

         if (!(external_color_usage & (1 << i))) {
            if (info->color[i]) {
               snprintf(name, 32, "gl_%s_FrontColor%i_dummy", mode_str, i);
               this->new_color[i] =
                  new (ctx) ir_variable(glsl_type::vec4_type, name,
                                        ir_var_temporary);
            }

            if (info->backcolor[i]) {
               snprintf(name, 32, "gl_%s_BackColor%i_dummy", mode_str, i);
               this->new_backcolor[i] =
                  new (ctx) ir_variable(glsl_type::vec4_type, name,
                                        ir_var_temporary);
            }
         }
      }

      if (!external_has_fog && !info->tfeedback_has_fog &&
          info->fog) {
         char name[32];

         snprintf(name, 32, "gl_%s_FogFragCoord_dummy", mode_str);
         this->new_fog = new (ctx) ir_variable(glsl_type::float_type, name,
                                               ir_var_temporary);
      }

      /* Now do the replacing. */
      visit_list_elements(this, ir);
   }

   virtual ir_visitor_status visit(ir_variable *var)
   {
      /* Remove the gl_TexCoord array. */
      if (this->info->lower_texcoord_array &&
          var == this->info->texcoord_array) {
         var->remove();
      }

      /* Replace set-but-unused color and fog outputs with dummy variables. */
      for (int i = 0; i < 2; i++) {
         if (var == this->info->color[i] && this->new_color[i]) {
            var->replace_with(this->new_color[i]);
         }
         if (var == this->info->backcolor[i] &&
             this->new_backcolor[i]) {
            var->replace_with(this->new_backcolor[i]);
         }
      }

      if (var == this->info->fog && this->new_fog) {
         var->replace_with(this->new_fog);
      }

      return visit_continue;
   }

   virtual void handle_rvalue(ir_rvalue **rvalue)
   {
      if (!*rvalue)
         return;

      void *ctx = ralloc_parent(*rvalue);

      /* Replace an array dereference gl_TexCoord[i] with a single
       * variable dereference representing gl_TexCoord[i].
       */
      if (this->info->lower_texcoord_array) {
         /* gl_TexCoord[i] occurence */
         ir_dereference_array *const da = (*rvalue)->as_dereference_array();

         if (da && da->variable_referenced() ==
             this->info->texcoord_array) {
            unsigned i = da->array_index->as_constant()->get_uint_component(0);

            *rvalue = new(ctx) ir_dereference_variable(this->new_texcoord[i]);
            return;
         }
      }

      /* Replace set-but-unused color and fog outputs with dummy variables. */
      ir_dereference_variable *const dv = (*rvalue)->as_dereference_variable();
      if (!dv)
         return;

      ir_variable *var = dv->variable_referenced();

      for (int i = 0; i < 2; i++) {
         if (var == this->info->color[i] && this->new_color[i]) {
            *rvalue = new(ctx) ir_dereference_variable(this->new_color[i]);
            return;
         }
         if (var == this->info->backcolor[i] &&
             this->new_backcolor[i]) {
            *rvalue = new(ctx) ir_dereference_variable(this->new_backcolor[i]);
            return;
         }
      }

      if (var == this->info->fog && this->new_fog) {
         *rvalue = new(ctx) ir_dereference_variable(this->new_fog);
      }
   }

   virtual ir_visitor_status visit_leave(ir_assignment *ir)
   {
      handle_rvalue(&ir->rhs);
      handle_rvalue(&ir->condition);

      /* We have to use set_lhs when changing the LHS of an assignment. */
      ir_rvalue *lhs = ir->lhs;

      handle_rvalue(&lhs);
      if (lhs != ir->lhs) {
         ir->set_lhs(lhs);
      }

      return visit_continue;
   }

private:
   const varying_info_visitor *info;
   struct ir_variable *new_texcoord[MAX_TEXTURE_COORD_UNITS];
   struct ir_variable *new_color[2];
   struct ir_variable *new_backcolor[2];
   struct ir_variable *new_fog;
};


static void
lower_texcoord_array(exec_list *ir, const varying_info_visitor *info)
{
   replace_varyings_visitor(ir, info,
                            (1 << MAX_TEXTURE_COORD_UNITS) - 1,
                            1 | 2, true);
}


void
do_dead_builtin_varyings(struct gl_context *ctx,
                         gl_shader *producer, gl_shader *consumer,
                         unsigned num_tfeedback_decls,
                         tfeedback_decl *tfeedback_decls)
{
   /* This optimization has no effect with the core context and GLES2, because
    * the built-in varyings we're eliminating here are not available there.
    *
    * EXT_separate_shader_objects doesn't allow this optimization,
    * because a program object can be bound partially (e.g. only one
    * stage of a program object can be bound).
    */
   if (ctx->API == API_OPENGL_CORE ||
       ctx->API == API_OPENGLES2 ||
       ctx->Extensions.EXT_separate_shader_objects) {
      return;
   }

   /* Information about built-in varyings. */
   varying_info_visitor producer_info(ir_var_shader_out);
   varying_info_visitor consumer_info(ir_var_shader_in);

   if (producer) {
      producer_info.get(producer->ir, num_tfeedback_decls, tfeedback_decls);

      if (!consumer) {
         /* At least eliminate unused gl_TexCoord elements. */
         if (producer_info.lower_texcoord_array) {
            lower_texcoord_array(producer->ir, &producer_info);
         }
         return;
      }
   }

   if (consumer) {
      consumer_info.get(consumer->ir, 0, NULL);

      if (!producer) {
         /* At least eliminate unused gl_TexCoord elements. */
         if (consumer_info.lower_texcoord_array) {
            lower_texcoord_array(consumer->ir, &consumer_info);
         }
         return;
      }
   }

   /* Eliminate the outputs unused by the consumer. */
   if (producer_info.lower_texcoord_array ||
       producer_info.color_usage ||
       producer_info.has_fog) {
      replace_varyings_visitor(producer->ir,
                               &producer_info,
                               consumer_info.texcoord_usage,
                               consumer_info.color_usage,
                               consumer_info.has_fog);
   }

   /* The gl_TexCoord fragment shader inputs can be initialized
    * by GL_COORD_REPLACE, so we can't eliminate them.
    *
    * This doesn't prevent elimination of the gl_TexCoord elements which
    * are not read by the fragment shader. We want to eliminate those anyway.
    */
   if (consumer->Type == GL_FRAGMENT_SHADER) {
      producer_info.texcoord_usage = (1 << MAX_TEXTURE_COORD_UNITS) - 1;
   }

   /* Eliminate the inputs uninitialized by the producer. */
   if (consumer_info.lower_texcoord_array ||
       consumer_info.color_usage ||
       consumer_info.has_fog) {
      replace_varyings_visitor(consumer->ir,
                               &consumer_info,
                               producer_info.texcoord_usage,
                               producer_info.color_usage,
                               producer_info.has_fog);
   }
}
