#include "bifrost_nir.h"

#include "nir.h"
#include "nir_builder.h"
#include "nir_search.h"
#include "nir_search_helpers.h"

/* What follows is NIR algebraic transform code for the following 6
 * transforms:
 *    ('fmul', 'a', 2.0) => ('fadd', 'a', 'a')
 *    ('fmin', ('fmax', 'a', -1.0), 1.0) => ('fsat_signed_mali', 'a')
 *    ('fmax', ('fmin', 'a', 1.0), -1.0) => ('fsat_signed_mali', 'a')
 *    ('fmax', 'a', 0.0) => ('fclamp_pos_mali', 'a')
 *    ('fabs', ('fddx', 'a')) => ('fabs', ('fddx_must_abs_mali', 'a'))
 *    ('fabs', ('fddy', 'b')) => ('fabs', ('fddy_must_abs_mali', 'b'))
 */


static const nir_search_value_union bifrost_nir_lower_algebraic_late_values[] = {
   /* ('fmul', 'a', 2.0) => ('fadd', 'a', 'a') */
   { .variable = {
      { nir_search_value_variable, -1 },
      0, /* a */
      false,
      nir_type_invalid,
      -1,
      {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15},
   } },
   { .constant = {
      { nir_search_value_constant, -1 },
      nir_type_float, { 0x4000000000000000 /* 2.0 */ },
   } },
   { .expression = {
      { nir_search_value_expression, -1 },
      false,
      false,
      false,
      nir_op_fmul,
      0, 1,
      { 0, 1 },
      -1,
   } },

   /* replace0_0 -> 0 in the cache */
   /* replace0_1 -> 0 in the cache */
   { .expression = {
      { nir_search_value_expression, -1 },
      false,
      false,
      false,
      nir_op_fadd,
      -1, 0,
      { 0, 0 },
      -1,
   } },

   /* ('fmin', ('fmax', 'a', -1.0), 1.0) => ('fsat_signed_mali', 'a') */
   /* search1_0_0 -> 0 in the cache */
   { .constant = {
      { nir_search_value_constant, -1 },
      nir_type_float, { 0xbff0000000000000 /* -1.0 */ },
   } },
   { .expression = {
      { nir_search_value_expression, -1 },
      false,
      false,
      false,
      nir_op_fmax,
      1, 1,
      { 0, 4 },
      -1,
   } },
   { .constant = {
      { nir_search_value_constant, -1 },
      nir_type_float, { 0x3ff0000000000000 /* 1.0 */ },
   } },
   { .expression = {
      { nir_search_value_expression, -1 },
      false,
      false,
      false,
      nir_op_fmin,
      0, 2,
      { 5, 6 },
      -1,
   } },

   /* replace1_0 -> 0 in the cache */
   { .expression = {
      { nir_search_value_expression, -1 },
      false,
      false,
      false,
      nir_op_fsat_signed_mali,
      -1, 0,
      { 0 },
      -1,
   } },

   /* ('fmax', ('fmin', 'a', 1.0), -1.0) => ('fsat_signed_mali', 'a') */
   /* search2_0_0 -> 0 in the cache */
   /* search2_0_1 -> 6 in the cache */
   { .expression = {
      { nir_search_value_expression, -1 },
      false,
      false,
      false,
      nir_op_fmin,
      1, 1,
      { 0, 6 },
      -1,
   } },
   /* search2_1 -> 4 in the cache */
   { .expression = {
      { nir_search_value_expression, -1 },
      false,
      false,
      false,
      nir_op_fmax,
      0, 2,
      { 9, 4 },
      -1,
   } },

   /* replace2_0 -> 0 in the cache */
   /* replace2 -> 8 in the cache */

   /* ('fmax', 'a', 0.0) => ('fclamp_pos_mali', 'a') */
   /* search3_0 -> 0 in the cache */
   { .constant = {
      { nir_search_value_constant, -1 },
      nir_type_float, { 0x0 /* 0.0 */ },
   } },
   { .expression = {
      { nir_search_value_expression, -1 },
      false,
      false,
      false,
      nir_op_fmax,
      0, 1,
      { 0, 11 },
      -1,
   } },

   /* replace3_0 -> 0 in the cache */
   { .expression = {
      { nir_search_value_expression, -1 },
      false,
      false,
      false,
      nir_op_fclamp_pos_mali,
      -1, 0,
      { 0 },
      -1,
   } },

   /* ('fabs', ('fddx', 'a')) => ('fabs', ('fddx_must_abs_mali', 'a')) */
   /* search4_0_0 -> 0 in the cache */
   { .expression = {
      { nir_search_value_expression, -1 },
      false,
      false,
      false,
      nir_op_fddx,
      -1, 0,
      { 0 },
      -1,
   } },
   { .expression = {
      { nir_search_value_expression, -1 },
      false,
      false,
      false,
      nir_op_fabs,
      -1, 0,
      { 14 },
      -1,
   } },

   /* replace4_0_0 -> 0 in the cache */
   { .expression = {
      { nir_search_value_expression, -1 },
      false,
      false,
      false,
      nir_op_fddx_must_abs_mali,
      -1, 0,
      { 0 },
      -1,
   } },
   { .expression = {
      { nir_search_value_expression, -1 },
      false,
      false,
      false,
      nir_op_fabs,
      -1, 0,
      { 16 },
      -1,
   } },

   /* ('fabs', ('fddy', 'b')) => ('fabs', ('fddy_must_abs_mali', 'b')) */
   { .variable = {
      { nir_search_value_variable, -1 },
      0, /* b */
      false,
      nir_type_invalid,
      -1,
      {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15},
   } },
   { .expression = {
      { nir_search_value_expression, -1 },
      false,
      false,
      false,
      nir_op_fddy,
      -1, 0,
      { 18 },
      -1,
   } },
   { .expression = {
      { nir_search_value_expression, -1 },
      false,
      false,
      false,
      nir_op_fabs,
      -1, 0,
      { 19 },
      -1,
   } },

   /* replace5_0_0 -> 18 in the cache */
   { .expression = {
      { nir_search_value_expression, -1 },
      false,
      false,
      false,
      nir_op_fddy_must_abs_mali,
      -1, 0,
      { 18 },
      -1,
   } },
   { .expression = {
      { nir_search_value_expression, -1 },
      false,
      false,
      false,
      nir_op_fabs,
      -1, 0,
      { 21 },
      -1,
   } },

};



static const struct transform bifrost_nir_lower_algebraic_late_transforms[] = {
   { ~0, ~0, ~0 }, /* Sentinel */

   { 2, 3, 0 },
   { ~0, ~0, ~0 }, /* Sentinel */

   { 12, 13, 0 },
   { ~0, ~0, ~0 }, /* Sentinel */

   { 10, 8, 0 },
   { 12, 13, 0 },
   { ~0, ~0, ~0 }, /* Sentinel */

   { 7, 8, 0 },
   { ~0, ~0, ~0 }, /* Sentinel */

   { 15, 17, 0 },
   { ~0, ~0, ~0 }, /* Sentinel */

   { 20, 22, 0 },
   { ~0, ~0, ~0 }, /* Sentinel */

};

static const struct per_op_table bifrost_nir_lower_algebraic_late_pass_op_table[nir_num_search_ops] = {
   [nir_op_fmul] = {
      .filter = (const uint16_t []) {
         0,
         1,
         0,
         0,
         0,
         0,
         0,
         0,
         0,
         0,
         0,
      },
      
      .num_filtered_states = 2,
      .table = (const uint16_t []) {
      
         0,
         2,
         2,
         2,
      },
   },
   [nir_op_fmin] = {
      .filter = (const uint16_t []) {
         0,
         1,
         0,
         0,
         2,
         0,
         0,
         2,
         0,
         0,
         0,
      },
      
      .num_filtered_states = 3,
      .table = (const uint16_t []) {
      
         0,
         3,
         0,
         3,
         3,
         8,
         0,
         8,
         0,
      },
   },
   [nir_op_fmax] = {
      .filter = (const uint16_t []) {
         0,
         1,
         0,
         2,
         0,
         0,
         0,
         0,
         2,
         0,
         0,
      },
      
      .num_filtered_states = 3,
      .table = (const uint16_t []) {
      
         0,
         4,
         0,
         4,
         4,
         7,
         0,
         7,
         0,
      },
   },
   [nir_op_fabs] = {
      .filter = (const uint16_t []) {
         0,
         0,
         0,
         0,
         0,
         1,
         2,
         0,
         0,
         0,
         0,
      },
      
      .num_filtered_states = 3,
      .table = (const uint16_t []) {
      
         0,
         9,
         10,
      },
   },
   [nir_op_fddx] = {
      .filter = NULL,
      
      .num_filtered_states = 1,
      .table = (const uint16_t []) {
      
         5,
      },
   },
   [nir_op_fddy] = {
      .filter = NULL,
      
      .num_filtered_states = 1,
      .table = (const uint16_t []) {
      
         6,
      },
   },
};

/* Mapping from state index to offset in transforms (0 being no transforms) */
static const uint16_t bifrost_nir_lower_algebraic_late_transform_offsets[] = {
   0,
   0,
   1,
   0,
   3,
   0,
   0,
   5,
   8,
   10,
   12,
};

static const nir_algebraic_table bifrost_nir_lower_algebraic_late_table = {
   .transforms = bifrost_nir_lower_algebraic_late_transforms,
   .transform_offsets = bifrost_nir_lower_algebraic_late_transform_offsets,
   .pass_op_table = bifrost_nir_lower_algebraic_late_pass_op_table,
   .values = bifrost_nir_lower_algebraic_late_values,
   .expression_cond = NULL,
   .variable_cond = NULL,
};

bool
bifrost_nir_lower_algebraic_late(nir_shader *shader)
{
   bool progress = false;
   bool condition_flags[1];
   const nir_shader_compiler_options *options = shader->options;
   const shader_info *info = &shader->info;
   (void) options;
   (void) info;

   STATIC_ASSERT(23 == ARRAY_SIZE(bifrost_nir_lower_algebraic_late_values));
   condition_flags[0] = true;

   nir_foreach_function(function, shader) {
      if (function->impl) {
         progress |= nir_algebraic_impl(function->impl, condition_flags,
                                        &bifrost_nir_lower_algebraic_late_table);
      }
   }

   return progress;
}

