Minor cleanup.

4 files changed, 24 insertions, 165 deletions

diff --git a/crates/buddy_allocator/src/lib.rs b/crates/buddy_allocator/src/lib.rs
index f0c8daa..c543ca6 100644
--- a/crates/buddy_allocator/src/lib.rs
+++ b/crates/buddy_allocator/src/lib.rs
@@ -48,8 +48,6 @@ mod bitvec;
 mod free_list;
 mod tree;
 
-// mod stripe;
-
 use crate::{
     bitvec::{Bitset, SubregionStatus},
     free_list::{FreeList, FreeListNode},
@@ -176,7 +174,7 @@ impl<
         while let Some((ptr, mut len_pages)) = free_list_alloc.pop() {
             while len_pages > 0 {
                 let size_class = (len_pages.trailing_zeros() as usize).min(SIZE_CLASS_COUNT - 1);
-                trees[i].base_ptr = ptr.cast().as_ptr();
+                trees[i].base_ptr = Some(ptr.cast());
                 trees[i].size_class = size_class;
                 i += 1;
                 len_pages -= 1 << size_class;
@@ -192,7 +190,7 @@ impl<
         // Slice the trees slice to fit how many trees actually remained, then sort them by
         // address.
         let trees = &mut trees[..i];
-        trees.sort_by_key(|tree| tree.base_ptr as usize);
+        trees.sort_by_key(|tree| tree.base_addr());
 
         // Compute the number of bits that belong to each tree and assign them.
         let mut bitset_offset = 0;
@@ -213,16 +211,12 @@ impl<
             unsafe { metadata.cast::<u8>().add(metadata.len()).as_ptr() } as usize;
         for tree in &mut *trees {
             // SAFETY: This is the one-past-the-end address.
-            let tree_addr_hi = unsafe { tree.base_ptr.add(1 << tree.size_class) } as usize;
+            let tree_addr_hi = tree.base_addr() + (PAGE_SIZE << tree.size_class);
             if metadata_addr_hi == tree_addr_hi {
                 // If the metadata was present inside the tree... TODO.
                 std::eprintln!("TODO");
             } else {
                 // Otherwise, we can take the whole range of the tree and add it to the free list.
-                //
-                // UNWRAP: The pointer must have been non-null when it was in the free-list
-                // allocator.
-                let ptr = NonNull::new(tree.base_ptr as *mut u8).unwrap();
 
                 // SAFETY: There are as many sentinels as size classes, so this will be in-bounds.
                 let sentinel = unsafe { free_list_sentinels.add(tree.size_class) };
@@ -233,7 +227,7 @@ impl<
                 // SAFETY: This region is not owned or borrowed by anything else (the only thing it
                 // could be would be the metadata, but that's not in this tree), is of the correct
                 // size, and is composed of untyped memory.
-                unsafe { free_list.push(ptr) };
+                unsafe { free_list.push(tree.base_ptr.unwrap().cast()) };
 
                 // Then, set a bit in the bitset to say that this region is present.
                 tree.bitset_mark_as_present(bitset, tree.size_class, 0);
@@ -276,7 +270,7 @@ 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>, mut size_class: usize) {
+    pub unsafe fn dealloc(&mut self, ptr: NonNull<u8>, 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];
@@ -288,12 +282,17 @@ impl<
         );
 
         // Start combining size classes.
+        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}");
         }
 
-        todo!("{tree:#?}")
+        // Mark the pointer as being in the tree.
+        tree.bitset_mark_as_present(&mut self.bitset, size_class, offset);
+
+        // Return the pointer to the appropriate free list.
+        self.free_list(size_class).push(ptr);
     }
 
     /// Returns the free list with the given size class.
@@ -313,7 +312,7 @@ impl<
         let addr = ptr.as_ptr() as usize;
         let tree_index = self
             .trees
-            .binary_search_by_key(&addr, |tree| tree.base_ptr as usize);
+            .binary_search_by_key(&addr, |tree| tree.base_addr());
         let tree_index = match tree_index {
             Ok(i) => i,
             Err(i) => {
@@ -321,7 +320,7 @@ impl<
                 i - 1
             }
         };
-        let offset = addr - self.trees[tree_index].base_ptr as usize;
+        let offset = addr - self.trees[tree_index].base_addr();
         (tree_index, offset)
     }
 }
diff --git a/crates/buddy_allocator/src/stripe.rs b/crates/buddy_allocator/src/stripe.rs
deleted file mode 100644
index 2ff686b..0000000
--- a/crates/buddy_allocator/src/stripe.rs
+++ /dev/null
@@ -1,148 +0,0 @@
-use crate::{bitvec::BitVec, PAGE_SIZE_BITS};
-use core::ptr::NonNull;
-
-/// A single size class for a single region of the allocator. See the comment on the
-/// `crate::allocators::buddy` module for more information.
-pub struct Stripe<'buddy> {
-    /// The base address of the tree this stripe is a part of.
-    pub base: *const (),
-
-    /// The order of the stripe. Order `n` means the subregions are `2ⁿ` pages in size.
-    pub order: usize,
-
-    /// The sentinel node of the free-list for this node. As an invariant of the type, there are no
-    /// live references to any node in this list.
-    pub free_list: NonNull<FreeListNode>,
-
-    /// The bitset used to track whether a given subregion is allocated or not. A `true` bit
-    /// corresponds to an allocated subregion.
-    pub bitset: &'buddy mut BitVec,
-
-    /// The offset from the start of the bitset to the region used by this stripe.
-    pub bitset_offset: usize,
-}
-
-impl<'buddy> Stripe<'buddy> {
-    /// Returns the buddy of the given pointer.
-    ///
-    /// ## Safety
-    ///
-    /// - The pointer must actually be part of this region.
-    unsafe fn buddy_of(&self, ptr: NonNull<FreeListNode>) -> NonNull<FreeListNode> {
-        let index = self.buddy_of_index(self.index_of(ptr));
-        let addr = self.base as usize + (index << (PAGE_SIZE_BITS + self.order));
-        NonNull::new_unchecked(addr as *mut FreeListNode)
-    }
-
-    /// Returns the buddy of the given index.
-    fn buddy_of_index(&self, index: usize) -> usize {
-        index ^ (1 << (PAGE_SIZE_BITS + self.order))
-    }
-
-    /// Returns the index the given pointer should have in the BitVec.
-    fn index_of(&self, ptr: NonNull<FreeListNode>) -> usize {
-        (ptr.as_ptr() as usize - self.base as usize) >> (PAGE_SIZE_BITS + self.order)
-    }
-
-    /// Pops a subregion from the free-list.
-    pub fn pop(&mut self) -> Option<NonNull<FreeListNode>> {
-        // SAFETY: The `free_list` is guaranteed to be valid by the invariants of the buddy
-        // allocator. Retrieving the next pointer doesn't, on its own, break aliasing rules.
-        let next = unsafe { self.free_list.read().next };
-
-        // If the sentinel and the next pointer refer to the same spot, the free-list was empty, so
-        // we can't pop from it.
-        if self.free_list == next {
-            return None;
-        }
-
-        // Otherwise, remove the node from the free-list.
-        unsafe {
-            FreeListNode::remove(next);
-        }
-
-        // Finally, mark the node as allocated in the bitvec.
-        let index = self.index_of(next);
-        assert!(self.bitset.get(self.bitset_offset + index));
-        self.bitset.set(self.bitset_offset + index, true);
-
-        Some(next)
-    }
-
-    /// Pushes a subregion back into the free-list.
-    ///
-    /// # Safety
-    ///
-    /// - There must be no live references to `subregion`.
-    /// - `subregion` must not be a member of any list.
-    pub unsafe fn push(&mut self, subregion: NonNull<FreeListNode>) {
-        // Insert the subregion as the first element of the free-list.
-        //
-        // SAFETY: The free-list is guaranteed to be valid by the invariants of the buddy
-        // allocator.
-        unsafe {
-            FreeListNode::insert(self.free_list, subregion);
-        }
-
-        // Mark the node as unallocated in the bitvec.
-        let index = self.index_of(subregion);
-        assert!(self.bitset.get(self.bitset_offset + index));
-        self.bitset.set(self.bitset_offset + index, false);
-    }
-
-    /// Pushes a subregion into the free-list for the first time.
-    ///
-    /// # Safety
-    ///
-    /// - There must be no live references to `subregion`.
-    /// - `subregion` must not be a member of any list.
-    pub unsafe fn push_initial(&mut self, subregion: NonNull<FreeListNode>) {
-        // Insert the subregion as the first element of the free-list.
-        //
-        // SAFETY: The free-list is guaranteed to be valid by the invariants of the buddy
-        // allocator.
-        unsafe {
-            FreeListNode::insert(self.free_list, subregion);
-        }
-
-        // Mark the node as unallocated in the bitvec.
-        let index = self.index_of(subregion);
-        assert!(self.bitset.get(self.bitset_offset + index));
-        self.bitset.set(self.bitset_offset + index, false);
-    }
-}
-
-pub struct FreeListNode {
-    next: NonNull<FreeListNode>,
-    prev: NonNull<FreeListNode>,
-}
-
-impl FreeListNode {
-    /// Inserts `new` after `prev`, initializing it. `prev` may be the sentinel.
-    ///
-    /// # Safety
-    ///
-    /// - There must be no live references to any node in the list that includes `prev`, including
-    ///   the sentinel.
-    /// - There must be no live references to `new`.
-    /// - `new` must not be a member of any list.
-    pub unsafe fn insert(prev: NonNull<FreeListNode>, new: NonNull<FreeListNode>) {
-        let next = prev.read().next;
-        (*prev.as_ptr()).next = new;
-        (*next.as_ptr()).prev = new;
-        new.write(FreeListNode { next, prev });
-    }
-
-    /// Removes this node from the free list it is a part of.
-    ///
-    /// # Safety
-    ///
-    /// - The pointer must point to a node that is part of a free list.
-    /// - There must be no live references to any node in the list, including the sentinel.
-    /// - The pointer must not point to the sentinel node.
-    pub unsafe fn remove(ptr: NonNull<FreeListNode>) {
-        let FreeListNode { next, prev } = ptr.read();
-        (*next.as_ptr()).prev = prev;
-        (*prev.as_ptr()).next = next;
-    }
-}
diff --git a/crates/buddy_allocator/src/tree.rs b/crates/buddy_allocator/src/tree.rs
index 7c67f3a..72ee466 100644
--- a/crates/buddy_allocator/src/tree.rs
+++ b/crates/buddy_allocator/src/tree.rs
@@ -1,5 +1,6 @@
 use crate::bitvec::{Bitset, SubregionStatus};
 use contracts::requires;
+use core::ptr::NonNull;
 
 /// A single region of the allocator. See the comment on the `crate::allocators::buddy` module for
 /// more information.
@@ -8,7 +9,7 @@ use contracts::requires;
 #[derive(Debug)]
 pub struct Tree<'buddy, const PAGE_SIZE: usize, const PAGE_SIZE_BITS: usize> {
     /// The base address of the tree.
-    pub base_ptr: *const [u8; PAGE_SIZE],
+    pub base_ptr: Option<NonNull<[u8; PAGE_SIZE]>>,
 
     /// The log2 of the number of pages in the region represented by the tree.
     pub size_class: usize,
@@ -22,6 +23,12 @@ pub struct Tree<'buddy, const PAGE_SIZE: usize, const PAGE_SIZE_BITS: usize> {
 impl<'buddy, const PAGE_SIZE: usize, const PAGE_SIZE_BITS: usize>
     Tree<'buddy, PAGE_SIZE, PAGE_SIZE_BITS>
 {
+    /// Returns the base address of the tree.
+    #[requires(self.base_ptr.is_some())]
+    pub fn base_addr(&self) -> usize {
+        self.base_ptr.unwrap().as_ptr() as usize
+    }
+
     /// Reads a bit from the bitset.
     #[requires(size_class <= self.size_class)]
     #[requires(offset_bytes <= (PAGE_SIZE << self.size_class))]
@@ -84,8 +91,9 @@ impl<'buddy, const PAGE_SIZE: usize, const PAGE_SIZE_BITS: usize>
     }
 
     /// Returns whether the tree contains an address.
+    #[requires(self.base_ptr.is_some())]
     pub fn contains(&self, addr: usize) -> bool {
-        let tree_addr_lo = self.base_ptr as usize;
+        let tree_addr_lo = self.base_addr();
         let tree_addr_hi = tree_addr_lo + (PAGE_SIZE << self.size_class);
         (tree_addr_lo..tree_addr_hi).contains(&addr)
     }
diff --git a/crates/buddy_allocator/tests/hosted_test.rs b/crates/buddy_allocator/tests/hosted_test.rs
index 89398ce..0e502b3 100644
--- a/crates/buddy_allocator/tests/hosted_test.rs
+++ b/crates/buddy_allocator/tests/hosted_test.rs
@@ -20,7 +20,7 @@ proptest! {
 #[test]
 fn test_simple_scenario() {
     let scenario = Scenario {
-        range_sizes: vec![5],
+        range_sizes: vec![7],
         actions: vec![
             Action::Debug,
             Action::Alloc {