refactor heap code and add heap debugging

5 files changed, 113 insertions, 27 deletions

diff --git a/stage3/heap.c b/stage3/heap.c
index 6addd6d..85bb8e9 100644
--- a/stage3/heap.c
+++ b/stage3/heap.c
@@ -1,34 +1,57 @@
 #include "halt.h"
 #include "heap.h"
 #include "memory.h"
+#include "font.h"
+#include "interrupts.h"
 
 #define PAGESIZE 0x1000
 #define MAGIC ((void *) 0x69)
 
-typedef struct __attribute__((packed)) Header {
-	struct Header *next;
-	usize size;
-} Header;
+static heap_header init_free_ptr;
+static heap_header *free_ptr = nil;
 
-static Header init_free_ptr;
-static Header *free_ptr = nil;
+#ifdef DEBUG
+void heap_check()
+{
+	heap_header *h = free_ptr;
+	for (;;) {
+		if ((u64) h < 0x100000) {
+			print(S("heap corruption\n"));
+			asm volatile("int $1");
+		}
+
+		h = h->next;
+		if (h == free_ptr)
+			break;
+	}
+}
+#endif
 
 void kfree(void *ptr)
 {
-	Header *h = ((Header *) ptr) - 1;
+	ISR_UNSAFE
 
+	heap_check();
+
+	heap_header *h = ((heap_header *) ptr) - 1;
 	if (h->next != MAGIC)
-		panic(S("kfree: invalid pointer"));
+		panic(S("kfree: invalid pointer\n"));
 
-	Header *next = free_ptr->next;
+	heap_header *next = free_ptr->next;
 	free_ptr->next = h;
 	h->next = next;
+
+	heap_check();
 }
 
 void *try_kmalloc(usize size)
 {
-	for (Header *prev = free_ptr;; prev = prev->next) {
-		Header *h = prev->next;
+	ISR_UNSAFE
+
+	heap_check();
+
+	for (heap_header *prev = free_ptr;; prev = prev->next) {
+		heap_header *h = prev->next;
 
 		if (h->size < size) {
 			if (h == free_ptr)
@@ -37,17 +60,20 @@ void *try_kmalloc(usize size)
 				continue;
 		}
 
-		if (h->size <= size + sizeof(Header)) {
+		if (h->size <= size + sizeof(heap_header)) {
 			prev->next = h->next;
 		} else {
 			// split
-			h->size -= size + sizeof(Header);
-			h = ((void *) h) + sizeof(Header) + h->size;
+			h->size -= size + sizeof(heap_header);
+			h = ((void *) h) + sizeof(heap_header) + h->size;
 			h->size = size;
 		}
 
 		h->next = MAGIC;
 		free_ptr = prev;
+
+		heap_check();
+
 		return h + 1;
 	}
 
@@ -56,29 +82,29 @@ void *try_kmalloc(usize size)
 
 void *kmalloc(usize size)
 {
-	void *p;
-
-	p = try_kmalloc(size);
-	if (p) return p;
-	panic(S("out of memory"));
+	ISR_UNSAFE
 
-	return nil;
+	void *ptr = try_kmalloc(size);
+	if (ptr == nil)
+		panic(S("kmalloc: out of memory\n"));
+	return ptr;
 }
 
 void *krealloc(void *ptr, usize size)
 {
+	ISR_UNSAFE
+
 	if (ptr == nil)
 		return kmalloc(size);
 
-	Header *h = ((Header *) ptr) - 1;
-
+	heap_header *h = ((heap_header *) ptr) - 1;
 	if (h->next != MAGIC)
-		panic(S("krealloc: invalid pointer"));
+		panic(S("krealloc: invalid pointer\n"));
 
 	void *new = kmalloc(size);
 
 	lmemcpy(new, ptr, h->size);
-	kfree(ptr);
+	kfree(h + 1);
 
 	return new;
 }
@@ -93,12 +119,12 @@ void heap_init()
 void heap_add(void *ptr, usize size)
 {
 	// discard blocks that are too small
-	if (size <= sizeof(Header))
+	if (size <= sizeof(heap_header))
 		return;
 
-	Header *h = ptr;
+	heap_header *h = ptr;
 	h->next = MAGIC;
-	h->size = size - sizeof(Header);
+	h->size = size - sizeof(heap_header);
 
 	kfree(h + 1);
 }
@@ -130,3 +156,8 @@ void heap_add_region(MemRegion *region)
 
 	region->used = -1; // just to be safe
 }
+
+heap_header *heap_get_free_ptr()
+{
+	return free_ptr;
+}
diff --git a/stage3/heap.h b/stage3/heap.h
index 50a3fcf..10fe261 100644
--- a/stage3/heap.h
+++ b/stage3/heap.h
@@ -14,4 +14,17 @@ void *kmalloc(usize siz);
 void kfree(void *ptr);
 void *krealloc(void *ptr, usize size);
 
+typedef struct __attribute__((packed)) heap_header {
+	struct heap_header *next;
+	usize size;
+} heap_header;
+
+heap_header *heap_get_free_ptr();
+
+#ifdef DEBUG
+void heap_check();
+#else
+#define heap_check()
+#endif
+
 #endif
diff --git a/stage3/interrupts.c b/stage3/interrupts.c
index 9c7bcc8..a13bd98 100644
--- a/stage3/interrupts.c
+++ b/stage3/interrupts.c
@@ -9,10 +9,14 @@
 #include "string.h"
 #include "debug.h"
 
+bool in_isr = false;
+
 extern u64 idt_entries[256]; // isr.asm
 
 void interrupt_handler(interrupt_frame *frame)
 {
+	in_isr = true;
+
 	if (frame->which < 32) {
 		debug_exception(frame);
 	} else if (frame->which-32 < 16) {
@@ -43,6 +47,16 @@ void interrupt_handler(interrupt_frame *frame)
 
 		ack_irq(irq);
 	}
+
+	in_isr = false;
+}
+
+void interrupts_unsafe(const char *func)
+{
+	if (in_isr) {
+		print(str_intro((char *) func));
+		panic(S(" called in ISR\n"));
+	}
 }
 
 typedef struct {
diff --git a/stage3/interrupts.h b/stage3/interrupts.h
index 595ed22..a73a68d 100644
--- a/stage3/interrupts.h
+++ b/stage3/interrupts.h
@@ -11,4 +11,13 @@ typedef struct __attribute__((packed)) {
 
 void interrupts_init();
 
+extern bool in_isr;
+
+#ifdef DEBUG
+void interrupts_unsafe(const char *function);
+#define ISR_UNSAFE interrupts_unsafe(__FUNCTION__);
+#else
+#define ISR_UNSAFE
+#endif
+
 #endif
diff --git a/stage3/shell.c b/stage3/shell.c
index 0d24859..7930130 100644
--- a/stage3/shell.c
+++ b/stage3/shell.c
@@ -317,6 +317,24 @@ static void cmd_choose(str arg)
 	}
 }
 
+static void cmd_heapdbg(str arg)
+{
+	(void) arg;
+	heap_header *free_ptr = heap_get_free_ptr();
+
+	heap_header *h = free_ptr;
+	for (;;) {
+		print_hex((u64) h);
+		print(S(" "));
+		print_hex(h->size);
+		print(S("\n"));
+
+		h = h->next;
+		if (h == free_ptr)
+			break;
+	}
+}
+
 typedef struct {
 	str name;
 	void (*fn)(str arg);
@@ -340,6 +358,7 @@ static command registry[] = {
 	{ S("watchdog"),  &cmd_watchdog  },
 	{ S("clocktest"), &cmd_clocktest },
 	{ S("choose"),    &cmd_choose    },
+	{ S("heapdbg"),   &cmd_heapdbg   },
 };
 
 void shell_run_cmd(str cmd)