pool.h

Go to the documentation of this file.

// Copyright (c) 2022 The Bitcoin Core developers
// Distributed under the MIT software license, see the accompanying
// file COPYING or http://www.opensource.org/licenses/mit-license.php.
 
#ifndef BITCOIN_SUPPORT_ALLOCATORS_POOL_H
#define BITCOIN_SUPPORT_ALLOCATORS_POOL_H
 
#include <array>
#include <cassert>
#include <cstddef>
#include <list>
#include <memory>
#include <new>
#include <type_traits>
#include <utility>
 
template <std::size_t MAX_BLOCK_SIZE_BYTES, std::size_t ALIGN_BYTES>
class PoolResource final {
    static_assert(ALIGN_BYTES > 0, "ALIGN_BYTES must be nonzero");
    static_assert((ALIGN_BYTES & (ALIGN_BYTES - 1)) == 0,
                  "ALIGN_BYTES must be a power of two");
 
    struct ListNode {
        ListNode *m_next;
 
        explicit ListNode(ListNode *next) : m_next(next) {}
    };
    static_assert(std::is_trivially_destructible_v<ListNode>,
                  "Make sure we don't need to manually call a destructor");
 
    static constexpr std::size_t ELEM_ALIGN_BYTES =
        std::max(alignof(ListNode), ALIGN_BYTES);
    static_assert((ELEM_ALIGN_BYTES & (ELEM_ALIGN_BYTES - 1)) == 0,
                  "ELEM_ALIGN_BYTES must be a power of two");
    static_assert(
        sizeof(ListNode) <= ELEM_ALIGN_BYTES,
        "Units of size ELEM_SIZE_ALIGN need to be able to store a ListNode");
    static_assert(
        (MAX_BLOCK_SIZE_BYTES & (ELEM_ALIGN_BYTES - 1)) == 0,
        "MAX_BLOCK_SIZE_BYTES needs to be a multiple of the alignment.");
 
    const size_t m_chunk_size_bytes;
 
    std::list<std::byte *> m_allocated_chunks{};
 
    std::array<ListNode *, MAX_BLOCK_SIZE_BYTES / ELEM_ALIGN_BYTES + 1>
        m_free_lists{};
 
    std::byte *m_available_memory_it = nullptr;
 
    std::byte *m_available_memory_end = nullptr;
 
    [[nodiscard]] static constexpr std::size_t
    NumElemAlignBytes(std::size_t bytes) {
        return (bytes + ELEM_ALIGN_BYTES - 1) / ELEM_ALIGN_BYTES + (bytes == 0);
    }
 
    [[nodiscard]] static constexpr bool
    IsFreeListUsable(std::size_t bytes, std::size_t alignment) {
        return alignment <= ELEM_ALIGN_BYTES && bytes <= MAX_BLOCK_SIZE_BYTES;
    }
 
    void PlacementAddToList(void *p, ListNode *&node) {
        node = new (p) ListNode{node};
    }
 
    void AllocateChunk() {
        // if there is still any available memory left, put it into the
        // freelist.
        size_t remaining_available_bytes =
            std::distance(m_available_memory_it, m_available_memory_end);
        if (0 != remaining_available_bytes) {
            PlacementAddToList(
                m_available_memory_it,
                m_free_lists[remaining_available_bytes / ELEM_ALIGN_BYTES]);
        }
 
        void *storage = ::operator new(m_chunk_size_bytes,
                                       std::align_val_t{ELEM_ALIGN_BYTES});
        m_available_memory_it = new (storage) std::byte[m_chunk_size_bytes];
        m_available_memory_end = m_available_memory_it + m_chunk_size_bytes;
        m_allocated_chunks.emplace_back(m_available_memory_it);
    }
 
    friend class PoolResourceTester;
 
public:
    explicit PoolResource(std::size_t chunk_size_bytes)
        : m_chunk_size_bytes(NumElemAlignBytes(chunk_size_bytes) *
                             ELEM_ALIGN_BYTES) {
        assert(m_chunk_size_bytes >= MAX_BLOCK_SIZE_BYTES);
        AllocateChunk();
    }
 
    PoolResource() : PoolResource(262144) {}
 
    PoolResource(const PoolResource &) = delete;
    PoolResource &operator=(const PoolResource &) = delete;
    PoolResource(PoolResource &&) = delete;
    PoolResource &operator=(PoolResource &&) = delete;
 
    ~PoolResource() {
        for (std::byte *chunk : m_allocated_chunks) {
            std::destroy(chunk, chunk + m_chunk_size_bytes);
            ::operator delete((void *)chunk,
                              std::align_val_t{ELEM_ALIGN_BYTES});
        }
    }
 
    void *Allocate(std::size_t bytes, std::size_t alignment) {
        if (IsFreeListUsable(bytes, alignment)) {
            const std::size_t num_alignments = NumElemAlignBytes(bytes);
            if (nullptr != m_free_lists[num_alignments]) {
                // we've already got data in the pool's freelist, unlink one
                // element and return the pointer to the unlinked memory. Since
                // FreeList is trivially destructible we can just treat it as
                // uninitialized memory.
                return std::exchange(m_free_lists[num_alignments],
                                     m_free_lists[num_alignments]->m_next);
            }
 
            // freelist is empty: get one allocation from allocated chunk
            // memory.
            const std::ptrdiff_t round_bytes =
                static_cast<std::ptrdiff_t>(num_alignments * ELEM_ALIGN_BYTES);
            if (round_bytes > m_available_memory_end - m_available_memory_it) {
                // slow path, only happens when a new chunk needs to be
                // allocated
                AllocateChunk();
            }
 
            // Make sure we use the right amount of bytes for that freelist
            // (might be rounded up),
            return std::exchange(m_available_memory_it,
                                 m_available_memory_it + round_bytes);
        }
 
        // Can't use the pool => use operator new()
        return ::operator new(bytes, std::align_val_t{alignment});
    }
 
    void Deallocate(void *p, std::size_t bytes,
                    std::size_t alignment) noexcept {
        if (IsFreeListUsable(bytes, alignment)) {
            const std::size_t num_alignments = NumElemAlignBytes(bytes);
            // put the memory block into the linked list. We can placement
            // construct the FreeList into the memory since we can be sure the
            // alignment is correct.
            PlacementAddToList(p, m_free_lists[num_alignments]);
        } else {
            // Can't use the pool => forward deallocation to ::operator
            // delete().
            ::operator delete(p, std::align_val_t{alignment});
        }
    }
 
    [[nodiscard]] std::size_t NumAllocatedChunks() const {
        return m_allocated_chunks.size();
    }
 
    [[nodiscard]] size_t ChunkSizeBytes() const { return m_chunk_size_bytes; }
};
 
template <class T, std::size_t MAX_BLOCK_SIZE_BYTES,
          std::size_t ALIGN_BYTES = alignof(T)>
class PoolAllocator {
    PoolResource<MAX_BLOCK_SIZE_BYTES, ALIGN_BYTES> *m_resource;
 
    template <typename U, std::size_t M, std::size_t A>
    friend class PoolAllocator;
 
public:
    using value_type = T;
    using ResourceType = PoolResource<MAX_BLOCK_SIZE_BYTES, ALIGN_BYTES>;
 
    PoolAllocator(ResourceType *resource) noexcept : m_resource(resource) {}
 
    PoolAllocator(const PoolAllocator &other) noexcept = default;
    PoolAllocator &operator=(const PoolAllocator &other) noexcept = default;
 
    template <class U>
    PoolAllocator(const PoolAllocator<U, MAX_BLOCK_SIZE_BYTES, ALIGN_BYTES>
                      &other) noexcept
        : m_resource(other.resource()) {}
 
    template <typename U> struct rebind {
        using other = PoolAllocator<U, MAX_BLOCK_SIZE_BYTES, ALIGN_BYTES>;
    };
 
    T *allocate(size_t n) {
        return static_cast<T *>(
            m_resource->Allocate(n * sizeof(T), alignof(T)));
    }
 
    void deallocate(T *p, size_t n) noexcept {
        m_resource->Deallocate(p, n * sizeof(T), alignof(T));
    }
 
    ResourceType *resource() const noexcept { return m_resource; }
};
 
template <class T1, class T2, std::size_t MAX_BLOCK_SIZE_BYTES,
          std::size_t ALIGN_BYTES>
bool operator==(
    const PoolAllocator<T1, MAX_BLOCK_SIZE_BYTES, ALIGN_BYTES> &a,
    const PoolAllocator<T2, MAX_BLOCK_SIZE_BYTES, ALIGN_BYTES> &b) noexcept {
    return a.resource() == b.resource();
}
 
template <class T1, class T2, std::size_t MAX_BLOCK_SIZE_BYTES,
          std::size_t ALIGN_BYTES>
bool operator!=(
    const PoolAllocator<T1, MAX_BLOCK_SIZE_BYTES, ALIGN_BYTES> &a,
    const PoolAllocator<T2, MAX_BLOCK_SIZE_BYTES, ALIGN_BYTES> &b) noexcept {
    return !(a == b);
}
 
#endif // BITCOIN_SUPPORT_ALLOCATORS_POOL_H