Grow or fail flag for realloc (#132)

rpmalloc/rpmalloc.c

@@ -1442,11 +1442,12 @@ _memory_aligned_allocate(heap_t* heap, size_t alignment, size_t size) {
 	size_t align_mask = alignment - 1;
 	if (alignment <= _memory_page_size) {
 		ptr = _memory_allocate(heap, size + alignment);
-		if ((uintptr_t)ptr & align_mask)
+		if ((uintptr_t)ptr & align_mask) {
 			ptr = (void*)(((uintptr_t)ptr & ~(uintptr_t)align_mask) + alignment);
-		//Mark as having aligned blocks
-		span_t* span = (span_t*)((uintptr_t)ptr & _memory_span_mask);
-		span->flags |= SPAN_FLAG_ALIGNED_BLOCKS;
+			//Mark as having aligned blocks
+			span_t* span = (span_t*)((uintptr_t)ptr & _memory_span_mask);
+			span->flags |= SPAN_FLAG_ALIGNED_BLOCKS;
+		}
 		return ptr;
 	}
 
@@ -1848,6 +1849,9 @@ _memory_reallocate(heap_t* heap, void* p, size_t size, size_t oldsize, unsigned
 		oldsize = 0;
 	}
 
+	if (!!(flags & RPMALLOC_GROW_OR_FAIL))
+		return 0;
+
 	//Size is greater than block size, need to allocate a new block and deallocate the old
 	//Avoid hysteresis by overallocating if increase is small (below 37%)
 	size_t lower_bound = oldsize + (oldsize >> 2) + (oldsize >> 3);
@@ -1863,27 +1867,26 @@ _memory_reallocate(heap_t* heap, void* p, size_t size, size_t oldsize, unsigned
 }
 
 static void*
-_memory_aligned_reallocate(heap_t* heap, void* p, size_t alignment, size_t size, size_t oldsize,
+_memory_aligned_reallocate(heap_t* heap, void* ptr, size_t alignment, size_t size, size_t oldsize,
                            unsigned int flags) {
 	if (alignment <= SMALL_GRANULARITY)
-		return _memory_reallocate(heap, p, size, oldsize, flags);
-
-	size_t usablesize = _memory_usable_size(p);
-	if ((usablesize >= size) && (size >= (usablesize / 2)) && !((uintptr_t)p & (alignment - 1)))
-		return p;
-
-	void* block = _memory_aligned_allocate(heap, alignment, size);
-	if (p) {
-		if (!oldsize)
-			oldsize = usablesize;
-		if (!(flags & RPMALLOC_NO_PRESERVE))
-			memcpy(block, p, oldsize < size ? oldsize : size);
-		_memory_deallocate(p);
-	}
-	if (block) {
-		//Mark as having aligned blocks
-		span_t* span = (span_t*)((uintptr_t)block & _memory_span_mask);
-		span->flags |= SPAN_FLAG_ALIGNED_BLOCKS;
+		return _memory_reallocate(heap, ptr, size, oldsize, flags);
+
+	int no_alloc = !!(flags & RPMALLOC_GROW_OR_FAIL);
+	size_t usablesize = _memory_usable_size(ptr);
+	if ((usablesize >= size) && !((uintptr_t)ptr & (alignment - 1))) {
+		if (no_alloc || (size >= (usablesize / 2)))
+			return ptr;
+	}
+	// Aligned alloc marks span as having aligned blocks
+	void* block = (!no_alloc ? _memory_aligned_allocate(heap, alignment, size) : 0);
+	if (EXPECTED(block)) {
+		if (!(flags & RPMALLOC_NO_PRESERVE) && ptr) {
+			if (!oldsize)
+				oldsize = usablesize;
+			memcpy(block, ptr, oldsize < size ? oldsize : size);
+		}
+		rpfree(ptr);
 	}
 	return block;
 }

rpmalloc/rpmalloc.h

@@ -59,6 +59,10 @@ extern "C" {
 
 //! Flag to rpaligned_realloc to not preserve content in reallocation
 #define RPMALLOC_NO_PRESERVE    1
+//! Flag to rpaligned_realloc to fail and return null pointer if grow cannot be done in-place,
+//  in which case the original pointer is still valid (just like a call to realloc which failes to allocate
+//  a new block).
+#define RPMALLOC_GROW_OR_FAIL   2
 
 typedef struct rpmalloc_global_statistics_t {
 	//! Current amount of virtual memory mapped, all of which might not have been committed (only if ENABLE_STATISTICS=1)

test/main.c

@@ -111,6 +111,8 @@ test_alloc(void) {
 			return test_fail("Bad usable size (aligned realloc)");
 		if (*((uintptr_t*)testptr) != 0x12345678)
 			return test_fail("Data not preserved on realloc");
+		if (rpaligned_realloc(testptr, 128, size * 1024 * 4, 0, RPMALLOC_GROW_OR_FAIL))
+			return test_fail("Realloc with grow-or-fail did not fail as expected");
 		void* unaligned = rprealloc(testptr, size);
 		if (unaligned != testptr) {
 			ptrdiff_t diff = pointer_diff(testptr, unaligned);
@@ -318,6 +320,10 @@ test_realloc(void) {
 	while (bigsize < 3 * 1024 * 1024) {
 		++bigsize;
 		bigptr = rprealloc(bigptr, bigsize);
+		if (rpaligned_realloc(bigptr, 0, bigsize * 32, 0, RPMALLOC_GROW_OR_FAIL))
+			return test_fail("Reallocation with grow-or-fail did not fail as expected");
+		if (rpaligned_realloc(bigptr, 128, bigsize * 32, 0, RPMALLOC_GROW_OR_FAIL))
+			return test_fail("Reallocation with aligned grow-or-fail did not fail as expected");
 	}
 	rpfree(bigptr);