project-aster/samples/03_model_render/gpu_resource_manager.h

// =============================================
//  Aster: gpu_resource_manager.h
//  Copyright (c) 2020-2024 Anish Bhobe
// =============================================

#pragma once

#include "aster/aster.h"

#include <EASTL/deque.h>
#include <EASTL/vector_map.h>

struct Device;
struct Texture;
struct StorageTexture;
struct StorageBuffer;

struct GpuResourceHandle
{
    constexpr static u32 INVALID_HANDLE = MaxValue<u32>;
    u32 m_Index = INVALID_HANDLE; // Default = invalid

    [[nodiscard]] bool
    IsInvalid() const
    {
        return m_Index == INVALID_HANDLE;
    }
};

struct BufferHandle : GpuResourceHandle
{
};

struct TextureHandle : GpuResourceHandle
{
};

struct StorageTextureHandle : GpuResourceHandle
{
};

struct SamplerHandle : GpuResourceHandle
{
};

struct TextureManager
{
    eastl::vector<Texture> m_Textures;
    u32 m_MaxCapacity;
    u32 m_FreeHead;

    void Init(u32 maxCapacity);
    TextureHandle Commit(Texture *texture);
    Texture *Fetch(TextureHandle handle);
    void Release(const Device *device, TextureHandle handle);
    void Destroy(const Device *device);
};

struct BufferManager
{
    eastl::vector<StorageBuffer> m_Buffers;
    u32 m_MaxCapacity;
    u32 m_FreeHead;

    void Init(u32 maxCapacity);
    BufferHandle Commit(StorageBuffer *buffer);
    StorageBuffer *Fetch(BufferHandle handle);
    void Release(const Device *device, BufferHandle handle);
    void Destroy(const Device *device);
};

struct StorageTextureManager : TextureManager
{
    StorageTextureHandle Commit(StorageTexture *texture);
    StorageTexture *Fetch(StorageTextureHandle handle);
    void Release(const Device *device, StorageTextureHandle handle);
};

struct SamplerManager
{
    // There can only be so many samplers.
    eastl::vector<vk::Sampler> m_Samplers;
    eastl::vector<usize> m_SamplerHashes;

    void Init(usize size);
    SamplerHandle Create(const Device *device, const vk::SamplerCreateInfo *createInfo);
    vk::Sampler Fetch(SamplerHandle handle);
    void Destroy(const Device *device);
};

struct GpuResourceManager
{
  private:
    union WriteInfo {
        vk::DescriptorBufferInfo uBufferInfo;
        vk::DescriptorImageInfo uImageInfo;
        vk::BufferView uBufferView;

        WriteInfo()
        {
        }

        explicit WriteInfo(vk::DescriptorBufferInfo info);
        explicit WriteInfo(vk::DescriptorImageInfo info);
        explicit WriteInfo(vk::BufferView info);
    };

    enum class HandleType
    {
        eBuffer,
        eTexture,
        eStorageTexture,
    };

    using WriteOwner = eastl::pair<HandleType, u32>;

    eastl::deque<WriteInfo> m_WriteInfos;
    eastl::vector<vk::WriteDescriptorSet> m_Writes;
    eastl::vector<WriteOwner> m_WriteOwner;

    vk::Sampler m_DefaultSampler;

    BufferManager m_BufferManager;
    TextureManager m_TextureManager;
    StorageTextureManager m_StorageTextureManager;
    SamplerManager m_SamplerManager;

    void EraseWrites(u32 handleIndex, HandleType handleType);

  public:
    Device *m_Device;

    constexpr static u32 BUFFER_BINDING_INDEX = 0;
    constexpr static u32 TEXTURE_BINDING_INDEX = 1;
    constexpr static u32 STORAGE_TEXTURE_BINDING_INDEX = 2;

    vk::SamplerCreateInfo m_DefaultSamplerCreateInfo;

    vk::DescriptorPool m_DescriptorPool;
    vk::DescriptorSetLayout m_SetLayout;
    vk::DescriptorSet m_DescriptorSet;

    BufferHandle Commit(StorageBuffer *storageBuffer);                           // Commit to GPU and take Ownership
    void Write(BufferHandle handle, usize offset, usize size, const void *data); // Write to buffer
    void Release(BufferHandle handle);                                           // Release and Destroy
    void Release(StorageBuffer *storageBuffer, BufferHandle handle);             // Release and Return

    TextureHandle CommitTexture(Texture *texture, SamplerHandle sampler = {}); // Commit to GPU and take Ownership
    void Release(TextureHandle handle);                                        // Release and Destroy
    void Release(Texture *texture, TextureHandle handle);                      // Release and Return

    StorageTextureHandle
    CommitStorageTexture(StorageTexture *storageTexture, SamplerHandle sampler = {}); // Commit to GPU and take Ownership
    void Release(StorageTextureHandle handle);                                   // Release and Destroy
    void Release(StorageTexture *texture, StorageTextureHandle handle);          // Release and Return

    SamplerHandle CreateSampler(const vk::SamplerCreateInfo *samplerCreateInfo);

    void Update(); // Update all the descriptors required.

    // Ctor/Dtor
    GpuResourceManager(Device *device, u16 maxSize);
    ~GpuResourceManager();

    GpuResourceManager(GpuResourceManager &&other) noexcept;
    GpuResourceManager &operator=(GpuResourceManager &&other) noexcept;

#if !defined(ASTER_NDEBUG)
    usize m_CommitedBufferCount = 0;
    usize m_CommitedTextureCount = 0;
    usize m_CommitedStorageTextureCount = 0;
#endif

    DISALLOW_COPY_AND_ASSIGN(GpuResourceManager);
};