object_pool.hpp

#pragma once
 
#include <atomic>
#include <cstddef>
#include <cstdint>
#include <limits>
#include <type_traits>
#include <utility>
 
#include "libxr_def.hpp"
 
namespace LibXR
{
template <typename Data>
class ObjectPool
{
 public:
  using ValueType = Data;      
  using IndexType = uint32_t;  
 
  class Slot
  {
   public:
    Slot()
      requires std::is_default_constructible_v<Data>
    = default;
 
    template <typename... Args>
      requires std::is_constructible_v<Data, Args...>
    explicit Slot(std::in_place_t, Args&&... args) : data_(std::forward<Args>(args)...)
    {
    }
 
    Slot(const Slot&) = delete;
    Slot& operator=(const Slot&) = delete;
    Slot(Slot&&) = delete;
    Slot& operator=(Slot&&) = delete;
 
   private:
    friend class ObjectPool;
    std::atomic<uint32_t> references_{0};
    IndexType next_ = 0;
    Data data_;
  };
 
 public:
  class ConstHandle;
 
  class Handle
  {
   public:
    Handle() = default;
 
    Handle(const Handle&) = delete;
    Handle& operator=(const Handle&) = delete;
 
    Handle(Handle&& other) noexcept
        : pool_(std::exchange(other.pool_, nullptr)),
          index_(std::exchange(other.index_, IndexType{}))
    {
    }
 
    Handle& operator=(Handle&& other) noexcept
    {
      if (this != &other)
      {
        Reset();
        pool_ = std::exchange(other.pool_, nullptr);
        index_ = std::exchange(other.index_, IndexType{});
      }
      return *this;
    }
 
    ~Handle() { Reset(); }
 
    [[nodiscard]] bool Valid() const { return pool_ != nullptr; }
 
    [[nodiscard]] Data& Get()
    {
      ASSERT(Valid());
      return pool_->slots_[index_].data_;
    }
 
    [[nodiscard]] const Data& Get() const
    {
      ASSERT(Valid());
      return pool_->slots_[index_].data_;
    }
 
    [[nodiscard]] Data* operator->() { return &Get(); }
    [[nodiscard]] const Data* operator->() const { return &Get(); }
    [[nodiscard]] Data& operator*() { return Get(); }
    [[nodiscard]] const Data& operator*() const { return Get(); }
 
    [[nodiscard]] IndexType Index() const
    {
      ASSERT(Valid());
      return index_;
    }
 
    void Reset()
    {
      ObjectPool* pool = std::exchange(pool_, nullptr);
      const IndexType index = std::exchange(index_, IndexType{});
      if (pool != nullptr)
      {
        pool->Release(index);
      }
    }
 
   private:
    friend class ObjectPool;
    friend class ConstHandle;
 
    Handle(ObjectPool* pool, IndexType index) : pool_(pool), index_(index) {}
 
    ObjectPool* pool_ = nullptr;
    IndexType index_ = {};
  };
 
  class ConstHandle
  {
   public:
    ConstHandle() = default;
 
    ConstHandle(const ConstHandle& other) noexcept
        : pool_(other.pool_), index_(other.index_)
    {
      if (pool_ != nullptr)
      {
        pool_->Retain(index_);
      }
    }
 
    ConstHandle(ConstHandle&& other) noexcept
        : pool_(std::exchange(other.pool_, nullptr)),
          index_(std::exchange(other.index_, IndexType{}))
    {
    }
 
    ConstHandle(Handle&& writer) noexcept
        : pool_(std::exchange(writer.pool_, nullptr)),
          index_(std::exchange(writer.index_, IndexType{}))
    {
    }
 
    ConstHandle& operator=(const ConstHandle& other) noexcept
    {
      if (pool_ != other.pool_ || index_ != other.index_)
      {
        ConstHandle copy(other);
        Swap(copy);
      }
      return *this;
    }
 
    ConstHandle& operator=(ConstHandle&& other) noexcept
    {
      if (this != &other)
      {
        ConstHandle moved(std::move(other));
        Swap(moved);
      }
      return *this;
    }
 
    ConstHandle& operator=(Handle&& writer) noexcept
    {
      ConstHandle moved(std::move(writer));
      Swap(moved);
      return *this;
    }
 
    ~ConstHandle() { Reset(); }
 
    [[nodiscard]] bool Valid() const { return pool_ != nullptr; }
 
    [[nodiscard]] const Data& Get() const
    {
      ASSERT(Valid());
      return pool_->slots_[index_].data_;
    }
 
    [[nodiscard]] const Data* operator->() const { return &Get(); }
    [[nodiscard]] const Data& operator*() const { return Get(); }
 
    [[nodiscard]] IndexType Index() const
    {
      ASSERT(Valid());
      return index_;
    }
 
    void Reset()
    {
      ObjectPool* pool = std::exchange(pool_, nullptr);
      const IndexType index = std::exchange(index_, IndexType{});
      if (pool != nullptr)
      {
        pool->Release(index);
      }
    }
 
   private:
    void Swap(ConstHandle& other) noexcept
    {
      std::swap(pool_, other.pool_);
      std::swap(index_, other.index_);
    }
 
    ObjectPool* pool_ = nullptr;
    IndexType index_ = {};
  };
 
  template <typename T = Data>
    requires std::is_default_constructible_v<T>
  explicit ObjectPool(size_t slot_count)
      : slots_(new Slot[slot_count]), slot_count_(slot_count), owns_slots_(true)
  {
    InitializeFreeStack();
  }
 
  ObjectPool(size_t slot_count, Slot* slots)
      : slots_(slots), slot_count_(slot_count), owns_slots_(false)
  {
    ASSERT(slots_ != nullptr);
    InitializeFreeStack();
  }
 
  ~ObjectPool()
  {
    ASSERT(EmptySize() == slot_count_);
 
    if (owns_slots_)
    {
      delete[] slots_;
    }
  }
 
  ObjectPool(const ObjectPool&) = delete;
  ObjectPool& operator=(const ObjectPool&) = delete;
  ObjectPool(ObjectPool&&) = delete;
  ObjectPool& operator=(ObjectPool&&) = delete;
 
  [[nodiscard]] ErrorCode Acquire(Handle& handle)
  {
    ASSERT(!handle.Valid());
#ifdef LIBXR_DEBUG_BUILD
    // 重叠申请会让空闲栈出现 ABA，同一槽位被两个申请方取得。
    // Overlapping acquisitions cause ABA on the free stack and hand one slot to two
    // acquirers.
    uint32_t idle = 0U;
    ASSERT(acquiring_.compare_exchange_strong(idle, 1U, std::memory_order_acquire,
                                              std::memory_order_relaxed));
#endif
 
    IndexType index = 0;
    const ErrorCode ec = PopFree(index);
    if (ec == ErrorCode::OK)
    {
      free_count_.fetch_sub(1U, std::memory_order_relaxed);
      DEV_ASSERT(slots_[index].references_.load(std::memory_order_relaxed) == 0U);
      slots_[index].references_.store(1U, std::memory_order_relaxed);
      handle = Handle(this, index);
    }
 
#ifdef LIBXR_DEBUG_BUILD
    acquiring_.store(0U, std::memory_order_release);
#endif
    return ec;
  }
 
  [[nodiscard]] size_t EmptySize() const
  {
    return free_count_.load(std::memory_order_relaxed);
  }
 
  [[nodiscard]] size_t Size() const { return slot_count_; }
 
  [[nodiscard]] Data& UnsafeAt(size_t index)
  {
    ASSERT(index < slot_count_);
    return slots_[index].data_;
  }
 
  [[nodiscard]] const Data& UnsafeAt(size_t index) const
  {
    ASSERT(index < slot_count_);
    return slots_[index].data_;
  }
 
 private:
  static constexpr IndexType EMPTY_INDEX = std::numeric_limits<IndexType>::max();
 
  void Retain(IndexType index) noexcept
  {
    DEV_ASSERT(static_cast<size_t>(index) < slot_count_);
    [[maybe_unused]] const uint32_t previous =
        slots_[index].references_.fetch_add(1U, std::memory_order_relaxed);
    DEV_ASSERT(previous > 0U);
    ASSERT(previous < std::numeric_limits<uint32_t>::max());
  }
 
  void Release(IndexType index)
  {
    DEV_ASSERT(static_cast<size_t>(index) < slot_count_);
    // acq_rel：其他持有者的负载访问先于最终归还。
    // acq_rel: other holders' payload accesses happen before the final return.
    const uint32_t previous =
        slots_[index].references_.fetch_sub(1U, std::memory_order_acq_rel);
    DEV_ASSERT(previous > 0U);
    if (previous == 1U)
    {
      free_count_.fetch_add(1U, std::memory_order_relaxed);
      PushFree(index);
    }
  }
 
  void InitializeFreeStack()
  {
    ASSERT(slot_count_ > 0);
    ASSERT(slot_count_ < static_cast<size_t>(EMPTY_INDEX));
    for (size_t index = 0; index < slot_count_; ++index)
    {
      slots_[index].next_ = static_cast<IndexType>(index + 1U);
    }
    slots_[slot_count_ - 1U].next_ = EMPTY_INDEX;
    free_head_.store(0U, std::memory_order_relaxed);
    free_count_.store(static_cast<uint32_t>(slot_count_), std::memory_order_relaxed);
  }
 
  ErrorCode PopFree(IndexType& index)
  {
    IndexType head = free_head_.load(std::memory_order_acquire);
    IndexType next = EMPTY_INDEX;
    do
    {
      if (head == EMPTY_INDEX)
      {
        return ErrorCode::EMPTY;
      }
      DEV_ASSERT(static_cast<size_t>(head) < slot_count_);
      // 只有申请方会取出槽位，head 在栈上时其 next_ 不会被改写，所以没有 ABA。
      // Only the acquirer removes slots, so next_ of a stacked head cannot change and
      // ABA cannot occur.
      next = slots_[head].next_;
    } while (!free_head_.compare_exchange_weak(head, next, std::memory_order_acquire,
                                               std::memory_order_acquire));
    index = head;
    return ErrorCode::OK;
  }
 
  void PushFree(IndexType index)
  {
    IndexType head = free_head_.load(std::memory_order_relaxed);
    do
    {
      // 槽位已不在栈上且只属于本次归还，可以直接写 next_。
      // The slot is off the stack and owned by this return, so next_ is private here.
      slots_[index].next_ = head;
    } while (!free_head_.compare_exchange_weak(head, index, std::memory_order_release,
                                               std::memory_order_relaxed));
  }
 
  Slot* slots_;              
  const size_t slot_count_;  
  const bool owns_slots_;    
  std::atomic<IndexType> free_head_{EMPTY_INDEX};  
  std::atomic<uint32_t> free_count_{0};            
  std::atomic<uint32_t> acquiring_{0};
};
}  // namespace LibXR