From f55f7f9275feb2347ecffc2b747f1e329f3587aa Mon Sep 17 00:00:00 2001
From: James Munns <james@onevariable.com>
Date: Fri, 1 Jul 2022 22:35:34 +0200
Subject: [PATCH] Fix free logic

---
 src/hole.rs | 51 +++++++++++++++++++++++++++++----------------------
 src/lib.rs  |  8 ++++++--
 2 files changed, 35 insertions(+), 24 deletions(-)

diff --git a/src/hole.rs b/src/hole.rs
index 5c25649..ad1b8bf 100644
--- a/src/hole.rs
+++ b/src/hole.rs
@@ -234,7 +234,6 @@ impl HoleList {
         }
     }
 
-    #[cfg(test)]
     #[allow(dead_code)]
     pub(crate) fn debug(&mut self) {
         if let Some(cursor) = self.cursor() {
@@ -421,36 +420,44 @@ impl Cursor {
     }
 
     fn try_insert_after(&mut self, mut node: NonNull<Hole>) -> Result<(), ()> {
-        if self.hole < node {
-            let node_u8 = node.as_ptr().cast::<u8>();
-            let node_size = unsafe { node.as_ref().size };
-            let hole_u8 = self.hole.as_ptr().cast::<u8>();
-            let hole_size = self.current().size;
+        let node_u8 = node.as_ptr().cast::<u8>();
+        let node_size = unsafe { node.as_ref().size };
 
-            // Does hole overlap node?
-            assert!(
-                hole_u8.wrapping_add(hole_size) <= node_u8,
-                "Freed node aliases existing hole! Bad free?",
-            );
-
-            // If we have a next, does the node overlap next?
-            if let Some(next) = self.current().next.as_ref() {
+        // If we have a next, does the node overlap next?
+        if let Some(next) = self.current().next.as_ref() {
+            if node < *next {
                 let node_u8 = node_u8 as *const u8;
                 assert!(
                     node_u8.wrapping_add(node_size) <= next.as_ptr().cast::<u8>(),
                     "Freed node aliases existing hole! Bad free?",
                 );
+            } else {
+                // The new hole isn't between current and next.
+                return Err(());
             }
+        }
 
-            // All good! Let's insert that after.
-            unsafe {
-                let maybe_next = self.hole.as_mut().next.replace(node);
-                node.as_mut().next = maybe_next;
-            }
-            Ok(())
-        } else {
-            Err(())
+        // At this point, we either have no "next" pointer, or the hole is
+        // between current and "next". The following assert can only trigger
+        // if we've gotten our list out of order.
+        debug_assert!(self.hole < node, "Hole list out of order?");
+
+        let hole_u8 = self.hole.as_ptr().cast::<u8>();
+        let hole_size = self.current().size;
+
+        // Does hole overlap node?
+        assert!(
+            hole_u8.wrapping_add(hole_size) <= node_u8,
+            "Freed node aliases existing hole! Bad free?",
+        );
+
+        // All good! Let's insert that after.
+        unsafe {
+            let maybe_next = self.hole.as_mut().next.replace(node);
+            node.as_mut().next = maybe_next;
         }
+
+        Ok(())
     }
 
     // Merge the current node with up to n following nodes
diff --git a/src/lib.rs b/src/lib.rs
index d257629..e7f0b7c 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -5,7 +5,6 @@
 )]
 #![no_std]
 
-#[cfg(test)]
 #[macro_use]
 extern crate std;
 
@@ -160,10 +159,14 @@ impl Heap {
     pub fn allocate_first_fit(&mut self, layout: Layout) -> Result<NonNull<u8>, ()> {
         match self.holes.allocate_first_fit(layout) {
             Ok((ptr, aligned_layout)) => {
+                self.holes.debug();
                 self.used += aligned_layout.size();
                 Ok(ptr)
             }
-            Err(err) => Err(err),
+            Err(err) => {
+                println!("ERR");
+                Err(err)
+            }
         }
     }
 
@@ -180,6 +183,7 @@ impl Heap {
     /// identical layout. Undefined behavior may occur for invalid arguments.
     pub unsafe fn deallocate(&mut self, ptr: NonNull<u8>, layout: Layout) {
         self.used -= self.holes.deallocate(ptr, layout).size();
+        self.holes.debug();
     }
 
     /// Returns the bottom address of the heap.