STLLocalAllocator.h

Go to the documentation of this file.

// Jolt Physics Library (https://github.com/jrouwe/JoltPhysics)
// SPDX-FileCopyrightText: 2025 Jorrit Rouwe
// SPDX-License-Identifier: MIT
 
#pragma once
 
#include <Jolt/Core/STLAllocator.h>
 
JPH_NAMESPACE_BEGIN
 
#ifndef JPH_DISABLE_CUSTOM_ALLOCATOR
 
template <typename T, size_t N>
class STLLocalAllocator : private STLAllocator<T>
{
    using Base = STLAllocator<T>;
 
public:
    using value_type = T;
    using pointer = T *;
    using const_pointer = const T *;
    using reference = T &;
    using const_reference = const T &;
    using size_type = size_t;
    using difference_type = ptrdiff_t;
 
    using is_always_equal = std::false_type;
 
    using propagate_on_container_copy_assignment = std::false_type;
    using propagate_on_container_move_assignment = std::false_type;
    using propagate_on_container_swap = std::false_type;
 
                            STLLocalAllocator() = default;
                            STLLocalAllocator(const STLLocalAllocator &) = delete; // Can't copy an allocator as the buffer is local to the original
                            STLLocalAllocator(STLLocalAllocator &&) = delete; // Can't move an allocator as the buffer is local to the original
    STLLocalAllocator &     operator = (const STLLocalAllocator &) = delete; // Can't copy an allocator as the buffer is local to the original
 
    template <class T2>
    explicit                STLLocalAllocator(const STLLocalAllocator<T2, N> &) : mNumElementsUsed(N) { }
 
    inline bool             is_local(const_pointer inPointer) const
    {
        ptrdiff_t diff = inPointer - reinterpret_cast<const_pointer>(mElements);
        return diff >= 0 && diff < ptrdiff_t(N);
    }
 
    inline pointer          allocate(size_type inN)
    {
        // If we allocate more than we have, fall back to the heap
        if (mNumElementsUsed + inN > N)
            return Base::allocate(inN);
 
        // Allocate from our local buffer
        pointer result = reinterpret_cast<pointer>(mElements) + mNumElementsUsed;
        mNumElementsUsed += inN;
        return result;
    }
 
    static constexpr bool   has_reallocate = true;
 
    inline pointer          reallocate(pointer inOldPointer, size_type inOldSize, size_type inNewSize)
    {
        JPH_ASSERT(inNewSize > 0); // Reallocating to zero size is implementation dependent, so we don't allow it
 
        // If there was no previous allocation, we can go through the regular allocate function
        if (inOldPointer == nullptr)
            return allocate(inNewSize);
 
        // If the pointer is outside our local buffer, fall back to the heap
        if (!is_local(inOldPointer))
        {
            if constexpr (AllocatorHasReallocate<Base>::sValue)
                return Base::reallocate(inOldPointer, inOldSize, inNewSize);
            else
                return ReallocateImpl(inOldPointer, inOldSize, inNewSize);
        }
 
        // If we happen to have space left, we only need to update our bookkeeping
        pointer base_ptr = reinterpret_cast<pointer>(mElements) + mNumElementsUsed - inOldSize;
        if (inOldPointer == base_ptr
            && mNumElementsUsed - inOldSize + inNewSize <= N)
        {
            mNumElementsUsed += inNewSize - inOldSize;
            return base_ptr;
        }
 
        // We can't reallocate in place, fall back to the heap
        return ReallocateImpl(inOldPointer, inOldSize, inNewSize);
    }
 
    inline void             deallocate(pointer inPointer, size_type inN)
    {
        // If the pointer is not in our local buffer, fall back to the heap
        if (!is_local(inPointer))
            return Base::deallocate(inPointer, inN);
 
        // Else we can only reclaim memory if it was the last allocation
        if (inPointer == reinterpret_cast<pointer>(mElements) + mNumElementsUsed - inN)
            mNumElementsUsed -= inN;
    }
 
    inline bool             operator == (const STLLocalAllocator<T, N> &inRHS) const
    {
        return this == &inRHS;
    }
 
    inline bool             operator != (const STLLocalAllocator<T, N> &inRHS) const
    {
        return this != &inRHS;
    }
 
    template <typename T2>
    struct rebind
    {
        using other = STLLocalAllocator<T2, N>;
    };
 
private:
    inline pointer          ReallocateImpl(pointer inOldPointer, size_type inOldSize, size_type inNewSize)
    {
        pointer new_pointer = Base::allocate(inNewSize);
        size_type n = min(inOldSize, inNewSize);
        if constexpr (std::is_trivially_copyable<T>())
        {
            // Can use mem copy
            memcpy(new_pointer, inOldPointer, n * sizeof(T));
        }
        else
        {
            // Need to actually move the elements
            for (size_t i = 0; i < n; ++i)
            {
                new (new_pointer + i) T(std::move(inOldPointer[i]));
                inOldPointer[i].~T();
            }
        }
        deallocate(inOldPointer, inOldSize);
        return new_pointer;
    }
 
    alignas(T) uint8        mElements[N * sizeof(T)];
    size_type               mNumElementsUsed = 0;
};
 
template <class T, size_t N> struct AllocatorHasReallocate<STLLocalAllocator<T, N>> { static constexpr bool sValue = STLLocalAllocator<T, N>::has_reallocate; };
 
#else
 
template <typename T, size_t N> using STLLocalAllocator = std::allocator<T>;
 
#endif // !JPH_DISABLE_CUSTOM_ALLOCATOR
 
JPH_NAMESPACE_END