Added STLLocalAllocator (#1435)

This allocator holds a fixed size buffer that it allocates from. It falls back to the heap when it runs out.

Jolt/Core/STLLocalAllocator.h

@@ -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

Jolt/Jolt.cmake

@@ -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

UnitTests/Core/STLLocalAllocatorTest.cpp

@@ -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>();
+	}
+}

UnitTests/UnitTests.cmake

@@ -12,6 +12,7 @@ set(UNIT_TESTS_SRC_FILES
 	${UNIT_TESTS_ROOT}/Core/LinearCurveTest.cpp
 	${UNIT_TESTS_ROOT}/Core/PreciseMathTest.cpp
 	${UNIT_TESTS_ROOT}/Core/ScopeExitTest.cpp
+	${UNIT_TESTS_ROOT}/Core/STLLocalAllocatorTest.cpp
 	${UNIT_TESTS_ROOT}/Core/StringToolsTest.cpp
 	${UNIT_TESTS_ROOT}/Core/QuickSortTest.cpp
 	${UNIT_TESTS_ROOT}/Core/UnorderedSetTest.cpp