multi_heap: introduce support for realloc()

Add support for realloc (and realloc_aligned) into the multi heap lib,
where the buffer sent in will either be reused (maybe shrinked),
or enlarged by allocating on any of the matching heaps of the multi heap.

Signed-off-by: Meir Komet <mskomet1@gmail.com>

doc/kernel/memory_management/heap.rst

@@ -127,6 +127,13 @@ application-provided callback is responsible for doing the underlying
 allocation from one of the managed heaps, and may use the
 configuration parameter in any way it likes to make that decision.
 
+For modifying the size of an allocated buffer (whether shrinking
+or enlarging it), you can use the
+:c:func:`sys_multi_heap_realloc` and
+:c:func:`sys_multi_heap_aligned_realloc` APIs.  If the buffer cannot be
+enlarged on the heap where it currently resides,
+any of the eligible heaps specified by the configuration parameter may be used.
+
 When unused, a multi heap may be freed via
 :c:func:`sys_multi_heap_free`.  The application does not need to pass
 a configuration parameter.  Memory allocated from any of the managed

include/zephyr/sys/multi_heap.h

@@ -168,6 +168,32 @@ const struct sys_multi_heap_rec *sys_multi_heap_get_heap(const struct sys_multi_
  */
 void sys_multi_heap_free(struct sys_multi_heap *mheap, void *block);
 
+/** @brief Expand the size of an existing allocation on the multi heap
+ *
+ * Returns a pointer to a new memory region with the same contents,
+ * but a different allocated size.  If the new allocation can be
+ * expanded in place, the pointer returned will be identical.
+ * Otherwise the data will be copies to a new block and the old one
+ * will be freed as per sys_heap_free().  If the specified size is
+ * smaller than the original, the block will be truncated in place and
+ * the remaining memory returned to the heap.  If the allocation of a
+ * new block fails, then NULL will be returned and the old block will
+ * not be freed or modified. If a new allocation is needed, the choice
+ * for the heap used will be bases on the cfg parameter (same as in sys_multi_heap_aligned_alloc).
+ *
+ * @param mheap Multi heap pointer
+ * @param cfg Opaque configuration parameter, as for sys_multi_heap_fn_t
+ * @param ptr Original pointer returned from a previous allocation
+ * @param align Alignment in bytes, must be a power of two
+ * @param bytes Number of bytes requested for the new block
+ * @return Pointer to memory the caller can now use, or NULL
+ */
+void *sys_multi_heap_aligned_realloc(struct sys_multi_heap *mheap, void *cfg,
+				     void *ptr, size_t align, size_t bytes);
+
+#define sys_multi_heap_realloc(mheap, cfg, ptr, bytes) \
+	sys_multi_heap_aligned_realloc(mheap, cfg, ptr, 0, bytes)
+
 /**
  * @}
  */

lib/heap/multi_heap.c

@@ -5,6 +5,7 @@
 #include <zephyr/sys/util.h>
 #include <zephyr/sys/sys_heap.h>
 #include <zephyr/sys/multi_heap.h>
+#include <string.h>
 
 void sys_multi_heap_init(struct sys_multi_heap *heap, sys_multi_heap_fn_t choice_fn)
 {
@@ -90,3 +91,38 @@ void sys_multi_heap_free(struct sys_multi_heap *mheap, void *block)
 		sys_heap_free(heap->heap, block);
 	}
 }
+
+void *sys_multi_heap_aligned_realloc(struct sys_multi_heap *mheap, void *cfg,
+				     void *ptr, size_t align, size_t bytes)
+{
+	/* special realloc semantics */
+	if (ptr == NULL) {
+		return sys_multi_heap_aligned_alloc(mheap, cfg, align, bytes);
+	}
+	if (bytes == 0) {
+		sys_multi_heap_free(mheap, ptr);
+		return NULL;
+	}
+
+	const struct sys_multi_heap_rec *rec = sys_multi_heap_get_heap(mheap, ptr);
+
+	__ASSERT_NO_MSG(rec);
+
+	/* Invoke the realloc function on the same heap, to try to reuse in place */
+	void *new_ptr = sys_heap_aligned_realloc(rec->heap, ptr, align, bytes);
+
+	if (new_ptr != NULL) {
+		return new_ptr;
+	}
+
+	size_t old_size = sys_heap_usable_size(rec->heap, ptr);
+
+	/* Otherwise, allocate a new block and copy the data */
+	new_ptr = sys_multi_heap_aligned_alloc(mheap, cfg, align, bytes);
+	if (new_ptr != NULL) {
+		memcpy(new_ptr, ptr, MIN(old_size, bytes));
+		sys_multi_heap_free(mheap, ptr);
+	}
+
+	return new_ptr;
+}

tests/lib/multi_heap/src/test_mheap_api.c

@@ -335,6 +335,11 @@ ZTEST(mheap_api, test_multi_heap)
 		zassert_true(blocks[i] >= &heap_mem[i][0] &&
 			     blocks[i] < &heap_mem[i+1][0],
 			     "allocation not in correct heap");
+
+		void *ptr = sys_multi_heap_realloc(&multi_heap, (void *)(long)i,
+			blocks[i], MHEAP_BYTES / 2);
+
+		zassert_equal(ptr, blocks[i], "realloc moved pointer");
 	}
 
 	/* Make sure all heaps fail to allocate another */
@@ -355,5 +360,25 @@ ZTEST(mheap_api, test_multi_heap)
 		blocks[i] = sys_multi_heap_alloc(&multi_heap, (void *)(long)i,
 						 MHEAP_BYTES / 2);
 		zassert_not_null(blocks[i], "final re-allocation failed");
+
+		/* Allocating smaller buffer should stay within */
+		void *ptr = sys_multi_heap_realloc(&multi_heap, (void *)(long)i,
+						   blocks[i], MHEAP_BYTES / 4);
+		zassert_equal(ptr, blocks[i], "realloc should return same value");
+
+		ptr = sys_multi_heap_alloc(&multi_heap, (void *)(long)i,
+					   MHEAP_BYTES / 4);
+		zassert_between_inclusive((uintptr_t)ptr, (uintptr_t)blocks[i] + MHEAP_BYTES / 4,
+			(uintptr_t)blocks[i] + MHEAP_BYTES / 2 - 1,
+			"realloc failed to shrink prev buffer");
 	}
+
+	/* Test realloc special cases */
+	void *ptr = sys_multi_heap_realloc(&multi_heap, (void *)0L,
+		blocks[0], /* size = */ 0);
+	zassert_is_null(ptr);
+
+	ptr = sys_multi_heap_realloc(&multi_heap, (void *)0L,
+		/* ptr = */ NULL, MHEAP_BYTES / 4);
+	zassert_not_null(ptr);
 }