refactor standard_allocator and fix realloc overflow bugs

Author: ZERICO2005

src/libc/allocator/allocator_standard.c

@@ -1,123 +1,124 @@
 #include <stddef.h>
-#include <stdlib.h>
 #include <stdint.h>
+#include <stdlib.h>
 #include <string.h>
 
-typedef struct __attribute__((packed)) block
-{
-    struct block *ptr;
-    size_t size;
+typedef struct __attribute__((packed)) block {
+    struct block *next;
+    size_t total_size;
+    uint8_t data[];
 } __attribute__((packed)) block_t;
 
 extern uint8_t __heap_low[];
 extern uint8_t __heap_high[];
 static uintptr_t heap_ptr = (uintptr_t)__heap_low;
-static block_t _alloc_base;
+// heap_base is the only node where total_size = 0
+static block_t heap_base = { .next = NULL, .total_size = 0 };
 
+#define BLOCK_HEADER_SIZE offsetof(block_t, data)
+
+// malloc(0) returns NULL
 void *malloc(size_t alloc_size)
 {
-    block_t *q;
-    block_t *r;
-
-    /* add size of block header to real size */
-    const size_t size = alloc_size + sizeof(block_t);
-    /* abort if alloc_size is 0 or size overflowed */
-    if (size <= alloc_size)
-    {
+    const size_t block_size = alloc_size + BLOCK_HEADER_SIZE;
+    if (block_size <= BLOCK_HEADER_SIZE) {
+        // reject if alloc_size is 0 or size overflowed
         return NULL;
     }
 
-    for (block_t *p = &_alloc_base; (q = p->ptr); p = q)
-    {
-        if (q->size >= size)
-        {
-            if (q->size <= size + sizeof(block_t))
-            {
-                p->ptr = q->ptr;
-            }
-            else
-            {
-                q->size -= size;
-                q = (block_t*)(((uint8_t*)q) + q->size);
-                q->size = size;
-            }
+    // search through the free-list for an open block (sorted by address)
+    block_t *previous_block = &heap_base;
+    while (previous_block->next != NULL) {
+        block_t *current_block = previous_block->next;
 
-            return q + 1;
+        if (current_block->total_size >= block_size) {
+            if (current_block->total_size - BLOCK_HEADER_SIZE <= block_size) {
+                // region is too small to split into two parts, so just claim the whole region
+                previous_block->next = current_block->next;
+                return current_block->data;
+            }
+            // split from the high end so the original free-list node stays valid
+            current_block->total_size -= block_size;
+            current_block = (block_t*)(((uint8_t*)current_block) + current_block->total_size);
+            current_block->total_size = block_size;
+            return current_block->data;
         }
+
+        previous_block = previous_block->next;
     }
 
-    /* compute next heap pointer */
-    if (heap_ptr + size < heap_ptr || heap_ptr + size >= (uintptr_t)__heap_high)
-    {
+    // no suitable free block exists, extend into fresh heap space
+    size_t heap_available = (uintptr_t)__heap_high - (uintptr_t)heap_ptr;
+    if (block_size > heap_available) {
         return NULL;
     }
 
-    r = (block_t*)heap_ptr;
-    r->size = size;
+    block_t *new_block = (block_t*)heap_ptr;
+    new_block->total_size = block_size;
+    heap_ptr = heap_ptr + block_size;
 
-    heap_ptr = heap_ptr + size;
-
-    return r + 1;
+    return new_block->data;
 }
 
 void free(void *ptr)
 {
-    if (ptr != NULL)
-    {
-        block_t *p;
-        block_t *q;
-
-        q = (block_t*)ptr - 1;
+    if (ptr == NULL) {
+        return;
+    }
 
-        for (p = &_alloc_base; p->ptr && p->ptr < q; p = p->ptr);
+    block_t *previous_block = &heap_base;
+    block_t *current_block = (block_t*)((uint8_t*)ptr - BLOCK_HEADER_SIZE);
+    while (
+        previous_block->next != NULL &&
+        (uintptr_t)previous_block->next < (uintptr_t)current_block
+    ) {
+        previous_block = previous_block->next;
+    }
 
-        if (p->ptr && (uint8_t*)p->ptr == ((uint8_t*)q) + q->size)
-        {
-            q->size += p->ptr->size;
-            q->ptr = p->ptr->ptr;
-        }
-        else
-        {
-            q->ptr = p->ptr;
-        }
+    // merge with the following free block if it is directly adjacent
+    if (
+        previous_block->next != NULL &&
+        (uint8_t*)previous_block->next == ((uint8_t*)current_block) + current_block->total_size
+    ) {
+        current_block->total_size += previous_block->next->total_size;
+        current_block->next = previous_block->next->next;
+    } else {
+        current_block->next = previous_block->next;
+    }
 
-        if (p->size && ((uint8_t*)p) + p->size == (uint8_t*)q)
-        {
-            p->size += q->size;
-            p->ptr = q->ptr;
-        }
-        else
-        {
-            p->ptr = q;
-        }
+    // merge with the preceding free block, unless that predecessor is heap_base
+    if (
+        previous_block->total_size != 0 &&
+        ((uint8_t*)previous_block) + previous_block->total_size == (uint8_t*)current_block
+    ) {
+        previous_block->total_size += current_block->total_size;
+        previous_block->next = current_block->next;
+    } else {
+        previous_block->next = current_block;
     }
 }
 
-void *realloc(void *ptr, size_t size)
+// realloc(ptr, 0) returns ptr and does nothing
+void *realloc(void *ptr, size_t alloc_size)
 {
-    block_t *h;
-    void *p;
-
-    if (ptr == NULL)
-    {
-        return malloc(size);
+    if (ptr == NULL) {
+        return malloc(alloc_size);
     }
 
-    h = (block_t*)ptr - 1;
-
-    if (h->size >= size + sizeof(block_t))
-    {
+    block_t *header = (block_t*)((uint8_t*)ptr - BLOCK_HEADER_SIZE);
+    if (header->total_size - BLOCK_HEADER_SIZE >= alloc_size) {
+        // realloc(ptr, 0) is undefined in C23 and returns here
         return ptr;
     }
 
-    p = malloc(size);
-    if (p == NULL)
-    {
+    // increase allocation size
+    void *new_ptr = malloc(alloc_size);
+    if (new_ptr == NULL) {
         return NULL;
     }
 
-    memcpy(p, ptr, h->size - sizeof(block_t));
+    memcpy(new_ptr, ptr, header->total_size - BLOCK_HEADER_SIZE);
     free(ptr);
 
-    return p;
+    return new_ptr;
 }