Crafter.Graphics/interfaces/Crafter.Graphics-ImageVulkan.cppm

/*
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  0215-1301 USA
*/

module;

#ifdef CRAFTER_GRAPHICS_VULKAN
#include <vulkan/vulkan.h>
#endif

export module Crafter.Graphics:ImageVulkan;
import std;
import :VulkanBuffer;

export namespace Crafter {
    #ifdef CRAFTER_GRAPHICS_VULKAN
    template <typename PixelType>
    class ImageVulkan {
    public:
        std::uint32_t width;
        std::uint32_t height;
        VkImage image;
        VkDeviceMemory imageMemory;
        VulkanBuffer<PixelType, true, false, false> buffer;
        VkImageView imageView;
        VkDescriptorImageInfo descriptor;

        void Create(std::uint32_t width, std::uint32_t height, VkCommandBuffer cmd, VkFormat format) {
            this->width = width;
            this->height = height;

            buffer.Create(VK_BUFFER_USAGE_TRANSFER_SRC_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 = 5;
            imageInfo.arrayLayers = 1;
            imageInfo.format = format;
            imageInfo.tiling = VK_IMAGE_TILING_OPTIMAL;
            imageInfo.initialLayout = VK_IMAGE_LAYOUT_UNDEFINED;
            imageInfo.usage = VK_IMAGE_USAGE_TRANSFER_SRC_BIT | VK_IMAGE_USAGE_TRANSFER_DST_BIT | VK_IMAGE_USAGE_SAMPLED_BIT; // Include TRANSFER_SRC_BIT
            imageInfo.samples = VK_SAMPLE_COUNT_1_BIT;
            imageInfo.sharingMode = VK_SHARING_MODE_EXCLUSIVE;
            
            VulkanDevice::CheckVkResult(vkCreateImage(VulkanDevice::device, &imageInfo, nullptr, &image));

            VkMemoryRequirements memRequirements;
            vkGetImageMemoryRequirements(VulkanDevice::device, image, &memRequirements);
            
            VkMemoryAllocateInfo allocInfo = {};
            allocInfo.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;
            allocInfo.allocationSize = memRequirements.size;
            allocInfo.memoryTypeIndex = VulkanDevice::GetMemoryType(memRequirements.memoryTypeBits, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT);
            
            VulkanDevice::CheckVkResult(vkAllocateMemory(VulkanDevice::device, &allocInfo, nullptr, &imageMemory));
            vkBindImageMemory(VulkanDevice::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 = 5;
            viewInfo.subresourceRange.baseArrayLayer = 0;
            viewInfo.subresourceRange.layerCount = 1;

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

            // Final transition to shader read-only layout
            TransitionImageLayout(cmd, image, VK_IMAGE_LAYOUT_UNDEFINED, VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL, 0, 5);
        }

        void Update(const PixelType* bufferdata, VkCommandBuffer cmd) {
            std::memcpy(buffer.value, bufferdata, height*width*sizeof(PixelType));
            buffer.FlushDevice(cmd, VK_ACCESS_MEMORY_READ_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT);
            TransitionImageLayout(cmd, image, VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, 0, 5);

            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
            );

            //Transition source and destination mip levels to correct layout
            TransitionImageLayout(cmd, image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL, 0, 1);

            // Mipmap generation (now includes the correct layout and offset management)
            for (std::uint32_t i = 1; i < 5; ++i) {
                std::uint32_t mipWidth = std::max(1u, width >> i);
                std::uint32_t mipHeight = std::max(1u, height >> i);

                std::uint32_t previousMipWidth = std::max(1u, width >> (i - 1));
                std::uint32_t previousMipHeight = std::max(1u, height >> (i - 1));

                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 }; // Ensure srcOffsets[0] is 0
                blit.srcOffsets[1] = { (int32_t)previousMipWidth, (int32_t)previousMipHeight, 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 }; // Ensure dstOffsets[0] is 0
                blit.dstOffsets[1] = { (int32_t)mipWidth, (int32_t)mipHeight, 1 };

                vkCmdBlitImage(cmd, image, VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL, image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, 1, &blit, VK_FILTER_LINEAR);

                // Transition back after blitting
                TransitionImageLayout(cmd, image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL, i, 1);
            }

            TransitionImageLayout(cmd, image, VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL, VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL, 0, 5);
        }

        private:
        void TransitionImageLayout(VkCommandBuffer cmd, VkImage image, VkImageLayout oldLayout, VkImageLayout newLayout, 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 = 0;
            barrier.dstAccessMask = 0;

            VkPipelineStageFlags sourceStage;
            VkPipelineStageFlags destinationStage;

            // Different cases for layout transition
            if (oldLayout == VK_IMAGE_LAYOUT_UNDEFINED && newLayout == VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL) {
                barrier.srcAccessMask = 0;
                barrier.dstAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT;
                sourceStage = VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT;
                destinationStage = VK_PIPELINE_STAGE_TRANSFER_BIT;
                // barrier.srcQueueFamilyIndex = VK_QUEUE_TRANSFER_BIT;
                // barrier.dstQueueFamilyIndex = VK_QUEUE_TRANSFER_BIT;
            } if (oldLayout == VK_IMAGE_LAYOUT_UNDEFINED && newLayout == VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL) {
                barrier.srcAccessMask = 0;
                barrier.dstAccessMask = VK_ACCESS_SHADER_READ_BIT;
                sourceStage = VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT;
                destinationStage = VK_PIPELINE_STAGE_RAY_TRACING_SHADER_BIT_KHR;
                // barrier.srcQueueFamilyIndex = VK_QUEUE_TRANSFER_BIT;
                // barrier.dstQueueFamilyIndex = VK_QUEUE_GRAPHICS_BIT;
            } else if (oldLayout == VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL && newLayout == VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL) {
                barrier.srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT;
                barrier.dstAccessMask = VK_ACCESS_SHADER_READ_BIT;
                sourceStage = VK_PIPELINE_STAGE_TRANSFER_BIT;
                destinationStage = VK_PIPELINE_STAGE_RAY_TRACING_SHADER_BIT_KHR;
                // barrier.srcQueueFamilyIndex = VK_QUEUE_TRANSFER_BIT;
                // barrier.dstQueueFamilyIndex = VK_QUEUE_GRAPHICS_BIT;
            } else if (oldLayout == VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL && newLayout == VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL) {
                barrier.srcAccessMask = VK_ACCESS_TRANSFER_READ_BIT;
                barrier.dstAccessMask = VK_ACCESS_SHADER_READ_BIT;
                sourceStage = VK_PIPELINE_STAGE_TRANSFER_BIT;
                destinationStage = VK_PIPELINE_STAGE_RAY_TRACING_SHADER_BIT_KHR;
                // barrier.srcQueueFamilyIndex = VK_QUEUE_TRANSFER_BIT;
                // barrier.dstQueueFamilyIndex = VK_QUEUE_GRAPHICS_BIT;
            } else if (oldLayout == VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL && newLayout == VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL) {
                barrier.srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT;
                barrier.dstAccessMask = VK_ACCESS_TRANSFER_READ_BIT;
                sourceStage = VK_PIPELINE_STAGE_TRANSFER_BIT;
                destinationStage = VK_PIPELINE_STAGE_TRANSFER_BIT;
                // barrier.srcQueueFamilyIndex = VK_QUEUE_TRANSFER_BIT;
                // barrier.dstQueueFamilyIndex = VK_QUEUE_TRANSFER_BIT;
            } else if (oldLayout == VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL && newLayout == VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL) {
                barrier.srcAccessMask = VK_ACCESS_TRANSFER_READ_BIT;
                barrier.dstAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT;
                sourceStage = VK_PIPELINE_STAGE_TRANSFER_BIT;
                destinationStage = VK_PIPELINE_STAGE_TRANSFER_BIT;
                // barrier.srcQueueFamilyIndex = VK_QUEUE_TRANSFER_BIT;
                // barrier.dstQueueFamilyIndex = VK_QUEUE_TRANSFER_BIT;
            } else if (oldLayout == VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL && newLayout == VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL) {
                barrier.srcAccessMask = VK_ACCESS_SHADER_READ_BIT;
                barrier.dstAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT;
                sourceStage = VK_PIPELINE_STAGE_RAY_TRACING_SHADER_BIT_KHR;
                destinationStage = VK_PIPELINE_STAGE_TRANSFER_BIT;
                // barrier.srcQueueFamilyIndex = VK_QUEUE_GRAPHICS_BIT;
                // barrier.dstQueueFamilyIndex = VK_QUEUE_TRANSFER_BIT;
            } else {
                std::cout << oldLayout << std::endl;
                std::cout << newLayout << std::endl;
                throw std::invalid_argument("unsupported layout transition!");
            }

            vkCmdPipelineBarrier(cmd, sourceStage, destinationStage, 0, 0, nullptr, 0, nullptr, 1, &barrier);

            descriptor = { .imageView = imageView, .imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL };
        }
    };
    #endif
}