Crafter.Graphics/interfaces/Crafter.Graphics-RenderingElement3D.cppm

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

module;
#ifndef CRAFTER_GRAPHICS_WINDOW_DOM
#include "vulkan/vulkan.h"
#endif // !CRAFTER_GRAPHICS_WINDOW_DOM
export module Crafter.Graphics:RenderingElement3D;
import :RT;
#ifndef CRAFTER_GRAPHICS_WINDOW_DOM
import std;
import :Mesh;
import :VulkanBuffer;
import Crafter.Math;
import :Window;

export namespace Crafter {
    struct TlasWithBuffer {
        VkDeviceAddress address = 0;
        VulkanBuffer<char, false> buffer;
        VkAccelerationStructureKHR accelerationStructure = VK_NULL_HANDLE;
        VulkanBuffer<VkAccelerationStructureInstanceKHR, true> instanceBuffer;
        VulkanBuffer<char, false> scratchBuffer;
        // Parallel to instanceBuffer, indexed by TLAS instance ID. Filled
        // from each element's userMetadata during BuildTLAS. Consumers
        // (e.g. ray-query collision) bind this in the descriptor heap and
        // look up via rayQueryGetIntersectionInstanceIdEXT to recover
        // application-side per-instance data without touching the
        // Vulkan-mandated instanceCustomIndex (which renderers may already
        // use for their own encoding).
        VulkanBuffer<std::uint32_t, true> metadataBuffer;
        // Last instance count this TLAS was built (not refit) for. When
        // elements.size() matches this, BuildTLAS does an in-place refit
        // (UPDATE mode) which is dramatically cheaper than a full rebuild
        // — refit walks the existing BVH and updates AABBs, while rebuild
        // reconstructs the topology from scratch. A change in count forces
        // a fresh rebuild, since a refit must keep the primitive count of
        // the build it updates.
        std::uint32_t builtInstanceCount = 0;
        // Build flags the current AS was created with — ALLOW_UPDATE (always
        // set, since BuildTLAS refits in place) plus the chosen
        // FAST_TRACE/FAST_BUILD preference bit. A refit (UPDATE mode) must
        // use flags identical to the originating build, so a change in the
        // requested preference forces a full rebuild rather than a refit.
        VkBuildAccelerationStructureFlagsKHR builtFlags = 0;
        // Per-slot record of the RenderingElement3D::hostDataVersion that
        // BuildTLAS last copied into this frame's instanceBuffer/metadataBuffer
        // — parallel to them, sized to the live instance count. The per-frame
        // copy re-uploads (and flushes) only the slots whose element advanced
        // past the recorded version; an element left at version 0 ("untracked")
        // reads dirty every frame, so callers that never opt in keep the
        // pre-#118 copy-every-frame behaviour. Reset to all-zero on every
        // topology change, since a rebuild reshuffles which element occupies
        // each slot (and may have reallocated the buffers). See
        // RenderingElement3D::hostDataVersion.
        std::vector<std::uint64_t> uploadedVersion;

        TlasWithBuffer() = default;
        // Moving hands the TLAS over: the moved-from TlasWithBuffer keeps no
        // acceleration structure, so only the last owner releases it, and its
        // next BuildTLAS builds a new one. Assigning queues the acceleration
        // structure this one held, as the destructor would; assign an empty
        // TlasWithBuffer to reset one.
        TlasWithBuffer(TlasWithBuffer&& other);
        TlasWithBuffer& operator=(TlasWithBuffer&& other);
        // Queues the acceleration structure for deferred deletion, ahead of
        // the storage it lives in (the buffers queue themselves): a command
        // buffer still being recorded or a frame in flight may still use it.
        ~TlasWithBuffer();
    };

   class RenderingElement3D {
        public:
        RTInstance instance;
        // Position in `elements`, maintained by Add/Remove for O(1) swap-and-pop.
        // Sentinel value = not currently registered.
        std::uint32_t indexInElements = std::numeric_limits<std::uint32_t>::max();
        // Application-defined per-instance tag, copied verbatim into
        // tlases[*].metadataBuffer at this element's TLAS instance ID
        // every BuildTLAS. Crafter doesn't interpret it.
        std::uint32_t userMetadata = 0;
        // When true, BuildTLAS skips copying instance.transform into the
        // TLAS instance buffer — the application's compute shader writes
        // the transform field directly into instanceBuffer at this
        // element's TLAS instance ID. Other instance fields (mask,
        // customIndex, SBT offset, BLAS reference) are still copied from
        // the CPU instance struct.
        //
        // Used to take per-frame transform updates off the CPU for bodies
        // whose transforms derive from GPU-side state (physics nodes that
        // already live on the GPU).
        bool transformOwnedByGpu = false;

        // Monotonic version of this element's host-authored TLAS data — the
        // instance fields the CPU writes (everything except a GPU-owned
        // transform; see transformOwnedByGpu) plus userMetadata. BuildTLAS
        // records, per frame, the version it last copied into each buffer slot
        // (TlasWithBuffer::uploadedVersion) and re-copies a slot only when its
        // element has advanced past the recorded version — so a TLAS dominated
        // by instances whose host fields are set once and then left alone (the
        // millions-of-GPU-driven-bodies target) pays no per-frame host copy or
        // flush after the first upload.
        //
        // 0 is the "untracked" sentinel: an element left at 0 is copied every
        // frame exactly as before this optimization, so code that mutates
        // instance/userMetadata without opting in stays correct. Opt in by
        // calling MarkHostDataDirty after every host-data change — the standard
        // dirty-flag contract (mark on change; the upload clears it until the
        // next change). The GPU-owned transform is exempt: BuildTLAS never
        // host-copies it, so changing it needs no mark.
        std::uint64_t hostDataVersion = 0;

        // Stamp this element's host data with a fresh global version so the next
        // BuildTLAS of each frame re-uploads (and flushes) its slot. Call after
        // changing any host-authored instance field or userMetadata. See
        // hostDataVersion.
        void MarkHostDataDirty() { hostDataVersion = ++hostDataVersionCounter; }

        // Source of globally-unique, monotonically-increasing host-data
        // versions. Global rather than per-element so a version value names a
        // unique (element, edit) pair: a per-slot recorded version then equals
        // the slot's current occupant only when that exact element's exact edit
        // was the last thing written there. That uniqueness is what makes the
        // relocation cases — swap-and-pop in Remove, or a remove+add that nets
        // the same instance count and so takes the refit path — fall out
        // correctly without tracking element identity: a relocated element's
        // version never collides with the (different) element's version recorded
        // for that slot. 64-bit, so it does not wrap in practice.
        inline static std::uint64_t hostDataVersionCounter = 0;

        static std::vector<RenderingElement3D*> elements;
        // Never destroyed, like Device::device. As a plain static its
        // destructor would run at process exit after the Vulkan driver and
        // layers, which are loaded during Device::Initialize (after this was
        // constructed), have torn down their own state, and destroying the
        // buffers would call into that freed state: the validation layer
        // segfaults on it. Exit reclaims them instead.
        inline static std::array<TlasWithBuffer, Window::numFrames>& tlases =
            *new std::array<TlasWithBuffer, Window::numFrames>();
        // Build (or in-place refit) the TLAS for frame `index`. `preference`
        // maps onto the FAST_TRACE/FAST_BUILD bits exactly as Mesh::Build
        // does for a BLAS: FastTrace (the default) spends more build time for
        // faster traversal — right for a TLAS rebuilt at most once per frame
        // and traced many times; FastBuild trades traversal speed for cheaper
        // builds, for scenes whose instance set churns hard every frame.
        // ALLOW_UPDATE is always set (refit is intrinsic to this path), so it
        // is not exposed as an option. Changing `preference` between frames
        // forces a full rebuild on the next call, since a refit must inherit
        // the original build's flags.
        //
        // The build writes tlases[index] in place: its instance and metadata
        // buffers on the host as it records, its acceleration structure and
        // scratch on the GPU. So call it only once every earlier submission
        // that used tlases[index] has completed. The TLASes are per frame in
        // flight, and Window::Render runs onUpdate after waiting for the fence
        // of its frame slot, so building tlases[currentFrame] there is safe.
        // Earlier builds, traces and compute ray queries in the same command
        // buffer are ordered before the build by the build itself.
        //
        // A rebuild for a new instance count or preference builds into the
        // same acceleration structure while it fits its storage, so
        // tlases[index].address stays the same. A TLAS that outgrows its
        // storage moves: read the address again after a build that grows the
        // scene.
        static void BuildTLAS(VkCommandBuffer cmd, std::uint32_t index, RTBuildPreference preference = RTBuildPreference::FastTrace);

        // Register / unregister with `elements`. Use these instead of touching
        // the vector directly: linear find+erase is O(n) and pathological at
        // the body counts physics targets (millions of braces).
        static void Add(RenderingElement3D* e);
        static void Remove(RenderingElement3D* e);
   };

    // A dependency BuildTLAS records around its build.
    struct TlasBarrier {
        VkPipelineStageFlags srcStage;
        VkPipelineStageFlags dstStage;
        VkMemoryBarrier memory;
    };

    // The dependency BuildTLAS records before its build, ordering the build
    // after the earlier work in the same command buffer that it must follow:
    // BLAS builds whose results it reads, an earlier build of this TLAS,
    // whose storage and scratch it writes again, and traces and compute ray
    // queries of the TLAS, which a refit or a rebuild in place would
    // otherwise overwrite while they read it.
    inline TlasBarrier BuildTlasWriteBarrier() {
        return {
            .srcStage = VK_PIPELINE_STAGE_ACCELERATION_STRUCTURE_BUILD_BIT_KHR | VK_PIPELINE_STAGE_RAY_TRACING_SHADER_BIT_KHR | VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
            .dstStage = VK_PIPELINE_STAGE_ACCELERATION_STRUCTURE_BUILD_BIT_KHR,
            .memory = {
                .sType = VK_STRUCTURE_TYPE_MEMORY_BARRIER,
                // Writes only: the reads need just the execution dependency
                // the stages give them. A read bit here does nothing, except
                // that synchronization validation (layers 1.4.363) then
                // reports the second of two BuildTLAS calls in one command
                // buffer as a write-after-write hazard.
                .srcAccessMask = VK_ACCESS_ACCELERATION_STRUCTURE_WRITE_BIT_KHR,
                .dstAccessMask = VK_ACCESS_ACCELERATION_STRUCTURE_READ_BIT_KHR | VK_ACCESS_ACCELERATION_STRUCTURE_WRITE_BIT_KHR,
            },
        };
    }

    // The dependency BuildTLAS records after its build, ordering it before
    // the shaders that read the TLAS later in the same command buffer:
    // traceRayEXT in a ray-tracing pass, rayQueryEXT in a compute pass. It
    // has to carry the acceleration-structure memory dependency; ordering
    // execution alone never makes the build's writes visible to those reads.
    inline TlasBarrier BuildTlasReadBarrier() {
        return {
            .srcStage = VK_PIPELINE_STAGE_ACCELERATION_STRUCTURE_BUILD_BIT_KHR,
            .dstStage = VK_PIPELINE_STAGE_RAY_TRACING_SHADER_BIT_KHR | VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
            .memory = {
                .sType = VK_STRUCTURE_TYPE_MEMORY_BARRIER,
                .srcAccessMask = VK_ACCESS_ACCELERATION_STRUCTURE_WRITE_BIT_KHR,
                .dstAccessMask = VK_ACCESS_ACCELERATION_STRUCTURE_READ_BIT_KHR,
            },
        };
    }
}
#endif // !CRAFTER_GRAPHICS_WINDOW_DOM

#ifdef CRAFTER_GRAPHICS_WINDOW_DOM
import std;
import :Mesh;
import :WebGPU;
import :WebGPUBuffer;
import :Window;

export namespace Crafter {
    // Per-frame TLAS storage. WebGPU has no real swapchain frame count
    // (Window::numFrames = 1 on DOM), so this is effectively a singleton —
    // the array form is kept for API symmetry with the Vulkan side so user
    // code that indexes `tlases[frameIdx]` ports unchanged.
    struct TlasWithBuffer {
        // Host-visible instance buffer holding RTInstance entries — same
        // layout as Vulkan's VkAccelerationStructureInstanceKHR, so user
        // code touching .instance.mask / .flags / .transform.matrix is
        // identical across backends. Also bound as a storage SSBO so
        // application compute shaders (e.g. physics-tlas-transform.comp.wgsl)
        // can write the .transform field directly when
        // RenderingElement3D::transformOwnedByGpu is set.
        WebGPUBuffer<RTInstance, true>     instanceBuffer;
        // Per-instance application metadata; parallel to instanceBuffer,
        // identical semantics to the Vulkan-side counterpart.
        WebGPUBuffer<std::uint32_t, true>  metadataBuffer;
        // GPU-built TLAS data: one TLASEntry per instance, written each
        // BuildTLAS by a compute pass on the JS bridge. Read by traceRay /
        // rayQuery as `@group(1) @binding(0) tlas: array<TLASEntry>`.
        // TLASEntry layout: 96 bytes — aabbMin (12) + maskHGoffset (4) +
        // aabbMax (12) + blasHandle (4) + invTransform 3x4 mat (48) +
        // customIndex (4) + _pad (12). Defined in the WGSL traversal
        // library; never directly read by C++.
        WebGPUBuffer<char, false>           buffer;
        // GPU LBVH support — see additional/dom-webgpu.js's TLAS-build
        // pipeline.
        //
        // entryOrder: per-frame permutation array of u32, indexing into
        // `buffer` (the TLASEntry[] array). Populated by the radix-sort
        // pass to spatially-coherent Morton order, then consumed by the
        // BVH construction + traversal passes.  In Stage 1 (this
        // baseline) it's the identity permutation written by
        // tlasBuildMain alongside the entries.
        WebGPUBuffer<char, false>           entryOrder;
        // mortonCodes: per-instance 32-bit Morton codes computed from the
        // world-AABB centroid, used as the radix-sort key. Written by
        // tlasBuildMain.
        WebGPUBuffer<char, false>           mortonCodes;
        // bvhNodes: 2N_PADDED - 1 sweep-tree BVH nodes built per frame
        // by the LBVH-build compute pass. Each node 32 bytes (aabbMin +
        // pad, aabbMax + pad). N_PADDED = 65536 (hardcoded in WGSL).
        // Internal nodes [0, N_PADDED-1); leaves [N_PADDED-1, 2*N_PADDED-1).
        // Node i's children are 2i+1, 2i+2 (implicit perfect binary
        // tree). Cap: 65536 instances per scene.
        WebGPUBuffer<char, false>           bvhNodes;
        // tlasBins: dead, kept allocated as a 64-byte placeholder so the
        // existing wgpuBuildTLAS C++ signature doesn't need a churn.
        // The pre-LBVH 64-bin partition was replaced by the full BVH.
        WebGPUBuffer<char, false>           tlasBins;
        // Sort ping-pong buffers for the radix sort. Each pass reads
        // from one and writes to the other, swapping role. Layout per
        // element: 1 u32 packed key = (morton16 << 16) | tlasIndex16.
        // Sized for N_PADDED.
        WebGPUBuffer<char, false>           sortTempA;
        WebGPUBuffer<char, false>           sortTempB;

        std::uint32_t builtInstanceCount = 0;
    };

    class RenderingElement3D {
        public:
        RTInstance instance{};
        std::uint32_t indexInElements = std::numeric_limits<std::uint32_t>::max();
        std::uint32_t userMetadata = 0;
        // Application compute shader writes the transform field of this
        // element's instanceBuffer slot directly — BuildTLAS preserves it.
        bool transformOwnedByGpu = false;

        // API-symmetric with the Vulkan side so portable code that opts its
        // instances into host-data dirty tracking compiles unchanged. The
        // WebGPU BuildTLAS re-uploads the whole CPU mirror every build (the
        // counts this path targets are small), so the version is not consulted
        // here — it exists purely for cross-backend source compatibility.
        std::uint64_t hostDataVersion = 0;
        void MarkHostDataDirty() { hostDataVersion = ++hostDataVersionCounter; }
        inline static std::uint64_t hostDataVersionCounter = 0;

        static std::vector<RenderingElement3D*> elements;
        inline static TlasWithBuffer tlases[Window::numFrames];

        // Repopulate the TLAS for frame `index`. WebGPU path always does
        // a fresh build (no refit) — the GPU build pass is cheap at the
        // ~10–100 instance counts the design targets; LBVH-for-TLAS is a
        // future optimization for larger scenes.
        //
        // BuildTLAS is now split into Upload + Build so a physics
        // compute pass (e.g. physics-tlas-transform) can run between the
        // CPU mirror upload and the GPU LBVH build. The compute pass
        // writes the per-instance transform bytes that BuildTLAS leaves
        // intact for elements flagged transformOwnedByGpu, and those
        // writes have to land before the LBVH reads them. The combined
        // BuildTLAS is kept as a convenience for callers that don't
        // interleave a compute pass (e.g. the ctor-time first build).
        static void BuildTLASUpload(WebGPUCommandEncoderRef cmd, std::uint32_t index);
        static void BuildTLASBuild(WebGPUCommandEncoderRef cmd, std::uint32_t index);
        // `preference` mirrors the Vulkan signature for portability but has
        // no effect here: the software LBVH-for-TLAS path has no hardware
        // FAST_TRACE/FAST_BUILD knob (same as Mesh::Build on WebGPU). The
        // parameter exists purely so portable code compiles unchanged.
        static void BuildTLAS(WebGPUCommandEncoderRef cmd, std::uint32_t index, RTBuildPreference preference = RTBuildPreference::FastTrace);

        static void Add(RenderingElement3D* e);
        static void Remove(RenderingElement3D* e);
    };
}
#endif // CRAFTER_GRAPHICS_WINDOW_DOM