Crafter.Graphics/implementations/Crafter.Graphics-Mesh.cpp

/*
Crafter®.Graphics
Copyright (C) 2026 Catcrafts®
catcrafts.net

This library is free software; you can redistribute it and/or
modify it under the terms of the GNU Lesser General Public
License version 3.0 as published by the Free Software Foundation;

This library is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
Lesser General Public License for more details.

You should have received a copy of the GNU Lesser General Public
License along with this library; if not, write to the Free Software
Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301 USA
*/

module;
#include "vulkan/vulkan.h"
module Crafter.Graphics:Mesh_impl;
import Crafter.Math;
import Crafter.Asset;
import :Mesh;
import :Device;
import :Decompress;
import :Types;
import std;

using namespace Crafter;

namespace {
    // Buffer-usage flag set shared by both Build paths. The compressed path
    // appends VK_BUFFER_USAGE_2_MEMORY_DECOMPRESSION_BIT_EXT.
    constexpr VkBufferUsageFlags2 kVertexUsageBase =
        VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT
        | VK_BUFFER_USAGE_2_ACCELERATION_STRUCTURE_BUILD_INPUT_READ_ONLY_BIT_KHR;
    constexpr VkBufferUsageFlags2 kIndexUsageBase =
        kVertexUsageBase | VK_BUFFER_USAGE_STORAGE_BUFFER_BIT;

    // Translate the portable RTBuildOptions to Vulkan build flags. The two
    // preferences are mutually exclusive; ALLOW_UPDATE is layered on top so
    // a later in-place refit (UPDATE mode) is permitted.
    VkBuildAccelerationStructureFlagsKHR BlasFlags(const RTBuildOptions& options) {
        VkBuildAccelerationStructureFlagsKHR flags =
            options.preference == RTBuildPreference::FastBuild
                ? VK_BUILD_ACCELERATION_STRUCTURE_PREFER_FAST_BUILD_BIT_KHR
                : VK_BUILD_ACCELERATION_STRUCTURE_PREFER_FAST_TRACE_BIT_KHR;
        if (options.allowUpdate)
            flags |= VK_BUILD_ACCELERATION_STRUCTURE_ALLOW_UPDATE_BIT_KHR;
        return flags;
    }

    // Shared BLAS sizing / creation / build-record tail — geometry-type
    // agnostic; both the triangle and the AABB (procedural) paths feed
    // their single geometry through here. `update` selects an in-place
    // refit (VK_BUILD_ACCELERATION_STRUCTURE_MODE_UPDATE_KHR, src == dst)
    // over a fresh build; the caller guarantees the geometry topology and
    // primitiveCount match the existing AS and that `flags` carried
    // ALLOW_UPDATE on the original build.
    void RecordBLASBuildFromGeometry(Mesh& self, const VkAccelerationStructureGeometryKHR& blasGeometry, std::uint32_t primitiveCount, VkBuildAccelerationStructureFlagsKHR flags, bool update, VkCommandBuffer cmd) {
        VkAccelerationStructureBuildGeometryInfoKHR blasBuildGeometryInfo{
            .sType         = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_BUILD_GEOMETRY_INFO_KHR,
            .type          = VK_ACCELERATION_STRUCTURE_TYPE_BOTTOM_LEVEL_KHR,
            .flags         = flags,
            .mode          = update
                                ? VK_BUILD_ACCELERATION_STRUCTURE_MODE_UPDATE_KHR
                                : VK_BUILD_ACCELERATION_STRUCTURE_MODE_BUILD_KHR,
            .geometryCount = 1,
            .pGeometries   = &blasGeometry,
        };

        if (!update) {
            // Fresh build: query sizes, (re)allocate scratch + AS storage,
            // create a new AS handle. The scratch buffer is sized for the
            // build (always >= the update scratch requirement), so it also
            // covers later in-place refits without resizing.
            VkAccelerationStructureBuildSizesInfoKHR blasBuildSizes = {
                .sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_BUILD_SIZES_INFO_KHR
            };
            Device::vkGetAccelerationStructureBuildSizesKHR(
                Device::device,
                VK_ACCELERATION_STRUCTURE_BUILD_TYPE_DEVICE_KHR,
                &blasBuildGeometryInfo,
                &primitiveCount,
                &blasBuildSizes
            );

            self.scratchBuffer.Resize(
                VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT,
                VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT,
                blasBuildSizes.buildScratchSize);

            self.blasBuffer.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,
                blasBuildSizes.accelerationStructureSize);

            // Re-Build of an existing Mesh: destroy the stale handle first
            // so the previous AS storage isn't leaked.
            if (self.accelerationStructure != VK_NULL_HANDLE) {
                Device::vkDestroyAccelerationStructureKHR(Device::device, self.accelerationStructure, nullptr);
                self.accelerationStructure = VK_NULL_HANDLE;
            }

            VkAccelerationStructureCreateInfoKHR blasCreateInfo{
                .sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_CREATE_INFO_KHR,
                .buffer = self.blasBuffer.buffer,
                .offset = 0,
                .size   = blasBuildSizes.accelerationStructureSize,
                .type   = VK_ACCELERATION_STRUCTURE_TYPE_BOTTOM_LEVEL_KHR,
            };
            Device::CheckVkResult(Device::vkCreateAccelerationStructureKHR(Device::device, &blasCreateInfo, nullptr, &self.accelerationStructure));

            VkAccelerationStructureDeviceAddressInfoKHR addrInfo {
                .sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_DEVICE_ADDRESS_INFO_KHR,
                .accelerationStructure = self.accelerationStructure
            };
            self.blasAddr = Device::vkGetAccelerationStructureDeviceAddressKHR(Device::device, &addrInfo);
        } else {
            // In-place refit: the AS may have been read by a prior TLAS
            // build / trace, so order that read before the build overwrites
            // it. Scratch reuse needs the same write-after-read guard.
            VkMemoryBarrier asBarrier {
                .sType = VK_STRUCTURE_TYPE_MEMORY_BARRIER,
                .srcAccessMask = VK_ACCESS_ACCELERATION_STRUCTURE_READ_BIT_KHR | 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_RAY_TRACING_SHADER_BIT_KHR,
                VK_PIPELINE_STAGE_ACCELERATION_STRUCTURE_BUILD_BIT_KHR,
                0, 1, &asBarrier, 0, nullptr, 0, nullptr);
        }

        blasBuildGeometryInfo.scratchData.deviceAddress = self.scratchBuffer.address;
        blasBuildGeometryInfo.dstAccelerationStructure  = self.accelerationStructure;
        // UPDATE refits in place (src == dst); a fresh build has no source.
        blasBuildGeometryInfo.srcAccelerationStructure  =
            update ? self.accelerationStructure : VK_NULL_HANDLE;

        VkAccelerationStructureBuildRangeInfoKHR blasRangeInfo {
            .primitiveCount  = primitiveCount,
            .primitiveOffset = 0,
            .firstVertex     = 0,
            .transformOffset = 0
        };
        VkAccelerationStructureBuildRangeInfoKHR* blasRangeInfoPP = &blasRangeInfo;
        Device::vkCmdBuildAccelerationStructuresKHR(cmd, 1, &blasBuildGeometryInfo, &blasRangeInfoPP);

        self.buildFlags          = flags;
        self.builtPrimitiveCount = primitiveCount;
    }

    void RecordBLASBuild(Mesh& self, std::uint32_t vertexCount, std::uint32_t indexCount, VkBuildAccelerationStructureFlagsKHR flags, bool update, VkCommandBuffer cmd) {
        VkDeviceOrHostAddressConstKHR vertexAddr;
        vertexAddr.deviceAddress = self.vertexBuffer.address;

        VkDeviceOrHostAddressConstKHR indexAddr;
        indexAddr.deviceAddress = self.indexBuffer.address;

        auto trianglesData = VkAccelerationStructureGeometryTrianglesDataKHR {
            .sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_TRIANGLES_DATA_KHR,
            .vertexFormat = VK_FORMAT_R32G32B32_SFLOAT,
            .vertexData = vertexAddr,
            .vertexStride = sizeof(Vector<float, 3, 3>),
            .maxVertex = vertexCount - 1,
            .indexType = VK_INDEX_TYPE_UINT32,
            .indexData = indexAddr,
            .transformData = {.deviceAddress = 0}
        };
        VkAccelerationStructureGeometryDataKHR geometryData(trianglesData);
        VkAccelerationStructureGeometryKHR blasGeometry {
            .sType = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_KHR,
            .geometryType = VK_GEOMETRY_TYPE_TRIANGLES_KHR,
            .geometry     = geometryData,
            .flags        = VK_GEOMETRY_OPAQUE_BIT_KHR
        };

        RecordBLASBuildFromGeometry(self, blasGeometry, indexCount / 3, flags, update, cmd);
    }
}

void Mesh::Build(std::span<Vector<float, 3, 3>> verticies, std::span<std::uint32_t> indicies, VkCommandBuffer cmd, RTBuildOptions options) {
    vertexBuffer.Resize(kVertexUsageBase, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT, verticies.size());
    indexBuffer.Resize(kIndexUsageBase, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT, indicies.size());

    std::memcpy(vertexBuffer.value, verticies.data(), verticies.size() * sizeof(Vector<float, 3, 3>));
    std::memcpy(indexBuffer.value, indicies.data(), indicies.size() * sizeof(std::uint32_t));

    vertexBuffer.FlushDevice(cmd, VK_ACCESS_MEMORY_READ_BIT, VK_PIPELINE_STAGE_ACCELERATION_STRUCTURE_BUILD_BIT_KHR);
    indexBuffer.FlushDevice(cmd, VK_ACCESS_MEMORY_READ_BIT, VK_PIPELINE_STAGE_ACCELERATION_STRUCTURE_BUILD_BIT_KHR);

    allowUpdate     = options.allowUpdate;
    builtInputCount = static_cast<std::uint32_t>(verticies.size());
    RecordBLASBuild(*this, static_cast<std::uint32_t>(verticies.size()), static_cast<std::uint32_t>(indicies.size()), BlasFlags(options), /*update*/ false, cmd);
}

void Mesh::Build(const CompressedMeshAsset& asset, VkCommandBuffer cmd, RTBuildOptions options) {
    if (!Device::memoryDecompressionSupported) {
        // CPU fallback: decompress into temporary host vectors, then take
        // the existing uncompressed path. The data region is decompressed
        // into a discard buffer (consumer is expected to handle data-stream
        // decoding via Compression::DecompressCPU on its own buffer).
        std::vector<Vector<float, 3, 3>> vertices(asset.vertexCount);
        std::vector<std::uint32_t> indices(asset.indexCount);
        std::vector<std::byte> dataDiscard(
            asset.blob.regions.size() >= 3 ? asset.blob.regions[2].decompressedSize : 0);
        std::array<std::span<std::byte>, 3> outputs = {
            std::as_writable_bytes(std::span(vertices)),
            std::as_writable_bytes(std::span(indices)),
            std::span<std::byte>(dataDiscard),
        };
        Compression::DecompressCPU(asset.blob, std::span(outputs).first(asset.blob.regions.size()));
        Build(vertices, indices, cmd, options);
        return;
    }

    vertexBuffer.Resize(
        kVertexUsageBase | VK_BUFFER_USAGE_2_MEMORY_DECOMPRESSION_BIT_EXT,
        VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT,
        asset.vertexCount);
    indexBuffer.Resize(
        kIndexUsageBase | VK_BUFFER_USAGE_2_MEMORY_DECOMPRESSION_BIT_EXT,
        VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT,
        asset.indexCount);

    compressedStaging.Resize(
        VK_BUFFER_USAGE_TRANSFER_SRC_BIT
        | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT
        | VK_BUFFER_USAGE_2_MEMORY_DECOMPRESSION_BIT_EXT,
        VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT,
        static_cast<std::uint32_t>(asset.blob.bytes.size()));
    std::memcpy(compressedStaging.value, asset.blob.bytes.data(), asset.blob.bytes.size());
    compressedStaging.FlushDevice();

    std::array<VkDeviceAddress, 3> dstAddresses = {
        vertexBuffer.address,
        indexBuffer.address,
        0,  // data region is not consumed by Mesh; caller handles it separately.
    };
    std::vector<VkDecompressMemoryRegionEXT> regions;
    for (std::size_t i = 0; i < asset.blob.regions.size() && i < 2; ++i) {
        const Compression::RegionMeta& r = asset.blob.regions[i];
        if (r.decompressedSize == 0) continue;
        std::span<const std::byte> streamBytes(
            asset.blob.bytes.data() + r.srcOffset,
            static_cast<std::size_t>(r.compressedSize));
        Decompress::ExpandStreamToTileRegions(
            streamBytes,
            compressedStaging.address + r.srcOffset,
            dstAddresses[i],
            regions);
    }

    Decompress::DecompressOnGPU(
        cmd,
        regions,
        VK_PIPELINE_STAGE_2_ACCELERATION_STRUCTURE_BUILD_BIT_KHR,
        VK_ACCESS_2_ACCELERATION_STRUCTURE_READ_BIT_KHR);

    allowUpdate     = options.allowUpdate;
    builtInputCount = asset.vertexCount;
    RecordBLASBuild(*this, asset.vertexCount, asset.indexCount, BlasFlags(options), /*update*/ false, cmd);
}

namespace {
    // Record an AABB BLAS build (fresh or in-place refit) from a device
    // address that already holds the VkAabbPositionsKHR-compatible boxes.
    // Geometry-source agnostic: the host-upload path (RecordProceduralBuild)
    // points this at self.aabbBuffer, while the device-buffer overloads
    // point it straight at the producer's buffer — no copy involved either
    // way. The geometry's opaque bit is read from self.opaque so an UPDATE
    // keeps the exact geometry description of the original build.
    void RecordProceduralBuildFromAddress(Mesh& self, VkDeviceAddress aabbAddress, std::uint32_t count, std::uint32_t stride, VkBuildAccelerationStructureFlagsKHR flags, bool update, VkCommandBuffer cmd) {
        VkAccelerationStructureGeometryAabbsDataKHR aabbsData {
            .sType  = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_AABBS_DATA_KHR,
            .data   = { .deviceAddress = aabbAddress },
            .stride = stride
        };
        VkAccelerationStructureGeometryDataKHR geometryData;
        geometryData.aabbs = aabbsData;
        VkAccelerationStructureGeometryKHR blasGeometry {
            .sType        = VK_STRUCTURE_TYPE_ACCELERATION_STRUCTURE_GEOMETRY_KHR,
            .geometryType = VK_GEOMETRY_TYPE_AABBS_KHR,
            .geometry     = geometryData,
            // Non-opaque by default (mirrors the WebGPU path) so any-hit
            // shaders run for procedural geometry unless the caller opts out.
            .flags        = self.opaque ? static_cast<VkGeometryFlagsKHR>(VK_GEOMETRY_OPAQUE_BIT_KHR) : VkGeometryFlagsKHR{}
        };

        RecordBLASBuildFromGeometry(self, blasGeometry, count, flags, update, cmd);
    }

    // Re-upload AABB build input from host memory and record an AABB BLAS
    // build (fresh or in-place refit). Shared by the host-span BuildProcedural
    // and RefitProcedural.
    void RecordProceduralBuild(Mesh& self, std::span<const RTAabb> aabbs, VkBuildAccelerationStructureFlagsKHR flags, bool update, VkCommandBuffer cmd) {
        // 24-byte-stride VkAabbPositionsKHR-compatible build input
        // (static_assert'd in the interface). Same usage set as the triangle
        // inputs: AS-build read-only + device address. A refit reuses the
        // existing same-sized buffer (count is unchanged), so the device
        // address — and the geometry it feeds — stays stable.
        if (!update) {
            self.aabbBuffer.Resize(kVertexUsageBase, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT, static_cast<std::uint32_t>(aabbs.size()));
        }
        std::memcpy(self.aabbBuffer.value, aabbs.data(), aabbs.size() * sizeof(RTAabb));
        self.aabbBuffer.FlushDevice(cmd, VK_ACCESS_MEMORY_READ_BIT, VK_PIPELINE_STAGE_ACCELERATION_STRUCTURE_BUILD_BIT_KHR);

        RecordProceduralBuildFromAddress(self, self.aabbBuffer.address, static_cast<std::uint32_t>(aabbs.size()), sizeof(RTAabb), flags, update, cmd);
    }
}

void Mesh::BuildProcedural(std::span<const RTAabb> aabbs, bool opaque, VkCommandBuffer cmd, RTBuildOptions options) {
    this->opaque    = opaque;
    allowUpdate     = options.allowUpdate;
    builtInputCount = static_cast<std::uint32_t>(aabbs.size());

    RecordProceduralBuild(*this, aabbs, BlasFlags(options), /*update*/ false, cmd);
}

void Mesh::BuildProcedural(VkDeviceAddress aabbAddress, std::uint32_t count, bool opaque, VkCommandBuffer cmd, RTBuildOptions options, std::uint32_t stride) {
    // Zero-copy procedural build: the AABBs already live in `aabbAddress`
    // (a device buffer a GPU compute pass wrote). Nothing touches
    // self.aabbBuffer — the build input is the producer's buffer directly.
    // The caller owns that buffer's lifetime + usage flags
    // (ACCELERATION_STRUCTURE_BUILD_INPUT_READ_ONLY + SHADER_DEVICE_ADDRESS)
    // and must barrier the producing writes before this build reads them.
    this->opaque    = opaque;
    allowUpdate     = options.allowUpdate;
    builtInputCount = count;

    RecordProceduralBuildFromAddress(*this, aabbAddress, count, stride, BlasFlags(options), /*update*/ false, cmd);
}

void Mesh::Refit(std::span<Vector<float, 3, 3>> verticies, std::span<std::uint32_t> indicies, VkCommandBuffer cmd) {
    // A hardware in-place UPDATE is only valid when the original build asked
    // for it, an AS already exists, and the topology is unchanged. Otherwise
    // fall back to a full rebuild (mode BUILD) of the same buffers — still
    // correct, just not the cheap path.
    const bool sameTopology =
        accelerationStructure != VK_NULL_HANDLE
        && verticies.size() == builtInputCount
        && indicies.size()  == static_cast<std::size_t>(builtPrimitiveCount) * 3;
    const bool update = allowUpdate && sameTopology;

    if (!update) {
        // Counts may have changed (or update wasn't permitted): take the
        // resizing build path, preserving the caller's original flags.
        Build(verticies, indicies, cmd, RTBuildOptions{
            .preference = (buildFlags & VK_BUILD_ACCELERATION_STRUCTURE_PREFER_FAST_BUILD_BIT_KHR)
                              ? RTBuildPreference::FastBuild : RTBuildPreference::FastTrace,
            .allowUpdate = allowUpdate,
        });
        return;
    }

    // Same-sized buffers: overwrite in place so the device addresses (and
    // thus the geometry the UPDATE reads) stay stable.
    std::memcpy(vertexBuffer.value, verticies.data(), verticies.size() * sizeof(Vector<float, 3, 3>));
    std::memcpy(indexBuffer.value, indicies.data(), indicies.size() * sizeof(std::uint32_t));
    vertexBuffer.FlushDevice(cmd, VK_ACCESS_MEMORY_READ_BIT, VK_PIPELINE_STAGE_ACCELERATION_STRUCTURE_BUILD_BIT_KHR);
    indexBuffer.FlushDevice(cmd, VK_ACCESS_MEMORY_READ_BIT, VK_PIPELINE_STAGE_ACCELERATION_STRUCTURE_BUILD_BIT_KHR);

    RecordBLASBuild(*this, static_cast<std::uint32_t>(verticies.size()), static_cast<std::uint32_t>(indicies.size()), buildFlags, /*update*/ true, cmd);
}

void Mesh::RefitProcedural(std::span<const RTAabb> aabbs, VkCommandBuffer cmd) {
    const bool sameTopology =
        accelerationStructure != VK_NULL_HANDLE
        && aabbs.size() == builtInputCount;
    const bool update = allowUpdate && sameTopology;

    if (!update && aabbs.size() != builtInputCount) {
        // Topology changed — a full rebuild needs the new count.
        builtInputCount = static_cast<std::uint32_t>(aabbs.size());
    }
    RecordProceduralBuild(*this, aabbs, buildFlags, update, cmd);
}

void Mesh::RefitProcedural(VkDeviceAddress aabbAddress, std::uint32_t count, VkCommandBuffer cmd, std::uint32_t stride) {
    // Zero-copy refit: the GPU producer rewrote the same device buffer in
    // place; re-record the build (UPDATE when allowed + count unchanged,
    // else a full rebuild) reading straight from `aabbAddress`. As with the
    // device-buffer BuildProcedural, nothing touches self.aabbBuffer and the
    // caller is responsible for barriering the producing writes.
    const bool sameTopology =
        accelerationStructure != VK_NULL_HANDLE
        && count == builtInputCount;
    const bool update = allowUpdate && sameTopology;

    if (!update && count != builtInputCount) {
        builtInputCount = count;
    }
    RecordProceduralBuildFromAddress(*this, aabbAddress, count, stride, buildFlags, update, cmd);
}