/*
 * Copyright © 2022 Valve Corporation
 *
 * 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

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

#include "build_interface.h"

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

void
main(void)
{
   uint32_t global_id = gl_GlobalInvocationID.x;
   uint32_t total_invocations = gl_NumWorkGroups.x * 64;

   uint64_t accel_struct_addr = args.mode == RADV_COPY_MODE_SERIALIZE ? args.src_addr : args.dst_addr;
   uint64_t serialized_addr = args.mode == RADV_COPY_MODE_SERIALIZE ? args.dst_addr : args.src_addr;

   uint64_t blas_addrs = serialized_addr + SIZEOF(radv_accel_struct_serialization_header);

   radv_accel_struct_serialization_header serialization_header =
      DEREF(REF(radv_accel_struct_serialization_header)(serialized_addr));

   radv_accel_struct_header header = DEREF(REF(radv_accel_struct_header)(accel_struct_addr));

   for (uint32_t i = global_id; i < serialization_header.instance_count; i += total_invocations) {
      uint64_t instance_offset_addr = accel_struct_addr + (header.leaf_node_offsets_offset + i * 4);
      uint64_t instance_addr = accel_struct_addr + (header.bvh_offset + DEREF(REF(uint32_t)(instance_offset_addr)));
      REF(radv_gfx12_instance_node) instance_node = REF(radv_gfx12_instance_node)(instance_addr);
      REF(radv_gfx12_instance_node_user_data) instance_data =
         REF(radv_gfx12_instance_node_user_data)(instance_addr + SIZEOF(radv_gfx12_instance_node));

      if (args.mode == RADV_COPY_MODE_SERIALIZE) {
         DEREF(INDEX(uint64_t, blas_addrs, i)) =
            node_to_addr(DEREF(instance_node).pointer_flags_bvh_addr) - DEREF(instance_data).bvh_offset;
      } else {
         uint32_t bvh_offset = DEREF(instance_data).bvh_offset;

         /* Replace the address while keeping the pointer flags. */
         uint64_t pointer_flags_bvh_addr = DEREF(instance_node).pointer_flags_bvh_addr;
         uint64_t blas_addr = DEREF(INDEX(uint64_t, blas_addrs, i));
         DEREF(instance_node).pointer_flags_bvh_addr =
            (pointer_flags_bvh_addr & 0xFFC0000000000000ul) | addr_to_node(blas_addr + bvh_offset);
      }
   }
}