Custom Allocators in C++: Pool, Arena & PMR Allocators Explained

#include <cassert>
#include <cstddef>
#include <cstdint>
#include <iostream>
#include <new>
#include <vector>

// ---------------------------------------------------------------------------
// PoolAllocator<T, BlockCount>
// Manages a fixed pool of BlockCount objects of type T.
// Allocation and deallocation are O(1) with no fragmentation.
// NOT thread-safe — use per-thread instances in concurrent scenarios.
// ---------------------------------------------------------------------------
template <typename T, std::size_t BlockCount>
class PoolResource {
    // Each free slot reuses its own memory to store a 'next' pointer.
    // This union lets us treat raw bytes as either a pointer or storage for T.
    union Slot {
        Slot*              next_free;   // when this slot is unused
        alignas(T) char    storage[sizeof(T)]; // when this slot is in use
    };

    Slot              pool_[BlockCount];  // the entire pool lives here (stack or member)
    Slot*             free_head_;        // points to the first available slot
    std::size_t       allocated_count_;  // for diagnostics

public:
    PoolResource() : allocated_count_(0) {
        // Chain all slots into the free-list at construction time.
        // Slot i's 'next' pointer points to slot i+1.
        for (std::size_t i = 0; i < BlockCount - 1; ++i) {
            pool_[i].next_free = &pool_[i + 1];
        }
        pool_[BlockCount - 1].next_free = nullptr; // last slot has no successor
        free_head_ = &pool_[0];                    // head starts at first slot
    }

    // Returns a pointer to raw storage for one T — does NOT construct T.
    T* allocate() {
        if (free_head_ == nullptr) {
            throw std::bad_alloc(); // pool exhausted
        }
        Slot* chosen_slot = free_head_;      // take the head slot
        free_head_ = chosen_slot->next_free; // advance the free-list head
        ++allocated_count_;
        return reinterpret_cast<T*>(chosen_slot->storage);
    }

    // Returns a slot to the free-list — does NOT destroy T.
    // Caller must call T's destructor manually before deallocating.
    void deallocate(T* object_ptr) noexcept {
        // Cast back to Slot so we can rewrite the next_free pointer.
        Slot* returned_slot = reinterpret_cast<Slot*>(object_ptr);
        returned_slot->next_free = free_head_; // push onto free-list
        free_head_ = returned_slot;
        --allocated_count_;
    }

    std::size_t allocated() const noexcept { return allocated_count_; }
    std::size_t capacity()  const noexcept { return BlockCount; }
};

// ---------------------------------------------------------------------------
// PoolAllocator<T> — STL-compatible wrapper around PoolResource.
// Satisfies the named Allocator requirement.
// ---------------------------------------------------------------------------
template <typename T, std::size_t BlockCount = 64>
struct PoolAllocator {
    using value_type = T;
    // Propagate allocator on container move — critical for stateful allocators.
    using propagate_on_container_move_assignment = std::true_type;
    using is_always_equal                        = std::false_type; // stateful!

    // Shared ownership of the resource via a raw pointer.
    // In production use std::shared_ptr<PoolResource<T, BlockCount>>.
    PoolResource<T, BlockCount>* resource_;

    explicit PoolAllocator(PoolResource<T, BlockCount>& res) noexcept
        : resource_(&res) {}

    template <typename U>
    PoolAllocator(const PoolAllocator<U, BlockCount>&) noexcept
        : resource_(nullptr) {
        // Rebind constructor — node allocators for list/map will arrive here.
        // Because block size changes with U, this is intentionally disabled:
        // use PMR for node containers instead.
        static_assert(sizeof(U) == sizeof(T),
            "PoolAllocator rebind is only valid for same-size types");
    }

    T* allocate(std::size_t n) {
        if (n != 1) throw std::bad_alloc(); // pool only handles single-object allocs
        return resource_->allocate();
    }

    void deallocate(T* ptr, std::size_t) noexcept {
        resource_->deallocate(ptr);
    }

    bool operator==(const PoolAllocator& other) const noexcept {
        return resource_ == other.resource_; // equal iff same resource
    }
    bool operator!=(const PoolAllocator& other) const noexcept {
        return !(*this == other);
    }
};

// ---------------------------------------------------------------------------
// Demonstration: a small event-like struct allocated from a pool.
// ---------------------------------------------------------------------------
struct NetworkPacket {
    uint32_t source_ip;
    uint32_t dest_ip;
    uint16_t payload_size;

    NetworkPacket(uint32_t src, uint32_t dst, uint16_t sz)
        : source_ip(src), dest_ip(dst), payload_size(sz) {}
};

int main() {
    constexpr std::size_t POOL_CAPACITY = 8;

    PoolResource<NetworkPacket, POOL_CAPACITY> packet_pool;

    std::cout << "Pool capacity: " << packet_pool.capacity() << " packets\n";

    // Allocate raw storage then placement-new to construct in-place.
    NetworkPacket* p1 = packet_pool.allocate();
    new (p1) NetworkPacket(0xC0A80001, 0xC0A80002, 512);

    NetworkPacket* p2 = packet_pool.allocate();
    new (p2) NetworkPacket(0xC0A80003, 0xC0A80004, 256);

    std::cout << "Allocated: " << packet_pool.allocated() << '\n';
    std::cout << "p1 payload: " << p1->payload_size << " bytes\n";
    std::cout << "p2 payload: " << p2->payload_size << " bytes\n";

    // Destroy then return to pool — ORDER matters.
    p1->~NetworkPacket();
    packet_pool.deallocate(p1);

    p2->~NetworkPacket();
    packet_pool.deallocate(p2);

    std::cout << "After deallocation, in-use: " << packet_pool.allocated() << '\n';
    return 0;
}

#include <array>
#include <cstddef>
#include <cstdint>
#include <iostream>
#include <memory_resource>  // C++17 — compile with -std=c++17
#include <string>
#include <vector>

// ---------------------------------------------------------------------------
// Part 1: Hand-rolled arena allocator (instructive, not for production)
// Shows the fundamental mechanic before PMR abstracts it.
// ---------------------------------------------------------------------------
class ArenaResource {
    std::byte* const buffer_start_;
    std::byte* const buffer_end_;
    std::byte*       bump_ptr_;       // next free byte

public:
    explicit ArenaResource(std::byte* buffer, std::size_t size) noexcept
        : buffer_start_(buffer)
        , buffer_end_(buffer + size)
        , bump_ptr_(buffer) {}

    void* allocate(std::size_t bytes, std::size_t alignment) {
        // Align bump_ptr_ up to the required alignment boundary.
        std::uintptr_t current_addr = reinterpret_cast<std::uintptr_t>(bump_ptr_);
        std::uintptr_t aligned_addr = (current_addr + alignment - 1) & ~(alignment - 1);
        std::byte*     aligned_ptr  = reinterpret_cast<std::byte*>(aligned_addr);

        if (aligned_ptr + bytes > buffer_end_) {
            throw std::bad_alloc(); // arena exhausted
        }

        bump_ptr_ = aligned_ptr + bytes; // advance past the allocation
        return aligned_ptr;
    }

    // Arena deallocation is intentionally a no-op.
    void deallocate(void*, std::size_t) noexcept {}

    // Reset the entire arena in O(1) — all previous allocations invalidated.
    void reset() noexcept { bump_ptr_ = buffer_start_; }

    std::size_t used()      const noexcept { return static_cast<std::size_t>(bump_ptr_ - buffer_start_); }
    std::size_t remaining() const noexcept { return static_cast<std::size_t>(buffer_end_ - bump_ptr_); }
};

// ---------------------------------------------------------------------------
// Part 2: C++17 PMR arena with std::pmr containers.
// monotonic_buffer_resource IS the production-grade arena.
// ---------------------------------------------------------------------------
void process_request_with_pmr(int request_id) {
    // Stack-allocate 4KB for this request's scratch memory.
    // Declared here so it outlives the resource and containers below.
    alignas(std::max_align_t) std::array<std::byte, 4096> scratch_buffer;

    // monotonic_buffer_resource: bump-pointer into scratch_buffer.
    // Falls back to std::pmr::new_delete_resource() if scratch fills up.
    std::pmr::monotonic_buffer_resource arena{
        scratch_buffer.data(),
        scratch_buffer.size(),
        std::pmr::new_delete_resource()  // upstream fallback
    };

    // polymorphic_allocator wraps the resource and adapts to any value_type.
    // std::pmr::vector IS std::vector<T, std::pmr::polymorphic_allocator<T>>.
    std::pmr::vector<std::pmr::string> log_lines{&arena};
    log_lines.reserve(16); // hits the arena, not the heap

    log_lines.push_back("Request received");
    log_lines.push_back("Validating parameters");
    log_lines.push_back("Processing payload");
    log_lines.push_back("Response dispatched");

    std::cout << "[Request " << request_id << "] arena used: "
              << arena.used_bytes()  // GCC extension; use custom tracking otherwise
              << " bytes (approx)\n";

    for (const auto& line : log_lines) {
        std::cout << "  > " << line << '\n';
    }

    // When this function returns:
    // 1. log_lines destructor runs — calls pmr::string destructors (NO heap frees).
    // 2. arena destructor runs — resets the bump pointer.
    // 3. scratch_buffer goes out of scope — no heap involved at all.
    // Total heap allocations for this call: ZERO (as long as scratch is sufficient).
}

// ---------------------------------------------------------------------------
// Part 3: Chaining resources — pool backed by an arena (common in game engines)
// ---------------------------------------------------------------------------
void demonstrate_chained_resources() {
    alignas(std::max_align_t) std::array<std::byte, 65536> big_arena;

    // The arena is the upstream for a pool resource.
    // Pool resource handles variable-size allocations efficiently within the arena.
    std::pmr::monotonic_buffer_resource upstream_arena{
        big_arena.data(), big_arena.size(), std::pmr::null_memory_resource()
        // null_memory_resource: throws if arena overflows — no hidden heap use.
    };
    std::pmr::unsynchronized_pool_resource pool{&upstream_arena};

    // Now allocate heterogeneous containers — all memory comes from big_arena.
    std::pmr::vector<int>         integers{&pool};
    std::pmr::vector<double>      doubles{&pool};
    std::pmr::vector<std::string> names{&pool}; // note: string's internal heap is NOT pooled

    integers.assign({10, 20, 30, 40, 50});
    doubles.assign({3.14, 2.71, 1.41});
    names.assign({"Alice", "Bob", "Carol"});

    std::cout << "\nChained resource demo:\n";
    std::cout << "Integers: ";
    for (int v : integers) std::cout << v << ' ';
    std::cout << '\n';

    std::cout << "Names: ";
    for (const auto& n : names) std::cout << n << ' ';
    std::cout << '\n';
}

int main() {
    // Part 2 demo
    process_request_with_pmr(42);
    std::cout << '\n';

    // Part 3 demo
    demonstrate_chained_resources();
    return 0;
}

#include <array>
#include <chrono>
#include <cstddef>
#include <iostream>
#include <memory_resource>
#include <string>
#include <vector>

// ---------------------------------------------------------------------------
// Micro-benchmark: heap allocator vs PMR monotonic arena
// for building a vector of strings inside a hot loop.
//
// Compile with optimisations: g++ -O2 -std=c++17 allocator_benchmark.cpp
// ---------------------------------------------------------------------------

constexpr int    ITERATION_COUNT     = 100'000;
constexpr int    STRINGS_PER_ITER    = 32;
constexpr size_t ARENA_SIZE_BYTES    = 32 * 1024; // 32KB per iteration

// Returns duration in microseconds.
long long benchmark_heap_allocator() {
    using Clock = std::chrono::high_resolution_clock;
    auto start  = Clock::now();

    for (int iter = 0; iter < ITERATION_COUNT; ++iter) {
        // Fresh heap vector every iteration — alloc + dealloc on every pass.
        std::vector<std::string> log_entries;
        log_entries.reserve(STRINGS_PER_ITER);

        for (int i = 0; i < STRINGS_PER_ITER; ++i) {
            // Each string that's > 15 chars triggers a heap allocation
            // (SSO threshold on most implementations).
            log_entries.push_back("event_log_entry_number_" + std::to_string(i));
        }
        // Vector destructs here: frees string buffers + vector buffer.
    }

    auto end = Clock::now();
    return std::chrono::duration_cast<std::chrono::microseconds>(end - start).count();
}

long long benchmark_arena_allocator() {
    using Clock = std::chrono::high_resolution_clock;

    // The backing buffer is declared ONCE outside the loop.
    // We reset the arena each iteration — O(1) reset, no dealloc overhead.
    alignas(std::max_align_t) std::array<std::byte, ARENA_SIZE_BYTES> backing_buffer;

    auto start = Clock::now();

    for (int iter = 0; iter < ITERATION_COUNT; ++iter) {
        // Reset resets the bump pointer — effectively "frees" all previous allocations.
        std::pmr::monotonic_buffer_resource arena{
            backing_buffer.data(),
            backing_buffer.size(),
            std::pmr::null_memory_resource() // crash if arena overflows — no hidden heap
        };

        // pmr::vector and pmr::string both allocate from 'arena'.
        std::pmr::vector<std::pmr::string> log_entries{&arena};
        log_entries.reserve(STRINGS_PER_ITER);

        for (int i = 0; i < STRINGS_PER_ITER; ++i) {
            log_entries.push_back("event_log_entry_number_" + std::to_string(i));
        }
        // Destructors run (pmr::string destructors call deallocate — a no-op on arena).
        // arena destructor runs — bump pointer reset in O(1).
    }

    auto end = Clock::now();
    return std::chrono::duration_cast<std::chrono::microseconds>(end - start).count();
}

int main() {
    std::cout << "Running " << ITERATION_COUNT << " iterations, "
              << STRINGS_PER_ITER << " strings each...\n\n";

    long long heap_us  = benchmark_heap_allocator();
    long long arena_us = benchmark_arena_allocator();

    std::cout << "Heap allocator total:  " << heap_us  << " us\n";
    std::cout << "Arena allocator total: " << arena_us << " us\n";
    std::cout << "Speedup:               "
              << (static_cast<double>(heap_us) / arena_us) << "x\n\n";

    std::cout << "Per-iteration averages:\n";
    std::cout << "  Heap:  " << (heap_us  * 1000.0 / ITERATION_COUNT) << " ns\n";
    std::cout << "  Arena: " << (arena_us * 1000.0 / ITERATION_COUNT) << " ns\n";
    return 0;
}