pool.hpp

// Copyright (c) Sleipnir contributors
 
#pragma once
 
#include <cstddef>
#include <memory>
 
#include <gch/small_vector.hpp>
 
#include "sleipnir/util/symbol_exports.hpp"
 
namespace slp {
 
class SLEIPNIR_DLLEXPORT PoolResource {
 public:
  explicit PoolResource(size_t blocks_per_chunk);
 
  ~PoolResource();
 
  PoolResource(const PoolResource&) = delete;
 
  PoolResource& operator=(const PoolResource&) = delete;
 
  PoolResource(PoolResource&&);
 
  PoolResource& operator=(PoolResource&&);
 
  [[nodiscard]]
  void* allocate(size_t bytes, size_t alignment = alignof(std::max_align_t));
 
  void deallocate(void* p, size_t bytes,
                  size_t alignment = alignof(std::max_align_t));
 
  bool is_equal(const PoolResource& other) const noexcept {
    return this == &other;
  }
 
  size_t blocks_in_use() const noexcept {
    size_t free_blocks = 0;
    for (const auto& pool : m_pools) {
      free_blocks += pool.free_list.size();
    }
    return m_chunks.size() * blocks_per_chunk - free_blocks;
  }
 
 private:
  struct Pool {
    size_t block_size;
 
    gch::small_vector<void*> free_list;
  };
 
  struct Chunk {
    std::byte* begin;
 
    size_t pool_index;
  };
 
  gch::small_vector<std::unique_ptr<std::byte[]>> m_buffer;
 
  gch::small_vector<Chunk> m_chunks;
 
  gch::small_vector<Pool> m_pools;
  size_t blocks_per_chunk;
 
  Pool& get_pool(size_t block_size);
 
  void add_chunk(Pool& pool);
};
 
template <typename T>
class PoolAllocator {
 public:
  using value_type = T;
 
  explicit constexpr PoolAllocator(PoolResource* r) : m_memory_resource{r} {}
 
  constexpr PoolAllocator(const PoolAllocator<T>&) = default;
 
  constexpr PoolAllocator<T>& operator=(const PoolAllocator<T>&) = default;
 
  [[nodiscard]]
  constexpr T* allocate(size_t n) {
    return static_cast<T*>(
        m_memory_resource->allocate(n * sizeof(T), alignof(T)));
  }
 
  constexpr void deallocate(T* p, size_t n) {
    m_memory_resource->deallocate(p, n * sizeof(T), alignof(T));
  }
 
 private:
  PoolResource* m_memory_resource;
};
 
SLEIPNIR_DLLEXPORT PoolResource& global_pool_resource();
 
template <typename T>
PoolAllocator<T> global_pool_allocator() {
  return PoolAllocator<T>{&global_pool_resource()};
}
 
}  // namespace slp