Implement size class splitting.

2 files changed, 57 insertions, 4 deletions

diff --git a/crates/alloc_buddy/src/lib.rs b/crates/alloc_buddy/src/lib.rs
index fc5d17e..0ace886 100644
--- a/crates/alloc_buddy/src/lib.rs
+++ b/crates/alloc_buddy/src/lib.rs
@@ -252,17 +252,69 @@ impl<
 
             // 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];
 
             // Mark the pointer as no longer being in the tree.
-            let Ok(()) = tree.bitset_mark_as_absent(self.bitset, size_class, offset) else {
+            let Ok(()) =
+                self.trees[tree_index].bitset_mark_as_absent(self.bitset, size_class, offset)
+            else {
                 panic!("the bit for tree {tree_index}, offset {offset} was not set in {self:#?}")
             };
 
             // Return the pointer.
             Ok(ptr)
+        } else if let Some(mut size_class_to_split) =
+            (size_class + 1..SIZE_CLASS_COUNT).find(|&i| !self.free_list(i).is_empty())
+        {
+            // If a larger size class had an available chunk, we'll need to split it.
+            //
+            // UNWRAP: We checked that the list was non-empty.
+            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);
+
+            // Mark the pointer as no longer being in the tree.
+            let Ok(()) = self.trees[tree_index].bitset_mark_as_absent(
+                self.bitset,
+                size_class_to_split,
+                offset,
+            ) else {
+                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).
+            while size_class_to_split != size_class {
+                // Move to the next size class.
+                size_class_to_split -= 1;
+
+                let offset_to_other_half = PAGE_SIZE << size_class_to_split;
+
+                // SAFETY: Thanks to our power-of-two policy, this must still fit within the bounds
+                // of the tree.
+                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.
+                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:#?}")
+                };
+
+                // Return the pointer to the appropriate free list.
+                //
+                // SAFETY: We know the higher half was valid, and we're semantically restricting
+                // the size of our ownership down to the lower half.
+                unsafe { self.free_list(size_class_to_split).push(other_half_ptr) };
+            }
+
+            // Return the pointer that finally remains.
+            Ok(ptr)
         } else {
-            // Try finding a larger chunk of memory in a larger size class.
+            // Otherwise, there was no way to serve the allocation.
             Err(AllocError)
         }
     }
diff --git a/crates/alloc_buddy/tests/hosted_test.rs b/crates/alloc_buddy/tests/hosted_test.rs
index 1afe263..d515f6e 100644
--- a/crates/alloc_buddy/tests/hosted_test.rs
+++ b/crates/alloc_buddy/tests/hosted_test.rs
@@ -20,7 +20,7 @@ proptest! {
 #[test]
 fn test_simple_scenario() {
     let scenario = Scenario {
-        range_sizes: vec![7],
+        range_sizes: vec![10],
         actions: vec![
             Action::Debug,
             Action::Alloc {
@@ -28,6 +28,7 @@ fn test_simple_scenario() {
                 size_class: 0,
             },
             Action::Debug,
+            Action::Dealloc { index: 0 },
         ],
     };
     scenario.run(false)