Implements block merging and fixes a provenance bug.

3 files changed, 143 insertions, 40 deletions

diff --git a/crates/alloc_buddy/src/free_list.rs b/crates/alloc_buddy/src/free_list.rs
index 9b470fd..973ad2c 100644
--- a/crates/alloc_buddy/src/free_list.rs
+++ b/crates/alloc_buddy/src/free_list.rs
@@ -56,16 +56,7 @@ impl FreeList {
             unsafe {
                 // UNWRAP: In a valid circular doubly-linked list, the pointers are all non-null.
                 let node = (*self.0.as_ptr()).next.unwrap();
-                let prev = (*node.as_ptr()).prev.unwrap();
-                let next = (*node.as_ptr()).next.unwrap();
-
-                (*prev.as_ptr()).next = Some(next);
-                (*next.as_ptr()).prev = Some(prev);
-
-                (*node.as_ptr()).next = None;
-                (*node.as_ptr()).prev = None;
-                (*node.as_ptr()).magic = USED_NODE_MAGIC;
-                Some(node.cast())
+                Some(FreeListNode::unlink(node))
             }
         }
     }
@@ -161,4 +152,25 @@ impl FreeListNode {
         });
         ptr
     }
+
+    /// Unlinks the node from the circular doubly-linked list.
+    ///
+    /// # Safety
+    ///
+    /// The pointer must be a non-sentinel node inside a valid free list.
+    #[requires(unsafe { (*node.as_ptr()).magic } == FREE_NODE_MAGIC)]
+    #[requires(unsafe { (*node.as_ptr()).self_addr } == node.as_ptr() as usize)]
+    pub unsafe fn unlink(node: NonNull<FreeListNode>) -> NonNull<u8> {
+        // UNWRAP: In a valid circular doubly-linked list, the pointers are all non-null.
+        let prev = (*node.as_ptr()).prev.unwrap();
+        let next = (*node.as_ptr()).next.unwrap();
+
+        (*prev.as_ptr()).next = Some(next);
+        (*next.as_ptr()).prev = Some(prev);
+
+        (*node.as_ptr()).next = None;
+        (*node.as_ptr()).prev = None;
+        (*node.as_ptr()).magic = USED_NODE_MAGIC;
+        node.cast()
+    }
 }
diff --git a/crates/alloc_buddy/src/lib.rs b/crates/alloc_buddy/src/lib.rs
index 0ace886..02d45f8 100644
--- a/crates/alloc_buddy/src/lib.rs
+++ b/crates/alloc_buddy/src/lib.rs
@@ -69,7 +69,7 @@ pub struct BuddyAllocator<
     free_list_sentinels: NonNull<FreeListNode>,
 
     /// The metadata associated with each tree.
-    trees: &'allocator mut [Tree<'allocator, PAGE_SIZE, PAGE_SIZE_BITS>],
+    trees: &'allocator [Tree<PAGE_SIZE, PAGE_SIZE_BITS>],
 
     /// The shared bitset.
     bitset: &'allocator mut Bitset,
@@ -251,7 +251,7 @@ impl<
             // Fast-path: try allocating from the right size class's free list.
 
             // Find the corresponding tree and the offset of the subregion in the tree.
-            let (tree_index, offset) = self.tree_index_and_offset(ptr);
+            let (tree_index, offset, ptr) = self.tree_index_offset_and_ptr(ptr);
 
             // Mark the pointer as no longer being in the tree.
             let Ok(()) =
@@ -271,7 +271,7 @@ impl<
             let ptr = self.free_list(size_class_to_split).pop().unwrap();
 
             // Find the corresponding tree and the offset of the subregion in the tree.
-            let (tree_index, offset) = self.tree_index_and_offset(ptr);
+            let (tree_index, offset, ptr) = self.tree_index_offset_and_ptr(ptr);
 
             // Mark the pointer as no longer being in the tree.
             let Ok(()) = self.trees[tree_index].bitset_mark_as_absent(
@@ -282,10 +282,12 @@ impl<
                 panic!("the bit for tree {tree_index}, offset {offset} was not set in {self:#?}")
             };
 
-            // For each size class between the one we found and the one we want, split the page and
-            // add the other half to the free list (and mark it as present in the bitset).
+            // For each size class between the one we found and the one we want, split the
+            // subregion and add the other half to the free list (and mark it as present in the
+            // bitset).
             while size_class_to_split != size_class {
-                // Move to the next size class.
+                // Move to the next size class, since that's the size class of the subregion we'll
+                // be putting back.
                 size_class_to_split -= 1;
 
                 let offset_to_other_half = PAGE_SIZE << size_class_to_split;
@@ -295,13 +297,13 @@ impl<
                 let other_half_ptr = unsafe { ptr.add(offset_to_other_half) };
                 let other_half_offset = offset + offset_to_other_half;
 
-                // Mark the pointer as being in the tree.
+                // Mark the pointer as being in the free list.
                 let Ok(()) = self.trees[tree_index].bitset_mark_as_present(
                     self.bitset,
                     size_class_to_split,
                     other_half_offset,
                 ) else {
-                    panic!("the bit for tree {tree_index}, offset {offset} was already set in {self:#?}")
+                    panic!("the bit for tree {tree_index}, offset {other_half_offset} was already set in {self:#?}")
                 };
 
                 // Return the pointer to the appropriate free list.
@@ -326,10 +328,10 @@ impl<
     /// - `ptr` must have been returned by a call to `alloc` with the same `size_class`.
     /// - `ptr` must convey ownership over the entire region.
     #[requires(size_class < SIZE_CLASS_COUNT)]
-    pub unsafe fn dealloc(&mut self, ptr: NonNull<u8>, size_class: usize) {
+    pub unsafe fn dealloc(&mut self, ptr: NonNull<u8>, mut size_class: usize) {
         // Find the corresponding tree and the offset of the subregion in the tree.
-        let (tree_index, offset) = self.tree_index_and_offset(ptr);
-        let tree = &mut self.trees[tree_index];
+        let (tree_index, mut offset, mut ptr) = self.tree_index_offset_and_ptr(ptr);
+        let tree = &self.trees[tree_index];
 
         // Ensure that the memory is marked as not being in the free list.
         assert_eq!(
@@ -341,7 +343,33 @@ impl<
         assert!(size_class <= tree.size_class);
         while size_class < tree.size_class {
             let buddy_offset = offset ^ (PAGE_SIZE << size_class);
-            todo!("{buddy_offset:#x} {offset:#x}");
+            match tree.bitset_get(self.bitset, size_class, buddy_offset) {
+                SubregionStatus::InFreeList => {
+                    // If the buddy was present, we can remove it from the bitset and free list and
+                    // keep merging.
+                    let buddy_node = tree
+                        .base_ptr
+                        .unwrap()
+                        .cast::<FreeListNode>()
+                        .byte_add(buddy_offset);
+                    let buddy_ptr = FreeListNode::unlink(buddy_node);
+
+                    let Ok(()) = tree.bitset_mark_as_absent(self.bitset, size_class, buddy_offset)
+                    else {
+                        panic!("the bit for tree {tree_index}, offset {buddy_offset} was not set in {self:#?}")
+                    };
+
+                    // Continue merging with the lower of the two pointers (so it will still have a
+                    // power-of-two-aligned offset).
+                    ptr = ptr.min(buddy_ptr);
+                    offset = offset.min(buddy_offset);
+                    size_class += 1;
+                }
+                SubregionStatus::NotInFreeList => {
+                    // If the buddy was absent, we won't be able to do any more merging.
+                    break;
+                }
+            }
         }
 
         // Mark the pointer as being in the tree.
@@ -363,10 +391,11 @@ impl<
         unsafe { FreeList::from_sentinel(sentinel) }
     }
 
-    /// Returns the index of the tree this pointer was allocated in.
+    /// Returns the index of the tree this pointer was allocated in, as well as the offset of the
+    /// pointer in the tree and a copy of the pointer with the provenance of the tree.
     #[ensures(self.trees[ret.0].contains(ptr.as_ptr() as usize))]
-    #[ensures(ret.1 < (1 << self.trees[ret.0].size_class))]
-    fn tree_index_and_offset(&self, ptr: NonNull<u8>) -> (usize, usize) {
+    #[ensures(ret.1 < (PAGE_SIZE << self.trees[ret.0].size_class))]
+    fn tree_index_offset_and_ptr(&self, ptr: NonNull<u8>) -> (usize, usize, NonNull<u8>) {
         let addr = ptr.as_ptr() as usize;
         let tree_index = self
             .trees
@@ -378,8 +407,10 @@ impl<
                 i - 1
             }
         };
-        let offset = addr - self.trees[tree_index].base_addr();
-        (tree_index, offset)
+        let tree = &self.trees[tree_index];
+        let offset = addr - tree.base_addr();
+        let ptr = unsafe { tree.base_ptr.unwrap().cast().add(offset) };
+        (tree_index, offset, ptr)
     }
 }
 
@@ -405,13 +436,28 @@ impl<
             }
         }
 
+        struct Bitsets<'allocator, const PAGE_SIZE: usize, const PAGE_SIZE_BITS: usize>(
+            &'allocator [Tree<PAGE_SIZE, PAGE_SIZE_BITS>],
+            &'allocator Bitset,
+        );
+
+        impl<'allocator, const PAGE_SIZE: usize, const PAGE_SIZE_BITS: usize> fmt::Debug
+            for Bitsets<'allocator, PAGE_SIZE, PAGE_SIZE_BITS>
+        {
+            fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
+                fmt.debug_list()
+                    .entries(self.0.iter().map(|tree| tree.debug_bitset(self.1)))
+                    .finish()
+            }
+        }
+
         fmt.debug_struct("BuddyAllocator")
             .field(
                 "free_lists",
                 &FreeLists::<SIZE_CLASS_COUNT>(self.free_list_sentinels),
             )
             .field("trees", &self.trees)
-            .field("bitset", &self.bitset)
+            .field("bitsets", &Bitsets(self.trees, self.bitset))
             .finish()
     }
 }
diff --git a/crates/alloc_buddy/src/tree.rs b/crates/alloc_buddy/src/tree.rs
index b713285..787af21 100644
--- a/crates/alloc_buddy/src/tree.rs
+++ b/crates/alloc_buddy/src/tree.rs
@@ -1,13 +1,13 @@
 use crate::bitset::{Bitset, SubregionStatus, UnexpectedBitsetState};
 use contracts::requires;
-use core::ptr::NonNull;
+use core::{fmt, ptr::NonNull};
 
 /// A single region of the allocator. See the comment on the `crate::allocators::buddy` module for
 /// more information.
 ///
 /// This type is valid when zero-initialized.
 #[derive(Debug)]
-pub struct Tree<'buddy, const PAGE_SIZE: usize, const PAGE_SIZE_BITS: usize> {
+pub struct Tree<const PAGE_SIZE: usize, const PAGE_SIZE_BITS: usize> {
     /// The base address of the tree.
     pub base_ptr: Option<NonNull<[u8; PAGE_SIZE]>>,
 
@@ -16,13 +16,9 @@ pub struct Tree<'buddy, const PAGE_SIZE: usize, const PAGE_SIZE_BITS: usize> {
 
     /// The offset in the bitset to the bits responsible for this tree's pages.
     pub bitset_offset: usize,
-
-    phantom: core::marker::PhantomData<&'buddy mut [u8; PAGE_SIZE]>,
 }
 
-impl<'buddy, const PAGE_SIZE: usize, const PAGE_SIZE_BITS: usize>
-    Tree<'buddy, PAGE_SIZE, PAGE_SIZE_BITS>
-{
+impl<const PAGE_SIZE: usize, const PAGE_SIZE_BITS: usize> Tree<PAGE_SIZE, PAGE_SIZE_BITS> {
     /// Returns the base address of the tree.
     #[requires(self.base_ptr.is_some())]
     pub fn base_addr(&self) -> usize {
@@ -31,11 +27,11 @@ impl<'buddy, const PAGE_SIZE: usize, const PAGE_SIZE_BITS: usize>
 
     /// Reads a bit from the bitset.
     #[requires(size_class <= self.size_class)]
-    #[requires(offset_bytes <= (PAGE_SIZE << self.size_class))]
+    #[requires(offset_bytes < (PAGE_SIZE << (self.size_class + 1)))]
     #[requires(offset_bytes.trailing_zeros() as usize >= PAGE_SIZE_BITS + size_class)]
     pub fn bitset_get(
         &self,
-        bitset: &mut Bitset,
+        bitset: &Bitset,
         size_class: usize,
         offset_bytes: usize,
     ) -> SubregionStatus {
@@ -45,7 +41,7 @@ impl<'buddy, const PAGE_SIZE: usize, const PAGE_SIZE_BITS: usize>
     /// Returns the index of a bit in the bitset that corresponds to the given size class and
     /// offset.
     #[requires(size_class <= self.size_class)]
-    #[requires(offset_bytes <= (PAGE_SIZE << self.size_class))]
+    #[requires(offset_bytes < (PAGE_SIZE << (self.size_class + 1)))]
     #[requires(offset_bytes.trailing_zeros() as usize >= PAGE_SIZE_BITS + size_class)]
     fn bitset_index(&self, size_class: usize, offset_bytes: usize) -> usize {
         // We store the largest size classes first in the bitset. Count how many we are away from
@@ -62,7 +58,7 @@ impl<'buddy, const PAGE_SIZE: usize, const PAGE_SIZE_BITS: usize>
 
     /// Changes a bit in the bitset from `InFreeList` to `NotInFreeList`.
     #[requires(size_class <= self.size_class)]
-    #[requires(offset_bytes <= (PAGE_SIZE << self.size_class))]
+    #[requires(offset_bytes < (PAGE_SIZE << (self.size_class + 1)))]
     #[requires(offset_bytes.trailing_zeros() as usize >= PAGE_SIZE_BITS + size_class)]
     pub fn bitset_mark_as_absent(
         &self,
@@ -79,7 +75,7 @@ impl<'buddy, const PAGE_SIZE: usize, const PAGE_SIZE_BITS: usize>
 
     /// Changes a bit in the bitset from `NotInFreeList` to `InFreeList`.
     #[requires(size_class <= self.size_class)]
-    #[requires(offset_bytes <= (PAGE_SIZE << self.size_class))]
+    #[requires(offset_bytes < (PAGE_SIZE << (self.size_class + 1)))]
     #[requires(offset_bytes.trailing_zeros() as usize >= PAGE_SIZE_BITS + size_class)]
     pub fn bitset_mark_as_present(
         &self,
@@ -101,4 +97,53 @@ impl<'buddy, const PAGE_SIZE: usize, const PAGE_SIZE_BITS: usize>
         let tree_addr_hi = tree_addr_lo + (PAGE_SIZE << self.size_class);
         (tree_addr_lo..tree_addr_hi).contains(&addr)
     }
+
+    /// Formats the region of the bitset corresponding to this tree.
+    pub fn debug_bitset<'a>(&'a self, bitset: &'a Bitset) -> impl 'a + fmt::Debug {
+        struct BitsetSlice<'a, const PAGE_SIZE: usize, const PAGE_SIZE_BITS: usize>(
+            &'a Tree<PAGE_SIZE, PAGE_SIZE_BITS>,
+            &'a Bitset,
+            usize,
+        );
+
+        impl<'a, const PAGE_SIZE: usize, const PAGE_SIZE_BITS: usize> fmt::Debug
+            for BitsetSlice<'a, PAGE_SIZE, PAGE_SIZE_BITS>
+        {
+            fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
+                for i in 0..(1 << (self.0.size_class - self.2)) {
+                    let offset_bytes = i << (PAGE_SIZE_BITS + self.2);
+                    let bit = match self.0.bitset_get(self.1, self.2, offset_bytes) {
+                        SubregionStatus::NotInFreeList => '0',
+                        SubregionStatus::InFreeList => '1',
+                    };
+                    write!(fmt, "{}", bit)?;
+                    for _ in 0..(1 << self.2) - 1 {
+                        write!(fmt, " ")?;
+                    }
+                }
+                Ok(())
+            }
+        }
+
+        struct BitsetTree<'a, const PAGE_SIZE: usize, const PAGE_SIZE_BITS: usize>(
+            &'a Tree<PAGE_SIZE, PAGE_SIZE_BITS>,
+            &'a Bitset,
+        );
+
+        impl<'a, const PAGE_SIZE: usize, const PAGE_SIZE_BITS: usize> fmt::Debug
+            for BitsetTree<'a, PAGE_SIZE, PAGE_SIZE_BITS>
+        {
+            fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
+                fmt.debug_list()
+                    .entries(
+                        (0..=self.0.size_class)
+                            .rev()
+                            .map(|size_class| BitsetSlice(self.0, self.1, size_class)),
+                    )
+                    .finish()
+            }
+        }
+
+        BitsetTree(self, bitset)
+    }
 }