From 40c1b55cc27a35e3e2eea81c50ab787a9f6b35d0 Mon Sep 17 00:00:00 2001
From: Andy Ross <andrew.j.ross@intel.com>
Date: Wed, 21 Oct 2020 13:24:32 -0700
Subject: [PATCH] lib/os/heap: Add sys_heap_realloc()

Add an optimized realloc() implementation that can successfully expand
allocations in place if there exists enough free memory after the
supplied block.

Signed-off-by: Andy Ross <andrew.j.ross@intel.com>
---
 include/sys/sys_heap.h | 24 +++++++++++++++++++++
 lib/os/heap.c          | 49 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 73 insertions(+)

diff --git a/include/sys/sys_heap.h b/include/sys/sys_heap.h
index fba61817cd6..66bcb600210 100644
--- a/include/sys/sys_heap.h
+++ b/include/sys/sys_heap.h
@@ -118,6 +118,30 @@ void *sys_heap_aligned_alloc(struct sys_heap *h, size_t align, size_t bytes);
  */
 void sys_heap_free(struct sys_heap *h, void *mem);
 
+/** @brief Expand the size of an existing allocation
+ *
+ * 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.
+ *
+ * @note The return of a NULL on failure is a different behavior than
+ * POSIX realloc(), which specifies that the original pointer will be
+ * returned (i.e. it is not possible to safely detect realloc()
+ * failure in POSIX, but it is here).
+ *
+ * @param heap Heap from which to allocate
+ * @param ptr Original pointer returned from a previous allocation
+ * @param bytes Number of bytes requested for the new block
+ * @return Pointer to memory the caller can now use, or NULL
+ */
+void *sys_heap_realloc(struct sys_heap *heap, void *ptr, size_t bytes);
+
 /** @brief Validate heap integrity
  *
  * Validates the internal integrity of a sys_heap.  Intended for unit
diff --git a/lib/os/heap.c b/lib/os/heap.c
index 5bc509a2e20..56129cc13a6 100644
--- a/lib/os/heap.c
+++ b/lib/os/heap.c
@@ -5,6 +5,7 @@
  */
 #include <sys/sys_heap.h>
 #include <kernel.h>
+#include <string.h>
 #include "heap.h"
 
 static void *chunk_mem(struct z_heap *h, chunkid_t c)
@@ -294,6 +295,54 @@ void *sys_heap_aligned_alloc(struct sys_heap *heap, size_t align, size_t bytes)
 	return mem;
 }
 
+void *sys_heap_realloc(struct sys_heap *heap, void *ptr, size_t bytes)
+{
+	struct z_heap *h = heap->heap;
+	chunkid_t c = mem_to_chunkid(h, ptr);
+	chunkid_t rc = right_chunk(h, c);
+	size_t chunks_need = bytes_to_chunksz(h, bytes);
+
+	if (chunk_size(h, c) > chunks_need) {
+		/* Shrink in place, split off and free unused suffix */
+		split_chunks(h, c, c + chunks_need);
+		set_chunk_used(h, c, true);
+		free_chunk(h, c + chunks_need);
+		return ptr;
+	} else if (!chunk_used(h, rc) &&
+		   (chunk_size(h, c) + chunk_size(h, rc) >= chunks_need)) {
+		/* Expand: split the right chunk and append */
+		chunkid_t split_size = chunks_need - chunk_size(h, c);
+
+		free_list_remove(h, rc);
+		if (split_size < chunk_size(h, rc)) {
+			split_chunks(h, rc, rc + split_size);
+			free_list_add(h, rc + split_size);
+		}
+
+		chunkid_t newsz = chunk_size(h, c) + split_size;
+
+		set_chunk_size(h, c, newsz);
+		set_chunk_used(h, c, true);
+		set_left_chunk_size(h, c + newsz, newsz);
+
+		CHECK(chunk_used(h, c));
+
+		return chunk_mem(h, c);
+	} else {
+		/* Reallocate and copy */
+		void *ptr2 = sys_heap_alloc(heap, bytes);
+
+		if (ptr2 == NULL) {
+			return NULL;
+		}
+
+		memcpy(ptr2, ptr,
+		       chunk_size(h, c) * CHUNK_UNIT - chunk_header_bytes(h));
+		sys_heap_free(heap, ptr);
+		return ptr2;
+	}
+}
+
 void sys_heap_init(struct sys_heap *heap, void *mem, size_t bytes)
 {
 	/* Must fit in a 32 bit count of HUNK_UNIT */