z_heap_aligned_alloc(): avoid memory wastage

The strategy used in z_heap_aligned_alloc() was to allocate an extra
align-sized memory block for storing a pointer to the memory heap.
This is wasteful in terms of memory usage when alignment is larger
than a pointer width. A loop is needed to find the initial memory
start when freeing it which isn't optimal either.

Instead, let's have sys_heap_aligned_alloc() rewind a pointer after
it is aligned to make just enough room for storing our heap reference.
This way the heap reference is always located immediately before the
aligned memory and any unused memory is returned to the heap.

The rewind and alignment values may coincide in which case only
the alignment is necessary anyway.

Signed-off-by: Nicolas Pitre <npitre@baylibre.com>

kernel/mempool.c

@@ -11,30 +11,33 @@
 
 static void *z_heap_aligned_alloc(struct k_heap *heap, size_t align, size_t size)
 {
-	uint8_t *mem;
+	void *mem;
 	struct k_heap **heap_ref;
-	size_t excess = MAX(sizeof(struct k_heap *), align);
+	size_t __align;
 
 	/*
-	 * get a block large enough to hold an initial (aligned and hidden) heap
+	 * Adjust the size to make room for our heap reference.
-	 * pointer, as well as the space the caller requested
+	 * Merge a rewind bit with align value (see sys_heap_aligned_alloc()).
+	 * This allows for storing the heap pointer right below the aligned
+	 * boundary without wasting any memory.
 	 */
-	if (size_add_overflow(size, excess, &size)) {
+	if (size_add_overflow(size, sizeof(heap_ref), &size)) {
 		return NULL;
 	}
+	__align = align | sizeof(heap_ref);
 
-	mem = k_heap_aligned_alloc(heap, align, size, K_NO_WAIT);
+	mem = k_heap_aligned_alloc(heap, __align, size, K_NO_WAIT);
 	if (mem == NULL) {
 		return NULL;
 	}
 
-	/* create (void *) values in the excess equal to (void *) -1 */
+	heap_ref = mem;
-	memset(mem, 0xff, excess);
-	heap_ref = (struct k_heap **)mem;
 	*heap_ref = heap;
+	mem = ++heap_ref;
+	__ASSERT(align == 0 || ((uintptr_t)mem & (align - 1)) == 0,
+		 "misaligned memory at %p (align = %zu)", mem, align);
 
-	/* return address of the user area part of the block to the caller */
+	return mem;
-	return mem + excess;
 }
 
 void k_free(void *ptr)
@@ -42,12 +45,8 @@ void k_free(void *ptr)
 	struct k_heap **heap_ref;
 
 	if (ptr != NULL) {
-		for (heap_ref = &((struct k_heap **)ptr)[-1];
+		heap_ref = ptr;
-			*heap_ref == (struct k_heap *)-1; --heap_ref) {
+		ptr = --heap_ref;
-			/* no-op */
-		}
-
-		ptr = (uint8_t *)heap_ref;
 		k_heap_free(*heap_ref, ptr);
 	}
 }

lib/os/heap.c

@@ -248,12 +248,28 @@ void *sys_heap_alloc(struct sys_heap *heap, size_t bytes)
 void *sys_heap_aligned_alloc(struct sys_heap *heap, size_t align, size_t bytes)
 {
 	struct z_heap *h = heap->heap;
+	size_t padded_sz, gap, rewind;
 
+	/*
+	 * Split align and rewind values (if any).
+	 * We allow for one bit of rewind in addition to the alignment
+	 * value to efficiently accommodate z_heap_aligned_alloc().
+	 * So if e.g. align = 0x28 (32 | 8) this means we align to a 32-byte
+	 * boundary and then rewind 8 bytes.
+	 */
+	rewind = align & -align;
+	if (align != rewind) {
+		align -= rewind;
+		gap = MIN(rewind, chunk_header_bytes(h));
+	} else {
+		if (align <= chunk_header_bytes(h)) {
+			return sys_heap_alloc(heap, bytes);
+		}
+		rewind = 0;
+		gap = chunk_header_bytes(h);
+	}
 	__ASSERT((align & (align - 1)) == 0, "align must be a power of 2");
 
-	if (align <= chunk_header_bytes(h)) {
-		return sys_heap_alloc(heap, bytes);
-	}
 	if (bytes == 0 || size_too_big(h, bytes)) {
 		return NULL;
 	}
@@ -263,16 +279,16 @@ void *sys_heap_aligned_alloc(struct sys_heap *heap, size_t align, size_t bytes)
 	 * We over-allocate to account for alignment and then free
 	 * the extra allocations afterwards.
 	 */
-	size_t padded_sz =
+	padded_sz = bytes_to_chunksz(h, bytes + align - gap);
-		bytes_to_chunksz(h, bytes + align - chunk_header_bytes(h));
 	chunkid_t c0 = alloc_chunk(h, padded_sz);
 
 	if (c0 == 0) {
 		return NULL;
 	}
+	uint8_t *mem = chunk_mem(h, c0);
 
 	/* Align allocated memory */
-	uint8_t *mem = (uint8_t *) ROUND_UP(chunk_mem(h, c0), align);
+	mem = (uint8_t *) ROUND_UP(mem + rewind, align) - rewind;
 	chunk_unit_t *end = (chunk_unit_t *) ROUND_UP(mem + bytes, CHUNK_UNIT);
 
 	/* Get corresponding chunks */