Implement heap_collect()

README.md

@@ -12,3 +12,5 @@ $ ./heap
 - The pointers to the heap can only be located in the heap and the stack.
 - No packed structs. All of the pointers should be aligned.
 - No tricks that obscure the pointers (like XOR Linked Lists).
+- Probably works only on x86_64
+- Probably works only when compiled with GCC

heap.c

@@ -1,7 +1,14 @@
 #include <stdio.h>
+#include <stdbool.h>
+#include <string.h>
 #include "./heap.h"
 
 uintptr_t heap[HEAP_CAP_WORDS] = {0};
+const uintptr_t *stack_base = 0;
+
+bool reachable_chunks[CHUNK_LIST_CAP] = {0};
+void *to_free[CHUNK_LIST_CAP] = {0};
+size_t to_free_count = 0;
 
 Chunk_List alloced_chunks = {0};
 Chunk_List freed_chunks = {
@@ -50,9 +57,9 @@ void chunk_list_merge(Chunk_List *dst, const Chunk_List *src)
     }
 }
 
-void chunk_list_dump(const Chunk_List *list)
+void chunk_list_dump(const Chunk_List *list, const char *name)
 {
-    printf("Chunks (%zu):\n", list->count);
+    printf("%s Chunks (%zu):\n", name, list->count);
     for (size_t i = 0; i < list->count; ++i) {
         printf("  start: %p, size: %zu\n",
                (void*) list->chunks[i].start,
@@ -121,7 +128,37 @@ void heap_free(void *ptr)
     }
 }
 
+static void mark_region(const uintptr_t *start, const uintptr_t *end)
+{
+    for (; start < end; start += 1) {
+        const uintptr_t *p = (const uintptr_t *) *start;
+        for (size_t i = 0; i < alloced_chunks.count; ++i) {
+            Chunk chunk = alloced_chunks.chunks[i];
+            if (chunk.start <= p && p < chunk.start + chunk.size) {
+                if (!reachable_chunks[i]) {
+                    reachable_chunks[i] = true;
+                    mark_region(chunk.start, chunk.start + chunk.size);
+                }
+            }
+        }
+    }
+}
+
 void heap_collect()
 {
-    UNIMPLEMENTED;
+    const uintptr_t *stack_start = (const uintptr_t*)__builtin_frame_address(0);
+    memset(reachable_chunks, 0, sizeof(reachable_chunks));
+    mark_region(stack_start, stack_base + 1);
+
+    to_free_count = 0;
+    for (size_t i = 0; i < alloced_chunks.count; ++i) {
+        if (!reachable_chunks[i]) {
+            assert(to_free_count < CHUNK_LIST_CAP);
+            to_free[to_free_count++] = alloced_chunks.chunks[i].start;
+        }
+    }
+
+    for (size_t i = 0; i < to_free_count; ++i) {
+        heap_free(to_free[i]);
+    }
 }

heap.h

@@ -19,6 +19,7 @@ static_assert(HEAP_CAP_BYTES % sizeof(uintptr_t) == 0,
 #define HEAP_CAP_WORDS (HEAP_CAP_BYTES / sizeof(uintptr_t))
 
 extern uintptr_t heap[HEAP_CAP_WORDS];
+extern const uintptr_t *stack_base;
 
 void *heap_alloc(size_t size_bytes);
 void heap_free(void *ptr);
@@ -42,7 +43,7 @@ extern Chunk_List tmp_chunks;
 
 void chunk_list_insert(Chunk_List *list, void *start, size_t size);
 void chunk_list_merge(Chunk_List *dst, const Chunk_List *src);
-void chunk_list_dump(const Chunk_List *list);
+void chunk_list_dump(const Chunk_List *list, const char *name);
 int chunk_list_find(const Chunk_List *list, uintptr_t *ptr);
 void chunk_list_remove(Chunk_List *list, size_t index);
 

main.c

@@ -57,8 +57,16 @@ void *ptrs[N] = {0};
 
 int main()
 {
+    stack_base = (const uintptr_t*)__builtin_frame_address(0);
+
+    for (size_t i = 0; i < 10; ++i) {
+        heap_alloc(i);
+    }
+
     Node *root = generate_tree(0, 3);
 
+    printf("root: %p\n", (void*)root);
+
     Jim jim = {
         .sink = stdout,
         .write = (Jim_Write) fwrite,
@@ -66,20 +74,15 @@ int main()
 
     print_tree(root, &jim);
 
+    printf("\n------------------------------\n");
+    heap_collect();
+    chunk_list_dump(&alloced_chunks, "Alloced");
+    chunk_list_dump(&freed_chunks, "Freed");
     printf("------------------------------\n");
-
-    size_t heap_ptrs_count = 0;
-    for (size_t i = 0; i < alloced_chunks.count; ++i) {
-        for (size_t j = 0; j < alloced_chunks.chunks[i].size; ++j) {
-            uintptr_t *p = (uintptr_t*) alloced_chunks.chunks[i].start[j];
-            if (heap <= p && p < heap + HEAP_CAP_WORDS) {
-                printf("DETECTED HEAP POINTER: %p\n", (void*) p);
-                heap_ptrs_count += 1;
-            }
-        }
-    }
-
-    printf("Detected %zu heap pointers\n", heap_ptrs_count);
+    root = NULL;
+    heap_collect();
+    chunk_list_dump(&alloced_chunks, "Alloced");
+    chunk_list_dump(&freed_chunks, "Freed");
 
     return 0;
 }