/*
 * Copyright © 2022 Konstantin Seurer
 *
 * SPDX-License-Identifier: MIT
 */

#version 460

#extension GL_GOOGLE_include_directive : require

#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require
#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require
#extension GL_EXT_shader_explicit_arithmetic_types_int32 : require
#extension GL_EXT_shader_explicit_arithmetic_types_int64 : require
#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require
#extension GL_EXT_scalar_block_layout : require
#extension GL_EXT_buffer_reference : require
#extension GL_EXT_buffer_reference2 : require
#extension GL_KHR_shader_subgroup_vote : require
#extension GL_KHR_shader_subgroup_arithmetic : require
#extension GL_KHR_shader_subgroup_ballot : require

layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in;

#include "build_interface.h"

layout(push_constant) uniform CONSTS {
   leaf_args args;
};

void
main(void)
{
   uint32_t global_id = gl_GlobalInvocationID.x;
   uint32_t primitive_id = args.geom_data.first_id + global_id;

   REF(key_id_pair) id_ptr = INDEX(key_id_pair, args.ids, primitive_id);
   uint32_t src_offset = global_id * args.geom_data.stride;

   uint32_t dst_stride;
   uint32_t node_type;
   if (args.geom_data.geometry_type == VK_GEOMETRY_TYPE_TRIANGLES_KHR) {
      dst_stride = SIZEOF(radv_bvh_triangle_node);
      node_type = radv_ir_node_triangle;
   } else if (args.geom_data.geometry_type == VK_GEOMETRY_TYPE_AABBS_KHR) {
      dst_stride = SIZEOF(radv_bvh_aabb_node);
      node_type = radv_ir_node_aabb;
   } else {
      dst_stride = SIZEOF(radv_bvh_instance_node);
      node_type = radv_ir_node_instance;
   }

   uint32_t dst_offset = primitive_id * dst_stride;
   VOID_REF dst_ptr = OFFSET(args.bvh, dst_offset);

   radv_aabb bounds;
   bool is_active;
   if (args.geom_data.geometry_type == VK_GEOMETRY_TYPE_TRIANGLES_KHR) {
      is_active = build_triangle(bounds, dst_ptr, args.geom_data, global_id);
   } else if (args.geom_data.geometry_type == VK_GEOMETRY_TYPE_AABBS_KHR) {
      VOID_REF src_ptr = OFFSET(args.geom_data.data, src_offset);
      is_active = build_aabb(bounds, src_ptr, dst_ptr, args.geom_data.geometry_id, global_id);
   } else {
      VOID_REF src_ptr = OFFSET(args.geom_data.data, src_offset);
      /* arrayOfPointers */
      if (args.geom_data.stride == 8) {
         src_ptr = DEREF(REF(VOID_REF)(src_ptr));
      }

      is_active = build_instance(bounds, src_ptr, dst_ptr, global_id);
   }

#if ALWAYS_ACTIVE
   if (!is_active && args.geom_data.geometry_type != VK_GEOMETRY_TYPE_INSTANCES_KHR) {
      bounds.min = vec3(0.0);
      bounds.max = vec3(0.0);
      is_active = true;
   }
#endif

   if (is_active) {
      REF(radv_ir_node) ir_node = INDEX(radv_ir_node, args.ir, primitive_id);
      DEREF(ir_node).aabb = bounds;
   }

   uint32_t ir_offset = primitive_id * SIZEOF(radv_ir_node);
   DEREF(id_ptr).id = is_active ? pack_ir_node_id(ir_offset, node_type) : RADV_BVH_INVALID_NODE;

   uvec4 ballot = subgroupBallot(is_active);
   if (subgroupElect())
      atomicAdd(DEREF(args.header).active_leaf_count, subgroupBallotBitCount(ballot));

   atomicMin(DEREF(args.header).min_bounds[0], to_emulated_float(bounds.min.x));
   atomicMin(DEREF(args.header).min_bounds[1], to_emulated_float(bounds.min.y));
   atomicMin(DEREF(args.header).min_bounds[2], to_emulated_float(bounds.min.z));
   atomicMax(DEREF(args.header).max_bounds[0], to_emulated_float(bounds.max.x));
   atomicMax(DEREF(args.header).max_bounds[1], to_emulated_float(bounds.max.y));
   atomicMax(DEREF(args.header).max_bounds[2], to_emulated_float(bounds.max.z));
}