/*
 * Copyright © 2023 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
#extension GL_KHR_memory_scope_semantics : require

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

#include "build_interface.h"

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

uint32_t fetch_parent_node(VOID_REF bvh, uint32_t node)
{
    uint64_t addr = bvh - node / 8 * 4 - 4;
    return DEREF(REF(uint32_t)(addr));
}

void main() {
    uint32_t bvh_offset = DEREF(args.src).bvh_offset;

    VOID_REF src_bvh = OFFSET(args.src, bvh_offset);
    VOID_REF dst_bvh = OFFSET(args.dst, bvh_offset);

    uint32_t leaf_node_size;
    if (args.geom_data.geometry_type == VK_GEOMETRY_TYPE_TRIANGLES_KHR)
        leaf_node_size = SIZEOF(radv_bvh_triangle_node);
    else if (args.geom_data.geometry_type == VK_GEOMETRY_TYPE_AABBS_KHR)
        leaf_node_size = SIZEOF(radv_bvh_aabb_node);
    else
        leaf_node_size = SIZEOF(radv_bvh_instance_node);

    uint32_t leaf_node_id = args.geom_data.first_id + gl_GlobalInvocationID.x;
    uint32_t first_leaf_offset = id_to_offset(RADV_BVH_ROOT_NODE) + SIZEOF(radv_bvh_box32_node);

    uint32_t dst_offset = leaf_node_id * leaf_node_size + first_leaf_offset;
    VOID_REF dst_ptr = OFFSET(dst_bvh, dst_offset);
    uint32_t src_offset = gl_GlobalInvocationID.x * args.geom_data.stride;

    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, gl_GlobalInvocationID.x);
    } else {
        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, gl_GlobalInvocationID.x);
    }

    if (!is_active)
        return;

    DEREF(INDEX(radv_aabb, args.leaf_bounds, leaf_node_id)) = bounds;
    memoryBarrier(gl_ScopeDevice,
        gl_StorageSemanticsBuffer,
        gl_SemanticsAcquireRelease | gl_SemanticsMakeAvailable | gl_SemanticsMakeVisible);

    uint32_t node_id = pack_node_id(dst_offset, 0);
    uint32_t parent_id = fetch_parent_node(src_bvh, node_id);
    uint32_t internal_nodes_offset = first_leaf_offset + args.leaf_node_count * leaf_node_size;
    while (parent_id != RADV_BVH_INVALID_NODE) {
        uint32_t offset = id_to_offset(parent_id);

        uint32_t parent_index = (offset - internal_nodes_offset) / SIZEOF(radv_bvh_box32_node) + 1;
        if (parent_id == RADV_BVH_ROOT_NODE)
            parent_index = 0;

        /* Make accesses to internal nodes in dst_bvh available and visible */
        memoryBarrier(gl_ScopeDevice,
                      gl_StorageSemanticsBuffer,
                      gl_SemanticsAcquireRelease | gl_SemanticsMakeAvailable | gl_SemanticsMakeVisible);

        REF(radv_bvh_box32_node) src_node = REF(radv_bvh_box32_node)OFFSET(src_bvh, offset);
        REF(radv_bvh_box32_node) dst_node = REF(radv_bvh_box32_node)OFFSET(dst_bvh, offset);
        uint32_t children[4];
        for (uint32_t i = 0; i < 4; ++i)
            children[i] = DEREF(src_node).children[i];

        uint32_t valid_child_count = 0;
        for (uint32_t i = 0; i < 4; ++valid_child_count, ++i)
            if (children[i] == RADV_BVH_INVALID_NODE)
                break;

        /* Check if all children have been processed. As this is an atomic the last path coming from
         * a child will pass here, while earlier paths break.
         */
        uint32_t ready_child_count = atomicAdd(
            DEREF(INDEX(uint32_t, args.internal_ready_count, parent_index)), 1, gl_ScopeDevice,
            gl_StorageSemanticsBuffer,
            gl_SemanticsAcquireRelease | gl_SemanticsMakeAvailable | gl_SemanticsMakeVisible);

        if (ready_child_count != valid_child_count - 1)
            break;

        for (uint32_t i = 0; i < 4; ++i)
            DEREF(dst_node).children[i] = children[i];

        for (uint32_t i = 0; i < valid_child_count; ++i) {
            uint32_t child_offset = id_to_offset(children[i]);
            radv_aabb child_bounds;
            if (child_offset == dst_offset)
                child_bounds = bounds;
            else if (child_offset >= internal_nodes_offset) {
                child_bounds = radv_aabb(vec3(INFINITY), vec3(-INFINITY));
                REF(radv_bvh_box32_node) child_node = REF(radv_bvh_box32_node)OFFSET(dst_bvh, child_offset);
                for (uint32_t j = 0; j < 4; ++j) {
                    if (DEREF(child_node).children[j] == RADV_BVH_INVALID_NODE)
                        break;
                    child_bounds.min = min(child_bounds.min, DEREF(child_node).coords[j].min);
                    child_bounds.max = max(child_bounds.max, DEREF(child_node).coords[j].max);
                }
            } else {
                uint32_t child_index = (child_offset - first_leaf_offset) / leaf_node_size;
                child_bounds = DEREF(INDEX(radv_aabb, args.leaf_bounds, child_index));
            }

            DEREF(dst_node).coords[i] = child_bounds;
        }

        if (parent_id == RADV_BVH_ROOT_NODE) {
            radv_aabb root_bounds = radv_aabb(vec3(INFINITY), vec3(-INFINITY));
            for (uint32_t i = 0; i < valid_child_count; ++i) {
                radv_aabb bounds = DEREF(dst_node).coords[i];
                root_bounds.min = min(root_bounds.min, bounds.min);
                root_bounds.max = max(root_bounds.max, bounds.max);
            }
            DEREF(args.dst).aabb = root_bounds;
        }

        parent_id = fetch_parent_node(src_bvh, parent_id);
    }
}