XCEngine/engine/include/XCEngine/RHI/D3D12/D3D12CommandList.h

#pragma once

#include <d3d12.h>
#include <wrl/client.h>
#include <vector>
#include <unordered_map>

#include "../RHICommandList.h"
#include "../RHIEnums.h"
#include "../RHITypes.h"
#include "D3D12Enums.h"

using Microsoft::WRL::ComPtr;

namespace XCEngine {
namespace RHI {

class RHIPipelineState;
class D3D12ResourceView;
class D3D12Shader;
class D3D12PipelineLayout;
class D3D12RenderPass;
class D3D12Framebuffer;

class D3D12CommandList : public RHICommandList {
public:
    D3D12CommandList();
    ~D3D12CommandList() override;

    bool Initialize(ID3D12Device* device, CommandQueueType type = CommandQueueType::Direct, ID3D12CommandAllocator* allocator = nullptr);
    void Shutdown() override;

    void Reset() override;
    void Close() override;

    ID3D12GraphicsCommandList* GetCommandList() const { return m_commandList.Get(); }

    void SetShader(RHIShader* shader) override;

    void SetUniformInt(const char* name, int value) override;
    void SetUniformFloat(const char* name, float value) override;
    void SetUniformVec3(const char* name, float x, float y, float z) override;
    void SetUniformVec4(const char* name, float x, float y, float z, float w) override;
    void SetUniformMat4(const char* name, const float* value) override;

    void SetGlobalInt(const char* name, int value) override;
    void SetGlobalFloat(const char* name, float value) override;
    void SetGlobalVec3(const char* name, float x, float y, float z) override;
    void SetGlobalVec4(const char* name, float x, float y, float z, float w) override;
    void SetGlobalMat4(const char* name, const float* value) override;
    void SetGlobalTexture(const char* name, RHIResourceView* texture) override;

    void TransitionBarrier(RHIResourceView* resource, ResourceStates stateBefore, ResourceStates stateAfter) override;
    void TransitionBarrier(ID3D12Resource* resource, ResourceStates stateBefore, ResourceStates stateAfter);
    void TransitionBarrierInternal(ID3D12Resource* resource, ResourceStates stateBefore, ResourceStates stateAfter, uint32_t subresource = D3D12_RESOURCE_BARRIER_ALL_SUBRESOURCES);
    void UAVBarrier(ID3D12Resource* resource = nullptr);
    void UAVBarrierInternal(ID3D12Resource* resource);
    void AliasBarrier(ID3D12Resource* beforeResource = nullptr, ID3D12Resource* afterResource = nullptr);
    void AliasBarrierInternal(ID3D12Resource* beforeResource, ID3D12Resource* afterResource);

    void BeginRenderPass(class RHIRenderPass* renderPass, class RHIFramebuffer* framebuffer,
                       const Rect& renderArea, uint32_t clearValueCount, const ClearValue* clearValues) override;
    void EndRenderPass() override;

    void SetPipelineState(RHIPipelineState* pso) override;
    void SetPipelineState(ID3D12PipelineState* pso);
    void SetPipelineStateInternal(ID3D12PipelineState* pso);
    void SetPipelineLayout(D3D12PipelineLayout* layout);
    void SetRootSignature(ID3D12RootSignature* signature);
    void SetViewport(const Viewport& viewport) override;
    void SetViewports(uint32_t count, const Viewport* viewports) override;
    void SetScissorRect(const Rect& rect) override;
    void SetScissorRects(uint32_t count, const Rect* rects) override;
    void SetPrimitiveTopology(PrimitiveTopology topology);
    void SetRenderTargets(uint32_t count, RHIResourceView** renderTargets, RHIResourceView* depthStencil = nullptr) override;
    void SetRenderTargetsInternal(uint32_t count, ID3D12Resource** renderTargets, ID3D12Resource* depthStencil = nullptr);
    void SetRenderTargetsHandle(uint32_t count, const D3D12_CPU_DESCRIPTOR_HANDLE* renderTargetHandles, const D3D12_CPU_DESCRIPTOR_HANDLE* depthStencilHandle = nullptr);

    void SetVertexBuffers(uint32_t startSlot, uint32_t count, RHIResourceView** buffers, const uint64_t* offsets, const uint32_t* strides) override;
    void SetVertexBuffersInternal(uint32_t startSlot, uint32_t count, const D3D12_VERTEX_BUFFER_VIEW* views);
    void SetIndexBuffer(RHIResourceView* buffer, uint64_t offset) override;
    void SetIndexBuffer(ID3D12Resource* buffer, uint64_t offset, Format format);
    void SetIndexBufferInternal(ID3D12Resource* buffer, uint64_t offset, Format indexFormat);

    void SetDescriptorHeap(ID3D12DescriptorHeap* heap);
    void SetDescriptorHeaps(uint32_t count, ID3D12DescriptorHeap** heaps);
    void SetGraphicsDescriptorTable(uint32_t rootParameterIndex, D3D12_GPU_DESCRIPTOR_HANDLE baseHandle);
    void SetComputeDescriptorTable(uint32_t rootParameterIndex, D3D12_GPU_DESCRIPTOR_HANDLE baseHandle);

    void SetGraphicsRootConstantBufferView(uint32_t rootParameterIndex, D3D12_GPU_VIRTUAL_ADDRESS bufferLocation);
    void SetGraphicsRoot32BitConstants(uint32_t rootParameterIndex, uint32_t num32BitValuesToSet, const void* pSrcData, uint32_t destOffsetIn32BitValues);
    void SetGraphicsRootDescriptorTable(uint32_t rootParameterIndex, D3D12_GPU_DESCRIPTOR_HANDLE baseDescriptor);
    void SetGraphicsRootShaderResourceView(uint32_t rootParameterIndex, D3D12_GPU_VIRTUAL_ADDRESS shaderResource);

    void SetStencilRef(uint8_t stencilRef) override;
    void SetDepthStencilState(const DepthStencilState& state) override;
    void SetBlendState(const BlendState& state) override;
    void SetBlendFactor(const float blendFactor[4]);
    void SetDepthBias(float depthBias, float slopeScaledDepthBias, float depthBiasClamp);

    void Draw(uint32_t vertexCount, uint32_t instanceCount = 1, uint32_t startVertex = 0, uint32_t startInstance = 0) override;
    void DrawIndexed(uint32_t indexCount, uint32_t instanceCount = 1, uint32_t startIndex = 0, int32_t baseVertex = 0, uint32_t startInstance = 0) override;
    void DrawInstancedIndirect(ID3D12Resource* argBuffer, uint64_t alignedByteOffset);
    void DrawInstancedIndirectInternal(ID3D12Resource* argBuffer, uint64_t alignedByteOffset);
    void DrawIndexedInstancedIndirect(ID3D12Resource* argBuffer, uint64_t alignedByteOffset);
    void DrawIndexedInstancedIndirectInternal(ID3D12Resource* argBuffer, uint64_t alignedByteOffset);

    void Clear(float r, float g, float b, float a, uint32_t buffers) override;
    void ClearRenderTarget(RHIResourceView* renderTarget, const float color[4]) override;
    void ClearRenderTargetView(ID3D12Resource* renderTarget, const float color[4], uint32_t rectCount = 0, const D3D12_RECT* rects = nullptr);
    void ClearRenderTargetView(D3D12_CPU_DESCRIPTOR_HANDLE renderTargetHandle, const float color[4], uint32_t rectCount = 0, const D3D12_RECT* rects = nullptr);
    void ClearDepthStencil(RHIResourceView* depthStencil, float depth, uint8_t stencil) override;
    void ClearDepthStencilView(ID3D12Resource* depthStencil, uint32_t clearFlags, float depth = 1.0f, uint8_t stencil = 0, uint32_t rectCount = 0, const D3D12_RECT* rects = nullptr);
    void ClearDepthStencilView(D3D12_CPU_DESCRIPTOR_HANDLE depthStencilHandle, uint32_t clearFlags, float depth = 1.0f, uint8_t stencil = 0, uint32_t rectCount = 0, const D3D12_RECT* rects = nullptr);
    void ClearUnorderedAccessView(D3D12_GPU_DESCRIPTOR_HANDLE viewHandle, D3D12_CPU_DESCRIPTOR_HANDLE resourceHandle, ID3D12Resource* unorderedAccess, const float values[4], uint32_t rectCount = 0, const D3D12_RECT* rects = nullptr);

    void CopyResource(RHIResourceView* dst, RHIResourceView* src) override;
    void CopyResourceInternal(ID3D12Resource* dst, ID3D12Resource* src);
    void CopyBuffer(ID3D12Resource* dst, uint64_t dstOffset, ID3D12Resource* src, uint64_t srcOffset, uint64_t size);
    void CopyTexture(ID3D12Resource* dst, const D3D12_TEXTURE_COPY_LOCATION& dstLocation, ID3D12Resource* src, const D3D12_TEXTURE_COPY_LOCATION& srcLocation);

    void BeginQuery(ID3D12QueryHeap* queryHeap, QueryType type, uint32_t index);
    void EndQuery(ID3D12QueryHeap* queryHeap, QueryType type, uint32_t index);
    void ResolveQueryData(ID3D12QueryHeap* queryHeap, QueryType type, uint32_t startIndex, uint32_t count, ID3D12Resource* resultBuffer, uint64_t resultOffset);

    void Dispatch(uint32_t x, uint32_t y, uint32_t z) override;
    void DispatchIndirect(ID3D12Resource* argBuffer, uint64_t alignedByteOffset);
    void DispatchIndirectInternal(ID3D12Resource* argBuffer, uint64_t alignedByteOffset);

    void* GetNativeHandle() override { return GetCommandList(); }

    void ExecuteBundle(ID3D12GraphicsCommandList* bundle);

    ResourceStates GetResourceState(ID3D12Resource* resource) const;
    void TrackResource(ID3D12Resource* resource);

private:
    ComPtr<ID3D12GraphicsCommandList> m_commandList;
    ComPtr<ID3D12CommandAllocator> m_commandAllocator;
    ComPtr<ID3D12DescriptorHeap> m_rtvHeap;
    ID3D12Device* m_device = nullptr;
    CommandQueueType m_type;

    std::unordered_map<ID3D12Resource*, ResourceStates> m_resourceStateMap;
    std::vector<ID3D12Resource*> m_trackedResources;

    D3D12_PRIMITIVE_TOPOLOGY m_currentTopology;
    ID3D12PipelineState* m_currentPipelineState;
    ID3D12RootSignature* m_currentRootSignature;
    ID3D12DescriptorHeap* m_currentDescriptorHeap;

    std::vector<D3D12_CPU_DESCRIPTOR_HANDLE> m_boundRenderTargets;
    D3D12_CPU_DESCRIPTOR_HANDLE m_boundDepthStencil;
    bool m_depthStencilBound = false;

    D3D12Shader* m_currentShader;
    class D3D12PipelineLayout* m_currentPipelineLayout;
    std::unordered_map<std::string, int> m_globalIntCache;
    std::unordered_map<std::string, float> m_globalFloatCache;
    std::unordered_map<std::string, std::vector<float>> m_globalVec3Cache;
    std::unordered_map<std::string, std::vector<float>> m_globalVec4Cache;
    std::unordered_map<std::string, std::vector<float>> m_globalMat4Cache;
    std::unordered_map<std::string, RHIResourceView*> m_globalTextureCache;
};

} // namespace RHI
} // namespace XCEngine