/*
 * Copyright © 2022 Bas Nieuwenhuizen
 *
 * 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 lanes = gl_NumWorkGroups.x * 64;
   uint32_t increment = lanes * 16;

   uint64_t copy_src_addr = args.src_addr;
   uint64_t copy_dst_addr = args.dst_addr;

   if (args.mode == RADV_COPY_MODE_DESERIALIZE) {
      copy_src_addr += SIZEOF(radv_accel_struct_serialization_header) +
                       DEREF(REF(radv_accel_struct_serialization_header)(args.src_addr)).instance_count * SIZEOF(uint64_t);

   }

   REF(radv_accel_struct_header) header = REF(radv_accel_struct_header)(copy_src_addr);

   uint64_t instance_base = args.src_addr + SIZEOF(radv_accel_struct_serialization_header);
   uint64_t node_offset = DEREF(header).instance_offset;
   uint64_t node_end = DEREF(header).instance_count * SIZEOF(radv_bvh_instance_node);
   if (node_end > 0)
      node_end += node_offset;

   if (args.mode == RADV_COPY_MODE_SERIALIZE) {
      copy_dst_addr += SIZEOF(radv_accel_struct_serialization_header) +
                       DEREF(REF(radv_accel_struct_header)(args.src_addr)).instance_count * SIZEOF(uint64_t);

      if (global_id == 0) {
         REF(radv_accel_struct_serialization_header) ser_header =
            REF(radv_accel_struct_serialization_header)(args.dst_addr);
         DEREF(ser_header).serialization_size = DEREF(header).serialization_size;
         DEREF(ser_header).compacted_size = DEREF(header).compacted_size;
         DEREF(ser_header).instance_count = DEREF(header).instance_count;
      }

      instance_base = args.dst_addr + SIZEOF(radv_accel_struct_serialization_header);
   } else if (args.mode == RADV_COPY_MODE_COPY)
      node_end = 0;

   uint64_t size = DEREF(header).compacted_size;
   for (uint64_t offset = global_id * 16; offset < size; offset += increment) {
      DEREF(REF(uvec4)(copy_dst_addr + offset)) =
         DEREF(REF(uvec4)(copy_src_addr + offset));

      /* Do the adjustment inline in the same invocation that copies the data so that we don't have
       * to synchronize. This is only possible on pre-GFX12 HW because leaf nodes have a different
       * order on GFX12.
       */
      if (!VK_BUILD_FLAG(RADV_BUILD_FLAG_BVH8) && offset < node_end && offset >= node_offset &&
          (offset - node_offset) % SIZEOF(radv_bvh_instance_node) == 0) {
         uint64_t idx = (offset - node_offset) / SIZEOF(radv_bvh_instance_node);

         uint32_t bvh_offset = DEREF(REF(radv_bvh_instance_node)(copy_src_addr + offset)).bvh_offset;
         if (args.mode == RADV_COPY_MODE_SERIALIZE) {
            DEREF(INDEX(uint64_t, instance_base, idx)) =
               node_to_addr(DEREF(REF(radv_bvh_instance_node)(copy_src_addr + offset)).bvh_ptr) - bvh_offset;
         } else { /* RADV_COPY_MODE_DESERIALIZE */
            uint64_t blas_addr = DEREF(INDEX(uint64_t, instance_base, idx));
            DEREF(REF(radv_bvh_instance_node)(copy_dst_addr + offset)).bvh_ptr = addr_to_node(blas_addr + bvh_offset);
         }
      }
   }
}