Improved allocator reallocate logic (#1094)

Implemented helper struct AllocatorHasReallocate that indicates if an allocater has the reallocate() function and using it to automatically fall back to allocate-copy-deallocate when reallocate is not available or the allocation needs an alignment that the allocator doesn't provide by default.

Jolt/Core/Array.h

@@ -72,6 +72,31 @@ private:
 		}
 	}
 
+	/// Reallocate the data block to inNewCapacity
+	inline void				reallocate(size_type inNewCapacity)
+	{
+		JPH_ASSERT(inNewCapacity > 0 && inNewCapacity >= mSize);
+
+		pointer pointer;
+		if constexpr (AllocatorHasReallocate<Allocator>::sValue)
+		{
+			// Reallocate data block
+			pointer = get_allocator().reallocate(mElements, mCapacity, inNewCapacity);
+		}
+		else
+		{
+			// Copy data to a new location
+			pointer = get_allocator().allocate(inNewCapacity);
+			if (mElements != nullptr)
+			{
+				move(pointer, mElements, mSize);
+				get_allocator().deallocate(mElements, mCapacity);
+			}
+		}
+		mElements = pointer;
+		mCapacity = inNewCapacity;
+	}
+
 	/// Destruct elements [inStart, inEnd - 1]
 	inline void				destruct(size_type inStart, size_type inEnd)
 	{
@@ -86,24 +111,7 @@ public:
 	inline void				reserve(size_type inNewSize)
 	{
 		if (mCapacity < inNewSize)
-		{
-			pointer pointer;
-			if constexpr (std::is_trivially_copyable<T>() && !std::is_same<Allocator, std::allocator<T>>())
-			{
-				pointer = get_allocator().reallocate(mElements, mCapacity, inNewSize);
-			}
-			else
-			{
-				pointer = get_allocator().allocate(inNewSize);
-				if (mElements != nullptr)
-				{
-					move(pointer, mElements, mSize);
-					get_allocator().deallocate(mElements, mCapacity);
-				}
-			}
-			mElements = pointer;
-			mCapacity = inNewSize;
-		}
+			reallocate(inNewSize);
 	}
 
 	/// Resize array to new length
@@ -150,15 +158,21 @@ private:
 		}
 	}
 
+	/// Free memory
+	inline void				free()
+	{
+		get_allocator().deallocate(mElements, mCapacity);
+		mElements = nullptr;
+		mCapacity = 0;
+	}
+
 	/// Destroy all elements and free memory
 	inline void				destroy()
 	{
 		if (mElements != nullptr)
 		{
 			clear();
-			get_allocator().deallocate(mElements, mCapacity);
-			mElements = nullptr;
-			mCapacity = 0;
+			free();
 		}
 	}
 
@@ -315,24 +329,12 @@ public:
 	/// Reduce the capacity of the array to match its size
 	void					shrink_to_fit()
 	{
-		if (mSize == 0)
-		{
-			// Free memory block
-			if (mElements != nullptr)
-			{
-				get_allocator().deallocate(mElements, mCapacity);
-				mElements = nullptr;
-				mCapacity = 0;
-			}
-		}
-		else if (mCapacity > mSize)
+		if (mElements != nullptr)
 		{
-			// Reallocate memory block
-			pointer pointer = get_allocator().allocate(mSize);
-			move(pointer, mElements, mSize);
-			get_allocator().deallocate(mElements, mCapacity);
-			mElements = pointer;
-			mCapacity = mSize;
+			if (mSize == 0)
+				free();
+			else if (mCapacity > mSize)
+				reallocate(mSize);
 		}
 	}
 

Jolt/Core/STLAlignedAllocator.h

@@ -44,18 +44,6 @@ public:
 		return (pointer)AlignedAllocate(inN * sizeof(value_type), N);
 	}
 
-	/// Reallocate memory
-	inline pointer			reallocate(pointer inOldPointer, size_type inOldSize, size_type inNewSize)
-	{
-		pointer new_pointer = allocate(inNewSize);
-		if (inOldPointer != nullptr)
-		{
-			memcpy(new_pointer, inOldPointer, inOldSize * sizeof(value_type));
-			deallocate(inOldPointer, inOldSize);
-		}
-		return new_pointer;
-	}
-
 	/// Free memory
 	inline void				deallocate(pointer inPointer, size_type)
 	{

Jolt/Core/STLAllocator.h

@@ -6,6 +6,9 @@
 
 JPH_NAMESPACE_BEGIN
 
+/// Default implementation of AllocatorHasReallocate which tells if an allocator has a reallocate function
+template <class T> struct AllocatorHasReallocate { static constexpr bool sValue = false; };
+
 #ifndef JPH_DISABLE_CUSTOM_ALLOCATOR
 
 /// STL allocator that forwards to our allocation functions
@@ -40,37 +43,33 @@ public:
 	template <typename T2>
 	inline					STLAllocator(const STLAllocator<T2> &) { }
 
+	/// If this allocator needs to fall back to aligned allocations because the type requires it
+	static constexpr bool	needs_aligned_allocate = alignof(T) > (JPH_CPU_ADDRESS_BITS == 32? 8 : 16);
+
 	/// Allocate memory
 	inline pointer			allocate(size_type inN)
 	{
-		if constexpr (alignof(T) > (JPH_CPU_ADDRESS_BITS == 32? 8 : 16))
+		if constexpr (needs_aligned_allocate)
 			return pointer(AlignedAllocate(inN * sizeof(value_type), alignof(T)));
 		else
 			return pointer(Allocate(inN * sizeof(value_type)));
 	}
 
+	/// Should we expose a reallocate function?
+	static constexpr bool	has_reallocate = std::is_trivially_copyable<T>() && !needs_aligned_allocate;
+
 	/// Reallocate memory
-	inline pointer			reallocate(pointer inOldPointer, size_type inOldSize, size_type inNewSize)
+	template <bool has_reallocate_v = has_reallocate, typename = std::enable_if_t<has_reallocate_v>>
+	inline pointer			reallocate(pointer inOldPointer, [[maybe_unused]] size_type inOldSize, size_type inNewSize)
 	{
-		if constexpr (alignof(T) > (JPH_CPU_ADDRESS_BITS == 32? 8 : 16))
-		{
-			// Can't reallocate aligned blocks
-			pointer new_pointer = allocate(inNewSize);
-			if (inOldPointer != nullptr)
-			{
-				memcpy(new_pointer, inOldPointer, inOldSize * sizeof(value_type));
-				deallocate(inOldPointer, inOldSize);
-			}
-			return new_pointer;
-		}
-		else
-			return pointer(Reallocate(inOldPointer, inNewSize * sizeof(value_type)));
+		JPH_ASSERT(inNewSize > 0); // Reallocating to zero size is implementation dependent, so we don't allow it
+		return pointer(Reallocate(inOldPointer, inNewSize * sizeof(value_type)));
 	}
 
 	/// Free memory
 	inline void				deallocate(pointer inPointer, size_type)
 	{
-		if constexpr (alignof(T) > (JPH_CPU_ADDRESS_BITS == 32? 8 : 16))
+		if constexpr (needs_aligned_allocate)
 			AlignedFree(inPointer);
 		else
 			Free(inPointer);
@@ -95,6 +94,9 @@ public:
 	};
 };
 
+/// The STLAllocator implements the reallocate function if the alignment of the class is smaller or equal to the default alignment for the platform
+template <class T> struct AllocatorHasReallocate<STLAllocator<T>> { static constexpr bool sValue = STLAllocator<T>::has_reallocate; };
+
 #else
 
 template <typename T> using STLAllocator = std::allocator<T>;

Jolt/Core/STLTempAllocator.h

@@ -43,18 +43,6 @@ public:
 		return pointer(mAllocator.Allocate(uint(inN * sizeof(value_type))));
 	}
 
-	/// Reallocate memory
-	inline pointer			reallocate(pointer inOldPointer, size_type inOldSize, size_type inNewSize)
-	{
-		pointer new_pointer = allocate(inNewSize);
-		if (inOldPointer != nullptr)
-		{
-			memcpy(new_pointer, inOldPointer, inOldSize * sizeof(value_type));
-			deallocate(inOldPointer, inOldSize);
-		}
-		return new_pointer;
-	}
-
 	/// Free memory
 	inline void				deallocate(pointer inPointer, size_type inN)
 	{