sparse_set.h

#pragma once
 
#include <algorithm>
#include <cstdint>
#include <vector>
#include <type_traits>
 
#include "l_assert.h"
 
namespace Vkm::Engine {

class ISparseSet {
    public:
        ISparseSet() = default;
        virtual ~ISparseSet() = default;
 
        ISparseSet(const ISparseSet& other) = delete;
        ISparseSet& operator=(const ISparseSet& other) = delete;
 
        ISparseSet(ISparseSet && other) = delete;
        ISparseSet& operator=(ISparseSet && other) = delete;

        virtual void removeIfPresent(uint32_t key) = 0;

        virtual void compact() = 0;
};

template<typename T>
class SparseSet : public ISparseSet {
    public:
        SparseSet() = default;
        ~SparseSet() override = default;
 
        SparseSet(const SparseSet& other) = delete;
        SparseSet& operator=(const SparseSet& other) = delete;
 
        SparseSet(SparseSet && other) = delete;
        SparseSet& operator=(SparseSet && other) = delete;
 
    public:
        T& add(uint32_t key, T && value)     { return addInternal(key, std::move(value)); }
        T& add(uint32_t key, const T& value) { return addInternal(key, value); }

        void remove(uint32_t key) {
            VKM_ASSERT(contains(key), "SparseSet::remove called with invalid key");
            // Guarded too: a build without asserts would index m_data with garbage.
            if (!contains(key)) return;
 
            uint32_t dataIdx = m_dataIndex[key];
            uint32_t lastIdx = static_cast<uint32_t>(m_data.size() - 1);
 
            if (dataIdx != lastIdx) {
                m_data[dataIdx] = std::move(m_data[lastIdx]);
 
                m_dataId[dataIdx]              = m_dataId[lastIdx];
                m_dataIndex[m_dataId[dataIdx]] = dataIdx;
            }
 
            m_data.pop_back();
            m_dataId.pop_back();
            m_dataIndex[key] = EMPTY;
        }

        void removeIfPresent(uint32_t key) override {
            if (contains(key)) remove(key);
        }

        bool contains(uint32_t key) const {
            return key < m_dataIndex.size() && m_dataIndex[key] != EMPTY;
        }

        T& get(uint32_t key) {
            VKM_ASSERT(contains(key), "SparseSet::get called with invalid key");
            return m_data[m_dataIndex[key]];
        }
 
        const T& get(uint32_t key) const {
            VKM_ASSERT(contains(key), "SparseSet::get called with invalid key");
            return m_data[m_dataIndex[key]];
        }
 
    public:
        template<typename Fn>
        void forEach(Fn&& fn) {
            for (uint32_t i = 0; i < m_data.size(); ++i) {
                fn(m_dataId[i], m_data[i]);
            }
        }
 
        template<typename Fn>
        void forEach(Fn&& fn) const {
            for (uint32_t i = 0; i < m_data.size(); ++i) {
                fn(m_dataId[i], m_data[i]);
            }
        }

        size_t size() const { return m_data.size(); }

        void clear() {
            m_data.clear();
            m_dataId.clear();
            std::fill(m_dataIndex.begin(), m_dataIndex.end(), EMPTY);
        }

        void compact() override {
            if (m_data.empty()) {
                m_dataIndex.clear();
            } else {
                uint32_t maxKey = 0;
                for (uint32_t i = 0; i < m_dataId.size(); ++i) {
                    if (m_dataId[i] > maxKey) maxKey = m_dataId[i];
                }
                if (maxKey + 1 < m_dataIndex.size()) m_dataIndex.resize(maxKey + 1);
            }
            m_dataIndex.shrink_to_fit();
        }

        uint32_t keyAt(uint32_t denseIndex) const { return m_dataId[denseIndex]; }

        T&       dataAt(uint32_t denseIndex)       { return m_data[denseIndex]; }
        const T& dataAt(uint32_t denseIndex) const { return m_data[denseIndex]; }
 
    private:
        static constexpr uint32_t EMPTY = UINT32_MAX;

        void ensureCapacity(uint32_t key) {
            if (key >= m_dataIndex.size())
                m_dataIndex.resize(key + 1, EMPTY);
        }

        template<typename... Args>
        T& addInternal(uint32_t key, Args&&... args) {
            VKM_ASSERT(key != EMPTY, "SparseSet::add key cannot be EMPTY sentinel");
            VKM_ASSERT(key != 0, "SparseSet::add key 0 is reserved");
            ensureCapacity(key);
            VKM_ASSERT(!contains(key), "SparseSet::add key already present");
            // Guarded too: without asserts the key would strand its first entry.
            if (contains(key)) return m_data[m_dataIndex[key]];
 
            uint32_t dataIdx = static_cast<uint32_t>(m_data.size());
 
            m_data.emplace_back(std::forward<Args>(args)...);
            m_dataId.push_back(key);
            m_dataIndex[key] = dataIdx;
 
            return m_data[dataIdx];
        }
 
    private:
        std::vector<uint32_t> m_dataIndex; 
        std::vector<uint32_t> m_dataId;    
        std::vector<T>        m_data;
};
 
} // namespace Vkm::Engine