XCEngine/engine/src/RHI/Vulkan/VulkanCommandList.cpp

#include "XCEngine/RHI/Vulkan/VulkanCommandList.h"

#include "XCEngine/RHI/Vulkan/VulkanBuffer.h"
#include "XCEngine/RHI/Vulkan/VulkanDescriptorSet.h"
#include "XCEngine/RHI/Vulkan/VulkanDevice.h"
#include "XCEngine/RHI/Vulkan/VulkanPipelineLayout.h"
#include "XCEngine/RHI/Vulkan/VulkanPipelineState.h"
#include "XCEngine/RHI/Vulkan/VulkanResourceView.h"
#include "XCEngine/RHI/Vulkan/VulkanTexture.h"

#include <vector>

namespace XCEngine {
namespace RHI {

namespace {

VkImageLayout ToVulkanImageLayout(ResourceStates state) {
    switch (state) {
        case ResourceStates::RenderTarget:
            return VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL;
        case ResourceStates::CopySrc:
            return VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL;
        case ResourceStates::CopyDst:
            return VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL;
        case ResourceStates::Present:
            return VK_IMAGE_LAYOUT_PRESENT_SRC_KHR;
        case ResourceStates::PixelShaderResource:
            return VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
        case ResourceStates::DepthWrite:
            return VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL;
        case ResourceStates::Common:
        default:
            return VK_IMAGE_LAYOUT_UNDEFINED;
    }
}

VkAccessFlags ToVulkanAccessMask(ResourceStates state) {
    switch (state) {
        case ResourceStates::RenderTarget:
            return VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT;
        case ResourceStates::CopySrc:
            return VK_ACCESS_TRANSFER_READ_BIT;
        case ResourceStates::CopyDst:
            return VK_ACCESS_TRANSFER_WRITE_BIT;
        case ResourceStates::PixelShaderResource:
            return VK_ACCESS_SHADER_READ_BIT;
        case ResourceStates::DepthWrite:
            return VK_ACCESS_DEPTH_STENCIL_ATTACHMENT_WRITE_BIT;
        case ResourceStates::Present:
        case ResourceStates::Common:
        default:
            return 0;
    }
}

VkPipelineStageFlags ToVulkanStageMask(ResourceStates state) {
    switch (state) {
        case ResourceStates::RenderTarget:
            return VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT;
        case ResourceStates::CopySrc:
        case ResourceStates::CopyDst:
            return VK_PIPELINE_STAGE_TRANSFER_BIT;
        case ResourceStates::PixelShaderResource:
            return VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT;
        case ResourceStates::DepthWrite:
            return VK_PIPELINE_STAGE_EARLY_FRAGMENT_TESTS_BIT | VK_PIPELINE_STAGE_LATE_FRAGMENT_TESTS_BIT;
        case ResourceStates::Present:
            return VK_PIPELINE_STAGE_BOTTOM_OF_PIPE_BIT;
        case ResourceStates::Common:
        default:
            return VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT;
    }
}

} // namespace

VulkanCommandList::~VulkanCommandList() {
    Shutdown();
}

bool VulkanCommandList::Initialize(VulkanDevice* device) {
    if (device == nullptr || device->GetDevice() == VK_NULL_HANDLE) {
        return false;
    }

    m_device = device;

    VkCommandPoolCreateInfo poolInfo = {};
    poolInfo.sType = VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO;
    poolInfo.flags = VK_COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT;
    poolInfo.queueFamilyIndex = device->GetGraphicsQueueFamilyIndex();
    if (vkCreateCommandPool(device->GetDevice(), &poolInfo, nullptr, &m_commandPool) != VK_SUCCESS) {
        return false;
    }

    VkCommandBufferAllocateInfo allocateInfo = {};
    allocateInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO;
    allocateInfo.commandPool = m_commandPool;
    allocateInfo.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;
    allocateInfo.commandBufferCount = 1;
    return vkAllocateCommandBuffers(device->GetDevice(), &allocateInfo, &m_commandBuffer) == VK_SUCCESS;
}

void VulkanCommandList::Shutdown() {
    EndActiveRenderPass();
    DestroyTransientFramebuffers();

    if (m_commandPool != VK_NULL_HANDLE && m_device != nullptr) {
        if (m_commandBuffer != VK_NULL_HANDLE) {
            vkFreeCommandBuffers(m_device->GetDevice(), m_commandPool, 1, &m_commandBuffer);
        }
        vkDestroyCommandPool(m_device->GetDevice(), m_commandPool, nullptr);
    }

    m_commandBuffer = VK_NULL_HANDLE;
    m_commandPool = VK_NULL_HANDLE;
    m_currentColorTarget = nullptr;
    m_currentDepthTarget = nullptr;
    m_currentPipelineState = nullptr;
    m_device = nullptr;
}

void VulkanCommandList::Reset() {
    if (m_device == nullptr || m_commandPool == VK_NULL_HANDLE || m_commandBuffer == VK_NULL_HANDLE) {
        return;
    }

    vkResetCommandPool(m_device->GetDevice(), m_commandPool, 0);
    DestroyTransientFramebuffers();

    m_currentPipelineState = nullptr;
    m_currentPrimitiveTopology = PrimitiveTopology::TriangleList;
    m_renderPassActive = false;
    m_hasViewport = false;
    m_hasScissor = false;

    VkCommandBufferBeginInfo beginInfo = {};
    beginInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;
    beginInfo.flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT;
    vkBeginCommandBuffer(m_commandBuffer, &beginInfo);
}

void VulkanCommandList::Close() {
    EndActiveRenderPass();

    auto* colorView = static_cast<VulkanResourceView*>(m_currentColorTarget);
    if (colorView != nullptr && colorView->GetTexture() != nullptr) {
        TransitionTexture(colorView->GetTexture(), ResourceStates::Present);
    }

    if (m_commandBuffer != VK_NULL_HANDLE) {
        vkEndCommandBuffer(m_commandBuffer);
    }
}

void VulkanCommandList::TransitionTexture(VulkanTexture* texture, ResourceStates newState) {
    if (texture == nullptr || m_commandBuffer == VK_NULL_HANDLE) {
        return;
    }

    if (texture->GetState() == newState) {
        return;
    }

    VkImageMemoryBarrier barrier = {};
    barrier.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER;
    barrier.oldLayout = ToVulkanImageLayout(texture->GetState());
    barrier.newLayout = ToVulkanImageLayout(newState);
    barrier.srcAccessMask = ToVulkanAccessMask(texture->GetState());
    barrier.dstAccessMask = ToVulkanAccessMask(newState);
    barrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
    barrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
    barrier.image = texture->GetImage();
    barrier.subresourceRange.aspectMask = GetImageAspectMask(texture->GetFormat());
    barrier.subresourceRange.baseMipLevel = 0;
    barrier.subresourceRange.levelCount = 1;
    barrier.subresourceRange.baseArrayLayer = 0;
    barrier.subresourceRange.layerCount = 1;

    vkCmdPipelineBarrier(
        m_commandBuffer,
        ToVulkanStageMask(texture->GetState()),
        ToVulkanStageMask(newState),
        0,
        0, nullptr,
        0, nullptr,
        1, &barrier);

    texture->SetState(newState);
}

void VulkanCommandList::TransitionBarrier(RHIResourceView* resource, ResourceStates stateBefore, ResourceStates stateAfter) {
    (void)stateBefore;
    auto* view = static_cast<VulkanResourceView*>(resource);
    if (view != nullptr && view->GetTexture() != nullptr) {
        TransitionTexture(view->GetTexture(), stateAfter);
    }
}

void VulkanCommandList::BeginRenderPass(class RHIRenderPass* renderPass, class RHIFramebuffer* framebuffer, const Rect& renderArea, uint32_t clearValueCount, const ClearValue* clearValues) {
    (void)renderPass;
    (void)framebuffer;
    (void)renderArea;
    (void)clearValueCount;
    (void)clearValues;
}

void VulkanCommandList::EndRenderPass() {
    EndActiveRenderPass();
}

void VulkanCommandList::SetShader(RHIShader* shader) {
    (void)shader;
}

void VulkanCommandList::SetPipelineState(RHIPipelineState* pso) {
    m_currentPipelineState = static_cast<VulkanPipelineState*>(pso);
}

void VulkanCommandList::SetGraphicsDescriptorSets(uint32_t firstSet, uint32_t count, RHIDescriptorSet** descriptorSets, RHIPipelineLayout* pipelineLayout) {
    if (count == 0 || descriptorSets == nullptr || m_commandBuffer == VK_NULL_HANDLE) {
        return;
    }

    VkPipelineLayout nativePipelineLayout = VK_NULL_HANDLE;
    if (pipelineLayout != nullptr) {
        nativePipelineLayout = static_cast<VulkanPipelineLayout*>(pipelineLayout)->GetPipelineLayout();
    } else if (m_currentPipelineState != nullptr) {
        nativePipelineLayout = m_currentPipelineState->GetPipelineLayout();
    }

    if (nativePipelineLayout == VK_NULL_HANDLE) {
        return;
    }

    std::vector<VkDescriptorSet> nativeSets;
    nativeSets.reserve(count);
    for (uint32_t i = 0; i < count; ++i) {
        auto* descriptorSet = static_cast<VulkanDescriptorSet*>(descriptorSets[i]);
        if (descriptorSet == nullptr || descriptorSet->GetDescriptorSet() == VK_NULL_HANDLE) {
            return;
        }
        nativeSets.push_back(descriptorSet->GetDescriptorSet());
    }

    vkCmdBindDescriptorSets(
        m_commandBuffer,
        VK_PIPELINE_BIND_POINT_GRAPHICS,
        nativePipelineLayout,
        firstSet,
        static_cast<uint32_t>(nativeSets.size()),
        nativeSets.data(),
        0,
        nullptr);
}

void VulkanCommandList::SetComputeDescriptorSets(uint32_t firstSet, uint32_t count, RHIDescriptorSet** descriptorSets, RHIPipelineLayout* pipelineLayout) {
    if (count == 0 || descriptorSets == nullptr || m_commandBuffer == VK_NULL_HANDLE || pipelineLayout == nullptr) {
        return;
    }

    std::vector<VkDescriptorSet> nativeSets;
    nativeSets.reserve(count);
    for (uint32_t i = 0; i < count; ++i) {
        auto* descriptorSet = static_cast<VulkanDescriptorSet*>(descriptorSets[i]);
        if (descriptorSet == nullptr || descriptorSet->GetDescriptorSet() == VK_NULL_HANDLE) {
            return;
        }
        nativeSets.push_back(descriptorSet->GetDescriptorSet());
    }

    vkCmdBindDescriptorSets(
        m_commandBuffer,
        VK_PIPELINE_BIND_POINT_COMPUTE,
        static_cast<VulkanPipelineLayout*>(pipelineLayout)->GetPipelineLayout(),
        firstSet,
        static_cast<uint32_t>(nativeSets.size()),
        nativeSets.data(),
        0,
        nullptr);
}

void VulkanCommandList::SetPrimitiveTopology(PrimitiveTopology topology) {
    m_currentPrimitiveTopology = topology;
}

void VulkanCommandList::SetViewport(const Viewport& viewport) {
    m_viewport.x = viewport.topLeftX;
    // Match the RHI viewport convention used by D3D12/OpenGL, where the origin is top-left.
    m_viewport.y = viewport.topLeftY + viewport.height;
    m_viewport.width = viewport.width;
    m_viewport.height = -viewport.height;
    m_viewport.minDepth = viewport.minDepth;
    m_viewport.maxDepth = viewport.maxDepth;
    m_hasViewport = true;
    vkCmdSetViewport(m_commandBuffer, 0, 1, &m_viewport);
}

void VulkanCommandList::SetViewports(uint32_t count, const Viewport* viewports) {
    if (count == 0 || viewports == nullptr) {
        return;
    }

    SetViewport(viewports[0]);
}

void VulkanCommandList::SetScissorRect(const Rect& rect) {
    m_scissor.offset.x = rect.left;
    m_scissor.offset.y = rect.top;
    m_scissor.extent.width = static_cast<uint32_t>(rect.right - rect.left);
    m_scissor.extent.height = static_cast<uint32_t>(rect.bottom - rect.top);
    m_hasScissor = true;
    vkCmdSetScissor(m_commandBuffer, 0, 1, &m_scissor);
}

void VulkanCommandList::SetScissorRects(uint32_t count, const Rect* rects) {
    if (count == 0 || rects == nullptr) {
        return;
    }

    SetScissorRect(rects[0]);
}

void VulkanCommandList::SetRenderTargets(uint32_t count, RHIResourceView** renderTargets, RHIResourceView* depthStencil) {
    m_currentColorTarget = (count > 0 && renderTargets != nullptr) ? renderTargets[0] : nullptr;
    m_currentDepthTarget = depthStencil;
}

void VulkanCommandList::SetStencilRef(uint8_t ref) {
    (void)ref;
}

void VulkanCommandList::SetBlendFactor(const float factor[4]) {
    (void)factor;
}

void VulkanCommandList::SetVertexBuffers(uint32_t startSlot, uint32_t count, RHIResourceView** buffers, const uint64_t* offsets, const uint32_t* strides) {
    if (count == 0 || buffers == nullptr) {
        return;
    }

    std::vector<VkBuffer> nativeBuffers;
    std::vector<VkDeviceSize> nativeOffsets;
    nativeBuffers.reserve(count);
    nativeOffsets.reserve(count);

    for (uint32_t i = 0; i < count; ++i) {
        auto* view = static_cast<VulkanResourceView*>(buffers[i]);
        if (view == nullptr || !view->IsValid()) {
            continue;
        }

        nativeBuffers.push_back(view->GetBuffer());
        nativeOffsets.push_back(view->GetBufferOffset() + (offsets != nullptr ? offsets[i] : 0));
        (void)strides;
    }

    if (!nativeBuffers.empty()) {
        vkCmdBindVertexBuffers(
            m_commandBuffer,
            startSlot,
            static_cast<uint32_t>(nativeBuffers.size()),
            nativeBuffers.data(),
            nativeOffsets.data());
    }
}

void VulkanCommandList::SetIndexBuffer(RHIResourceView* buffer, uint64_t offset) {
    auto* view = static_cast<VulkanResourceView*>(buffer);
    if (view == nullptr || !view->IsValid()) {
        return;
    }

    vkCmdBindIndexBuffer(
        m_commandBuffer,
        view->GetBuffer(),
        view->GetBufferOffset() + offset,
        ToVulkanIndexType(view->GetFormat()));
}

void VulkanCommandList::Draw(uint32_t vertexCount, uint32_t instanceCount, uint32_t startVertex, uint32_t startInstance) {
    if (!EnsureGraphicsRenderPass()) {
        return;
    }

    vkCmdBindPipeline(m_commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, m_currentPipelineState->GetPipeline());
    vkCmdDraw(m_commandBuffer, vertexCount, instanceCount, startVertex, startInstance);
}

void VulkanCommandList::DrawIndexed(uint32_t indexCount, uint32_t instanceCount, uint32_t startIndex, int32_t baseVertex, uint32_t startInstance) {
    if (!EnsureGraphicsRenderPass()) {
        return;
    }

    vkCmdBindPipeline(m_commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, m_currentPipelineState->GetPipeline());
    vkCmdDrawIndexed(m_commandBuffer, indexCount, instanceCount, startIndex, baseVertex, startInstance);
}

void VulkanCommandList::Clear(float r, float g, float b, float a, uint32_t buffers) {
    (void)buffers;

    auto* colorView = static_cast<VulkanResourceView*>(m_currentColorTarget);
    if (colorView == nullptr || colorView->GetTexture() == nullptr) {
        return;
    }

    EndActiveRenderPass();

    VulkanTexture* texture = colorView->GetTexture();
    TransitionTexture(texture, ResourceStates::CopyDst);

    VkClearColorValue clearColor = {};
    clearColor.float32[0] = r;
    clearColor.float32[1] = g;
    clearColor.float32[2] = b;
    clearColor.float32[3] = a;

    VkImageSubresourceRange range = {};
    range.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
    range.baseMipLevel = 0;
    range.levelCount = 1;
    range.baseArrayLayer = 0;
    range.layerCount = 1;

    vkCmdClearColorImage(
        m_commandBuffer,
        texture->GetImage(),
        VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
        &clearColor,
        1,
        &range);

    TransitionTexture(texture, ResourceStates::RenderTarget);
}

void VulkanCommandList::ClearRenderTarget(RHIResourceView* renderTarget, const float color[4]) {
    m_currentColorTarget = renderTarget;
    Clear(color[0], color[1], color[2], color[3], 1);
}

void VulkanCommandList::ClearDepthStencil(RHIResourceView* depthStencil, float depth, uint8_t stencil) {
    (void)depthStencil;
    (void)depth;
    (void)stencil;
}

void VulkanCommandList::CopyResource(RHIResourceView* dst, RHIResourceView* src) {
    (void)dst;
    (void)src;
}

void VulkanCommandList::Dispatch(uint32_t x, uint32_t y, uint32_t z) {
    (void)x;
    (void)y;
    (void)z;
}

bool VulkanCommandList::EnsureGraphicsRenderPass() {
    if (m_renderPassActive) {
        return true;
    }

    if (m_currentPipelineState == nullptr || !m_currentPipelineState->IsValid()) {
        return false;
    }

    auto* colorView = static_cast<VulkanResourceView*>(m_currentColorTarget);
    if (colorView == nullptr || colorView->GetTexture() == nullptr || colorView->GetImageView() == VK_NULL_HANDLE) {
        return false;
    }

    VulkanTexture* texture = colorView->GetTexture();
    TransitionTexture(texture, ResourceStates::RenderTarget);

    VkFramebufferCreateInfo framebufferInfo = {};
    framebufferInfo.sType = VK_STRUCTURE_TYPE_FRAMEBUFFER_CREATE_INFO;
    framebufferInfo.renderPass = m_currentPipelineState->GetRenderPass();
    framebufferInfo.attachmentCount = 1;
    const VkImageView attachment = colorView->GetImageView();
    framebufferInfo.pAttachments = &attachment;
    framebufferInfo.width = texture->GetWidth();
    framebufferInfo.height = texture->GetHeight();
    framebufferInfo.layers = 1;

    VkFramebuffer framebuffer = VK_NULL_HANDLE;
    if (vkCreateFramebuffer(m_device->GetDevice(), &framebufferInfo, nullptr, &framebuffer) != VK_SUCCESS) {
        return false;
    }
    m_transientFramebuffers.push_back(framebuffer);

    VkRenderPassBeginInfo renderPassInfo = {};
    renderPassInfo.sType = VK_STRUCTURE_TYPE_RENDER_PASS_BEGIN_INFO;
    renderPassInfo.renderPass = m_currentPipelineState->GetRenderPass();
    renderPassInfo.framebuffer = framebuffer;
    renderPassInfo.renderArea.offset = { 0, 0 };
    renderPassInfo.renderArea.extent = { texture->GetWidth(), texture->GetHeight() };
    renderPassInfo.clearValueCount = 0;
    renderPassInfo.pClearValues = nullptr;

    vkCmdBeginRenderPass(m_commandBuffer, &renderPassInfo, VK_SUBPASS_CONTENTS_INLINE);

    if (m_hasViewport) {
        vkCmdSetViewport(m_commandBuffer, 0, 1, &m_viewport);
    }
    if (m_hasScissor) {
        vkCmdSetScissor(m_commandBuffer, 0, 1, &m_scissor);
    }

    m_renderPassActive = true;
    return true;
}

void VulkanCommandList::EndActiveRenderPass() {
    if (m_renderPassActive && m_commandBuffer != VK_NULL_HANDLE) {
        vkCmdEndRenderPass(m_commandBuffer);
        m_renderPassActive = false;
    }
}

void VulkanCommandList::DestroyTransientFramebuffers() {
    if (m_device == nullptr) {
        m_transientFramebuffers.clear();
        return;
    }

    for (VkFramebuffer framebuffer : m_transientFramebuffers) {
        if (framebuffer != VK_NULL_HANDLE) {
            vkDestroyFramebuffer(m_device->GetDevice(), framebuffer, nullptr);
        }
    }
    m_transientFramebuffers.clear();
}

} // namespace RHI
} // namespace XCEngine