From 135ffaff74458c6f0e9a814e6698e2c3e1ee30f9 Mon Sep 17 00:00:00 2001
From: Christopher Friedt <chrisfriedt@gmail.com>
Date: Thu, 26 Nov 2020 08:19:10 -0500
Subject: [PATCH] kernel/k_malloc: add k_aligned_alloc

This change adds z_heap_aligned_alloc() and k_aligned_alloc()
and changes z_heap_malloc() and k_malloc() to be small wrappers around
the aligned variants.

Fixes #29519

Signed-off-by: Christopher Friedt <chrisfriedt@gmail.com>
---
 include/kernel.h | 27 ++++++++++++++++++++++--
 kernel/mempool.c | 53 ++++++++++++++++++++++++++++++++++--------------
 2 files changed, 63 insertions(+), 17 deletions(-)

diff --git a/include/kernel.h b/include/kernel.h
index 378f1d8fea0..62d0822ec0e 100644
--- a/include/kernel.h
+++ b/include/kernel.h
@@ -4540,7 +4540,27 @@ extern void z_mem_pool_free_id(struct k_mem_block_id *id);
  */
 
 /**
- * @brief Allocate memory from heap.
+ * @brief Allocate memory from the heap with a specified alignment.
+ *
+ * This routine provides semantics similar to aligned_alloc(); memory is
+ * allocated from the heap with a specified alignment. However, one minor
+ * difference is that k_aligned_alloc() accepts any non-zero @p size,
+ * wherase aligned_alloc() only accepts a @p size that is an integral
+ * multiple of @p align.
+ *
+ * Above, aligned_alloc() refers to:
+ * C11 standard (ISO/IEC 9899:2011): 7.22.3.1
+ * The aligned_alloc function (p: 347-348)
+ *
+ * @param align Alignment of memory requested (in bytes).
+ * @param size Amount of memory requested (in bytes).
+ *
+ * @return Address of the allocated memory if successful; otherwise NULL.
+ */
+extern void *k_aligned_alloc(size_t align, size_t size);
+
+/**
+ * @brief Allocate memory from the heap.
  *
  * This routine provides traditional malloc() semantics. Memory is
  * allocated from the heap memory pool.
@@ -4549,7 +4569,10 @@ extern void z_mem_pool_free_id(struct k_mem_block_id *id);
  *
  * @return Address of the allocated memory if successful; otherwise NULL.
  */
-extern void *k_malloc(size_t size);
+static inline void *k_malloc(size_t size)
+{
+	return k_aligned_alloc(sizeof(void *), size);
+}
 
 /**
  * @brief Free memory allocated from heap.
diff --git a/kernel/mempool.c b/kernel/mempool.c
index 301e8202dd5..f67fde3585f 100644
--- a/kernel/mempool.c
+++ b/kernel/mempool.c
@@ -7,37 +7,53 @@
 #include <kernel.h>
 #include <string.h>
 #include <sys/math_extras.h>
+#include <sys/util.h>
 
-void *z_heap_malloc(struct k_heap *heap, size_t size)
+static void *z_heap_aligned_alloc(struct k_heap *heap, size_t align, size_t size)
 {
+	uint8_t *mem;
+	struct k_heap **heap_ref;
+	size_t excess = MAX(sizeof(struct k_heap *), align);
+
 	/*
-	 * get a block large enough to hold an initial (hidden) heap
+	 * get a block large enough to hold an initial (aligned and hidden) heap
 	 * pointer, as well as the space the caller requested
 	 */
-	if (size_add_overflow(size, sizeof(struct k_heap *),
-			      &size)) {
+	if (size_add_overflow(size, excess, &size)) {
 		return NULL;
 	}
 
-	struct k_heap **blk = k_heap_alloc(heap, size, K_NO_WAIT);
-
-	if (blk == NULL) {
+	mem = k_heap_aligned_alloc(heap, align, size, K_NO_WAIT);
+	if (mem == NULL) {
 		return NULL;
 	}
 
-	blk[0] = heap;
+	/* create (void *) values in the excess equal to (void *) -1 */
+	memset(mem, 0xff, excess);
+	heap_ref = (struct k_heap **)mem;
+	*heap_ref = heap;
 
 	/* return address of the user area part of the block to the caller */
-	return (char *)&blk[1];
+	return mem + excess;
+}
+
+static void *z_heap_malloc(struct k_heap *heap, size_t size)
+{
+	return z_heap_aligned_alloc(heap, sizeof(void *), size);
 }
 
 void k_free(void *ptr)
 {
-	if (ptr != NULL) {
-		struct k_heap **blk = &((struct k_heap **)ptr)[-1];
-		struct k_heap *heap = *blk;
+	struct k_heap **heap_ref;
 
-		k_heap_free(heap, blk);
+	if (ptr != NULL) {
+		for (heap_ref = &((struct k_heap **)ptr)[-1];
+			*heap_ref == (struct k_heap *)-1; --heap_ref) {
+			/* no-op */
+		}
+
+		ptr = (uint8_t *)heap_ref;
+		k_heap_free(*heap_ref, ptr);
 	}
 }
 
@@ -46,9 +62,16 @@ void k_free(void *ptr)
 K_HEAP_DEFINE(_system_heap, CONFIG_HEAP_MEM_POOL_SIZE);
 #define _SYSTEM_HEAP (&_system_heap)
 
-void *k_malloc(size_t size)
+void *k_aligned_alloc(size_t align, size_t size)
 {
-	return z_heap_malloc(_SYSTEM_HEAP, size);
+	__ASSERT(align / sizeof(void *) >= 1
+		&& (align % sizeof(void *)) == 0,
+		"align must be a multiple of sizeof(void *)");
+
+	__ASSERT((align & (align - 1)) == 0,
+		"align must be a power of 2");
+
+	return z_heap_aligned_alloc(_SYSTEM_HEAP, align, size);
 }
 
 void *k_calloc(size_t nmemb, size_t size)