Fixed diasterous bug with hfree.

Before this fix, hfree would neglect to set the prev pointer in the next
used block and such was leaving the prev pointers invalid after
coalescing frees.

3 files changed, 236 insertions, 43 deletions

diff --git a/02-usart/include/kern/mem.h b/02-usart/include/kern/mem.h
index d150744..c0999f5 100644
--- a/02-usart/include/kern/mem.h
+++ b/02-usart/include/kern/mem.h
@@ -24,6 +24,8 @@ void* debug_halloc_get_prev_ptr(void* ptr);
 
 int debug_halloc_assert_consistency(char* error, size_t len);
 
+void debug_print_blocks();
+
 #endif
 
 
diff --git a/02-usart/src/kern/mem.c b/02-usart/src/kern/mem.c
index 79bcabf..5234fff 100644
--- a/02-usart/src/kern/mem.c
+++ b/02-usart/src/kern/mem.c
@@ -26,6 +26,8 @@ void* memset(void* dest, int c, size_t n);
 
 typedef uint16_t halloc_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.
@@ -34,19 +36,25 @@ typedef struct HALLOC_NODE {
     uint32_t header;
     struct {
       /* Is this memory block currently in use (hasn't been hfree'd) */
-      bool     used:1;
+      uint8_t     used:1;
       /* Number of words allocated. Does not include the header. */
-      uint16_t size:15;
+      uint16_t size:12;
       /* The location of the previous block (in WORDS from offest) */
-      halloc_off_t prev;
+      halloc_off_t prev:12;
+      uint8_t canary:7;
     } PACKED;
   };
 
   uint8_t mem[]; /* The memory to use. */
 } halloc_node_t;
 
+static_assert(offsetof(halloc_node_t, mem) == 4, "Offset check failed.");
+
 halloc_node_t* halloc_start;
 
+#define halloc_node_out_of_range(node) \
+   ((uint8_t*) (node) == ((uint8_t*)&DATA_SEGMENT_STOP) + MAX_HEAP_SIZE)
+
 #define halloc_node_next(cur) \
   ((halloc_node_t*)(((uint8_t*)(cur)) + (((cur)->size + 1) * 4)))
 
@@ -58,6 +66,9 @@ halloc_node_t* halloc_start;
 #define halloc_node_get_off(node) \
   (((uint32_t)(((uint8_t*)(node)) - ((uint8_t*)(halloc_start)))) / 4)
 
+#define get_halloc_node(mem) \
+  ((halloc_node_t*)(((uint8_t*)mem) - 4))
+
 #define size_for(n) \
   (((n) / 4) + ((n) % 4 != 0))
 
@@ -67,6 +78,7 @@ void* halloc(size_t size)
     halloc_start = (halloc_node_t*) DATA_SEGMENT_STOP_ADDR;
     memset(halloc_start, 0, sizeof(halloc_node_t));
     halloc_start->size = (MAX_HEAP_SIZE / 4) - 1;
+    halloc_start->canary = CANARY;
   }
 
   size_t realsz = size_for(size); /* Clip the size to the nearest word. */
@@ -85,6 +97,7 @@ void* halloc(size_t size)
         next->used = 0;
         next->size = orig_size - realsz - sizeof(halloc_node_t) / 4;
         next->prev = offset;
+        next->canary = CANARY;
 
         halloc_node_t* nextnext = halloc_node_next(next);
         if (halloc_node_get_off(nextnext) < (MAX_HEAP_SIZE / 4)) {
@@ -101,36 +114,64 @@ void* halloc(size_t size)
   return NULL;
 }
 
+
 /* Joins this node with the previous and next nodes if they're free. */
 static void coalesce(halloc_node_t* cur)
 {
-  if (cur->used) return;
+  halloc_node_t* orig = cur;
+  halloc_node_t* last_freed;
+  halloc_node_t* next_used;
 
-  halloc_node_t* next = halloc_node_next(cur);
-  if (!next->used) {
-    cur->size += next->size + sizeof(halloc_node_t) / 4;
+  /* Find the earliest contiguous free'd block. */
+  while (!cur->used && cur != halloc_start) {
+    cur = halloc_node_prev(cur);
   }
 
-  if (cur != halloc_start) {
-    coalesce(halloc_node_prev(cur));
+  if (cur == halloc_start && !cur->used) {
+    last_freed = cur;
+  } else {
+    last_freed = halloc_node_next(cur);
   }
-}
 
-void hfree(void* mem)
-{
-  /* TODO this should be a critical section. */
-  if (!mem) {
-    /* Consistent with C, freeing a NULL pointer does nothing. */
-    return;
+  /* Find the next used block. */
+  cur = orig;
+  while (!cur->used && !halloc_node_out_of_range(cur)) {
+    cur = halloc_node_next(cur);
   }
 
-  halloc_node_t* header = ((halloc_node_t*) mem) - 1;
+  next_used = cur;
+
+  if (!halloc_node_out_of_range(next_used)) {
+    next_used->prev = halloc_node_get_off(last_freed);
+  }
+
+  last_freed->size = ((uint8_t*) next_used - (last_freed->mem)) / 4;
+}
 
-  if (!header->used) {
 #ifdef FOR_TESTING
-    assert(header->used);
+#include <assert.h>
+#include <stdio.h>
+void panic(const char* x)
+{
+  fprintf(stderr, "%s\n", x);
+  assert(0);
+}
+#else
+void panic(const char* x)
+{
+  for(;;);
+}
 #endif
-    return; // TODO make this fail on ARM.
+
+void hfree(void* mem)
+{
+  /* Like normal free(), do nothing on free'ing NULL */
+  if (!mem) return;
+
+  halloc_node_t* header = get_halloc_node(mem);
+  if (!header->used) {
+    panic("Heap double free or corruption!\n");
+    return;
   }
 
   header->used = 0;
@@ -157,6 +198,18 @@ void* debug_halloc_get_prev_ptr(void* ptr)
   return prev->mem;
 }
 
+void debug_print_blocks()
+{
+  printf("------ Print Blocks -------\n");
+  halloc_node_t* cur = halloc_node_at_off(0);
+
+  while (!halloc_node_out_of_range(cur)) {
+    printf("header (%04x) {used=%d, size=%5d, prev=%04x, canary=%02x}\n",
+        halloc_node_get_off(cur), cur->used, cur->size, cur->prev, cur->canary);
+    cur = halloc_node_next(cur);
+  }
+}
+
 /* Tests that we can walk up and down the allocated blocks and that they
  * are properly aligned. */
 int debug_halloc_assert_consistency(char* error, size_t len)
@@ -164,8 +217,16 @@ int debug_halloc_assert_consistency(char* error, size_t len)
   halloc_node_t* cur = halloc_node_at_off(0);
   size_t total_size = 0;
   size_t loop_check = 0;
+  size_t n_blocks = 1;
+  size_t n_blocks_back = 1;
 
   while(1) {
+    if (cur->canary != CANARY) {
+      snprintf(error, len, "Node has corrupted canary. %02x vs expected %02x\n",
+          cur->canary, CANARY);
+      return 1;
+    }
+
     total_size += cur->size + 1;
 
     halloc_node_t* next = halloc_node_next(cur);
@@ -178,7 +239,9 @@ int debug_halloc_assert_consistency(char* error, size_t len)
           (void*)(DATA_SEGMENT_STOP_ADDR + MAX_HEAP_SIZE));
       return 1;
     }
+
     cur = next;
+    ++ n_blocks;
   }
 
   if (total_size * 4 != MAX_HEAP_SIZE) {
@@ -188,10 +251,21 @@ int debug_halloc_assert_consistency(char* error, size_t len)
     return 1;
   }
 
+  if (cur == halloc_start) {
+    return 0;
+  }
+
   while (loop_check < 10000) {
     halloc_node_t* prev = halloc_node_prev(cur);
+    ++ n_blocks_back;
 
     if (prev == halloc_start)  {
+      if (n_blocks != n_blocks_back) {
+        snprintf(
+            error, len, "Different number of blocks found on the way back. Found %lu on the way back vs %lu up.\n",
+            n_blocks_back, n_blocks);
+        return 1;
+      }
       return 0;
     }
 
diff --git a/02-usart/tests/test_memory.c b/02-usart/tests/test_memory.c
index 2272f20..04e9289 100644
--- a/02-usart/tests/test_memory.c
+++ b/02-usart/tests/test_memory.c
@@ -9,13 +9,6 @@
 #include "kern/common.h"
 #include "kern/mem.h"
 
-extern uint32_t* halloc_start;
-static void wipeout_halloc()
-{
-  memset(halloc_start, 0, 1024);
-  halloc_start = NULL;
-}
-
 struct TEST_STRUCT {
   uint32_t array[3];
 };
@@ -24,6 +17,32 @@ struct TEST_STRUCT2 {
   uint32_t array[10];
 };
 
+/* Copy of the node structure. */
+typedef struct HALLOC_NODE {
+  union {
+    uint32_t header;
+    struct {
+      /* Is this memory block currently in use (hasn't been hfree'd) */
+      bool     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) */
+      uint16_t prev:12;
+      uint8_t canary:7;
+    } PACKED;
+  };
+
+  uint8_t mem[]; /* The memory to use. */
+} halloc_node_t;
+
+extern halloc_node_t* halloc_start;
+
+static void wipeout_halloc()
+{
+  memset(halloc_start, 0, 1024);
+  halloc_start = NULL;
+}
+
 
 static struct TEST_STRUCT* new_test_struct()
 {
@@ -71,6 +90,18 @@ TEST(memory, halloc)
   ASSERT_CHAIN(test3, test4);
   ASSERT_CHAIN(test4, test5);
 
+
+  char buf[1024];
+  if (debug_halloc_assert_consistency(buf, 1024)) {
+    fprintf(
+        stderr,
+        "Consistency check failed. (%s:%d)\n",
+        __FILE__,
+        __LINE__);
+    fprintf(stderr, buf);
+    ASSERT_TRUE(false);
+  }
+
   wipeout_halloc();
 
   return 0;
@@ -98,6 +129,10 @@ struct UNEVEN_STRUCT* new_uneven_struct()
 
 TEST(memory, uneven_halloc)
 {
+  if (halloc_start) {
+    wipeout_halloc();
+  }
+
   struct UNEVEN_STRUCT* test1 = new_uneven_struct();
   struct UNEVEN_STRUCT* test2 = new_uneven_struct();
 
@@ -110,6 +145,10 @@ TEST(memory, uneven_halloc)
 
 TEST(memory, halloc_free)
 {
+  if (halloc_start) {
+    wipeout_halloc();
+  }
+
   struct TEST_STRUCT* test1 = new_test_struct();
   struct TEST_STRUCT2* test2 = new_test_struct2();
   struct TEST_STRUCT* test3 = new_test_struct();
@@ -122,7 +161,7 @@ TEST(memory, halloc_free)
   hfree(test1);
   hfree(test5);
 
-  ASSERT_EQ((*halloc_start >> 1) * 4, MAX_HEAP_SIZE - 4);
+  ASSERT_EQ((int) halloc_start->size * 4, MAX_HEAP_SIZE - 4);
 
   test1 = new_test_struct();
   test2 = new_test_struct2();
@@ -136,7 +175,7 @@ TEST(memory, halloc_free)
   hfree(test4);
   hfree(test5);
 
-  ASSERT_EQ((*halloc_start >> 1) * 4, MAX_HEAP_SIZE - 4);
+  ASSERT_EQ((int) halloc_start->size * 4, MAX_HEAP_SIZE - 4);
 
   test1 = new_test_struct();
   test2 = new_test_struct2();
@@ -150,7 +189,7 @@ TEST(memory, halloc_free)
   hfree(test2);
   hfree(test5);
 
-  ASSERT_EQ((*halloc_start >> 1) * 4, MAX_HEAP_SIZE - 4);
+  ASSERT_EQ((int) halloc_start->size * 4, MAX_HEAP_SIZE - 4);
 
   wipeout_halloc();
 
@@ -159,6 +198,10 @@ TEST(memory, halloc_free)
 
 TEST(memory, halloc_free_alloc2)
 {
+  if (halloc_start) {
+    wipeout_halloc();
+  }
+
   struct TEST_STRUCT2* test1 = new_test_struct2();
   struct TEST_STRUCT2* test2 = new_test_struct2();
 
@@ -189,53 +232,127 @@ TEST(memory, halloc_free_alloc2)
   return 0;
 }
 
+TEST(memory, relink_backref_after_free)
+{
+  if (halloc_start) {
+    wipeout_halloc();
+  }
+
+  struct TEST_STRUCT* test2 = new_test_struct();
+  struct TEST_STRUCT* test3 = new_test_struct();
+
+  hfree(test2);
+  hfree(test3);
+
+  char buf[1024];
+  if (debug_halloc_assert_consistency(buf, 1024)) {
+    fprintf(stderr, "Consistency check failed.\n");
+    fprintf(stderr, buf);
+    ASSERT_TRUE(false);
+  }
+
+  return 0;
+}
+
 TEST(memory, consistency_stress)
 {
+#define NRUNS 500
   if (halloc_start) {
     wipeout_halloc();
   }
 
   int i;
-  void* allocd[500] = { 0 };
+  void* allocd[NRUNS] = { 0 };
   char buf[1024];
 
-  for (i = 0; i < 500; ++ i) {
+  for (i = 0; i < NRUNS; ++ i) {
     size_t nalloc = rand() % 20;
     allocd[i] = halloc(nalloc);
 
+    if (debug_halloc_assert_consistency(buf, 1024)) {
+      fprintf(
+          stderr,
+          "Consistency check failed. (At index=%d, %s:%d)\n",
+          i,
+          __FILE__,
+          __LINE__);
+      fprintf(stderr, buf);
+      ASSERT_TRUE(false);
+    }
+
     ASSERT_TRUE(allocd[i]);
-    size_t idx = rand() % 500;
+
+    memset(allocd[i], 0xFF, nalloc);
+    size_t idx = rand() % NRUNS;
+
+    if (debug_halloc_assert_consistency(buf, 1024)) {
+      fprintf(
+          stderr,
+          "Consistency check failed. (At index=%d, %s:%d)\n",
+          i,
+          __FILE__,
+          __LINE__);
+      fprintf(stderr, buf);
+      ASSERT_TRUE(false);
+    }
 
     hfree(allocd[idx]);
     allocd[idx] = NULL;
 
-    idx = rand() % 500;
+    if (debug_halloc_assert_consistency(buf, 1024)) {
+      fprintf(
+          stderr,
+          "Consistency check failed. (At index=%d, %s:%d)\n",
+          i,
+          __FILE__,
+          __LINE__);
+      fprintf(stderr, buf);
+      ASSERT_TRUE(false);
+    }
+
+    idx = rand() % NRUNS;
     hfree(allocd[idx]);
     allocd[idx] = NULL;
 
     if (debug_halloc_assert_consistency(buf, 1024)) {
-      fprintf(stderr, "Consistency check failed. (At index %d) nalloc=%lu\n", i, nalloc);
+      fprintf(
+          stderr,
+          "Consistency check failed. (At index=%d, %s:%d)\n",
+          i,
+          __FILE__,
+          __LINE__);
       fprintf(stderr, buf);
       ASSERT_TRUE(false);
     }
   }
 
-  for(i = 0; i < 500; ++ i) {
-    hfree(allocd[i]);
-  }
+  for(i = 0; i < NRUNS; ++ i) {
+    if (allocd[i]) {
+      hfree(allocd[i]);
+    }
 
-  if (debug_halloc_assert_consistency(buf, 1024)) {
-    fprintf(stderr, "Consistency check failed. (Final)\n");
-    fprintf(stderr, buf);
-    ASSERT_TRUE(false);
+    if (debug_halloc_assert_consistency(buf, 1024)) {
+        fprintf(
+            stderr,
+            "Consistency check failed. (At index=%d, %s:%d)\n",
+            i,
+            __FILE__,
+            __LINE__);
+      fprintf(stderr, buf);
+      ASSERT_TRUE(false);
+    }
   }
-  ASSERT_EQ((*halloc_start >> 1) * 4, MAX_HEAP_SIZE - 4);
+  ASSERT_EQ((int) halloc_start->size * 4, MAX_HEAP_SIZE - 4);
 
   return 0;
 }
 
 TEST(memory, halloc_free_alloc)
 {
+  if (halloc_start) {
+    wipeout_halloc();
+  }
+
   new_test_struct();
   struct TEST_STRUCT2* test2 = new_test_struct2();
   new_test_struct();