From b5d3417572729534e643f7d78b8b0858872a37fd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Alexander=20Neuh=C3=A4user?= <alexander.neuhaeuser@tuta.io>
Date: Sat, 20 Dec 2025 16:51:18 +0100
Subject: [PATCH] Fix first fit implementation of heap

---
 heap/src/lib.rs   |  10 +++--
 heap/src/tests.rs | 102 ++++++++++++++++++++++++++++++++--------------
 2 files changed, 77 insertions(+), 35 deletions(-)

diff --git a/heap/src/lib.rs b/heap/src/lib.rs
index e728c1a..843c871 100644
--- a/heap/src/lib.rs
+++ b/heap/src/lib.rs
@@ -61,7 +61,7 @@ impl Heap {
     unsafe fn find_first_fit(&self, size: usize) -> Result<*mut HeapHeader, NovaError> {
         let mut current = self.start_address;
         while !fits(size, current) {
-            if let Some(next) = (*self.start_address).next {
+            if let Some(next) = (*current).next {
                 current = next;
             } else {
                 return Err(NovaError::HeapFull);
@@ -122,9 +122,11 @@ impl Heap {
                 },
             )
         };
-        (*current).next = Some(new_address);
-        (*current).free = false;
-        (*current).size = size;
+        unsafe {
+            (*current).next = Some(new_address);
+            (*current).free = false;
+            (*current).size = size;
+        }
     }
 
     pub fn free(&self, pointer: *mut u8) -> Result<(), NovaError> {
diff --git a/heap/src/tests.rs b/heap/src/tests.rs
index 7d82076..13f3326 100644
--- a/heap/src/tests.rs
+++ b/heap/src/tests.rs
@@ -88,38 +88,78 @@ fn test_freeing_root() {
         assert_eq!((*root_header).size, root_header_start_size);
         assert!((*root_header).next.is_none());
     }
+}
 
-    fn test_merging_free_sections() {
-        let heap_vector = Box::new([0u8; HEAP_SIZE]);
-        let mut heap = Heap::empty();
-        heap.init(
-            &heap_vector[0] as *const u8 as usize,
-            &heap_vector[HEAP_SIZE - 1] as *const u8 as usize,
+#[test]
+fn test_merging_free_sections() {
+    let heap_vector = Box::new([0u8; HEAP_SIZE]);
+    let mut heap = Heap::empty();
+    heap.init(
+        &heap_vector[0] as *const u8 as usize,
+        &heap_vector[HEAP_SIZE - 1] as *const u8 as usize,
+    );
+
+    let root_header = heap.start_address;
+    let root_header_start_size = unsafe { (*root_header).size };
+
+    let malloc1 = heap.malloc(MIN_BLOCK_SIZE).unwrap();
+    let malloc_header_before = unsafe { *Heap::get_header_ref_from_data_pointer(malloc1) };
+    let malloc2 = heap.malloc(MIN_BLOCK_SIZE).unwrap();
+    let _ = heap.malloc(MIN_BLOCK_SIZE).unwrap();
+
+    unsafe {
+        assert!(heap.free(malloc1).is_ok());
+
+        let malloc_header_free = *Heap::get_header_ref_from_data_pointer(malloc1);
+        assert_ne!(malloc_header_before.free, malloc_header_free.free);
+        assert_eq!(malloc_header_before.size, malloc_header_free.size);
+
+        assert!(heap.free(malloc2).is_ok());
+        let malloc_header_merge = *Heap::get_header_ref_from_data_pointer(malloc1);
+
+        assert!(malloc_header_merge.free);
+        assert_eq!(
+            malloc_header_merge.size,
+            malloc_header_free.size + MIN_BLOCK_SIZE + HEAP_HEADER_SIZE
         );
-
-        let root_header = heap.start_address;
-        let root_header_start_size = unsafe { (*root_header).size };
-
-        let malloc1 = heap.malloc(MIN_BLOCK_SIZE).unwrap();
-        let malloc_header_before = unsafe { *Heap::get_header_ref_from_data_pointer(malloc1) };
-        let malloc2 = heap.malloc(MIN_BLOCK_SIZE).unwrap();
-        let malloc3 = heap.malloc(MIN_BLOCK_SIZE).unwrap();
-
-        unsafe {
-            assert!(heap.free(malloc1).is_ok());
-
-            let malloc_header_free = *Heap::get_header_ref_from_data_pointer(malloc1);
-            assert_ne!(malloc_header_before.free, malloc_header_free.free);
-            assert_eq!(malloc_header_before.size, malloc_header_free.size);
-
-            assert!(heap.free(malloc2).is_ok());
-            let malloc_header_merge = *Heap::get_header_ref_from_data_pointer(malloc1);
-
-            assert!(malloc_header_merge.free);
-            assert_eq!(
-                malloc_header_merge.size,
-                malloc_header_free.size + MIN_BLOCK_SIZE + HEAP_HEADER_SIZE
-            );
-        }
+    }
+}
+
+#[test]
+fn test_first_fit() {
+    let heap_vector = Box::new([0u8; HEAP_SIZE]);
+    let mut heap = Heap::empty();
+    heap.init(
+        &heap_vector[0] as *const u8 as usize,
+        &heap_vector[HEAP_SIZE - 1] as *const u8 as usize,
+    );
+
+    let root_header = heap.start_address;
+    let root_header_start_size = unsafe { (*root_header).size };
+
+    let malloc1 = heap.malloc(MIN_BLOCK_SIZE).unwrap();
+    let malloc2 = heap.malloc(MIN_BLOCK_SIZE).unwrap();
+    let malloc3 = heap.malloc(MIN_BLOCK_SIZE * 3).unwrap();
+    let malloc4 = heap.malloc(MIN_BLOCK_SIZE).unwrap();
+
+    unsafe {
+        assert!(heap.free(malloc1).is_ok());
+        assert!(heap.free(malloc3).is_ok());
+        let malloc5 = heap.malloc(MIN_BLOCK_SIZE * 2).unwrap();
+        let malloc1_header = unsafe { *Heap::get_header_ref_from_data_pointer(malloc1) };
+
+        // First free block stays empty
+        assert!(malloc1_header.free);
+
+        // New allocation takes the first fit aka. malloc3
+        assert_eq!(malloc5, malloc3);
+
+        // If no free slot could be found, append to the end
+        let malloc6 = heap.malloc(MIN_BLOCK_SIZE * 2).unwrap();
+        assert!(malloc6 > malloc4);
+
+        // Malloc7 takes slot of Malloc1
+        let malloc7 = heap.malloc(MIN_BLOCK_SIZE).unwrap();
+        assert_eq!(malloc1, malloc7);
     }
 }