object_pool.hpp

#pragma once
 
#include <cstddef>
#include <cstdint>
#include <limits>
#include <new>
#include <type_traits>
#include <utility>
 
#include "libxr_def.hpp"
#include "queue.hpp"
 
namespace LibXR
{
template <typename QueueType>
concept PoolIndexQueue =
    requires(QueueType queue, const typename QueueType::ValueType& index_const,
             typename QueueType::ValueType& index_mut) {
      typename QueueType::ValueType;
      { queue.Push(index_const) } -> std::same_as<ErrorCode>;
      { queue.Pop(index_mut) } -> std::same_as<ErrorCode>;
      { queue.Size() } -> std::convertible_to<size_t>;
    };
 
template <typename Data, PoolIndexQueue FreeQueue>
class BasicObjectPool
{
 public:
  using ValueType = Data;       
  using QueueType = FreeQueue;  
  using IndexType = typename FreeQueue::ValueType;  
 
  static_assert(std::is_integral_v<IndexType>,
                "BasicObjectPool requires an integral index queue");
  static_assert(std::is_unsigned_v<IndexType>,
                "BasicObjectPool requires an unsigned index queue");
 
  class Handle
  {
   public:
    Handle() = default;
 
    Handle(BasicObjectPool* pool, IndexType index) : pool_(pool), index_(index) {}
 
    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)
      {
        return *this;
      }
 
      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(pool_ != nullptr);
      return pool_->slots_[index_];
    }
 
    [[nodiscard]] const Data& Get() const
    {
      ASSERT(pool_ != nullptr);
      return pool_->slots_[index_];
    }
 
    [[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(pool_ != nullptr);
      return index_;
    }
 
    void Reset()
    {
      if (pool_ == nullptr)
      {
        return;
      }
 
      const ErrorCode ec = pool_->Release(index_);
      ASSERT(ec == ErrorCode::OK);
      pool_ = nullptr;
      index_ = IndexType{};
    }
 
   private:
    BasicObjectPool* pool_ = nullptr;  
    IndexType index_ = {};             
  };
 
  template <typename T = Data>
    requires std::is_default_constructible_v<T>
  explicit BasicObjectPool(size_t slot_count) : slot_count_(slot_count)
  {
    ASSERT(slot_count_ > 0);
    ConstructOwnedQueue(slot_count_);
    AllocateOwnedSlots();
    InitializeFreeQueue();
  }
 
  BasicObjectPool(size_t slot_count, Data* slots) : slot_count_(slot_count), slots_(slots)
  {
    ASSERT(slot_count_ > 0);
    ASSERT(slots_ != nullptr);
    ConstructOwnedQueue(slot_count_);
    InitializeFreeQueue();
  }
 
  template <typename T = Data>
    requires std::is_default_constructible_v<T>
  BasicObjectPool(FreeQueue& free_queue, size_t slot_count)
      : free_queue_(&free_queue), slot_count_(slot_count)
  {
    ASSERT(slot_count_ > 0);
    ASSERT(free_queue.Size() == 0);
    AllocateOwnedSlots();
    InitializeFreeQueue();
  }
 
  BasicObjectPool(FreeQueue& free_queue, size_t slot_count, Data* slots)
      : free_queue_(&free_queue), slot_count_(slot_count), slots_(slots)
  {
    ASSERT(slot_count_ > 0);
    ASSERT(slots_ != nullptr);
    ASSERT(free_queue.Size() == 0);
    InitializeFreeQueue();
  }
 
  ~BasicObjectPool()
  {
    ASSERT(EmptySize() == slot_count_);
 
    if (owns_slots_)
    {
      delete[] slots_;
    }
 
    if (owns_free_queue_)
    {
      free_queue_->~FreeQueue();
    }
  }
 
  BasicObjectPool(const BasicObjectPool&) = delete;
  BasicObjectPool& operator=(const BasicObjectPool&) = delete;
  BasicObjectPool(BasicObjectPool&&) = delete;
  BasicObjectPool& operator=(BasicObjectPool&&) = delete;
 
  [[nodiscard]] ErrorCode Acquire(Handle& handle)
  {
    ASSERT(!handle.Valid());
 
    IndexType index = 0;
    const ErrorCode ec = free_queue_->Pop(index);
    if (ec != ErrorCode::OK)
    {
      return ec;
    }
 
    ASSERT(static_cast<size_t>(index) < slot_count_);
    handle = Handle(this, index);
    return ErrorCode::OK;
  }
 
  [[nodiscard]] size_t EmptySize() const { return free_queue_->Size(); }
 
  [[nodiscard]] size_t Size() const { return slot_count_; }
 
  [[nodiscard]] Data& UnsafeAt(size_t index)
  {
    ASSERT(index < slot_count_);
    return slots_[index];
  }
 
  [[nodiscard]] const Data& UnsafeAt(size_t index) const
  {
    ASSERT(index < slot_count_);
    return slots_[index];
  }
 
 private:
  ErrorCode Release(IndexType index)
  {
    ASSERT(static_cast<size_t>(index) < slot_count_);
    return free_queue_->Push(index);
  }
 
  void ConstructOwnedQueue(size_t slot_count)
  {
    free_queue_ = new (free_queue_storage_) FreeQueue(slot_count);
    owns_free_queue_ = true;
  }
 
  void AllocateOwnedSlots()
  {
    slots_ = new Data[slot_count_];
    owns_slots_ = true;
  }
 
  void InitializeFreeQueue()
  {
    ASSERT(slot_count_ > 0);
    ASSERT(slot_count_ - 1 <= static_cast<size_t>(std::numeric_limits<IndexType>::max()));
    for (size_t index = 0; index < slot_count_; ++index)
    {
      const ErrorCode ec = free_queue_->Push(static_cast<IndexType>(index));
      ASSERT(ec == ErrorCode::OK);
    }
  }
 
  alignas(FreeQueue) std::byte free_queue_storage_[sizeof(FreeQueue)] = {};
  FreeQueue* free_queue_ =
      nullptr;               
  const size_t slot_count_;  
  Data* slots_ = nullptr;    
  bool owns_free_queue_ =
      false;  
  bool owns_slots_ =
      false;  
};
 
template <typename Data, typename IndexType = uint32_t>
using ObjectPool = BasicObjectPool<Data, Queue<IndexType>>;
 
template <typename Data, typename IndexType = uint32_t>
using SPSCObjectPool = BasicObjectPool<Data, SPSCQueue<IndexType>>;
 
template <typename Data, typename IndexType = uint32_t>
using MPMCObjectPool = BasicObjectPool<Data, MPMCQueue<IndexType>>;
}  // namespace LibXR