RecyclableIdAllocator.hpp

// SPDX-FileCopyrightText: 2025 Vincent Leroy
// SPDX-License-Identifier: MIT
//
// This file is part of dfx.
//
// Licensed under the MIT License. See the LICENSE file in the project root
// for full license information.
 
#pragma once
 
// Standard includes
#include <concepts>
#include <cstdint>
#include <queue>
#include <unordered_set>
 
// Project includes
#include "Exception.hpp"
 
namespace dfx::Utils
{
template<std::integral T>
class RecyclableIdAllocator
{
public:
    using value_type = T;
 
private:
    using FreeIdList = std::priority_queue<value_type, std::vector<value_type>, std::greater<value_type>>;
 
public:
    value_type next() noexcept
    {
        value_type id;
 
        if (_freeIds.empty())
            id = _nextId++;
        else
        {
            id = _freeIds.top();
            _freeIds.pop();
            _freedIds.erase(id);
        }
 
        ++_activeCount;
        return id;
    }

    void release(value_type id)
    {
        // Ignore invalid frees (id was never allocated)
        if (id >= _nextId)
            return ;
 
        // Prevent double-free
        if (_freedIds.contains(id))
            return ;
 
        --_activeCount;
 
        if (_activeCount == 0) // Everything has been released — full reset
            return reset();
 
        if (_nextId == id + 1)
            --_nextId; // Reclaim end-of-sequence ID without heap use
        else
        {
            _freeIds.push(id);
            _freedIds.insert(id);
        }
    }

    void reserve(value_type id)
    {
        if (id >= _nextId)
        {
            // Fill the gap from _nextId to id
            for (auto i = _nextId; i < id; ++i)
            {
                _freeIds.push(i);
                _freedIds.insert(i);
            }
            _nextId = id + 1;
        }
        else
        {
            B_ASSERT(_freedIds.contains(id), "ID {} is already in use", id);
 
            // Remove from free lists
            _freedIds.erase(id);
 
            // _freeIds has no direct erase, so we rebuild the heap
            FreeIdList tmp;
            while (!_freeIds.empty())
            {
                auto const v = _freeIds.top();
                _freeIds.pop();
                if (v != id)
                    tmp.push(v);
            }
 
            _freeIds.swap(tmp);
        }
 
        ++_activeCount;
    }

    void reset() noexcept
    {
        _nextId = value_type{};
        _activeCount = 0;
        _freeIds = {};
        _freedIds.clear();
    }
 
private:
    value_type                     _nextId = value_type{};
    uint64_t                       _activeCount = 0;
    FreeIdList                     _freeIds;
    std::unordered_set<value_type> _freedIds;
};
} // !namespace dfx::Utils