/*
 * Copyright © 2021 Google
 * Copyright © 2023 Valve Corporation
 * SPDX-License-Identifier: MIT
 */

#ifndef LVP_NIR_RAY_TRACING_H
#define LVP_NIR_RAY_TRACING_H

#include "nir/nir.h"
#include "nir/nir_builder.h"

nir_def *lvp_mul_vec3_mat(nir_builder *b, nir_def *vec, nir_def *matrix[], bool translation);

void lvp_load_wto_matrix(nir_builder *b, nir_def *instance_addr, nir_def **node_data, nir_def **out);

nir_def *lvp_load_vertex_position(nir_builder *b, nir_def *primitive_addr, uint32_t index);

struct lvp_ray_traversal_args;

struct lvp_ray_flags {
   nir_def *force_opaque;
   nir_def *force_not_opaque;
   nir_def *terminate_on_first_hit;
   nir_def *no_cull_front;
   nir_def *no_cull_back;
   nir_def *no_cull_opaque;
   nir_def *no_cull_no_opaque;
   nir_def *no_skip_triangles;
   nir_def *no_skip_aabbs;
};

struct lvp_leaf_intersection {
   nir_def *node_addr;
   nir_def *primitive_id;
   nir_def *geometry_id_and_flags;
   nir_def *opaque;
};

typedef void (*lvp_aabb_intersection_cb)(nir_builder *b, struct lvp_leaf_intersection *intersection,
                                         const struct lvp_ray_traversal_args *args,
                                         const struct lvp_ray_flags *ray_flags);

struct lvp_triangle_intersection {
   struct lvp_leaf_intersection base;

   nir_def *t;
   nir_def *frontface;
   nir_def *barycentrics;
};

typedef void (*lvp_triangle_intersection_cb)(nir_builder *b,
                                             struct lvp_triangle_intersection *intersection,
                                             const struct lvp_ray_traversal_args *args,
                                             const struct lvp_ray_flags *ray_flags);

struct lvp_ray_traversal_vars {
   /* For each accepted hit, tmax will be set to the t value. This allows for automatic intersection
    * culling.
    */
   nir_deref_instr *tmax;

   /* Those variables change when entering and exiting BLASes. */
   nir_deref_instr *origin;
   nir_deref_instr *dir;
   nir_deref_instr *inv_dir;

   /* The base address of the current TLAS/BLAS. */
   nir_deref_instr *bvh_base;

   nir_deref_instr *current_node;

   nir_deref_instr *stack_base;
   nir_deref_instr *stack_ptr;
   nir_deref_instr *stack;

   /* Information about the current instance used for culling. */
   nir_deref_instr *instance_addr;
   nir_deref_instr *sbt_offset_and_flags;
};

struct lvp_ray_traversal_args {
   nir_def *root_bvh_base;
   nir_def *flags;
   nir_def *cull_mask;
   nir_def *origin;
   nir_def *tmin;
   nir_def *dir;

   struct lvp_ray_traversal_vars vars;

   lvp_aabb_intersection_cb aabb_cb;
   lvp_triangle_intersection_cb triangle_cb;

   void *data;
};

/* Builds the ray traversal loop and returns whether traversal is incomplete, similar to
 * rayQueryProceedEXT. Traversal will only be considered incomplete, if one of the specified
 * callbacks breaks out of the traversal loop.
 */
nir_def *lvp_build_ray_traversal(nir_builder *b, const struct lvp_ray_traversal_args *args);

struct lvp_device;
struct lvp_pipeline;
struct lvp_pipeline_layout;
struct vk_pipeline_robustness_state;

bool lvp_lower_exec_graph(struct lvp_pipeline *pipeline, nir_shader *nir);

bool lvp_lower_input_attachments(nir_shader *shader, bool use_fragcoord_sysval);

void lvp_lower_pipeline_layout(const struct lvp_device *device,
                               struct lvp_pipeline_layout *layout,
                               nir_shader *shader);

bool lvp_nir_lower_ray_queries(struct nir_shader *shader);

bool lvp_nir_lower_sparse_residency(struct nir_shader *shader);

bool lvp_nir_opt_robustness(struct nir_shader *shader, struct lvp_device *device,
                            struct vk_pipeline_robustness_state *robustness);

#endif