DirectXDescriptorHeap.h

#pragma once
 
// Standard.
#include <variant>
#include <memory>
#include <unordered_set>
#include <mutex>
#include <atomic>
#include <queue>
#include <functional>
 
// Custom.
#include "render/directx/descriptors/DirectXDescriptor.h"
#include "render/general/resource/GpuResource.h"
#include "DirectXDescriptorType.hpp"
 
// External.
#include "directx/d3dx12.h"
#include "misc/Error.h"
 
// OS.
#include <wrl.h>
 
namespace ne {
    using namespace Microsoft::WRL;
 
    class DirectXRenderer;
    class DirectXResource;
    class DirectXDescriptorHeap;
 
    enum class DescriptorHeapType : int {
        RTV = 0,
        DSV,
        CBV_SRV_UAV,
        SAMPLER,
    };
 
    class ContinuousDirectXDescriptorRange {
        // This class is fully managed by the heap.
        friend class DirectXDescriptorHeap;
 
    public:
        ContinuousDirectXDescriptorRange() = delete;
 
        /* Notifies the heap. */
        ~ContinuousDirectXDescriptorRange();
 
        ContinuousDirectXDescriptorRange(const ContinuousDirectXDescriptorRange& other) = delete;
        ContinuousDirectXDescriptorRange& operator=(const ContinuousDirectXDescriptorRange& other) = delete;
 
        // Intentionally disable `move` because heap will store a raw pointer to the range
        // and we don't want to accidentally `move` the range which will make the heap's raw pointer invalid.
        ContinuousDirectXDescriptorRange(ContinuousDirectXDescriptorRange&& other) noexcept = delete;
        ContinuousDirectXDescriptorRange&
        operator=(ContinuousDirectXDescriptorRange&& other) noexcept = delete;
 
        static constexpr INT getRangeGrowSize() { return iRangeGrowSize; }
 
        size_t getRangeSize();
 
        size_t getRangeCapacity();
 
        INT getRangeStartInHeap();
 
        D3D12_GPU_DESCRIPTOR_HANDLE getGpuDescriptorHandleToRangeStart() const;
 
        std::string getRangeName() const { return sRangeName; }
 
    private:
        struct InternalData {
            std::unordered_set<DirectXDescriptor*> allocatedDescriptors;
 
            std::queue<INT> noLongerUsedDescriptorIndices;
 
            INT iRangeStartInHeap = -1;
 
            INT iRangeCapacity = 0;
 
            INT iNextFreeIndexInRange = 0;
        };
 
        ContinuousDirectXDescriptorRange(
            DirectXDescriptorHeap* pHeap,
            const std::function<void()>& onRangeIndicesChanged,
            const std::string& sRangeName);
 
        [[nodiscard]] std::optional<Error> markDescriptorAsUnused(DirectXDescriptor* pDescriptor);
 
        std::variant<std::optional<INT>, Error> tryReserveFreeHeapIndexToCreateDescriptor();
 
        std::pair<std::recursive_mutex, InternalData> mtxInternalData;
 
        const std::function<void()> onRangeIndicesChanged;
 
        const std::string sRangeName;
 
        DirectXDescriptorHeap* const pHeap = nullptr;
 
        static constexpr INT iRangeGrowSize = 50; // NOLINT: see static asserts
    };
 
    class DirectXDescriptorHeap {
        // Notifies the heap about descriptor being destroyed.
        friend class DirectXDescriptor;
 
        // Notifies the heap about range being destroyed.
        friend class ContinuousDirectXDescriptorRange;
 
        // Only resource can request descriptors.
        friend class DirectXResource;
 
    public:
        struct InternalData {
            ComPtr<ID3D12DescriptorHeap> pHeap;
 
            std::unordered_set<ContinuousDirectXDescriptorRange*> continuousDescriptorRanges;
 
            INT iHeapCapacity = 0;
 
            INT iHeapSize = 0;
 
            INT iNextFreeHeapIndex = 0;
 
            std::queue<INT> noLongerUsedSingleDescriptorIndices;
 
            std::unordered_set<DirectXDescriptor*> bindedSingleDescriptors;
        };
 
        DirectXDescriptorHeap() = delete;
        DirectXDescriptorHeap(const DirectXDescriptorHeap&) = delete;
        DirectXDescriptorHeap& operator=(const DirectXDescriptorHeap&) = delete;
 
        ~DirectXDescriptorHeap();
 
        static constexpr INT getHeapGrowSize() { return iHeapGrowSize; }
 
        static std::variant<std::unique_ptr<DirectXDescriptorHeap>, Error>
        create(DirectXRenderer* pRenderer, DescriptorHeapType heapType);
 
        std::variant<std::shared_ptr<ContinuousDirectXDescriptorRange>, Error>
        allocateContinuousDescriptorRange(
            const std::string& sRangeName, const std::function<void()>& onRangeIndicesChanged);
 
        INT getHeapCapacity();
 
        INT getHeapSize();
 
        size_t getNoLongerUsedDescriptorCount();
 
        UINT getDescriptorSize() const { return iDescriptorSize; }
 
        ID3D12DescriptorHeap* getInternalHeap() const { return mtxInternalData.second.pHeap.Get(); }
 
        std::pair<std::recursive_mutex, InternalData>* getInternalData() { return &mtxInternalData; }
 
    protected:
        static std::string convertHeapTypeToString(DescriptorHeapType heapType);
 
        DirectXDescriptorHeap(DirectXRenderer* pRenderer, DescriptorHeapType heapType);
 
        void onDescriptorBeingDestroyed(
            DirectXDescriptor* pDescriptor, ContinuousDirectXDescriptorRange* pRange = nullptr);
 
        void onDescriptorRangeBeingDestroyed(ContinuousDirectXDescriptorRange* pRange);
 
        void createView(
            CD3DX12_CPU_DESCRIPTOR_HANDLE heapHandle,
            const DirectXResource* pResource,
            DirectXDescriptorType descriptorType,
            std::optional<size_t> cubemapFaceIndex) const;
 
        [[nodiscard]] std::optional<Error> expandHeap(ContinuousDirectXDescriptorRange* pChangedRange);
 
        [[nodiscard]] std::variant<bool, Error>
        shrinkHeapIfPossible(ContinuousDirectXDescriptorRange* pChangedRange);
 
        [[nodiscard]] std::optional<Error>
        createHeap(INT iCapacity, ContinuousDirectXDescriptorRange* pChangedRange);
 
        std::vector<DirectXDescriptorType> getDescriptorTypesHandledByThisHeap() const;
 
        [[nodiscard]] std::optional<Error> rebindViewsUpdateIndices();
 
    private:
        [[nodiscard]] std::optional<Error> assignDescriptor(
            DirectXResource* pResource,
            DirectXDescriptorType descriptorType,
            const std::shared_ptr<ContinuousDirectXDescriptorRange>& pRange = nullptr,
            bool bBindDescriptorsToCubemapFaces = true);
 
        static bool isShrinkingPossible(INT iSize, INT iCapacity, INT iGrowSize);
 
        [[nodiscard]] std::optional<Error> expandRange(ContinuousDirectXDescriptorRange* pRange);
 
        DirectXRenderer* pRenderer;
 
        std::pair<std::recursive_mutex, InternalData> mtxInternalData;
 
        UINT iDescriptorSize = 0;
 
        DescriptorHeapType heapType;
 
        std::string sHeapType;
 
        D3D12_DESCRIPTOR_HEAP_TYPE d3dHeapType;
 
        static constexpr INT iHeapGrowSize = 300; // NOLINT: don't recreate heap too often
    };
} // namespace ne