SlotPool.hpp

#pragma once
 
#include <atomic>
#include <cstddef>
#include <vector>
 
namespace wmtk {
 
inline constexpr size_t INVALID_SLOT = static_cast<size_t>(-1);
 
template <class T>
class SlotPool
{
public:
    using value_type = T;
 
    SlotPool() = default;
    SlotPool(const SlotPool&) = delete;
    SlotPool& operator=(const SlotPool&) = delete;
 
    size_t capacity() const { return m_data.size(); }
 
    size_t live() const { return m_live.load(std::memory_order_relaxed); }
 
    void set_live(size_t n) { m_live.store(n, std::memory_order_relaxed); }
 
    void resize(size_t n) { m_data.resize(n); }
 
    void shrink_to_fit() { m_data.shrink_to_fit(); }
 
    void clear()
    {
        m_data.clear();
        set_live(0);
    }
 
    bool refused() const { return m_refused.load(std::memory_order_relaxed); }
    void clear_refused() { m_refused.store(false, std::memory_order_relaxed); }
    static size_t taken_by_this_thread() { return t_taken; }
 
    T& operator[](size_t i) { return m_data[i]; }
    const T& operator[](size_t i) const { return m_data[i]; }
 
    auto begin() { return m_data.begin(); }
    auto end() { return m_data.end(); }
    auto begin() const { return m_data.begin(); }
    auto end() const { return m_data.end(); }
    size_t request(size_t n)
    {
        size_t first = m_live.load(std::memory_order_relaxed);
        if (n == 0) {
            return first;
        }
        const size_t cap = m_data.size();
        do {
            // Phrased as a subtraction rather than `first + n > cap` so that it cannot
            // overflow; `first <= cap` always holds, so `cap - first` is well defined.
            if (first > cap || n > cap - first) {
                m_refused.store(true, std::memory_order_relaxed);
                return INVALID_SLOT;
            }
        } while (!m_live.compare_exchange_weak(
            first,
            first + n,
            std::memory_order_acq_rel,
            std::memory_order_relaxed));
 
        // Reset the handed-out slots. The storage keeps whatever a previous generation left in
        // the spare region, so a fresh slot has to be cleared here rather than being fresh by
        // construction -- which is what the tbb::concurrent_vector this replaced used to give.
        for (size_t i = first; i < first + n; ++i) {
            m_data[i] = T{};
        }
        t_taken += n;
        return first;
    }
 
private:
    std::vector<T> m_data;
    std::atomic<size_t> m_live{0};
    static inline thread_local size_t t_taken = 0;
    std::atomic<bool> m_refused{false};
};
 
} // namespace wmtk