From af8df2bd6fd3eabd4b7015174c49ce7746bdfa89 Mon Sep 17 00:00:00 2001
From: Pabloader <pabloader@pabloader.ru>
Date: Thu, 8 May 2025 19:10:18 +0000
Subject: [PATCH] Proper memory allocators

---
 build/assets/stdlib.c | 118 ++++++++++++++++++++++++++++++++++++++++--
 build/assets/stdlib.h |   8 ++-
 2 files changed, 121 insertions(+), 5 deletions(-)

diff --git a/build/assets/stdlib.c b/build/assets/stdlib.c
index ecdbf15..78a5ebd 100644
--- a/build/assets/stdlib.c
+++ b/build/assets/stdlib.c
@@ -3,6 +3,12 @@
 #define BLOCK_SIZE (64 * 1024)
 #define BLOCK_1MB 16
 
+typedef struct header {
+    size_t size : 31;
+    uint8_t is_free : 1;
+} header_t;
+static header_t* head;
+
 static uintptr_t bump_pointer = (uintptr_t)&__heap_base;
 static uintptr_t heap_end = BLOCK_1MB * BLOCK_SIZE;
 IMPORT(grow) void grow(uint32_t blocks);
@@ -19,11 +25,117 @@ static void* bump_alloc(uintptr_t n) {
     return (void*)r;
 }
 
-void* malloc(uintptr_t n) {
-    return bump_alloc(n);
+static header_t* get_next_header(header_t* block) {
+    if (!block) {
+        return NULL;
+    }
+    uint8_t* data = (uint8_t*)(block + 1);
+    uint8_t* next_header = data + block->size;
+    if ((uintptr_t)next_header >= bump_pointer) {
+        return NULL;
+    }
+    return (header_t*)next_header;
 }
 
-void free(void* p) { (void)p; }
+static header_t* get_prev_header(header_t* block) {
+    header_t* curr = head;
+    while (curr) {
+        header_t* next = get_next_header(curr);
+        if (next == block) {
+            return curr;
+        }
+        curr = next;
+    }
+    return NULL;
+}
+
+static header_t* get_free_block(size_t size) {
+    header_t* curr = head;
+    while (curr) {
+        if (curr->is_free && curr->size >= size) {
+            return curr;
+        }
+        curr = get_next_header(curr);
+    }
+    return NULL;
+}
+
+void* malloc(size_t size) {
+    if (!size) {
+        return NULL;
+    }
+
+    header_t* header = get_free_block(size);
+    if (header) {
+        header->is_free = 0;
+        if (header->size > size + sizeof(header_t) * 2) {
+            size_t old_size = header->size;
+            header->size = size;
+
+            header_t* next = get_next_header(header);
+            next->size = old_size - size - sizeof(header_t);
+            next->is_free = 1;
+        }
+        return (void*)(header + 1);
+    }
+    size_t total_size = sizeof(header_t) + size;
+
+    void* block = bump_alloc(total_size);
+    if (!block) {
+        return NULL;
+    }
+
+    header = block;
+    header->size = size;
+    header->is_free = 0;
+
+    if (!head) {
+        head = header;
+    }
+
+    return (void*)(header + 1);
+}
+
+void free(void* block) {
+    if (!block) {
+        return;
+    }
+
+    header_t* header = ((header_t*)block) - 1;
+    if (header->is_free) {
+        return;
+    }
+    header->is_free = 1;
+
+    header_t* next_header = get_next_header(header);
+    if (next_header && next_header->is_free) {
+        header->size += next_header->size + sizeof(header_t);
+    }
+
+    header_t* prev_header = get_prev_header(header);
+    if (prev_header && prev_header->is_free) {
+        prev_header->size += header->size + sizeof(header_t);
+    }
+}
+
+void* realloc(void* block, size_t size) {
+    header_t* header;
+    void* ret;
+
+    if (!block || !size)
+        return malloc(size);
+
+    header = (header_t*)block - 1;
+    if (header->size >= size)
+        return block;
+
+    ret = malloc(size);
+    if (ret) {
+        memcpy(ret, block, header->size);
+        free(block);
+    }
+    return ret;
+}
 
 void* memset(void* s, uint8_t c, uint32_t n) {
     uint8_t* p = (uint8_t*)s;
diff --git a/build/assets/stdlib.h b/build/assets/stdlib.h
index 3639e49..c3516a5 100644
--- a/build/assets/stdlib.h
+++ b/build/assets/stdlib.h
@@ -6,8 +6,12 @@ extern unsigned char __heap_base;
 #define IMPORT(name) __attribute__((import_module("env"), import_name(#name)))
 #define EXPORT(name) __attribute__((export_name(#name)))
 
-EXPORT(malloc) void* malloc(uintptr_t);
-void free(void*);
+#define NULL ((void*)0)
+
+typedef uint32_t size_t;
+
+EXPORT(malloc) void* malloc(size_t);
+EXPORT(free) void free(void*);
 
 void* memset(void* s, uint8_t c, uint32_t n);
 void* memcpy(void* dest, const void* src, uint32_t n);