adding kmalloc and memory management

src/kernel/kernel.c

@@ -11,10 +11,12 @@ void kmain() {
     knewline();
     page_init();
 
-    kprint("Stress testing memory...\n");
-    while(1) {
-        void *p = page_alloc();
-    }
+    //kprint("Stress testing memory...\n");
+    //while(1) {
+    //    void *p = page_alloc();
+    //}
+
+    test_memory_integrity();
 
     poweroff();
 }

src/kernel/memory.c

@@ -4,20 +4,8 @@
 #include "memory.h"
 #include "panic.h"
 
-#define PAGE_SIZE 4096
-
-struct Page {
-	struct Page *next;
-};
-
-struct HeapHeader {
-	size_t size;
-	int is_free;
-	struct HeapHeader *next;
-};
-
-static struct Page *free_list = NULL;
 
+HeapHeader* heap_free_list;
 extern uint8_t _heap_start[]; // named in the linker script
 
 void page_init() {
@@ -34,7 +22,6 @@ void page_free(void *addr) {
 	if (addr == NULL) return;
 
 	struct Page *p = (struct Page *)addr;
-
 	p->next = free_list;
 	free_list = p;
 }
@@ -47,14 +34,111 @@ void *page_alloc() {
 	free_list = free_list->next;
 
 	//zero out the page
-	for (int i = 0; i < (PAGE_SIZE/8); i++) {
+	for (int i = 0; i < (PAGE_SIZE); i++) {
 		((uint8_t *)p)[i] = 0;
 	}
 
 	return (void *)p;
 }
 
-void kmalloc(size_t size) {
-	//check for existing heap
+void *kmalloc(size_t size) {
+	if (size == 0) return NULL;
 
+	// Align size to 8 bytes
+	size = (size + 7) & ~7;
+
+	if (heap_free_list == NULL) {
+		HeapHeader* header = (HeapHeader *)page_alloc();
+		header->size = PAGE_SIZE - sizeof(HeapHeader);
+		header->is_free = 1;
+		header->next = NULL;
+		
+		heap_free_list = header;
+	}
+
+	HeapHeader *current = heap_free_list;
+	HeapHeader *prev = NULL;
+
+	while (current != NULL) {
+		if (current->is_free && current->size >= size) {
+			break;
+		}
+		prev = current;
+		current = current->next;
+	}
+
+	if(current == NULL) {
+		// No suitable block found, allocate a new page
+		HeapHeader* header = (HeapHeader *)page_alloc();
+		header->size = PAGE_SIZE - sizeof(HeapHeader);
+		header->is_free = 1;
+		header->next = NULL;
+
+		if (prev) {
+			prev->next = header;
+		} else {
+			heap_free_list = header;
+		}
+		current = header;
+	}
+
+	// Now, current is a block that can be used
+    if (!current) {
+        return NULL;
+    }
+
+	current->is_free = 0;
+
+	// If the block is larger than needed, split it
+	size_t min_split_size = sizeof(HeapHeader) + 16;
+	if (current->size >= size + min_split_size) {
+		HeapHeader *new_header = (HeapHeader *)((uint8_t *)current + sizeof(HeapHeader) + size);
+		new_header->size = current->size - size - sizeof(HeapHeader);
+		new_header->is_free = 1;
+		new_header->next = current->next;
+
+		current->size = size;
+		current->next = new_header;
+	}
+
+    return (void*)((char*)current + sizeof(HeapHeader));
+}
+
+void kfree(void *ptr) {
+	if (ptr == NULL) return;
+
+	HeapHeader *header = (HeapHeader *)((uint8_t *)ptr - sizeof(HeapHeader));
+	if(header->is_free) {
+		kpanic("Double free detected!");
+	}
+	header->is_free = 1;
+
+	kcoalesce(header);
+}
+
+void kcoalesce(HeapHeader *header) {
+	if(!header || !header->next || !header->is_free) return;
+	uintptr_t current_end = (uintptr_t)header + sizeof(HeapHeader) + header->size;
+	if (current_end == (uintptr_t)header->next && header->next->is_free) {
+		header->size += sizeof(HeapHeader) + header->next->size;
+		header->next = header->next->next;
+		kcoalesce(header);
+	}
+}
+
+void test_memory_integrity() {
+    kprint("Running Integrity Test...\n");
+    uint64_t *a = (uint64_t*)kmalloc(16);
+    uint64_t *b = (uint64_t*)kmalloc(16);
+    
+    *a = 0x1122334455667788;
+    *b = 0x99AABBCCDDEEFF00;
+
+    if (*a == 0x1122334455667788) {
+        kprint("Integrity Pass!\n");
+    } else {
+        kprint("CORRUPTION DETECTED!\n");
+    }
+    kfree(a);
+    kfree(b);
 }

src/kernel/memory.h

@@ -1,8 +1,26 @@
 #ifndef KMEMORY_H
 #define KMEMORY_H
 
+#define PAGE_SIZE 4096
+
+typedef struct Page {
+	struct Page *next;
+} Page;
+
+typedef struct HeapHeader {
+	size_t size;
+	int is_free;
+	struct HeapHeader *next;
+} HeapHeader;
+
+static struct Page *free_list = NULL;
+
 void page_init();
 void page_free(void *addr);
 void *page_alloc();
+void *kmalloc(size_t size);
+void kfree(void *ptr);
+void kcoalesce(HeapHeader *header);
+void test_memory_integrity();
 
 #endif