slot_allocator.h

#pragma once
 
#include <algorithm>
#include <cstdint>
#include <vector>
 
#include "l_assert.h"
#include "core/memory/types.h"
 
namespace Vkm::Engine {

class SlotAllocator {
    public:
        static constexpr uint32_t MAX_CLAIMED_INDEX = 1u << 22;
 
    public:
        SlotAllocator() : m_generation({GenerationIndex{}}) {}
        ~SlotAllocator() = default;
 
        SlotAllocator(const SlotAllocator& other) = delete;
        SlotAllocator& operator=(const SlotAllocator& other) = delete;
 
        SlotAllocator(SlotAllocator && other) = delete;
        SlotAllocator& operator=(SlotAllocator && other) = delete;
 
    public:
        StorageIndex allocate(bool* recycled = nullptr) {
            bool reused = false;
            const uint32_t idx = allocateSlot(reused);
            if (recycled) *recycled = reused;
            m_generation[idx].setAlive(true);
            ++m_liveCount;
            return StorageIndex{idx, m_generation[idx].generation()};
        }

        void free(StorageIndex id) {
            VKM_ASSERT(has(id), "SlotAllocator::free called with invalid handle");
            if (!has(id)) return;
 
            m_generation[id.index].setAlive(false);
            m_generation[id.index].bumpGeneration();
            m_freeList.push_back(id.index);
            --m_liveCount;
        }

        bool has(StorageIndex id) const {
            return id
                && id.index < m_generation.size()
                && m_generation[id.index].alive()
                && m_generation[id.index].generation() == id.generation;
        }

        StorageIndex handleAt(uint32_t index) const {
            if (index >= m_generation.size()) return {};
            return StorageIndex{index, m_generation[index].generation()};
        }

        bool isAliveAtIndex(uint32_t index) const {
            return index > 0
                && index < m_generation.size()
                && m_generation[index].alive();
        }

        StorageIndex allocateAt(uint32_t index) {
            if (index == 0 || index > MAX_CLAIMED_INDEX) return {};
 
            // Highest first, so popping from the back hands the gap out lowest first.
            const uint32_t oldSize = static_cast<uint32_t>(m_generation.size());
            while (m_generation.size() <= index) m_generation.push_back({});
            for (uint32_t gap = index; gap-- > oldSize;) m_freeList.push_back(gap);
 
            VKM_ASSERT(
                !m_generation[index].alive(),
                "SlotAllocator::allocateAt: slot %u already alive",
                index
            );
            // Guarded too, so without asserts two owners never share a slot.
            if (m_generation[index].alive()) return {};
 
            m_generation[index].setAlive(true);
            ++m_liveCount;
            return StorageIndex{index, m_generation[index].generation()};
        }

        template<typename Fn>
        void forEach(Fn&& fn) const {
            for (uint32_t i = 1; i < m_generation.size(); ++i) {
                if (m_generation[i].alive()) fn(i);
            }
        }

        void clear() {
            m_freeList.clear();
            m_freeList.reserve(m_generation.size());
            // Highest first, so slot 1 comes out next, as from a new allocator.
            for (uint32_t i = static_cast<uint32_t>(m_generation.size()); i-- > 1;) {
                if (m_generation[i].alive()) {
                    m_generation[i].setAlive(false);
                    m_generation[i].bumpGeneration();
                }
                m_freeList.push_back(i);
            }
            m_liveCount = 0;
        }

        void swap(SlotAllocator& other) noexcept {
            using std::swap;
            swap(m_generation, other.m_generation);
            swap(m_freeList,   other.m_freeList);
            swap(m_liveCount,  other.m_liveCount);
        }

        size_t size() const { return m_liveCount; }

        size_t extent() const { return m_generation.size(); }
 
    private:
        uint32_t allocateSlot(bool& reused) {
            // Skip slots allocateAt claimed; each is discarded once, so amortised O(1).
            while (!m_freeList.empty()) {
                const uint32_t idx = m_freeList.back();
                m_freeList.pop_back();
                if (!m_generation[idx].alive()) {
                    reused = true;
                    return idx;
                }
            }
 
            reused = false;
            const uint32_t idx = static_cast<uint32_t>(m_generation.size());
            m_generation.push_back({});
            return idx;
        }
 
    private:
        std::vector<GenerationIndex> m_generation;
        std::vector<uint32_t> m_freeList;   
        size_t m_liveCount = 0;
};
 
} // namespace Vkm::Engine