Crafter.Graphics/interfaces/Crafter.Graphics-ImageVulkan.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:ImageVulkan;
#ifndef CRAFTER_GRAPHICS_WINDOW_DOM
import std;
import Crafter.Asset;
import :Decompress;
import :VulkanBuffer;

export namespace Crafter {
    // Builds the batched final layout transition for a freshly-generated mip
    // chain (mipLevels >= 2). After the blit loop, levels [0, mipLevels-1) sit
    // in TRANSFER_SRC_OPTIMAL — each was read as a blit source — while the last
    // level is still in TRANSFER_DST_OPTIMAL: it was written by the final blit
    // and is never read, so it never needed a DST->SRC barrier of its own.
    // Both groups move to the consumer `layout` in one vkCmdPipelineBarrier
    // (two VkImageMemoryBarrier entries sharing src=TRANSFER / dst=consumer
    // stage), shaving one barrier call off the previous N+1-per-chain count.
    // The interleaved per-level barriers stay one-at-a-time — each is mandated
    // because the next blit reads the level it transitions. Writes the entries
    // into `out` and returns the count (always 2).
    std::uint32_t BuildMipChainFinalBarriers(VkImage image, VkImageLayout layout, std::uint32_t mipLevels, std::array<VkImageMemoryBarrier, 2>& out) {
        auto fill = [&](VkImageMemoryBarrier& b, VkImageLayout oldLayout, std::uint32_t baseMip, std::uint32_t count, VkAccessFlags srcAccess) {
            b = {};
            b.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER;
            b.oldLayout = oldLayout;
            b.newLayout = layout;
            b.image = image;
            b.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
            b.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
            b.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
            b.subresourceRange.baseMipLevel = baseMip;
            b.subresourceRange.levelCount = count;
            b.subresourceRange.baseArrayLayer = 0;
            b.subresourceRange.layerCount = 1;
            b.srcAccessMask = srcAccess;
            b.dstAccessMask = VK_ACCESS_SHADER_READ_BIT;
        };
        // Levels [0, mipLevels-1): blit sources, currently TRANSFER_SRC.
        fill(out[0], VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL, 0, mipLevels - 1u, VK_ACCESS_TRANSFER_READ_BIT);
        // Final level: the last blit's destination, still TRANSFER_DST.
        fill(out[1], VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, mipLevels - 1u, 1u, VK_ACCESS_TRANSFER_WRITE_BIT);
        return 2;
    }

    template <typename PixelType>
    class ImageVulkan {
    public:
        std::uint16_t width;
        std::uint16_t height;
        std::uint8_t mipLevels;
        VkImage image = VK_NULL_HANDLE;
        VkDeviceMemory imageMemory = VK_NULL_HANDLE;
        VulkanBuffer<PixelType, true> buffer;
        // Transient host-visible staging for the compressed Update path. Same
        // lifetime contract as Mesh::compressedStaging: the compressed Update
        // releases it via DeferredClear() right after recording the decompress,
        // so the fence-keyed deletion queue (#101/#102) frees it once that
        // submit's frame has cleared instead of pinning it for the image's life.
        // Between Updates the handle is null; the next Update's Resize re-creates it.
        VulkanBuffer<std::byte, true> compressedStaging;
        VkImageView imageView = VK_NULL_HANDLE;
        VkDescriptorImageInfo descriptor;
        // Pipeline stage that samples this image after upload. Set once at
        // Create and reused by every Update so the upload barrier's dst (and
        // the next-upload entry barrier's src) scope actually covers the real
        // consumer. Defaults to RT — the font atlas overrides it to COMPUTE,
        // since UI text is rendered by a compute shader, not an RT pipeline.
        VkPipelineStageFlags consumerStage = VK_PIPELINE_STAGE_RAY_TRACING_SHADER_BIT_KHR;
        // Whether this image is continuously re-uploaded from a persistent
        // CPU-side staging map. False (the default) marks a static texture
        // uploaded once — e.g. Sponza albedo: the first Update's buffer→image
        // copy is the only reader of the host-visible staging `buffer`, so it
        // is released to the fence-keyed deletion queue (#101/#102) right after
        // the copy instead of pinning HOST_VISIBLE / small-BAR memory for the
        // image's whole life (issue #114). True keeps `buffer` alive: a streamed
        // image (the FontAtlas) re-fills buffer.value on the CPU and re-uploads
        // every frame, so the persistent map must survive — its uploads go
        // through UpdateRegion, which never releases the staging.
        bool streamed = false;

        void Create(std::uint16_t width, std::uint16_t height, std::uint8_t mipLevels, VkCommandBuffer cmd, VkFormat format, VkImageCreateFlags flags, VkImageLayout layout, VkPipelineStageFlags consumerStage = VK_PIPELINE_STAGE_RAY_TRACING_SHADER_BIT_KHR, bool streamed = false) {
            this->width = width;
            this->height = height;
            this->mipLevels = mipLevels;
            // The RT default is not a valid pipeline stage on a device without
            // ray tracing, where the only shader that can sample an image is a
            // compute shader. Resolved here rather than in the member
            // initializer, which can run before Device::Initialize. The
            // parameter is corrected, not only the member: the transition at
            // the end of this function is recorded with it.
            if (!Device::rayTracingSupported && consumerStage == VK_PIPELINE_STAGE_RAY_TRACING_SHADER_BIT_KHR) {
                consumerStage = VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT;
            }
            this->consumerStage = consumerStage;
            this->streamed = streamed;
            buffer.Create(
                VK_BUFFER_USAGE_TRANSFER_SRC_BIT | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT,
                VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT,
                width * height
            );
            
            VkImageCreateInfo imageInfo = {};
            imageInfo.sType = VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO;
            imageInfo.imageType = VK_IMAGE_TYPE_2D;
            imageInfo.extent.width = width;
            imageInfo.extent.height = height;
            imageInfo.extent.depth = 1;
            imageInfo.mipLevels = mipLevels;
            imageInfo.arrayLayers = 1;
            imageInfo.format = format;
            imageInfo.tiling = VK_IMAGE_TILING_OPTIMAL;
            imageInfo.initialLayout = VK_IMAGE_LAYOUT_UNDEFINED;
            imageInfo.usage = flags;
            imageInfo.samples = VK_SAMPLE_COUNT_1_BIT;
            imageInfo.sharingMode = VK_SHARING_MODE_EXCLUSIVE;
            
            Device::CheckVkResult(vkCreateImage(Device::device, &imageInfo, nullptr, &image));

            VkMemoryRequirements memRequirements;
            vkGetImageMemoryRequirements(Device::device, image, &memRequirements);
            
            VkMemoryAllocateInfo allocInfo = {};
            allocInfo.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;
            allocInfo.allocationSize = memRequirements.size;
            allocInfo.memoryTypeIndex = Device::GetMemoryType(memRequirements.memoryTypeBits, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT);
            
            Device::CheckVkResult(vkAllocateMemory(Device::device, &allocInfo, nullptr, &imageMemory));
            vkBindImageMemory(Device::device, image, imageMemory, 0);

            VkImageViewCreateInfo viewInfo = {};
            viewInfo.sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO;
            viewInfo.image = image;
            viewInfo.viewType = VK_IMAGE_VIEW_TYPE_2D;
            viewInfo.format = format;
            viewInfo.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
            viewInfo.subresourceRange.baseMipLevel = 0;
            viewInfo.subresourceRange.levelCount = mipLevels;
            viewInfo.subresourceRange.baseArrayLayer = 0;
            viewInfo.subresourceRange.layerCount = 1;

            Device::CheckVkResult(vkCreateImageView(Device::device, &viewInfo, nullptr, &imageView));

            // Final transition to shader read-only layout
            TransitionImageLayout(cmd, image, VK_IMAGE_LAYOUT_UNDEFINED, layout, VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT, consumerStage, 0, VK_ACCESS_SHADER_READ_BIT, 0, mipLevels);

            descriptor = { .imageView = imageView, .imageLayout = layout };
        }

        void Update(VkCommandBuffer cmd, VkImageLayout layout) {
            buffer.FlushDevice(cmd, VK_ACCESS_MEMORY_READ_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT);
            TransitionImageLayout(cmd, image, layout, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, consumerStage, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_ACCESS_SHADER_READ_BIT, VK_ACCESS_TRANSFER_WRITE_BIT, 0, mipLevels);

            VkBufferImageCopy region{};
            region.bufferOffset = 0;
            region.bufferRowLength = 0;
            region.bufferImageHeight = 0;
            region.imageSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
            region.imageSubresource.mipLevel = 0;
            region.imageSubresource.baseArrayLayer = 0;
            region.imageSubresource.layerCount = 1;
            region.imageOffset = {0, 0, 0};
            region.imageExtent = { width, height, 1};
            vkCmdCopyBufferToImage(
                cmd,
                buffer.buffer,
                image,
                VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
                1,
                &region
            );

            BuildMipChain(cmd, layout);
            ReleaseStaging();
        }

        // Upload only the sub-rectangle [x, x+w) × [y, y+h) of the staging
        // buffer into the image, instead of the whole extent. The staging
        // buffer keeps the full image row stride, so bufferRowLength stays
        // `width` and bufferOffset just points at the rect's first texel —
        // imageOffset/imageExtent then carve out the rect on the GPU side.
        // Layout transitions stay whole-image (cheap); only the copy extent
        // shrinks. Single mip level only: the partial-upload path is for
        // CPU-streamed atlases (FontAtlas) which carry no mips, so there is
        // no blit chain to regenerate.
        void UpdateRegion(VkCommandBuffer cmd, VkImageLayout layout, std::uint32_t x, std::uint32_t y, std::uint32_t w, std::uint32_t h) {
            buffer.FlushDevice(cmd, VK_ACCESS_MEMORY_READ_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT);
            TransitionImageLayout(cmd, image, layout, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, consumerStage, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_ACCESS_SHADER_READ_BIT, VK_ACCESS_TRANSFER_WRITE_BIT, 0, mipLevels);

            VkBufferImageCopy region{};
            region.bufferOffset = (static_cast<VkDeviceSize>(y) * width + x) * sizeof(PixelType);
            region.bufferRowLength = width;
            region.bufferImageHeight = 0;
            region.imageSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
            region.imageSubresource.mipLevel = 0;
            region.imageSubresource.baseArrayLayer = 0;
            region.imageSubresource.layerCount = 1;
            region.imageOffset = { static_cast<std::int32_t>(x), static_cast<std::int32_t>(y), 0 };
            region.imageExtent = { w, h, 1 };
            vkCmdCopyBufferToImage(cmd, buffer.buffer, image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, 1, &region);

            TransitionImageLayout(cmd, image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, layout, VK_PIPELINE_STAGE_TRANSFER_BIT, consumerStage, VK_ACCESS_TRANSFER_WRITE_BIT, VK_ACCESS_SHADER_READ_BIT, 0, mipLevels);
        }

        // GPU compressed-asset Update: stage compressed bytes, decompress
        // into `buffer` via VK_EXT_memory_decompression, then copy buffer→image
        // and transition to `layout`. Falls back to CPU decode + the existing
        // Update path when Device::memoryDecompressionSupported is false.
        // Caller is responsible for the dimensions matching: asset.sizeX/sizeY
        // must equal this->width/height (set by Create), and asset.pixelStride
        // must equal sizeof(PixelType).
        void Update(const CompressedTextureAsset& asset, VkCommandBuffer cmd, VkImageLayout layout) {
            if (asset.pixelStride != sizeof(PixelType)) {
                throw std::runtime_error("ImageVulkan::Update(compressed): pixel stride mismatch");
            }
            if (!Device::memoryDecompressionSupported) {
                std::span<PixelType> dst{ buffer.value, static_cast<std::size_t>(width) * height };
                std::array<std::span<std::byte>, 1> outputs = {
                    std::as_writable_bytes(dst),
                };
                Compression::DecompressCPU(asset.blob, outputs);
                Update(cmd, layout);
                return;
            }

            // Re-create the staging-into-image buffer with MEMORY_DECOMPRESSION
            // permission so the GPU codec can write into it. Keeps it
            // HOST_VISIBLE (matches the existing path) — on UMA / ReBAR that's
            // a fast path, on older systems the decompress writes traverse
            // PCIe but correctness is unchanged.
            buffer.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>(width) * height);

            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::vector<VkDecompressMemoryRegionEXT> regions;
            for (const Compression::RegionMeta& r : asset.blob.regions) {
                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,
                    buffer.address,
                    regions);
            }
            Decompress::DecompressOnGPU(
                cmd,
                regions,
                VK_PIPELINE_STAGE_2_COPY_BIT,
                VK_ACCESS_2_TRANSFER_READ_BIT);

            // Compressed staging is read only by the decompress recorded above;
            // the buffer→image copy below reads `buffer` (the decompress dst),
            // never this. Release it to the fence-keyed deletion queue
            // (#101/#102) now instead of pinning host-visible memory for the
            // image's whole life. The recorded vkCmdDecompressMemoryEXT still
            // references compressedStaging.address, so it must outlive this
            // submit — which the queue guarantees by retiring the allocation
            // only after framesInFlight frames (i.e. after the submit's fence).
            compressedStaging.DeferredClear();

            // Continue with the existing buffer→image upload + layout transitions.
            // We've already inserted the decompress→transfer-read barrier,
            // so we skip the FlushDevice host-write barrier the regular Update
            // would emit (no host write happened).
            TransitionImageLayout(cmd, image, layout, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, consumerStage, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_ACCESS_SHADER_READ_BIT, VK_ACCESS_TRANSFER_WRITE_BIT, 0, mipLevels);

            VkBufferImageCopy region{};
            region.bufferOffset = 0;
            region.bufferRowLength = 0;
            region.bufferImageHeight = 0;
            region.imageSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
            region.imageSubresource.mipLevel = 0;
            region.imageSubresource.baseArrayLayer = 0;
            region.imageSubresource.layerCount = 1;
            region.imageOffset = {0, 0, 0};
            region.imageExtent = { width, height, 1 };
            vkCmdCopyBufferToImage(cmd, buffer.buffer, image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, 1, &region);

            BuildMipChain(cmd, layout);
            ReleaseStaging();
        }

        // After level 0 was copied in (in TRANSFER_DST): blits each level from
        // the one above it, then takes the whole image to `layout` for
        // consumerStage. A level is max(1, size >> level) on each side, as
        // Vulkan sizes it: a non-square image's shorter side reaches 1 before
        // the chain ends, and a blit to 0 would write nothing.
        void BuildMipChain(VkCommandBuffer cmd, VkImageLayout layout) {
            if (mipLevels <= 1) {
                TransitionImageLayout(cmd, image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, layout, VK_PIPELINE_STAGE_TRANSFER_BIT, consumerStage, VK_ACCESS_TRANSFER_WRITE_BIT, VK_ACCESS_SHADER_READ_BIT, 0, mipLevels);
                return;
            }
            auto levelSize = [](std::uint16_t size, std::uint32_t level) {
                return std::max<std::int32_t>(1, static_cast<std::int32_t>(size >> level));
            };
            TransitionImageLayout(cmd, image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_ACCESS_TRANSFER_WRITE_BIT, VK_ACCESS_TRANSFER_READ_BIT, 0, 1);
            for (std::uint32_t i = 1; i < mipLevels; ++i) {
                VkImageBlit blit = {};
                blit.srcSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
                blit.srcSubresource.mipLevel = i - 1;
                blit.srcSubresource.baseArrayLayer = 0;
                blit.srcSubresource.layerCount = 1;
                blit.srcOffsets[0] = { 0, 0, 0 };
                blit.srcOffsets[1] = { levelSize(width, i - 1), levelSize(height, i - 1), 1 };
                blit.dstSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
                blit.dstSubresource.mipLevel = i;
                blit.dstSubresource.baseArrayLayer = 0;
                blit.dstSubresource.layerCount = 1;
                blit.dstOffsets[0] = { 0, 0, 0 };
                blit.dstOffsets[1] = { levelSize(width, i), levelSize(height, i), 1 };
                vkCmdBlitImage(cmd, image, VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL, image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, 1, &blit, VK_FILTER_LINEAR);
                // The final blit's destination is never read again, so it
                // skips the DST->SRC barrier and is taken straight to the
                // consumer layout by the batched final transition below.
                if (i + 1 < mipLevels) {
                    TransitionImageLayout(cmd, image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_ACCESS_TRANSFER_WRITE_BIT, VK_ACCESS_TRANSFER_READ_BIT, i, 1);
                }
            }
            std::array<VkImageMemoryBarrier, 2> finalBarriers;
            std::uint32_t finalCount = BuildMipChainFinalBarriers(image, layout, mipLevels, finalBarriers);
            vkCmdPipelineBarrier(cmd, VK_PIPELINE_STAGE_TRANSFER_BIT, consumerStage, 0, 0, nullptr, 0, nullptr, finalCount, finalBarriers.data());
        }

        void Destroy() {
            vkDestroyImageView(Device::device, imageView, nullptr);
            vkDestroyImage(Device::device, image, nullptr);
            vkFreeMemory(Device::device, imageMemory, nullptr);
            imageView = VK_NULL_HANDLE;
            image = VK_NULL_HANDLE;
            imageMemory = VK_NULL_HANDLE;
            // Free any staging this image still owns: a streamed image's
            // persistent `buffer`, or either staging buffer on an image
            // destroyed before its first upload released them (issue #114).
            // ReleaseStaging()/Update already null released handles, so these
            // are no-ops then; Clear() is gated on a live handle so a static
            // texture (whose `buffer` was deferred-cleared) can't double-free.
            if (buffer.buffer != VK_NULL_HANDLE) buffer.Clear();
            if (compressedStaging.buffer != VK_NULL_HANDLE) compressedStaging.Clear();
        }

        private:
        // Release the host-visible staging `buffer` after a one-shot upload,
        // unless this is a streamed image (FontAtlas) whose persistent map must
        // survive. The recorded buffer→image copy still references buffer.buffer,
        // so route through the fence-keyed deletion queue (#101/#102) instead of
        // an immediate Clear(): DeferredClear nulls the handle (so re-entry and
        // Destroy treat the staging as already gone) while the allocation
        // outlives this submit's frame. Issue #114.
        void ReleaseStaging() {
            if (!streamed) {
                buffer.DeferredClear();
            }
        }

        void TransitionImageLayout(VkCommandBuffer cmd, VkImage image, VkImageLayout oldLayout, VkImageLayout newLayout, VkPipelineStageFlags sourceStage, VkPipelineStageFlags destinationStage, VkAccessFlags srcAccessMask, VkAccessFlags dstAccessMask, std::uint32_t mipLevel, std::uint32_t count) {
            VkImageMemoryBarrier barrier = {};
            barrier.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER;
            barrier.oldLayout = oldLayout;
            barrier.newLayout = newLayout;
            barrier.image = image;
            barrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
            barrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
            barrier.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
            barrier.subresourceRange.baseMipLevel = mipLevel;
            barrier.subresourceRange.levelCount = count;
            barrier.subresourceRange.baseArrayLayer = 0;
            barrier.subresourceRange.layerCount = 1;
            barrier.srcAccessMask = srcAccessMask;
            barrier.dstAccessMask = dstAccessMask;

            vkCmdPipelineBarrier(cmd, sourceStage, destinationStage, 0, 0, nullptr, 0, nullptr, 1, &barrier);
        }
    };
}
#endif // !CRAFTER_GRAPHICS_WINDOW_DOM