Skip to content

Commit

Permalink
Added STLLocalAllocator (#1435)
Browse files Browse the repository at this point in the history
This allocator holds a fixed size buffer that it allocates from. It falls back to the heap when it runs out.
  • Loading branch information
jrouwe authored Jan 4, 2025
1 parent cf1f2d5 commit 06793ab
Show file tree
Hide file tree
Showing 4 changed files with 386 additions and 0 deletions.
169 changes: 169 additions & 0 deletions Jolt/Core/STLLocalAllocator.h
Original file line number Diff line number Diff line change
@@ -0,0 +1,169 @@
// 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

/// STL allocator that keeps N elements in a local buffer before falling back to regular allocations
template <typename T, size_t N>
class STLLocalAllocator : private STLAllocator<T>
{
using Base = STLAllocator<T>;

public:
/// General properties
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;

/// The allocator is not stateless (has local buffer)
using is_always_equal = std::false_type;

/// We cannot copy, move or swap allocators
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;

/// Constructor
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

/// Constructor used when rebinding to another type. This expects the allocator to use the original memory pool from the first allocator,
/// but in our case we cannot use the local buffer of the original allocator as it has different size and alignment rules.
/// To solve this we make this allocator fall back to the heap immediately.
template <class T2> STLLocalAllocator(const STLLocalAllocator<T2, N> &) : mNumElementsUsed(N) { }

/// Check if inPointer is in the local buffer
inline bool is_local(const_pointer inPointer) const
{
ptrdiff_t diff = inPointer - reinterpret_cast<const_pointer>(mElements);
return diff >= 0 && diff < ptrdiff_t(N);
}

/// Allocate memory
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;
}

/// Always implements a reallocate function as we can often reallocate in place
static constexpr bool has_reallocate = true;

/// Reallocate memory
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);
}

/// Free memory
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;
}

/// Allocators are not-stateless, assume if allocator address matches that the allocators are the same
inline bool operator == (const STLLocalAllocator<T, N> &inRHS) const
{
return this == &inRHS;
}

inline bool operator != (const STLLocalAllocator<T, N> &inRHS) const
{
return this != &inRHS;
}

/// Converting to allocator for other type
template <typename T2>
struct rebind
{
using other = STLLocalAllocator<T2, N>;
};

private:
/// Implements reallocate when the base class doesn't or when we go from local buffer to heap
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;
};

/// The STLLocalAllocator always implements a reallocate function as it can often reallocate in place
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
1 change: 1 addition & 0 deletions Jolt/Jolt.cmake
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,7 @@ set(JOLT_PHYSICS_SRC_FILES
${JOLT_PHYSICS_ROOT}/Core/StaticArray.h
${JOLT_PHYSICS_ROOT}/Core/STLAlignedAllocator.h
${JOLT_PHYSICS_ROOT}/Core/STLAllocator.h
${JOLT_PHYSICS_ROOT}/Core/STLLocalAllocator.h
${JOLT_PHYSICS_ROOT}/Core/STLTempAllocator.h
${JOLT_PHYSICS_ROOT}/Core/StreamIn.h
${JOLT_PHYSICS_ROOT}/Core/StreamOut.h
Expand Down
215 changes: 215 additions & 0 deletions UnitTests/Core/STLLocalAllocatorTest.cpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,215 @@
// Jolt Physics Library (https://github.com/jrouwe/JoltPhysics)
// SPDX-FileCopyrightText: 2025 Jorrit Rouwe
// SPDX-License-Identifier: MIT

#include "UnitTestFramework.h"

#include <Jolt/Core/STLLocalAllocator.h>

TEST_SUITE("STLLocalAllocatorTest")
{
/// The number of elements in the local buffer
static constexpr size_t N = 20;

#ifndef JPH_DISABLE_CUSTOM_ALLOCATOR
template <class ArrayType>
static bool sIsLocal(ArrayType &inArray)
{
#ifdef JPH_USE_STD_VECTOR
// Check that the data pointer is within the array.
// Note that when using std::vector we cannot use get_allocator as that makes a copy of the allocator internally
// and we've disabled the copy constructor since our allocator cannot be copied.
const uint8 *data = reinterpret_cast<const uint8 *>(inArray.data());
const uint8 *array = reinterpret_cast<const uint8 *>(&inArray);
return data >= array && data < array + sizeof(inArray);
#else
return inArray.get_allocator().is_local(inArray.data());
#endif
}
#endif

template <class ArrayType, bool NonTrivial>
static void sTestArray()
{
// Allocate so that we will run out of local memory and reallocate from heap at least once
ArrayType arr;
for (int i = 0; i < 64; ++i)
arr.push_back(i);
CHECK(arr.size() == 64);
for (int i = 0; i < 64; ++i)
{
CHECK(arr[i] == i);
#if !defined(JPH_USE_STD_VECTOR) && !defined(JPH_DISABLE_CUSTOM_ALLOCATOR)
// We only have to move elements once we run out of the local buffer, this happens as we resize
// from 16 to 32 elements, we'll reallocate again at 32 and 64
if constexpr (NonTrivial)
CHECK(arr[i].GetNonTriv() == (i < 16? 3 : (i < 32? 2 : 1)));
#endif
}
CHECK(IsAligned(arr.data(), alignof(typename ArrayType::value_type)));
#ifndef JPH_DISABLE_CUSTOM_ALLOCATOR
CHECK(!sIsLocal(arr));
#endif

// Check that we can copy the array to another array
ArrayType arr2;
arr2 = arr;
for (int i = 0; i < 64; ++i)
{
CHECK(arr2[i] == i);
if constexpr (NonTrivial)
CHECK(arr2[i].GetNonTriv() == -999);
}
CHECK(IsAligned(arr2.data(), alignof(typename ArrayType::value_type)));
#ifndef JPH_DISABLE_CUSTOM_ALLOCATOR
CHECK(!sIsLocal(arr2));
#endif

// Clear the array
arr.clear();
arr.shrink_to_fit();
CHECK(arr.size() == 0);
#ifndef JPH_USE_STD_VECTOR // Some implementations of std::vector ignore shrink_to_fit
CHECK(arr.capacity() == 0);
CHECK(arr.data() == nullptr);
#endif

// Allocate so we stay within the local buffer
for (int i = 0; i < 10; ++i)
arr.push_back(i);
CHECK(arr.size() == 10);
for (int i = 0; i < 10; ++i)
{
CHECK(arr[i] == i);
#if !defined(JPH_USE_STD_VECTOR) && !defined(JPH_DISABLE_CUSTOM_ALLOCATOR)
// We never need to move elements as they stay within the local buffer
if constexpr (NonTrivial)
CHECK(arr[i].GetNonTriv() == 1);
#endif
}
CHECK(IsAligned(arr.data(), alignof(typename ArrayType::value_type)));
#if !defined(JPH_USE_STD_VECTOR) && !defined(JPH_DISABLE_CUSTOM_ALLOCATOR) // Doesn't work with std::vector since it doesn't use the reallocate function and runs out of space
CHECK(sIsLocal(arr));
#endif

// Check that we can copy the array to the local buffer
ArrayType arr3;
arr3 = arr;
CHECK(arr3.size() == 10);
for (int i = 0; i < 10; ++i)
{
CHECK(arr3[i] == i);
if constexpr (NonTrivial)
CHECK(arr3[i].GetNonTriv() == -999);
}
CHECK(IsAligned(arr3.data(), alignof(typename ArrayType::value_type)));
#ifndef JPH_DISABLE_CUSTOM_ALLOCATOR
CHECK(sIsLocal(arr3));
#endif

// Check that if we reserve the memory, that we can fully fill the array in local memory
ArrayType arr4;
arr4.reserve(N);
for (int i = 0; i < int(N); ++i)
arr4.push_back(i);
CHECK(arr4.size() == N);
CHECK(arr4.capacity() == N);
for (int i = 0; i < int(N); ++i)
{
CHECK(arr4[i] == i);
if constexpr (NonTrivial)
CHECK(arr4[i].GetNonTriv() == 1);
}
CHECK(IsAligned(arr4.data(), alignof(typename ArrayType::value_type)));
#ifndef JPH_DISABLE_CUSTOM_ALLOCATOR
CHECK(sIsLocal(arr4));
#endif
}

TEST_CASE("TestAllocation")
{
using Allocator = STLLocalAllocator<int, N>;
using ArrayType = Array<int, Allocator>;
#ifndef JPH_DISABLE_CUSTOM_ALLOCATOR
static_assert(AllocatorHasReallocate<Allocator>::sValue);
#endif

sTestArray<ArrayType, false>();
}

TEST_CASE("TestAllocationAligned")
{
// Force the need for an aligned allocation
struct alignas(64) Aligned
{
Aligned(int inValue) : mValue(inValue) { }
operator int() const { return mValue; }

private:
int mValue;
};
static_assert(std::is_trivially_copyable<Aligned>());

using Allocator = STLLocalAllocator<Aligned, N>;
using ArrayType = Array<Aligned, Allocator>;
#ifndef JPH_DISABLE_CUSTOM_ALLOCATOR
static_assert(AllocatorHasReallocate<Allocator>::sValue);
#endif

sTestArray<ArrayType, false>();
}

TEST_CASE("TestAllocationNonTrivial")
{
// Force non trivial copy constructor
struct NonTriv
{
NonTriv(int inValue) : mValue(inValue) { }
NonTriv(const NonTriv &inRHS) : mValue(inRHS.mValue), mMakeNonTriv(-999) { }
NonTriv(NonTriv &&inRHS) : mValue(inRHS.mValue), mMakeNonTriv(inRHS.mMakeNonTriv + 1) { }
NonTriv & operator = (const NonTriv &inRHS) { mValue = inRHS.mValue; mMakeNonTriv = -9999; return *this; }
operator int() const { return mValue; }
int GetNonTriv() const { return mMakeNonTriv; }

private:
int mValue;
int mMakeNonTriv = 0;
};
static_assert(!std::is_trivially_copyable<NonTriv>());

using Allocator = STLLocalAllocator<NonTriv, N>;
using ArrayType = Array<NonTriv, Allocator>;
#ifndef JPH_DISABLE_CUSTOM_ALLOCATOR
static_assert(AllocatorHasReallocate<Allocator>::sValue);
#endif

sTestArray<ArrayType, true>();
}

TEST_CASE("TestAllocationAlignedNonTrivial")
{
// Force non trivial copy constructor
struct alignas(64) AlNonTriv
{
AlNonTriv(int inValue) : mValue(inValue) { }
AlNonTriv(const AlNonTriv &inRHS) : mValue(inRHS.mValue), mMakeNonTriv(-999) { }
AlNonTriv(AlNonTriv &&inRHS) : mValue(inRHS.mValue), mMakeNonTriv(inRHS.mMakeNonTriv + 1) { }
AlNonTriv & operator = (const AlNonTriv &inRHS) { mValue = inRHS.mValue; mMakeNonTriv = -9999; return *this; }
operator int() const { return mValue; }
int GetNonTriv() const { return mMakeNonTriv; }

private:
int mValue;
int mMakeNonTriv = 0;
};
static_assert(!std::is_trivially_copyable<AlNonTriv>());

using Allocator = STLLocalAllocator<AlNonTriv, N>;
using ArrayType = Array<AlNonTriv, Allocator>;
#ifndef JPH_DISABLE_CUSTOM_ALLOCATOR
static_assert(AllocatorHasReallocate<Allocator>::sValue);
#endif

sTestArray<ArrayType, true>();
}
}
Loading

0 comments on commit 06793ab

Please sign in to comment.