project-aster/aster/include/aster/systems/commit_manager.h

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

#pragma once

#include "aster/aster.h"
#include "buffer_manager.h"
#include "image_manager.h"
#include "sampler_manager.h"

#include "aster/util/intrusive_slist.h"

#include "EASTL/deque.h"
#include "EASTL/intrusive_hash_map.h"
#include "EASTL/bonus/fixed_ring_buffer.h"

#include "resource.h"

namespace systems
{

class CommitManager
{
  private:
    template <concepts::Manageable T>
    struct HandleMapper
    {
        using Type = T;
        using Manager = Manager<Type>;
        using Handle = typename Manager::Handle;
        using Resource = ResId<Type>;

        struct Entry : public eastl::intrusive_hash_node_key<Handle>
        {
            std::atomic<u32> m_CommitCount;

            void
            AddRef()
            {
                const auto rc = ++m_CommitCount;
                assert(rc > 0);
            }

            void
            Release()
            {
                const auto rc = --m_CommitCount;
                assert(rc < MaxValue<u32>);
            }

            bool
            IsReferenced() const
            {
                return m_CommitCount;
            }

            bool
            operator==(const Entry &other) const
            {
                return this->mKey == other.mKey;
            }

            Entry *
            Next()
            {
                return Recast<Entry *>(this->mpNext);
            }

            void
            SetNext(Entry &entry)
            {
                this->mpNext = &entry;
            }

            struct Hash
            {
                usize
                operator()(const Handle &e)
                {
                    return eastl::hash<Type *>()(e.get());
                }
            };
        };
        static_assert(sizeof(Entry) == 24);

        eastl::vector<Entry> m_Data;
        IntrusiveStack<Entry> m_FreeList;
        eastl::intrusive_hash_map<typename Entry::key_type, Entry, 31, typename Entry::Hash> m_InUse;
        std::array<IntrusiveStack<Entry>, 4> m_ToDelete;
        u8 m_ToDeleteIndex = 0;

        explicit HandleMapper(const u32 maxCount)
            : m_Data{maxCount}
        {
            // Setup freelist
            for (auto it = m_Data.rbegin(); it != m_Data.rend(); ++it)
            {
                m_FreeList.Push(*it);
            }
        }

        ~HandleMapper()
        {
            for (auto & toDelete : m_ToDelete)
            {
                ClearEntries(toDelete);
            }
        }

        PIN_MEMORY(HandleMapper);

        /// Returns a commit, and a bool signifying if it is a new commit.
        std::tuple<Resource, bool>
        Create(const Handle &object)
        {
            // Get-from freelist
            assert(!m_FreeList.Empty());

            auto it = m_InUse.find(object);
            if (it != m_InUse.end())
            {
                it->AddRef();
                auto i = GetIndex(*it);
                return {Resource{i}, false};
            }

            Entry &data = m_FreeList.Pop();

            data.mKey = object;
            data.m_CommitCount = 1;

            m_InUse.insert(data);

            auto i = GetIndex(data);

            return {Resource{i}, true};
        }

        Handle
        GetHandle(const Resource &res)
        {
            return m_Data[res.m_Index].mKey;
        }

        void
        AddRef(const Resource &commit)
        {
            m_Data.at(commit.m_Index).AddRef();
        }

        void
        Release(const Resource &commit)
        {
            auto &entry = m_Data.at(commit.m_Index);
            entry.Release();
            if (!entry.IsReferenced())
            {
                QueueDelete(entry);
            }
        }

        /**
         * Sweeps through the delete queue.
         * All freed items are cleared. (With a 3 frame delay)
         */
        void
        Update()
        {
            m_ToDeleteIndex = (m_ToDeleteIndex + 1) % m_ToDelete.size();
            auto &list = m_ToDelete[m_ToDeleteIndex];
            ClearEntries(list);
        }

      private:
        u32
        GetIndex(const Entry &entry)
        {
            return Cast<u32>(&entry - m_Data.begin());
        }

        void
        QueueDelete(Entry &entry)
        {
            m_InUse.remove(entry);
            m_ToDelete[m_ToDeleteIndex].Push(entry);
        }

        void
        ClearEntries(IntrusiveStack<Entry>& entries)
        {
            while (auto item = entries.TryPop())
            {
                Entry &entry = item.value();
                entry.mKey.reset();
                entry.m_CommitCount = 0;
            }
        }
    };

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

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

    using WriteCommand = vk::WriteDescriptorSet;

    // using WriteOwner = std::variant<Handle<Buffer>, Handle<Image>>;

  public:
    const Device *m_Device;

    CommitManager(const Device *device, u32 maxBuffers, u32 maxImages, u32 maxStorageImages,
                  Ref<Sampler> defaultSampler);
    ~CommitManager();

    PIN_MEMORY(CommitManager);

    // Commit Buffer
    ResId<Buffer> CommitBuffer(const Ref<Buffer> &buffer);

    // Commit Storage Images
    ResId<StorageImage> CommitStorageImage(const concepts::StorageImageRef auto &image)
    {
        return CommitStorageImage(CastImage<StorageImage>(image));
    }
    ResId<StorageImage> CommitStorageImage(const Ref<StorageImage> &image);

    // Sampled Images
    ResId<Texture>
    CommitTexture(const concepts::SampledImageRef auto &image, const Ref<Sampler> &sampler)
    {
        return CommitTexture(CastImage<Texture>(image), sampler);
    }

    ResId<Texture>
    CommitTexture(const concepts::SampledImageRef auto &image)
    {
        return CommitTexture(CastImage<Texture>(image));
    }

    ResId<Texture> CommitTexture(const Ref<Texture> &handle);
    ResId<Texture> CommitTexture(const Ref<Texture> &image, const Ref<Sampler> &sampler);

    void Update();

    Ref<Buffer>
    FetchHandle(const ResId<Buffer> &id)
    {
        return m_Buffers.GetHandle(id);
    }

    Ref<Texture>
    FetchHandle(const ResId<Texture> &id)
    {
        return m_Images.GetHandle(id);
    }

    Ref<StorageImage>
    FetchHandle(const ResId<StorageImage> &id)
    {
        return m_StorageImages.GetHandle(id);
    }

    [[nodiscard]] const vk::DescriptorSetLayout &
    GetDescriptorSetLayout() const
    {
        return m_SetLayout;
    }

    [[nodiscard]] const vk::DescriptorSet &
    GetDescriptorSet() const
    {
        return m_DescriptorSet;
    }

    static CommitManager &Instance()
    {
        assert(m_Instance);
        return *m_Instance;
    }

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

    constexpr static u8 BUFFER_BINDING_INDEX = 0x0;
    constexpr static u8 IMAGE_BINDING_INDEX = 0x1;
    constexpr static u8 STORAGE_IMAGE_BINDING_INDEX = 0x2;

    HandleMapper<Buffer> m_Buffers;
    HandleMapper<Texture> m_Images;
    HandleMapper<StorageImage> m_StorageImages;

    Ref<Sampler> m_DefaultSampler;

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

    static CommitManager *m_Instance;
    friend ResId<Buffer>;
    friend ResId<Texture>;
    friend ResId<StorageImage>;

    void
    AddRef(const ResId<Buffer> &handle)
    {
        m_Buffers.AddRef(handle);
    }
    void
    AddRef(const ResId<Texture> &handle)
    {
        m_Images.AddRef(handle);
    }
    void
    AddRef(const ResId<StorageImage> &handle)
    {
        m_StorageImages.AddRef(handle);
    }

    void
    Release(const ResId<Buffer> &handle)
    {
        m_Buffers.Release(handle);
    }
    void
    Release(const ResId<Texture> &handle)
    {
        m_Images.Release(handle);
    }
    void
    Release(const ResId<StorageImage> &handle)
    {
        m_StorageImages.Release(handle);
    }
};

template <concepts::Manageable T>
void
ResId<T>::AddRef() const
{
    CommitManager::Instance().AddRef(*this);
}

template <concepts::Manageable T>
void
ResId<T>::Release() const
{
    CommitManager::Instance().Release(*this);
}

} // namespace systems