/*
 * Copyright 2024 Valve Corpoation
 * SPDX-License-Identifier: MIT
 */

#include "nir.h"
#include "nir_builder.h"
#include "nir_builder_opcodes.h"

/**
 * If load_frag_coord.xy is only used by conversions to integer,
 * replace it with load_pixel_coord.
 */

/* Return true if all uses of component_mask components ignore the fractional
 * part.
 *
 * Return false if any uses of components outside component_mask exist.
 */
bool
nir_all_uses_of_float_are_integer(nir_def *def, unsigned component_mask)
{
   /* Check if xy are only used by casts to integers. */
   nir_foreach_use(use, def) {
      /* Ifs aren't expected to source float. */
      assert(!nir_src_is_if(use));

      unsigned mask = nir_src_components_read(use);
      if (!(mask & component_mask))
         continue;

      /* Reject instructions that use more components than expected.
       *
       * (e.g. frag_coord that has zw uses is always pessimistically treated
       * as having no integer use)
       */
      if (mask & ~component_mask)
         return false;

      nir_instr *use_instr = nir_src_use_instr(use);

      if (use_instr->type != nir_instr_type_alu)
         return false;

      switch (nir_instr_as_alu(use_instr)->op) {
      case nir_op_f2i8:
      case nir_op_f2i16:
      case nir_op_f2i32:
      case nir_op_f2i64:
      case nir_op_f2u8:
      case nir_op_f2u16:
      case nir_op_f2u32:
      case nir_op_f2u64:
      case nir_op_ftrunc:
      case nir_op_ffloor:
         continue;
      default:
         return false;
      }
   }

   return true;
}

static bool
opt_frag_pos(nir_builder *b, nir_intrinsic_instr *intr, UNUSED void *data)
{
   if (intr->intrinsic != nir_intrinsic_load_frag_coord &&
       intr->intrinsic != nir_intrinsic_load_frag_coord_xy)
      return false;

   /* Don't increase precision. */
   if (intr->def.bit_size != 32)
      return false;

   if (!nir_all_uses_of_float_are_integer(&intr->def, 0x3))
      return false;

   b->cursor = nir_before_instr(&intr->instr);
   nir_def *pixel_coord = nir_load_pixel_coord(b);

   nir_foreach_use_safe(use, &intr->def) {
      if (intr->intrinsic == nir_intrinsic_load_frag_coord) {
         unsigned mask = nir_src_components_read(use);

         if (!(mask & 0x3))
            continue;

         /* This is implied by nir_all_uses_of_float_are_integer. */
         assert(!(mask & ~0x3));
      }

      nir_src_rewrite(use, pixel_coord);

      nir_alu_instr *use_instr = nir_instr_as_alu(nir_src_use_instr(use));

      /* load_frag_coord is always positive, so we should never sign extend here. */
      bool needs_float = use_instr->op == nir_op_ffloor || use_instr->op == nir_op_ftrunc;
      nir_alu_type dst_type = (needs_float ? nir_type_float : nir_type_uint) | use_instr->def.bit_size;
      use_instr->op = nir_type_conversion_op(nir_type_uint16, dst_type, nir_rounding_mode_undef);
      use_instr->fp_math_ctrl = nir_op_valid_fp_math_ctrl(use_instr->op, use_instr->fp_math_ctrl);
   }

   return true;
}

bool
nir_opt_frag_coord_to_pixel_coord(nir_shader *shader)
{
   return nir_shader_intrinsics_pass(shader, opt_frag_pos,
                                     nir_metadata_control_flow,
                                     NULL);
}
