1 files changed, 67 insertions, 50 deletions

diff --git a/src/kern/mem.c b/src/kern/mem.c
index 533b394..31756e5 100644
--- a/src/kern/mem.c
+++ b/src/kern/mem.c
@@ -28,79 +28,90 @@ void* memset(void* dest, int c, size_t n);
 
 typedef uint16_t kalloc_off_t;
 
-#define CANARY 0x5a
-
 // The sizes will count the number of WORDS allocated.
 // Since there's a max size of 16k, only 12 bits will be
 // needed for this.
 typedef struct KALLOC_NODE {
-  union {
-    uint32_t header;
-    struct {
-      /* Is this memory block currently in use (hasn't been kfree'd) */
-      uint8_t used : 1;
-      /* Number of words allocated. Does not include the header. */
-      uint16_t size : 12;
-      /* The location of the previous block (in WORDS from offest) */
-      kalloc_off_t prev : 12;
-      uint8_t canary : 7;
-    } PACKED;
-  };
+  uint16_t size_words;
+  kalloc_off_t prev;
+  /*
+   * LSB is whether this block is used.
+   *
+   * Rest of the bits must equal (0xdeadbeee >> 1)
+   */
+  uint32_t used_and_canary;
 
   uint8_t mem[]; /* The memory to use. */
-} kalloc_node_t;
-
-static_assert(offsetof(kalloc_node_t, mem) == 4, "Offset check failed.");
+} PACKED kalloc_node_t;
 
 kalloc_node_t* kalloc_start;
 
+#define CANARY 0xdeadbeee
+#define kalloc_node_in_use(node) ((node)->used_and_canary & 1)
+#define kalloc_node_get_canary(node) ((node)->used_and_canary & (~1))
+#define WORD_SIZE (sizeof(uint32_t))
+#define SIZEOF_KALLOC_NODE_WORDS (sizeof(kalloc_node_t) / WORD_SIZE)
 #define HEAP_START_ADDR ((ptrdiff_t)&HEAP_START)
 #define REAL_HEAP_START \
   (*((unsigned char*)((HEAP_START_ADDR & (~3)) + (HEAP_START_ADDR % 4 != 0))))
 #define MAX_HEAP_SIZE ((&HEAP_STOP - &REAL_HEAP_START))
+#define MAX_HEAP_SIZE_WORDS (MAX_HEAP_SIZE / WORD_SIZE)
 #define kalloc_node_out_of_range(node) ((void*)(node) >= (void*)&HEAP_STOP)
 
 #define kalloc_node_next(cur) \
-  ((kalloc_node_t*)(((uint8_t*)(cur)) + (((cur)->size + 1) * 4)))
+  ((kalloc_node_t*)(((uint8_t*)(cur)) + (((cur)->size_words + SIZEOF_KALLOC_NODE_WORDS) * WORD_SIZE)))
 
 #define kalloc_node_prev(cur) kalloc_node_at_off(cur->prev)
 
 #define kalloc_node_at_off(offset) \
-  ((kalloc_node_t*)(((uint8_t*)kalloc_start) + (offset)*4))
+  ((kalloc_node_t*)(((uint8_t*)kalloc_start) + (offset)*WORD_SIZE))
 
 #define kalloc_node_get_off(node) \
-  (((uint32_t)(((uint8_t*)(node)) - ((uint8_t*)(kalloc_start)))) / 4)
+  ((uint32_t)(((((uint8_t*)(node)) - ((uint8_t*)(kalloc_start)))) / WORD_SIZE))
 
-#define get_kalloc_node(mem) ((kalloc_node_t*)(((uint8_t*)mem) - 4))
+#define get_kalloc_node(mem) \
+  ((kalloc_node_t*)(((uint8_t*)mem) - SIZEOF_KALLOC_NODE_WORDS * WORD_SIZE))
 
 #define size_for(n) (((n) / 4) + ((n) % 4 != 0))
 
+void kalloc_init()
+{
+  kalloc_start = (kalloc_node_t*)&REAL_HEAP_START;
+  memset(kalloc_start, 0, sizeof(kalloc_node_t));
+  kalloc_start->size_words = MAX_HEAP_SIZE_WORDS - SIZEOF_KALLOC_NODE_WORDS;
+  kalloc_start->used_and_canary = CANARY;
+}
+
 void* kalloc(size_t size)
 {
   if (!kalloc_start) {
-    kalloc_start = (kalloc_node_t*)&REAL_HEAP_START;
-    memset(kalloc_start, 0, sizeof(kalloc_node_t));
-    kalloc_start->size = (MAX_HEAP_SIZE / 4) - 1;
-    kalloc_start->canary = CANARY;
+    kalloc_init();
   }
 
-  size_t realsz = size_for(size); /* Clip the size to the nearest word. */
+  size_t realsz_words = size_for(size); /* Clip the size to the nearest word. */
   kalloc_off_t offset = 0;
-  while (offset < (MAX_HEAP_SIZE / 4)) {
+  while (offset < (MAX_HEAP_SIZE_WORDS)) {
     kalloc_node_t* cur = kalloc_node_at_off(offset);
 
-    if (!cur->used && (cur->size >= realsz)) {
-      cur->used = true;
-      size_t orig_size = cur->size;
-      cur->size = realsz;
+    if (!kalloc_node_in_use(cur) && (cur->size_words >= realsz_words)) {
+      cur->used_and_canary |= 1;
+      size_t orig_size_words = cur->size_words;
+
+      if (orig_size_words < realsz_words + SIZEOF_KALLOC_NODE_WORDS * 2) {
+        /* If the original size is only slightly larger than the size to
+         * allocate, there's no point in saving the leftover space. */
+        realsz_words = orig_size_words;
+      }
+
+      cur->size_words = realsz_words;
 
-      if (orig_size > realsz) {
+      if (orig_size_words > realsz_words) {
         /* This kalloc node needs to split into two blocks. */
         kalloc_node_t* next = kalloc_node_next(cur);
-        next->used = 0;
-        next->size = orig_size - realsz - sizeof(kalloc_node_t) / 4;
+        next->used_and_canary = CANARY;
+        next->size_words =
+            orig_size_words - realsz_words - SIZEOF_KALLOC_NODE_WORDS;
         next->prev = offset;
-        next->canary = CANARY;
 
         kalloc_node_t* nextnext = kalloc_node_next(next);
         if (kalloc_node_get_off(nextnext) < (MAX_HEAP_SIZE / 4)) {
@@ -111,7 +122,7 @@ void* kalloc(size_t size)
       return (void*)cur->mem;
     }
 
-    offset += (sizeof(kalloc_node_t) / 4) + cur->size;
+    offset += SIZEOF_KALLOC_NODE_WORDS + cur->size_words;
   }
 
   return NULL;
@@ -125,11 +136,11 @@ static void coalesce(kalloc_node_t* cur)
   kalloc_node_t* next_used;
 
   /* Find the earliest contiguous free'd block. */
-  while (!cur->used && cur != kalloc_start) {
+  while (!kalloc_node_in_use(cur) && cur != kalloc_start) {
     cur = kalloc_node_prev(cur);
   }
 
-  if (cur == kalloc_start && !cur->used) {
+  if (cur == kalloc_start && !kalloc_node_in_use(cur)) {
     last_freed = cur;
   } else {
     last_freed = kalloc_node_next(cur);
@@ -137,7 +148,7 @@ static void coalesce(kalloc_node_t* cur)
 
   /* Find the next used block. */
   cur = orig;
-  while (!cur->used && !kalloc_node_out_of_range(cur)) {
+  while (!kalloc_node_in_use(cur) && !kalloc_node_out_of_range(cur)) {
     cur = kalloc_node_next(cur);
   }
 
@@ -147,7 +158,7 @@ static void coalesce(kalloc_node_t* cur)
     next_used->prev = kalloc_node_get_off(last_freed);
   }
 
-  last_freed->size = ((uint8_t*)next_used - (last_freed->mem)) / 4;
+  last_freed->size_words = ((uint8_t*)next_used - (last_freed->mem)) / WORD_SIZE;
 }
 
 void kfree(void* mem)
@@ -156,12 +167,12 @@ void kfree(void* mem)
   if (!mem) return;
 
   kalloc_node_t* header = get_kalloc_node(mem);
-  if (!header->used) {
+  if (!kalloc_node_in_use(header)) {
     panic("Heap double free or corruption!\n");
     return;
   }
 
-  header->used = 0;
+  header->used_and_canary &= ~1;
   coalesce(header);
 }
 
@@ -188,16 +199,22 @@ void* debug_kalloc_get_prev_ptr(void* ptr)
 void debug_print_blocks()
 {
   printf("------ Print Blocks -------\n");
+  printf("MAX_HEAP_SIZE      : %ld\n", MAX_HEAP_SIZE);
+  printf("MAX_HEAP_SIZE_WORDS: %ld\n", MAX_HEAP_SIZE_WORDS);
+  printf("HEAP_START         : %p\n", &HEAP_START);
+  printf("HEAP_STOP          : %p\n", &HEAP_STOP);
+  printf("---------------------------\n");
   kalloc_node_t* cur = kalloc_node_at_off(0);
 
   while (!kalloc_node_out_of_range(cur)) {
     printf(
-        "header (%04x) {used=%d, size=%5d, prev=%04x, canary=%02x}\n",
+        "header (%04x)@%p {used=%d, size=%5d, prev=%04x, canary=%04x}\n",
         kalloc_node_get_off(cur),
-        cur->used,
-        cur->size,
+        cur->mem,
+        kalloc_node_in_use(cur),
+        cur->size_words,
         cur->prev,
-        cur->canary);
+        kalloc_node_get_canary(cur));
     cur = kalloc_node_next(cur);
   }
 }
@@ -213,17 +230,17 @@ int debug_kalloc_assert_consistency(char* error, size_t len)
   size_t n_blocks_back = 1;
 
   while (1) {
-    if (cur->canary != CANARY) {
+    if (kalloc_node_get_canary(cur) != CANARY) {
       snprintf(
           error,
           len,
           "Node has corrupted canary. %02x vs expected %02x\n",
-          cur->canary,
+          kalloc_node_get_canary(cur),
           CANARY);
       return 1;
     }
 
-    total_size += cur->size + 1;
+    total_size += cur->size_words + SIZEOF_KALLOC_NODE_WORDS;
 
     kalloc_node_t* next = kalloc_node_next(cur);
     if ((void*)next == (void*)&HEAP_STOP) {
@@ -247,7 +264,7 @@ int debug_kalloc_assert_consistency(char* error, size_t len)
         error,
         len,
         "Total recorded size is inconsistent. %lu vs %lu\n",
-        total_size * 4,
+        total_size * WORD_SIZE,
         MAX_HEAP_SIZE);
     return 1;
   }