Crafter.Graphics/implementations/Crafter.Graphics-RenderingElement3D.cpp

//SPDX-License-Identifier: LGPL-3.0-only
//SPDX-FileCopyrightText: Copyright (C) 2026 Catcrafts®

module;
#include <vulkan/vulkan_core.h>
module Crafter.Graphics:RenderingElement3D_impl;
import :RenderingElement3D;
import std;

using namespace Crafter;


std::vector<RenderingElement3D*> RenderingElement3D::elements;

void RenderingElement3D::Add(RenderingElement3D* e) {
    e->indexInElements = static_cast<std::uint32_t>(elements.size());
    elements.push_back(e);
}

void RenderingElement3D::Remove(RenderingElement3D* e) {
    // Idempotent: callers like Builder ghost flow toggle elements in/out
    // and may try to remove an already-removed element.
    std::uint32_t idx = e->indexInElements;
    if (idx == std::numeric_limits<std::uint32_t>::max()) return;
    std::uint32_t last = static_cast<std::uint32_t>(elements.size() - 1);
    if (idx != last) {
        elements[idx] = elements[last];
        elements[idx]->indexInElements = idx;
    }
    elements.pop_back();
    e->indexInElements = std::numeric_limits<std::uint32_t>::max();
}

void RenderingElement3D::BuildTLAS(VkCommandBuffer cmd, std::uint32_t index, RTBuildPreference preference) {
    Device::RequireRayTracing();
    auto& tlas = tlases[index];
    const std::uint32_t primitiveCount = static_cast<std::uint32_t>(elements.size());

    // ALLOW_UPDATE is always set — BuildTLAS refits in place on later
    // frames. The FAST_TRACE/FAST_BUILD preference is layered on top; the
    // two preference bits are mutually exclusive, so exactly one is chosen.
    // Both bits are baked into the AS at BUILD time and must be identical on
    // every subsequent UPDATE, so a change in `preference` is treated as a
    // topology change below to force a fresh rebuild with the new flags.
    const VkBuildAccelerationStructureFlagsKHR buildFlags =
        VK_BUILD_ACCELERATION_STRUCTURE_ALLOW_UPDATE_BIT_KHR
        | (preference == RTBuildPreference::FastBuild
            ? VK_BUILD_ACCELERATION_STRUCTURE_PREFER_FAST_BUILD_BIT_KHR
            : VK_BUILD_ACCELERATION_STRUCTURE_PREFER_FAST_TRACE_BIT_KHR);

    // Refit (UPDATE) is allowed when the count matches the count this AS
    // was last built for and the build flags are unchanged. A change forces
    // a full rebuild because the AS storage and instance buffer were sized
    // for the old count (and a refit must inherit the original build flags).
    // Refit is dramatically cheaper at scale (millions of instances) — it
    // walks the existing BVH and updates AABBs rather than reconstructing
    // topology.
    const bool topologyChanged =
        tlas.accelerationStructure == VK_NULL_HANDLE
        || primitiveCount != tlas.builtInstanceCount
        || buildFlags != tlas.builtFlags;

    {
        VkMemoryBarrier asBarrier {
            .sType = VK_STRUCTURE_TYPE_MEMORY_BARRIER,
            .srcAccessMask = VK_ACCESS_ACCELERATION_STRUCTURE_WRITE_BIT_KHR,
            .dstAccessMask = VK_ACCESS_ACCELERATION_STRUCTURE_READ_BIT_KHR | VK_ACCESS_ACCELERATION_STRUCTURE_WRITE_BIT_KHR
        };
        vkCmdPipelineBarrier(cmd,
            VK_PIPELINE_STAGE_ACCELERATION_STRUCTURE_BUILD_BIT_KHR,
            VK_PIPELINE_STAGE_ACCELERATION_STRUCTURE_BUILD_BIT_KHR,
            0, 1, &asBarrier, 0, nullptr, 0, nullptr);
    }

    if (topologyChanged) {
        // Grow the host-visible inputs on a high-water mark rather than
        // resizing to the exact count every topology change. These buffers
        // only need to hold *at least* primitiveCount entries — the AS build
        // reads exactly primitiveCount of them via tlasRangeInfo, and the
        // copy loop below writes [0, primitiveCount) — so an allocation left
        // over from a larger earlier frame is reused as-is. This drops the
        // two host-buffer reallocations on a count decrease (and on an
        // increase that still fits the previous high-water capacity),
        // leaving only the AS storage + scratch rebuild below, which is tied
        // to the AS itself and dominates this path regardless.
        //
        // instanceBuffer and metadataBuffer grow in lockstep (always resized
        // together to the same entry count), so one capacity check covers
        // both. size is only meaningful once buffer is non-null, so the
        // null check must short-circuit before the division.
        // STORAGE_BUFFER_BIT is required because the application's compute
        // shaders bind these buffers as storage SSBOs (e.g. to write
        // per-instance transforms directly into the TLAS instance data).
        //
        // instanceBuffer prefers HOST_VISIBLE | DEVICE_LOCAL (BAR/VRAM) so the
        // two per-frame GPU accesses stay on-chip instead of crossing PCIe: the
        // compute shader writes the transform field in place, and the AS build
        // reads the whole instance. The CPU still writes the host-authored
        // fields below — those are write-only, sequential, never-read-back
        // writes, the ideal write-combined BAR workload. DEVICE_LOCAL is a
        // best-effort preference: GetMemoryType falls back to plain
        // HOST_VISIBLE (the previous behaviour) when no combined type exists
        // (no resizable BAR). HOST_COHERENT is intentionally dropped from the
        // required set — the combined type may not be coherent, and the
        // FlushDevice(cmd, ...) below already establishes host->build
        // visibility, gating its flush-skip on the *chosen* type's flags. The
        // single shared buffer is preserved: a whole-buffer copy would clobber
        // the GPU-written transforms, so staging is deliberately avoided.
        //
        // metadataBuffer (#75) gets the same HOST_VISIBLE | prefer-DEVICE_LOCAL
        // upgrade: it is CPU-written every frame (the copy loop below) and read
        // by the ray shaders as a STORAGE_BUFFER every frame, so it has the same
        // per-frame-shader-read-over-PCIe cost. Unlike instanceBuffer there is
        // no GPU co-write, so the direct upgrade is unconditional — but staging
        // is still wrong here: a per-frame-rewritten buffer would pay a stage+
        // copy+barrier every frame, defeating the point. HOST_COHERENT is
        // likewise dropped; the FlushDevice() after the copy loop makes the host
        // writes available before the consuming submit on a non-coherent type
        // (no-op when the chosen type is coherent — the previous behaviour),
        // and the queue submit carries the host-write -> shader-read ordering as
        // it always has.
        if (tlas.instanceBuffer.buffer == VK_NULL_HANDLE
            || primitiveCount > tlas.instanceBuffer.size / sizeof(VkAccelerationStructureInstanceKHR)) {
            tlas.instanceBuffer.Resize(VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_ACCELERATION_STRUCTURE_BUILD_INPUT_READ_ONLY_BIT_KHR | VK_BUFFER_USAGE_STORAGE_BUFFER_BIT, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT, primitiveCount, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT);
            tlas.metadataBuffer.Resize(VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT, primitiveCount, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT);
        }

        // A topology change reshuffles which element occupies each slot (a count
        // change, and Remove's swap-and-pop within it) and may have reallocated
        // the host buffers just above — so every slot's prior contents are
        // stale. Clear the per-slot version record (sized to the live count) so
        // the copy loop below re-uploads all primitiveCount slots this frame.
        tlas.uploadedVersion.assign(primitiveCount, 0);
    }

    // Copy only the slots whose host-authored data changed since this frame's
    // buffers last saw them (#118), tracking the first/last dirty index so the
    // flush below covers just that [first, last] envelope rather than the whole
    // high-water capacity (the envelope may span a few clean slots between
    // sparse dirty ones — flushing a couple of extra entries is far cheaper than
    // the per-slot vkFlushMappedMemoryRanges calls true sub-ranges would need).
    // A slot is dirty when its element advanced past the version
    // recorded for it; an element left at version 0 (never MarkHostDataDirty'd)
    // reads dirty every frame, preserving the pre-#118 copy-every-frame
    // behaviour for callers that don't opt into dirty tracking. Globally-unique
    // versions make this correct under relocation: a slot now holding a
    // different element never matches the version recorded for its previous
    // occupant. The GPU-owned transform is never copied here regardless — the
    // application's compute shader writes it in place earlier in this submission.
    bool anyDirty = false;
    std::uint32_t dirtyFirst = 0;
    std::uint32_t dirtyLast  = 0;
    for(std::uint32_t i = 0; i < primitiveCount; i++) {
        const std::uint64_t version = elements[i]->hostDataVersion;
        if (version != 0 && version == tlas.uploadedVersion[i]) {
            continue; // slot already holds this element's current host data
        }
        if (elements[i]->transformOwnedByGpu) {
            // Skip the transform field — the application's compute shader
            // writes it earlier in this submission. Copy everything else.
            auto& dst = tlas.instanceBuffer.value[i];
            const auto& src = elements[i]->instance;
            dst.instanceCustomIndex                    = src.instanceCustomIndex;
            dst.mask                                   = src.mask;
            dst.instanceShaderBindingTableRecordOffset = src.instanceShaderBindingTableRecordOffset;
            dst.flags                                  = src.flags;
            dst.accelerationStructureReference         = src.accelerationStructureReference;
        } else {
            tlas.instanceBuffer.value[i] = elements[i]->instance;
        }
        tlas.metadataBuffer.value[i] = elements[i]->userMetadata;
        tlas.uploadedVersion[i] = version;
        if (!anyDirty) { dirtyFirst = i; anyDirty = true; }
        dirtyLast = i;
    }

    if (anyDirty) {
        // Flush (and, for the instance buffer, barrier) only the [dirtyFirst,
        // dirtyLast] span actually rewritten this frame. When nothing is dirty
        // there were no host writes at all, so both flushes — and the
        // instanceBuffer HOST->build barrier — are skipped: the AS build reads
        // memory already made visible by the submit that last wrote it, and any
        // GPU-owned transform write is ordered by the application's own
        // compute->build barrier, not this host barrier.
        const std::uint32_t count = dirtyLast - dirtyFirst + 1;
        constexpr VkDeviceSize instStride = sizeof(VkAccelerationStructureInstanceKHR);
        constexpr VkDeviceSize metaStride = sizeof(std::uint32_t);

        tlas.instanceBuffer.FlushDevice(
            cmd, VK_ACCESS_MEMORY_READ_BIT, VK_PIPELINE_STAGE_ACCELERATION_STRUCTURE_BUILD_BIT_KHR,
            static_cast<VkDeviceSize>(dirtyFirst) * instStride,
            static_cast<VkDeviceSize>(count) * instStride);

        // Make the per-frame metadata host writes available before the consuming
        // submit. metadataBuffer is not an AS build input (the build never reads
        // it) and has no GPU co-write — only the ray shaders read it, in a later
        // pass — so it needs neither the build-stage barrier instanceBuffer takes
        // above nor a HOST->shader command barrier: the queue submit already
        // orders host writes made available before it against every command in
        // the submission, exactly as it did when this buffer was HOST_COHERENT.
        // This ranged host flush is what now makes those writes available on a
        // non-coherent (BAR/VRAM) type; it self-gates to a no-op on a coherent
        // type (#60), preserving the previous behaviour where the type ends up
        // coherent.
        tlas.metadataBuffer.FlushDevice(
            static_cast<VkDeviceSize>(dirtyFirst) * metaStride,
            static_cast<VkDeviceSize>(count) * metaStride);
    }

    VkAccelerationStructureGeometryInstancesDataKHR instancesData {
        .sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_INSTANCES_DATA_KHR,
        .arrayOfPointers = VK_FALSE,
        .data            = {tlas.instanceBuffer.address}
    };

    VkAccelerationStructureGeometryDataKHR geometryData;
    geometryData.instances = instancesData;

    VkAccelerationStructureGeometryKHR tlasGeometry {
        .sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_KHR,
        .geometryType = VK_GEOMETRY_TYPE_INSTANCES_KHR,
        .geometry     = geometryData
    };

    VkAccelerationStructureBuildGeometryInfoKHR tlasBuildGeometryInfo {
        .sType         = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_BUILD_GEOMETRY_INFO_KHR,
        .type          = VK_ACCELERATION_STRUCTURE_TYPE_TOP_LEVEL_KHR,
        // ALLOW_UPDATE (required for any subsequent UPDATE-mode refit) plus
        // the caller's FAST_TRACE/FAST_BUILD preference — see buildFlags above.
        .flags         = buildFlags,
        .mode          = topologyChanged
                            ? VK_BUILD_ACCELERATION_STRUCTURE_MODE_BUILD_KHR
                            : VK_BUILD_ACCELERATION_STRUCTURE_MODE_UPDATE_KHR,
        .geometryCount = 1,
        .pGeometries   = &tlasGeometry
    };

    if (topologyChanged) {
        // Query sizes for the fresh build, allocate AS storage + scratch.
        VkAccelerationStructureBuildSizesInfoKHR tlasBuildSizes {
            .sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_BUILD_SIZES_INFO_KHR
        };
        Device::vkGetAccelerationStructureBuildSizesKHR(
            Device::device,
            VK_ACCELERATION_STRUCTURE_BUILD_TYPE_DEVICE_KHR,
            &tlasBuildGeometryInfo,
            &primitiveCount,
            &tlasBuildSizes
        );

        // Scratch buffer must hold at least max(buildScratchSize, updateScratchSize).
        // Sizing for buildScratchSize covers both — refit is always smaller.
        tlas.scratchBuffer.Resize(VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, tlasBuildSizes.buildScratchSize);
        tlas.buffer.Resize(VK_BUFFER_USAGE_ACCELERATION_STRUCTURE_STORAGE_BIT_KHR | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT | VK_BUFFER_USAGE_ACCELERATION_STRUCTURE_BUILD_INPUT_READ_ONLY_BIT_KHR, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, tlasBuildSizes.accelerationStructureSize);

        // Destroy the previous AS handle before creating a new one — the
        // pre-refit path leaked here on every frame.
        if (tlas.accelerationStructure != VK_NULL_HANDLE) {
            Device::vkDestroyAccelerationStructureKHR(Device::device, tlas.accelerationStructure, nullptr);
            tlas.accelerationStructure = VK_NULL_HANDLE;
        }

        VkAccelerationStructureCreateInfoKHR tlasCreateInfo {
            .sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_CREATE_INFO_KHR,
            .buffer = tlas.buffer.buffer,
            .offset = 0,
            .size   = tlasBuildSizes.accelerationStructureSize,
            .type   = VK_ACCELERATION_STRUCTURE_TYPE_TOP_LEVEL_KHR,
        };
        Device::CheckVkResult(Device::vkCreateAccelerationStructureKHR(Device::device, &tlasCreateInfo, nullptr, &tlas.accelerationStructure));

        VkAccelerationStructureDeviceAddressInfoKHR addrInfo {
            .sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_DEVICE_ADDRESS_INFO_KHR,
            .accelerationStructure = tlas.accelerationStructure
        };
        tlas.address = Device::vkGetAccelerationStructureDeviceAddressKHR(Device::device, &addrInfo);

        tlas.builtInstanceCount = primitiveCount;
        tlas.builtFlags         = buildFlags;
    }

    // For UPDATE mode, src == dst (in-place refit). For BUILD, src is
    // VK_NULL_HANDLE and dst is the freshly-created handle.
    tlasBuildGeometryInfo.scratchData.deviceAddress = tlas.scratchBuffer.address;
    tlasBuildGeometryInfo.dstAccelerationStructure  = tlas.accelerationStructure;
    tlasBuildGeometryInfo.srcAccelerationStructure  =
        topologyChanged ? VK_NULL_HANDLE : tlas.accelerationStructure;

    VkAccelerationStructureBuildRangeInfoKHR tlasRangeInfo {
        .primitiveCount  = primitiveCount,
        .primitiveOffset = 0,
        .firstVertex     = 0,
        .transformOffset = 0
    };
    VkAccelerationStructureBuildRangeInfoKHR* tlasRangeInfoPP = &tlasRangeInfo;
    Device::vkCmdBuildAccelerationStructuresKHR(cmd, 1, &tlasBuildGeometryInfo, &tlasRangeInfoPP);

    vkCmdPipelineBarrier(
        cmd,
        VK_PIPELINE_STAGE_ACCELERATION_STRUCTURE_BUILD_BIT_KHR,
        VK_PIPELINE_STAGE_RAY_TRACING_SHADER_BIT_KHR | VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
        0,
        0, nullptr,
        0, nullptr,
        0, nullptr
    );
}