/*
 * Copyright © 2024 Valve Corporation
 * SPDX-License-Identifier: MIT
 */

#include "compiler/spirv/spirv.h"
#include "nir/nir_builder.h"
#include "nir/nir_control_flow.h"
#include "nir/nir_deref.h"

#include "bvh/tu_build_interface.h"
#include "tu_shader.h"

enum rq_intersection_var_index {
   rq_intersection_primitive_id,
   rq_intersection_geometry_id,
   rq_intersection_origin,
   rq_intersection_direction,
   rq_intersection_instance,
   rq_intersection_type_flags,
   rq_intersection_sbt_offset,
   rq_intersection_barycentrics,
   rq_intersection_t,
};

static const glsl_type *
get_rq_intersection_type(void)
{
   struct glsl_struct_field fields[] = {
#define FIELD(_type, _name) \
      [rq_intersection_##_name] = glsl_struct_field(_type, #_name),
      FIELD(glsl_uint_type(), primitive_id)
      FIELD(glsl_uint_type(), geometry_id)
      FIELD(glsl_vec_type(3), origin)
      FIELD(glsl_vec_type(3), direction)
      FIELD(glsl_uint_type(), instance)
      FIELD(glsl_uint_type(), type_flags)
      FIELD(glsl_uint_type(), sbt_offset)
      FIELD(glsl_vec2_type(), barycentrics)
      FIELD(glsl_float_type(), t)
#undef FIELD
   };

   return glsl_struct_type(fields, ARRAY_SIZE(fields), "ray_query_intersection", false);
}

enum rq_var_index {
   rq_index_accel_struct_base,
   rq_index_root_bvh_base,
   rq_index_bvh_base,
   rq_index_flags,
   rq_index_tmin,
   rq_index_world_origin,
   rq_index_world_direction,
   rq_index_incomplete,
   rq_index_closest,
   rq_index_candidate,
   rq_index_stack_ptr,
   rq_index_top_stack,
   rq_index_stack_low_watermark,
   rq_index_current_node,
   rq_index_previous_node,
   rq_index_instance_top_node,
   rq_index_instance_bottom_node,
   rq_index_stack,
};

/* Driver-internal flag to indicate that we haven't found an intersection */
#define TU_INTERSECTION_TYPE_NO_INTERSECTION (1u << 0)

#define MAX_STACK_DEPTH 8

static const glsl_type *
get_rq_type(void)
{
   const glsl_type *intersection_type = get_rq_intersection_type();

   struct glsl_struct_field fields[] = {
#define FIELD(_type, _name) \
      [rq_index_##_name] = glsl_struct_field(_type, #_name),
      FIELD(glsl_uvec2_type(), accel_struct_base)
      FIELD(glsl_uvec2_type(), root_bvh_base)
      FIELD(glsl_uvec2_type(), bvh_base)
      FIELD(glsl_uint_type(), flags)
      FIELD(glsl_float_type(), tmin)
      FIELD(glsl_vec_type(3), world_origin)
      FIELD(glsl_vec_type(3), world_direction)
      FIELD(glsl_bool_type(), incomplete)
      FIELD(intersection_type, closest)
      FIELD(intersection_type, candidate)
      FIELD(glsl_uint_type(), stack_ptr)
      FIELD(glsl_uint_type(), top_stack)
      FIELD(glsl_uint_type(), stack_low_watermark)
      FIELD(glsl_uint_type(), current_node)
      FIELD(glsl_uint_type(), previous_node)
      FIELD(glsl_uint_type(), instance_top_node)
      FIELD(glsl_uint_type(), instance_bottom_node)
      FIELD(glsl_array_type(glsl_uvec2_type(), MAX_STACK_DEPTH, 0), stack)
#undef FIELD
   };

   return glsl_struct_type(fields, ARRAY_SIZE(fields), "ray_query", false);
}

struct rq_var {
   nir_variable *rq;

   nir_intrinsic_instr *initialization;
   nir_def *uav_index;
};

static void
lower_ray_query(nir_shader *shader, nir_function_impl *impl,
                nir_variable *ray_query, struct hash_table *ht)
{
   struct rq_var *var = rzalloc(ht, struct rq_var);
   const glsl_type *type = ray_query->type;

   const glsl_type *rq_type = glsl_type_wrap_in_arrays(get_rq_type(), type);

   if (impl)
      var->rq = nir_local_variable_create(impl, rq_type, "ray_query");
   else
      var->rq = nir_variable_create(shader, nir_var_shader_temp, rq_type, "ray_query");

   _mesa_hash_table_insert(ht, ray_query, var);
}

static nir_deref_instr *
get_rq_deref(nir_builder *b, struct hash_table *ht, nir_def *def,
             struct rq_var **rq_var_out)
{
   nir_deref_instr *deref = nir_def_as_deref(def);

   nir_deref_path path;
   nir_deref_path_init(&path, deref, NULL);
   assert(path.path[0]->deref_type == nir_deref_type_var);

   nir_variable *opaque_var = nir_deref_instr_get_variable(path.path[0]);
   struct hash_entry *entry = _mesa_hash_table_search(ht, opaque_var);
   assert(entry);

   struct rq_var *rq = (struct rq_var *)entry->data;

   nir_deref_instr *out_deref = nir_build_deref_var(b, rq->rq);

   if (glsl_type_is_array(opaque_var->type)) {
      nir_deref_instr **p = &path.path[1];
      for (; *p; p++) {
         if ((*p)->deref_type == nir_deref_type_array) {
            nir_def *index = (*p)->arr.index.ssa;

            out_deref = nir_build_deref_array(b, out_deref, index);
         } else {
            UNREACHABLE("Unsupported deref type");
         }
      }
   }

   nir_deref_path_finish(&path);

   if (rq_var_out)
      *rq_var_out = rq;

   return out_deref;
}

static nir_def *
get_rq_initialize_uav_index(nir_intrinsic_instr *intr, struct rq_var *var)
{
   if (nir_src_is_intrinsic(intr->src[1]) &&
       nir_def_as_intrinsic(intr->src[1].ssa)->intrinsic ==
       nir_intrinsic_load_vulkan_descriptor) {
      return intr->src[1].ssa;
   } else {
      return NULL;
   }
}

/* Before we modify control flow, walk the shader and determine ray query
 * instructions for which we know the ray query has been initialized via a
 * descriptor instead of a pointer, and record the UAV descriptor.
 */
static void
calc_uav_index(nir_function_impl *impl, struct hash_table *ht)
{
   nir_metadata_require(impl, nir_metadata_dominance);

   nir_foreach_block (block, impl) {
      nir_foreach_instr (instr, block) {
         if (instr->type != nir_instr_type_intrinsic)
            continue;

         nir_intrinsic_instr *intr = nir_instr_as_intrinsic(instr);

         nir_def *rq_def;
         switch (intr->intrinsic) {
         case nir_intrinsic_rq_initialize:
         case nir_intrinsic_rq_load:
         case nir_intrinsic_rq_proceed:
            rq_def = intr->src[0].ssa;
            break;
         default:
            continue;
         }

         nir_deref_instr *deref = nir_def_as_deref(rq_def);

         if (deref->deref_type != nir_deref_type_var)
            continue;

         nir_variable *opaque_var = deref->var;
         struct hash_entry *entry = _mesa_hash_table_search(ht, opaque_var);
         assert(entry);

         struct rq_var *rq = (struct rq_var *)entry->data;

         if (intr->intrinsic == nir_intrinsic_rq_initialize) {
            rq->initialization = intr;
            rq->uav_index = get_rq_initialize_uav_index(intr, rq);
         } else {
            if (rq->initialization &&
                nir_block_dominates(rq->initialization->instr.block,
                                    block) && rq->uav_index) {
               _mesa_hash_table_insert(ht, instr, rq->uav_index);
            }
         }
      }
   }
}

/* Return a pointer to the TLAS descriptor, which is actually a UAV
 * descriptor, if we know that the ray query has been initialized via a
 * descriptor and not a pointer. If not known, returns NULL.
 */
static nir_def *
get_uav_index(nir_instr *cur_instr, struct hash_table *ht)
{
   struct hash_entry *entry = _mesa_hash_table_search(ht, cur_instr);
   if (entry)
      return (nir_def *)entry->data;
   return NULL;
}

/* Load some data from the TLAS header or instance descriptors. This uses the
 * UAV descriptor when available, via "uav_index" which should be obtained
 * from get_uav_index().
 */
static nir_def *
load_tlas(nir_builder *b, nir_def *tlas,
          nir_def *uav_index, nir_def *index,
          unsigned offset, unsigned components)
{
   if (uav_index) {
      return nir_load_uav_ir3(b, components, 32, uav_index,
                              nir_vec2(b, index, nir_imm_int(b, offset / 4)),
                              .access = (gl_access_qualifier)(
                                          ACCESS_NON_WRITEABLE |
                                          ACCESS_CAN_REORDER),
                              .align_mul = AS_RECORD_SIZE,
                              .align_offset = offset);
   } else {
      return nir_load_global_ir3(b, components, 32,
                                 nir_pack_64_2x32(b, tlas),
                                 nir_iadd_imm(b, nir_imul_imm(b, index, AS_RECORD_SIZE / 4),
                                              offset / 4),
                                 /* The required alignment of the
                                  * user-specified base from the Vulkan spec.
                                  */
                                 .align_mul = 256,
                                 .align_offset = 0);
   }
}

/* The first record is the TLAS header and the rest of the records are
 * instances, so we need to add 1 to the instance ID when reading data in an
 * instance.
 */
#define load_instance_offset(b, tlas, uav_index, instance, \
                             field, offset, components) \
   load_tlas(b, tlas, uav_index, nir_iadd_imm(b, instance, 1), \
             offsetof(struct tu_instance_descriptor, field) + offset, \
             components)

#define load_instance(b, tlas, uav_index, instance, field, components) \
   load_instance_offset(b, tlas, uav_index, instance, field, 0, components)

#define rq_deref(b, _rq, name) nir_build_deref_struct(b, _rq, rq_index_##name)
#define rq_load(b, _rq, name) nir_load_deref(b, rq_deref(b, _rq, name))
#define rq_store(b, _rq, name, val, wrmask) \
   nir_store_deref(b, rq_deref(b, _rq, name), val, wrmask)
#define rqi_deref(b, _rq, name) nir_build_deref_struct(b, _rq, rq_intersection_##name)
#define rqi_load(b, _rq, name) nir_load_deref(b, rqi_deref(b, _rq, name))
#define rqi_store(b, _rq, name, val, wrmask) \
   nir_store_deref(b, rqi_deref(b, _rq, name), val, wrmask)

static void
lower_rq_initialize(nir_builder *b, struct hash_table *ht,
                    nir_intrinsic_instr *intr)
{
   struct rq_var *var;
   nir_deref_instr *rq = get_rq_deref(b, ht, intr->src[0].ssa, &var);

   if (nir_def_as_deref(intr->src[0].ssa)->deref_type ==
       nir_deref_type_var) {
      var->initialization = intr;
   } else {
      var->initialization = NULL;
   }

   nir_def *uav_index = get_rq_initialize_uav_index(intr, var);

   nir_def *tlas = intr->src[1].ssa;
   nir_def *flags = intr->src[2].ssa;
   nir_def *cull_mask = intr->src[3].ssa;
   nir_def *origin = intr->src[4].ssa;
   nir_def *tmin = intr->src[5].ssa;
   nir_def *direction = intr->src[6].ssa;
   nir_def *tmax = intr->src[7].ssa;

   nir_def *tlas_base;
   if (uav_index) {
      tlas_base = load_tlas(b, NULL, uav_index, nir_imm_int(b, 0),
                            offsetof(struct tu_accel_struct_header,
                                     self_ptr), 2);
   } else {
      tlas_base = nir_unpack_64_2x32(b, tlas);
   }

   rq_store(b, rq, accel_struct_base, tlas_base, 0x3);

   nir_def *root_bvh_base = load_tlas(b, tlas_base, uav_index, nir_imm_int(b, 0),
                                      offsetof(struct tu_accel_struct_header,
                                               bvh_ptr), 2);

   nir_deref_instr *closest = rq_deref(b, rq, closest);
   nir_deref_instr *candidate = rq_deref(b, rq, candidate);

   rq_store(b, rq, flags,
            /* Fill out initial fourth src of ray_intersection */
            nir_ior_imm(b,
                        nir_ior(b, nir_ishl_imm(b, flags, 4),
                                nir_ishl_imm(b, cull_mask, 16)),
                        0b1111), 0x1);

   rqi_store(b, candidate, origin, origin, 0x7);
   rqi_store(b, candidate, direction, direction, 0x7);

   rq_store(b, rq, tmin, tmin, 0x1);
   rq_store(b, rq, world_origin, origin, 0x7);
   rq_store(b, rq, world_direction, direction, 0x7);

   rqi_store(b, closest, t, tmax, 0x1);
   rqi_store(b, closest, type_flags, nir_imm_int(b, TU_INTERSECTION_TYPE_NO_INTERSECTION), 0x1);

   /* Make sure that instance data loads don't hang in case of a miss by setting a valid initial instance. */
   rqi_store(b, closest, instance, nir_imm_int(b, 0), 0x1);
   rqi_store(b, candidate, instance, nir_imm_int(b, 0), 0x1);

   rq_store(b, rq, root_bvh_base, root_bvh_base, 0x3);
   rq_store(b, rq, bvh_base, root_bvh_base, 0x3);

   rq_store(b, rq, stack_ptr, nir_imm_int(b, 0), 0x1);
   rq_store(b, rq, top_stack, nir_imm_int(b, -1), 0x1);
   rq_store(b, rq, stack_low_watermark, nir_imm_int(b, 0), 0x1);
   rq_store(b, rq, current_node, nir_imm_int(b, 0), 0x1);
   rq_store(b, rq, previous_node, nir_imm_int(b, VK_BVH_INVALID_NODE), 0x1);
   rq_store(b, rq, instance_top_node, nir_imm_int(b, VK_BVH_INVALID_NODE), 0x1);
   rq_store(b, rq, instance_bottom_node, nir_imm_int(b, VK_BVH_INVALID_NODE), 0x1);

   rq_store(b, rq, incomplete, nir_imm_true(b), 0x1);
}

static void
insert_terminate_on_first_hit(nir_builder *b, nir_deref_instr *rq)
{
   nir_def *terminate_on_first_hit =
      nir_test_mask(b, rq_load(b, rq, flags),
                    SpvRayFlagsTerminateOnFirstHitKHRMask << 4);
   nir_push_if(b, terminate_on_first_hit);
   {
      rq_store(b, rq, incomplete, nir_imm_false(b), 0x1);
   }
   nir_pop_if(b, NULL);
}

static void
lower_rq_confirm_intersection(nir_builder *b, struct hash_table *ht, nir_intrinsic_instr *intr)
{
   nir_deref_instr *rq = get_rq_deref(b, ht, intr->src[0].ssa, NULL);
   nir_copy_deref(b, rq_deref(b, rq, closest), rq_deref(b, rq, candidate));
   insert_terminate_on_first_hit(b, rq);
}

static void
lower_rq_generate_intersection(nir_builder *b, struct hash_table *ht, nir_intrinsic_instr *intr)
{
   nir_deref_instr *rq = get_rq_deref(b, ht, intr->src[0].ssa, NULL);
   nir_deref_instr *closest = rq_deref(b, rq, closest);
   nir_deref_instr *candidate = rq_deref(b, rq, candidate);

   nir_push_if(b, nir_iand(b, nir_fge(b, rqi_load(b, closest, t),
                                      intr->src[1].ssa),
                           nir_fge(b, intr->src[1].ssa,
                                   rq_load(b, rq, tmin))));
   {
      nir_copy_deref(b, closest, candidate);
      insert_terminate_on_first_hit(b, rq);
      rqi_store(b, closest, t, intr->src[1].ssa, 0x1);
   }
   nir_pop_if(b, NULL);
}

static void
lower_rq_terminate(nir_builder *b, struct hash_table *ht, nir_intrinsic_instr *intr)
{
   nir_deref_instr *rq = get_rq_deref(b, ht, intr->src[0].ssa, NULL);
   rq_store(b, rq, incomplete, nir_imm_false(b), 0x1);
}

static nir_def *
lower_rq_load(nir_builder *b, struct hash_table *ht, nir_intrinsic_instr *intr)
{
   struct rq_var *var;
   nir_deref_instr *rq = get_rq_deref(b, ht, intr->src[0].ssa, &var);
   nir_def *uav_index = get_uav_index(&intr->instr, ht);
   nir_def *tlas = rq_load(b, rq, accel_struct_base);
   nir_deref_instr *closest = rq_deref(b, rq, closest);
   nir_deref_instr *candidate = rq_deref(b, rq, candidate);
   bool committed = nir_intrinsic_committed(intr);
   nir_deref_instr *intersection = committed ? closest : candidate;

   uint32_t column = nir_intrinsic_column(intr);

   nir_ray_query_value value = nir_intrinsic_ray_query_value(intr);
   switch (value) {
   case nir_ray_query_value_flags: {
      nir_def *flags = rq_load(b, rq, flags);
      return nir_ubitfield_extract(b, flags, nir_imm_int(b, 4),
                                   nir_imm_int(b, 12));
   }
   case nir_ray_query_value_intersection_barycentrics:
      return rqi_load(b, intersection, barycentrics);
   case nir_ray_query_value_intersection_candidate_aabb_opaque:
      return nir_ieq_imm(b, nir_iand_imm(b, rqi_load(b, candidate, type_flags),
                                         TU_INTERSECTION_TYPE_AABB |
                                         TU_INTERSECTION_TYPE_NONOPAQUE |
                                         TU_INTERSECTION_TYPE_NO_INTERSECTION),
                         TU_INTERSECTION_TYPE_AABB);
   case nir_ray_query_value_intersection_front_face:
      return nir_inot(b, nir_test_mask(b, rqi_load(b, intersection, type_flags),
                                       TU_INTERSECTION_BACK_FACE));
   case nir_ray_query_value_intersection_geometry_index:
      return rqi_load(b, intersection, geometry_id);
   case nir_ray_query_value_intersection_instance_custom_index: {
      nir_def *instance = rqi_load(b, intersection, instance);
      return load_instance(b, tlas, uav_index, instance, custom_instance_index, 1);
   }
   case nir_ray_query_value_intersection_instance_id:
      return rqi_load(b, intersection, instance);
   case nir_ray_query_value_intersection_instance_sbt_index:
      return rqi_load(b, intersection, sbt_offset);
   case nir_ray_query_value_intersection_object_ray_direction:
      return rqi_load(b, intersection, direction);
   case nir_ray_query_value_intersection_object_ray_origin:
      return rqi_load(b, intersection, origin);
   case nir_ray_query_value_intersection_object_to_world: {
      nir_def *instance = rqi_load(b, intersection, instance);
      nir_def *rows[3];
      for (unsigned r = 0; r < 3; ++r)
         rows[r] = load_instance_offset(b, tlas, uav_index, instance,
                                        otw_matrix.values,
                                        r * 16, 4);

      return nir_vec3(b, nir_channel(b, rows[0], column), nir_channel(b, rows[1], column),
                      nir_channel(b, rows[2], column));
   }
   case nir_ray_query_value_intersection_primitive_index:
      return rqi_load(b, intersection, primitive_id);
   case nir_ray_query_value_intersection_t:
      return rqi_load(b, intersection, t);
   case nir_ray_query_value_intersection_type: {
      nir_def *intersection_type =
         nir_iand_imm(b, nir_ishr_imm(b, rqi_load(b, intersection, type_flags),
                                      util_logbase2(TU_INTERSECTION_TYPE_AABB)), 1);
      if (committed) {
         nir_def *has_intersection =
            nir_inot(b,
                     nir_test_mask(b, rqi_load(b, intersection, type_flags),
                                   TU_INTERSECTION_TYPE_NO_INTERSECTION));
         intersection_type = nir_iadd(b, intersection_type,
                                      nir_b2i32(b, has_intersection));
      }
      return intersection_type;
   }
   case nir_ray_query_value_intersection_world_to_object: {
      nir_def *instance = rqi_load(b, intersection, instance);
      nir_def *rows[3];
      for (unsigned r = 0; r < 3; ++r)
         rows[r] = load_instance_offset(b, tlas, uav_index, instance,
                                        wto_matrix.values, r * 16, 4);

      return nir_vec3(b, nir_channel(b, rows[0], column), nir_channel(b, rows[1], column),
                      nir_channel(b, rows[2], column));
   }
   case nir_ray_query_value_tmin:
      return rq_load(b, rq, tmin);
   case nir_ray_query_value_world_ray_direction:
      return rq_load(b, rq, world_direction);
   case nir_ray_query_value_world_ray_origin:
      return rq_load(b, rq, world_origin);
   default:
      UNREACHABLE("Invalid nir_ray_query_value!");
   }

   return NULL;
}

/* For the initialization of instance_bottom_node. Explicitly different than
 * VK_BVH_INVALID_NODE or any real node, to ensure we never exit an instance
 * when we're not in one.
 */
#define TU_BVH_NO_INSTANCE_ROOT 0xfffffffeu

nir_def *
nir_build_vec3_mat_mult(nir_builder *b, nir_def *vec, nir_def *matrix[], bool translation)
{
   nir_def *result_components[3] = {
      nir_channel(b, matrix[0], 3),
      nir_channel(b, matrix[1], 3),
      nir_channel(b, matrix[2], 3),
   };
   for (unsigned i = 0; i < 3; ++i) {
      for (unsigned j = 0; j < 3; ++j) {
         nir_def *v = nir_fmul(b, nir_channels(b, vec, 1 << j), nir_channels(b, matrix[i], 1 << j));
         result_components[i] = (translation || j) ? nir_fadd(b, result_components[i], v) : v;
      }
   }
   return nir_vec(b, result_components, 3);
}

static nir_def *
fetch_parent_node(nir_builder *b, nir_def *bvh, nir_def *node)
{
   nir_def *offset = nir_iadd_imm(b, nir_imul_imm(b, node, 4), 4);

   return nir_load_global(
      b, 1, 32, nir_isub(b, nir_pack_64_2x32(b, bvh), nir_u2u64(b, offset)),
      .align_mul = 4);
}

static nir_def *
build_ray_traversal(nir_builder *b, nir_deref_instr *rq,
                    nir_def *tlas, nir_def *uav_index)
{
   nir_deref_instr *candidate = rq_deref(b, rq, candidate);
   nir_deref_instr *closest = rq_deref(b, rq, closest);

   nir_variable *incomplete = nir_local_variable_create(b->impl, glsl_bool_type(), "incomplete");
   nir_store_var(b, incomplete, nir_imm_true(b), 0x1);

   nir_loop *loop = nir_push_loop(b);
   {
      nir_loop_add_continue_construct(loop);

      /* Go up the stack if current_node == VK_BVH_INVALID_NODE */
      nir_push_if(b, nir_ieq_imm(b, rq_load(b, rq, current_node), VK_BVH_INVALID_NODE));
      {
         /* Early exit if we never overflowed the stack, to avoid having to backtrack to
          * the root for no reason. */
         nir_push_if(b, nir_ilt_imm(b, rq_load(b, rq, stack_ptr), 1));
         {
            nir_store_var(b, incomplete, nir_imm_false(b), 0x1);
            nir_jump(b, nir_jump_break);
         }
         nir_pop_if(b, NULL);

         nir_def *stack_instance_exit =
            nir_ige(b, rq_load(b, rq, top_stack), rq_load(b, rq, stack_ptr));
         nir_def *root_instance_exit =
            nir_ieq(b, rq_load(b, rq, previous_node), rq_load(b, rq, instance_bottom_node));
         nir_if *instance_exit = nir_push_if(b, nir_ior(b, stack_instance_exit, root_instance_exit));
         instance_exit->control = nir_selection_control_dont_flatten;
         {
            rq_store(b, rq, top_stack, nir_imm_int(b, -1), 1);
            rq_store(b, rq, previous_node, rq_load(b, rq, instance_top_node), 1);
            rq_store(b, rq, instance_bottom_node, nir_imm_int(b, TU_BVH_NO_INSTANCE_ROOT), 1);

            rq_store(b, rq, bvh_base, rq_load(b, rq, root_bvh_base), 3);
            rqi_store(b, candidate, origin, rq_load(b, rq, world_origin), 7);
            rqi_store(b, candidate, direction, rq_load(b, rq, world_direction), 7);
         }
         nir_pop_if(b, NULL);

         nir_push_if(
            b, nir_ige(b, rq_load(b, rq, stack_low_watermark), rq_load(b, rq, stack_ptr)));
         {
            /* Get the parent of the previous node using the parent pointers.
             * We will re-intersect the parent and figure out what index the
             * previous node was below.
             */
            nir_def *prev = rq_load(b, rq, previous_node);
            nir_def *bvh_addr = rq_load(b, rq, bvh_base);

            nir_def *parent = fetch_parent_node(b, bvh_addr, prev);
            nir_push_if(b, nir_ieq_imm(b, parent, VK_BVH_INVALID_NODE));
            {
               nir_store_var(b, incomplete, nir_imm_false(b), 0x1);
               nir_jump(b, nir_jump_break);
            }
            nir_pop_if(b, NULL);
            rq_store(b, rq, current_node, parent, 0x1);
         }
         nir_push_else(b, NULL);
         {
            /* Go up the stack and get the next child of the parent. */
            nir_def *stack_ptr = nir_iadd_imm(b, rq_load(b, rq, stack_ptr), -1);

            nir_def *stack_idx =
               nir_umod_imm(b, stack_ptr, MAX_STACK_DEPTH);
            nir_deref_instr *stack_deref =
               nir_build_deref_array(b, rq_deref(b, rq, stack), stack_idx);
            nir_def *stack_entry = nir_load_deref(b, stack_deref);
            nir_def *children_base = nir_channel(b, stack_entry, 0);
            nir_def *children = nir_channel(b, stack_entry, 1);

            nir_def *next_child_idx =
               nir_iadd_imm(b, nir_iand_imm(b, children, 0x1f), -3);
            
            nir_def *child_offset =
               nir_iand_imm(b, nir_ishr(b, children, next_child_idx), 0x7);
            nir_def *bvh_node = nir_iadd(b, children_base, child_offset);

            nir_push_if(b, nir_ieq_imm(b, next_child_idx, 8));
            {
               rq_store(b, rq, stack_ptr, stack_ptr, 1);
            }
            nir_push_else(b, NULL);
            {
               children = nir_bitfield_insert(b, children, next_child_idx,
                                              nir_imm_int(b, 0),
                                              nir_imm_int(b, 5));
               nir_store_deref(b, stack_deref,
                               nir_vec2(b, nir_undef(b, 1, 32), children),
                               0x2);
            }
            nir_pop_if(b, NULL);

            rq_store(b, rq, current_node, bvh_node, 0x1);
            /* We don't need previous_node when we have the stack. Indicate to
             * the internal intersection handling below that this isn't the
             * underflow case.
             */
            rq_store(b, rq, previous_node, nir_imm_int(b, VK_BVH_INVALID_NODE), 0x1);
         }
         nir_pop_if(b, NULL);
      }
      nir_push_else(b, NULL);
      {
         rq_store(b, rq, previous_node, nir_imm_int(b, VK_BVH_INVALID_NODE), 0x1);
      }
      nir_pop_if(b, NULL);

      nir_def *bvh_node = rq_load(b, rq, current_node);
      nir_def *bvh_base = rq_load(b, rq, bvh_base);

      nir_def *prev_node = rq_load(b, rq, previous_node);
      rq_store(b, rq, previous_node, bvh_node, 0x1);
      rq_store(b, rq, current_node, nir_imm_int(b, VK_BVH_INVALID_NODE), 0x1);

      nir_def *origin = rqi_load(b, candidate, origin);
      nir_def *tmin = rq_load(b, rq, tmin);
      nir_def *direction = rqi_load(b, candidate, direction);
      nir_def *tmax = rqi_load(b, closest, t);

      nir_def *intrinsic_result =
         nir_ray_intersection_ir3(b, 32, bvh_base, bvh_node,
                                  nir_vec8(b,
                                           nir_channel(b, origin, 0),
                                           nir_channel(b, origin, 1),
                                           nir_channel(b, origin, 2),
                                           tmin,
                                           nir_channel(b, direction, 0),
                                           nir_channel(b, direction, 1),
                                           nir_channel(b, direction, 2),
                                           tmax),
                                  rq_load(b, rq, flags));

      nir_def *intersection_flags = nir_channel(b, intrinsic_result, 0);
      nir_def *intersection_count =
         nir_ubitfield_extract_imm(b, intersection_flags, 4, 4);
      nir_def *intersection_id = nir_channel(b, intrinsic_result, 1);

      nir_push_if(b, nir_test_mask(b, intersection_flags,
                                   TU_INTERSECTION_TYPE_LEAF));
      {
         nir_def *processed_mask = nir_iand_imm(b, intersection_flags, 0xf);

         /* Keep processing the current node if the mask isn't yet 0 */
         rq_store(b, rq, current_node,
                  nir_bcsel(b, nir_ieq_imm(b, processed_mask, 0),
                            nir_imm_int(b, VK_BVH_INVALID_NODE),
                            bvh_node), 1);

         /* If the mask is 0, replace with the initial 0xf for the next
          * intersection.
          */
         processed_mask =
            nir_bcsel(b, nir_ieq_imm(b, processed_mask, 0),
                      nir_imm_int(b, 0xf), processed_mask);
         
         /* Replace the mask in the flags. */
         rq_store(b, rq, flags,
                  nir_bitfield_insert(b, rq_load(b, rq, flags),
                                      processed_mask, nir_imm_int(b, 0),
                                      nir_imm_int(b, 4)), 1);

         nir_push_if(b, nir_ieq_imm(b, intersection_count, 0));
         {
            nir_jump(b, nir_jump_continue);
         }
         nir_pop_if(b, NULL);

         nir_push_if(b, nir_test_mask(b, intersection_flags,
                                      TU_INTERSECTION_TYPE_TLAS));
         {
            /* instance */
            rqi_store(b, candidate, instance, intersection_id, 1);

            nir_def *wto_matrix[3];
            for (unsigned i = 0; i < 3; i++)
               wto_matrix[i] = load_instance_offset(b, tlas, uav_index,
                                                    intersection_id,
                                                    wto_matrix.values,
                                                    i * 16, 4);

            nir_def *sbt_offset_and_flags =
               load_instance(b, tlas, uav_index, intersection_id,
                             sbt_offset_and_flags, 1);
            nir_def *blas_bvh =
               load_instance(b, tlas, uav_index, intersection_id,
                             bvh_ptr, 2);

            nir_def *instance_flags = nir_iand_imm(b, sbt_offset_and_flags,
                                                   0xff000000);
            nir_def *sbt_offset = nir_iand_imm(b, sbt_offset_and_flags,
                                               0x00ffffff);
            nir_def *flags = rq_load(b, rq, flags);
            flags = nir_ior(b, nir_iand_imm(b, flags, 0x00ffffff),
                            instance_flags);
            rq_store(b, rq, flags, flags, 1);

            rqi_store(b, candidate, sbt_offset, sbt_offset, 1);

            rq_store(b, rq, top_stack, rq_load(b, rq, stack_ptr), 1);
            rq_store(b, rq, bvh_base, blas_bvh, 3);

            /* Push the instance root node onto the stack */
            rq_store(b, rq, current_node, nir_imm_int(b, 0), 0x1);
            rq_store(b, rq, instance_bottom_node, nir_imm_int(b, 0), 1);
            rq_store(b, rq, instance_top_node, bvh_node, 1);

            /* Transform the ray into object space */
            rqi_store(b, candidate, origin,
                      nir_build_vec3_mat_mult(b, rq_load(b, rq, world_origin),
                                              wto_matrix, true), 7);
            rqi_store(b, candidate, direction,
                      nir_build_vec3_mat_mult(b, rq_load(b, rq, world_direction),
                                              wto_matrix, false), 7);
         }
         nir_push_else(b, NULL);
         {
            /* AABB & triangle */
            rqi_store(b, candidate, type_flags,
                      nir_iand_imm(b, intersection_flags,
                                   TU_INTERSECTION_TYPE_AABB |
                                   TU_INTERSECTION_TYPE_NONOPAQUE |
                                   TU_INTERSECTION_BACK_FACE), 1);

            rqi_store(b, candidate, primitive_id, intersection_id, 1);

            /* TODO: Implement optimization to try to combine these into 1
             * 32-bit ID, for compressed nodes.
             *
             * load_global_ir3 doesn't have the required range so we have to
             * do the offset math ourselves.
             */
            nir_def *offset =
               nir_ior_imm(b, nir_imul_imm(b, nir_u2u64(b, bvh_node),
                                            sizeof(tu_leaf_node)),
                           offsetof(struct tu_leaf_node, geometry_id));
            nir_def *geometry_id_ptr = nir_iadd(b, nir_pack_64_2x32(b, bvh_base),
                                                offset);
            nir_def *geometry_id = nir_load_global(
               b, 1, 32, geometry_id_ptr, .access = ACCESS_NON_WRITEABLE,
               .align_mul = sizeof(struct tu_leaf_node),
               .align_offset = offsetof(struct tu_leaf_node, geometry_id));
            rqi_store(b, candidate, geometry_id, geometry_id, 1);

            nir_push_if(b, nir_test_mask(b, intersection_flags,
                                         TU_INTERSECTION_TYPE_AABB));
            {
               nir_jump(b, nir_jump_break);
            }
            nir_push_else(b, NULL);
            {
               rqi_store(b, candidate, barycentrics,
                         nir_vec2(b, nir_channel(b, intrinsic_result, 3),
                                  nir_channel(b, intrinsic_result, 4)), 0x3);
               rqi_store(b, candidate, t, nir_channel(b, intrinsic_result,
                                                      2), 0x1);
               nir_push_if(b, nir_test_mask(b, intersection_flags,
                                            TU_INTERSECTION_TYPE_NONOPAQUE));
               {
                  nir_jump(b, nir_jump_break);
               }
               nir_push_else(b, NULL);
               {
                  nir_copy_deref(b, closest, candidate);
                  nir_def *terminate_on_first_hit =
                     nir_test_mask(b, rq_load(b, rq, flags),
                                   SpvRayFlagsTerminateOnFirstHitKHRMask << 4);
                  nir_push_if(b, terminate_on_first_hit);
                  {
                     nir_store_var(b, incomplete, nir_imm_false(b), 0x1);
                     nir_jump(b, nir_jump_break);
                  }
                  nir_pop_if(b, NULL);
               }
               nir_pop_if(b, NULL);
            }
            nir_pop_if(b, NULL);
         }
         nir_pop_if(b, NULL);
      }
      nir_push_else(b, NULL);
      {
         /* internal */
         nir_push_if(b, nir_ine_imm(b, intersection_count, 0));
         {
            nir_def *children = nir_channel(b, intrinsic_result, 3);

            nir_push_if(b, nir_ieq_imm(b, prev_node, VK_BVH_INVALID_NODE));
            {
               /* The children array returned by the HW is specially set up so
                * that we can do this to get the first child.
                */
               nir_def *first_child_offset =
                  nir_iand_imm(b, nir_ishr(b, children, children), 0x7);

               rq_store(b, rq, current_node,
                        nir_iadd(b, intersection_id, first_child_offset),
                        0x1);

               nir_push_if(b, nir_igt_imm(b, intersection_count, 1));
               {
                  nir_def *stack_ptr = rq_load(b, rq, stack_ptr);
                  nir_def *stack_idx = nir_umod_imm(b, stack_ptr,
                                                    MAX_STACK_DEPTH);
                  nir_def *stack_entry =
                     nir_vec2(b, intersection_id, children);
                  nir_store_deref(b,
                                  nir_build_deref_array(b, rq_deref(b, rq, stack),
                                                        stack_idx),
                                  stack_entry, 0x7);
                  rq_store(b, rq, stack_ptr,
                           nir_iadd_imm(b, rq_load(b, rq, stack_ptr), 1), 0x1);

                  nir_def *new_watermark =
                     nir_iadd_imm(b, rq_load(b, rq, stack_ptr),
                                  -MAX_STACK_DEPTH);
                  new_watermark = nir_imax(b, rq_load(b, rq,
                                                      stack_low_watermark),
                                           new_watermark);
                  rq_store(b, rq, stack_low_watermark, new_watermark, 0x1);
               }
               nir_pop_if(b, NULL);
            }
            nir_push_else(b, NULL);
            {
               /* The underflow case. We have the previous_node and an array
                * of intersecting children of its parent, and we need to find
                * its position in the array so that we can return the next
                * child in the array or VK_BVH_INVALID_NODE if it's the last
                * child.
                */
               nir_def *prev_offset =
                  nir_isub(b, prev_node, intersection_id);

               /* A bit-pattern with ones at the LSB of each child's
                * position.
                */
               uint32_t ones = 0b1001001001001001001001 << 8;

               /* Replicate prev_offset into the position of each child. */
               nir_def *prev_offset_repl =
                  nir_imul_imm(b, prev_offset, ones);

               /* a == b <=> a ^ b == 0. Reduce the problem to finding the
                * child whose bits are 0.
                */
               nir_def *diff = nir_ixor(b, prev_offset_repl, children);

               /* This magic formula comes from Hacker's Delight, section 6.1
                * "Find First 0-byte", adapted for 3-bit "bytes". The first
                * zero byte will be the lowest byte with 1 set in the highest
                * position (i.e. bit 2). We need to then subtract 2 to get the
                * current position and 5 to get the next position.
                */
               diff = nir_iand_imm(b, nir_iand(b, nir_iadd_imm(b, diff, -ones),
                                               nir_inot(b, diff)),
                                   ones << 2);
               diff = nir_find_lsb(b, diff);

               nir_def *next_offset =
                  nir_iand_imm(b, nir_ishr(b, children,
                                           nir_iadd_imm(b, diff, -5)),
                               0x7);

               nir_def *next =
                  nir_bcsel(b, nir_ieq_imm(b, diff, 8 + 2),
                            nir_imm_int(b, VK_BVH_INVALID_NODE),
                            nir_iadd(b, next_offset, intersection_id));
               rq_store(b, rq, current_node, next, 0x1);
            }
            nir_pop_if(b, NULL);
         }
         nir_pop_if(b, NULL);
      }
      nir_pop_if(b, NULL);
   }
   nir_pop_loop(b, loop);

   return nir_load_var(b, incomplete);
}

static nir_def *
lower_rq_proceed(nir_builder *b, struct hash_table *ht, nir_intrinsic_instr *intr)
{
   struct rq_var *var;
   nir_deref_instr *rq = get_rq_deref(b, ht, intr->src[0].ssa, &var);
   nir_def *uav_index = get_uav_index(&intr->instr, ht);
   nir_def *tlas = rq_load(b, rq, accel_struct_base);

   nir_push_if(b, nir_load_deref(b, rq_deref(b, rq, incomplete)));
   {
      nir_def *incomplete = build_ray_traversal(b, rq, tlas, uav_index);
      nir_store_deref(b, rq_deref(b, rq, incomplete), incomplete, 0x1);
   }
   nir_pop_if(b, NULL);

   return nir_load_deref(b, rq_deref(b, rq, incomplete));
}

bool
tu_nir_lower_ray_queries(nir_shader *shader)
{
   bool progress = false;
   struct hash_table *query_ht = _mesa_pointer_hash_table_create(NULL);

   nir_foreach_variable_in_list (var, &shader->variables) {
      if (!var->data.ray_query)
         continue;

      lower_ray_query(shader, NULL, var, query_ht);

      progress = true;
   }

   nir_foreach_function (function, shader) {
      if (!function->impl)
         continue;

      nir_builder builder = nir_builder_create(function->impl);

      nir_foreach_variable_in_list (var, &function->impl->locals) {
         if (!var->data.ray_query)
            continue;

         lower_ray_query(shader, function->impl, var, query_ht);

         progress = true;
      }

      calc_uav_index(function->impl, query_ht);

      nir_foreach_block (block, function->impl) {
         nir_foreach_instr_safe (instr, block) {
            if (instr->type != nir_instr_type_intrinsic)
               continue;

            nir_intrinsic_instr *intrinsic = nir_instr_as_intrinsic(instr);

            if (!nir_intrinsic_is_ray_query(intrinsic->intrinsic))
               continue;

            builder.cursor = nir_before_instr(instr);

            nir_def *new_dest = NULL;

            switch (intrinsic->intrinsic) {
            case nir_intrinsic_rq_confirm_intersection:
               lower_rq_confirm_intersection(&builder, query_ht, intrinsic);
               break;
            case nir_intrinsic_rq_generate_intersection:
               lower_rq_generate_intersection(&builder, query_ht, intrinsic);
               break;
            case nir_intrinsic_rq_initialize:
               lower_rq_initialize(&builder, query_ht, intrinsic);
               break;
            case nir_intrinsic_rq_load:
               new_dest = lower_rq_load(&builder, query_ht, intrinsic);
               break;
            case nir_intrinsic_rq_proceed:
               new_dest = lower_rq_proceed(&builder, query_ht, intrinsic);
               break;
            case nir_intrinsic_rq_terminate:
               lower_rq_terminate(&builder, query_ht, intrinsic);
               break;
            default:
               UNREACHABLE("Unsupported ray query intrinsic!");
            }

            if (new_dest)
               nir_def_rewrite_uses(&intrinsic->def, new_dest);

            nir_instr_remove(instr);
            nir_instr_free(instr);

            progress = true;
         }
      }

      nir_progress(true, function->impl, nir_metadata_none);
   }

   ralloc_free(query_ht);

   if (progress)
      nir_lower_continue_constructs(shader);

   return progress;
}

