Juliet/Juliet/src/Graphics/D3D12/D3D12DescriptorHeap.cpp

#include <Core/Logging/LogManager.h>
#include <Core/Logging/LogTypes.h>
#include <Core/Memory/Allocator.h>
#include <Graphics/D3D12/D3D12DescriptorHeap.h>
#include <Graphics/D3D12/D3D12GraphicsDevice.h>
#include <Graphics/D3D12/D3D12Utils.h>

namespace Juliet::D3D12::Internal
{
    void CreateDescriptorHeapPool(NonNullPtr<D3D12Driver> driver, D3D12DescriptorHeapPool& heapPool,
                                  D3D12_DESCRIPTOR_HEAP_TYPE type, uint32 count)
    {
        // Heap pool is just single linked list of free elements
        constexpr size_t kInitialCapacity = 4;
        heapPool.Arena                    = ArenaAllocate();
        heapPool.FirstFreeDescriptorHeap  = nullptr;

        // Pre allocate 4
        for (uint32 i = 0; i < kInitialCapacity; ++i)
        {
            D3D12DescriptorHeap* descriptorHeap = CreateDescriptorHeap(driver, heapPool.Arena, type, count, false);
            descriptorHeap->Next                = heapPool.FirstFreeDescriptorHeap;
            heapPool.FirstFreeDescriptorHeap    = descriptorHeap;
        }
    }

    void DestroyDescriptorHeapPool(D3D12DescriptorHeapPool& heapPool)
    {
        D3D12DescriptorHeap* current = heapPool.FirstFreeDescriptorHeap;
        while (current != nullptr)
        {
            D3D12DescriptorHeap* next = current->Next;
            DestroyDescriptorHeap(current);
            current = next;
        }
        ArenaRelease(heapPool.Arena);
    }

    D3D12DescriptorHeap* CreateDescriptorHeap(NonNullPtr<D3D12Driver> driver, NonNullPtr<Arena> arena,
                                              D3D12_DESCRIPTOR_HEAP_TYPE type, uint32 count, bool isStaging)
    {
        D3D12DescriptorHeap* heap = ArenaPushStruct<D3D12DescriptorHeap>(arena.Get());
        Assert(heap);

        heap->CurrentDescriptorIndex = 0;

        heap->FreeIndices.Create();
        heap->FreeIndices.Resize(16);
        heap->CurrentFreeIndex = 0;

        D3D12_DESCRIPTOR_HEAP_DESC heapDesc;
        heapDesc.NumDescriptors = count;
        heapDesc.Type           = type;
        heapDesc.Flags    = isStaging ? D3D12_DESCRIPTOR_HEAP_FLAG_NONE : D3D12_DESCRIPTOR_HEAP_FLAG_SHADER_VISIBLE;
        heapDesc.NodeMask = 0;

        ID3D12DescriptorHeap* handle;
        HRESULT               result =
            driver->D3D12Device->CreateDescriptorHeap(&heapDesc, IID_ID3D12DescriptorHeap, reinterpret_cast<void**>(&handle));
        if (FAILED(result))
        {
            LogError(driver->D3D12Device, "Failed to create descriptor heap!", result);
            DestroyDescriptorHeap(heap);
            return nullptr;
        }

        heap->Handle                 = handle;
        heap->HeapType               = type;
        heap->MaxDescriptors         = count;
        heap->Staging                = isStaging;
        heap->DescriptorSize         = driver->D3D12Device->GetDescriptorHandleIncrementSize(type);
        heap->DescriptorHeapCPUStart = handle->GetCPUDescriptorHandleForHeapStart();
        if (!isStaging)
        {
            heap->DescriptorHeapGPUStart = handle->GetGPUDescriptorHandleForHeapStart();
        }

        return heap;
    }

    void DestroyDescriptorHeap(NonNullPtr<D3D12DescriptorHeap> heap)
    {
        heap->FreeIndices.Destroy();
        if (heap->Handle)
        {
            heap->Handle->Release();
        }
    }

    bool AssignDescriptor(D3D12DescriptorHeap* heap, D3D12Descriptor& outDescriptor)
    {
        uint32 index = UINT32_MAX;

        if (heap->CurrentFreeIndex > 0)
        {
            heap->CurrentFreeIndex -= 1;
            index = heap->FreeIndices[heap->CurrentFreeIndex];
        }
        else if (heap->CurrentDescriptorIndex < heap->MaxDescriptors)
        {
            index = heap->CurrentDescriptorIndex;
            heap->CurrentDescriptorIndex++;
        }
        else
        {
            Assert(false, "Descriptor Heap Full!");
            return false;
        }

        outDescriptor.Heap      = heap;
        outDescriptor.Index     = index;
        outDescriptor.CpuHandle = heap->DescriptorHeapCPUStart;
        outDescriptor.CpuHandle.ptr += heap->DescriptorSize * index;
        outDescriptor.GpuHandle = heap->DescriptorHeapGPUStart;
        outDescriptor.GpuHandle.ptr += heap->DescriptorSize * index;

        return true;
    }

    void ReleaseDescriptor(const D3D12Descriptor& descriptor)
    {
        if (descriptor.Index == UINT32_MAX || descriptor.Heap == nullptr)
        {
            return;
        }

        D3D12DescriptorHeap* heap = descriptor.Heap;

        if (heap->CurrentFreeIndex >= heap->FreeIndices.Count)
        {
            heap->FreeIndices.PushBack(descriptor.Index);
        }
        else
        {
            heap->FreeIndices[heap->CurrentFreeIndex] = descriptor.Index;
            heap->CurrentFreeIndex++;
        }
    }

    D3D12DescriptorHeap* AcquireSamplerHeapFromPool(NonNullPtr<D3D12Driver> d3d12Driver)
    {
        D3D12DescriptorHeapPool& pool = d3d12Driver->SamplerHeapPool;

        D3D12DescriptorHeap* result = pool.FirstFreeDescriptorHeap;
        if (result)
        {
            pool.FirstFreeDescriptorHeap = pool.FirstFreeDescriptorHeap->Next;
        }
        else
        {
            result = CreateDescriptorHeap(d3d12Driver, pool.Arena, D3D12_DESCRIPTOR_HEAP_TYPE_SAMPLER,
                                          GPUDriver::kSampler_HeapDescriptorCount, false);
        }

        return result;
    }

    void ReturnSamplerHeapToPool(NonNullPtr<D3D12Driver> d3d12Driver, NonNullPtr<D3D12DescriptorHeap> heap)
    {
        D3D12DescriptorHeapPool& pool = d3d12Driver->SamplerHeapPool;

        heap->CurrentDescriptorIndex = 0;

        heap->Next                   = pool.FirstFreeDescriptorHeap;
        pool.FirstFreeDescriptorHeap = heap;
    }
} // namespace Juliet::D3D12::Internal