1 files changed, 64 insertions, 102 deletions

diff --git a/crates/buddy_allocator/src/tree.rs b/crates/buddy_allocator/src/tree.rs
index 1673e26..7c67f3a 100644
--- a/crates/buddy_allocator/src/tree.rs
+++ b/crates/buddy_allocator/src/tree.rs
@@ -1,3 +1,6 @@
+use crate::bitvec::{Bitset, SubregionStatus};
+use contracts::requires;
+
 /// A single region of the allocator. See the comment on the `crate::allocators::buddy` module for
 /// more information.
 ///
@@ -13,118 +16,77 @@ 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,
 
-    // /// The bitset used to track whether subregion are allocated or not in this tree.
-    // pub bitset: &'buddy mut BitVec,
     phantom: core::marker::PhantomData<&'buddy mut [u8; PAGE_SIZE]>,
 }
 
-/*
-use crate::{
-    bitvec::BitVec,
-    stripe::{FreeListNode, Stripe},
-    MAX_ORDER, PAGE_SIZE_BITS,
-};
-use contracts::requires;
-
-/// A single region of the allocator. See the comment on the `crate::allocators::buddy` module for
-/// more information.
-pub struct Tree<'buddy> {
-    /// The base address of the tree.
-    pub base: *const (),
-
-    /// The log2 of the number of pages in the region represented by the tree.
-    pub log2_page_count: usize,
-
-    /// The array of sentinel nodes of the free-list for this node. As an invariant of the type,
-    /// there are no live references to any node in any list in this array, and there are
-    /// MAX_ORDER + 1 nodes.
-    pub free_lists: NonNull<FreeListNode>,
-
-    /// The bitset used to track whether subregion are allocated or not in this tree.
-    pub bitset: &'buddy mut BitVec,
-}
-
-impl<'buddy> Tree<'buddy> {
-    /// Tries to allocate a subregion with the given order, possibly splitting a larger subregion
-    /// in order to so so.
-    #[requires(order <= MAX_ORDER)]
-    pub fn alloc(&mut self, order: usize) -> Option<NonNull<FreeListNode>> {
-        if let Some(ptr) = self.stripe(order).pop() {
-            Some(ptr)
-        } else if order == MAX_ORDER {
-            None
-        } else {
-            // Get a larger region.
-            let ptr = self.alloc(order + 1)?;
-
-            // Get a pointer to the higher half.
-            //
-            // SAFETY: This has to be in-bounds, it's part of the same allocation!
-            let higher_half = unsafe { ptr.byte_add(1 << (PAGE_SIZE_BITS + order)) };
-
-            // Put the higher half in the current buddy's stripe.
-            //
-            // SAFETY: The higher half is from this region, not in the higher stripe, and of the
-            // right size.
-            unsafe {
-                self.stripe(order).push(higher_half);
-            }
-
-            // Return the pointer.
-            Some(ptr)
-        }
+impl<'buddy, const PAGE_SIZE: usize, const PAGE_SIZE_BITS: usize>
+    Tree<'buddy, PAGE_SIZE, PAGE_SIZE_BITS>
+{
+    /// Reads a bit from the bitset.
+    #[requires(size_class <= self.size_class)]
+    #[requires(offset_bytes <= (PAGE_SIZE << self.size_class))]
+    #[requires(offset_bytes.trailing_zeros() as usize >= PAGE_SIZE_BITS + size_class)]
+    pub fn bitset_get(
+        &self,
+        bitset: &mut Bitset,
+        size_class: usize,
+        offset_bytes: usize,
+    ) -> SubregionStatus {
+        bitset.get(self.bitset_index(size_class, offset_bytes))
     }
 
-    /// Returns the stripe with the given order.
-    #[requires(order <= MAX_ORDER)]
-    fn stripe<'stripe>(&'stripe mut self, order: usize) -> Stripe<'stripe> {
-        // TODO: There should be some smart bitwise-math version of this...
-        let mut bitset_offset = 0;
-        for i in 0..order {
-            bitset_offset += (1 << self.log2_page_count) >> i;
-        }
-
-        Stripe {
-            base: self.base,
-            order,
-            // SAFETY: order is constrained to be in-bounds.
-            free_list: unsafe { self.free_lists.add(order) },
-            bitset: self.bitset,
-            bitset_offset,
-        }
+    /// 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.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
+        // the largest.
+        let skipped_size_classes = self.size_class - size_class;
+        let bits_skipped_for_size_class = (1 << skipped_size_classes) - 1;
+
+        // Next, find our index in the size class.
+        let bits_skipped_for_index = offset_bytes >> (PAGE_SIZE_BITS + size_class);
+
+        // The sum of those two is simply our index.
+        bits_skipped_for_size_class + bits_skipped_for_index
     }
-}
 
-/// Evil bitwise version of the reasonable loop to compute the `bitset_offset` of a stripe.
-#[requires(log2_page_count < usize::BITS as usize)]
-#[requires(order < usize::BITS as usize)]
-fn compute_bitset_offset(log2_page_count: usize, order: usize) -> usize {
-    let ones = |i: usize| !(usize::MAX << i);
-
-    if order > log2_page_count + 1 {
-        ones(log2_page_count + 1)
-    } else {
-        ones(order).wrapping_shl((log2_page_count + 1 - order) as u32)
+    /// 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.trailing_zeros() as usize >= PAGE_SIZE_BITS + size_class)]
+    pub fn bitset_mark_as_absent(
+        &self,
+        bitset: &mut Bitset,
+        size_class: usize,
+        offset_bytes: usize,
+    ) {
+        let index = self.bitset_index(size_class, offset_bytes);
+        assert_eq!(bitset.get(index), SubregionStatus::InFreeList);
+        bitset.set(index, SubregionStatus::NotInFreeList)
     }
-}
 
-#[test]
-fn compute_bitset_offset_test() {
-    fn compute_bitset_offset_loop(log2_page_count: usize, order: usize) -> usize {
-        let mut bitset_offset = 0;
-        for i in 0..order {
-            bitset_offset += (1 << log2_page_count) >> i;
-        }
-        bitset_offset
+    /// 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.trailing_zeros() as usize >= PAGE_SIZE_BITS + size_class)]
+    pub fn bitset_mark_as_present(
+        &self,
+        bitset: &mut Bitset,
+        size_class: usize,
+        offset_bytes: usize,
+    ) {
+        let index = self.bitset_index(size_class, offset_bytes);
+        assert_eq!(bitset.get(index), SubregionStatus::NotInFreeList);
+        bitset.set(index, SubregionStatus::InFreeList)
     }
 
-    for log2_page_count in 0..64 {
-        for order in 0..64 {
-            assert_eq!(
-                compute_bitset_offset(log2_page_count, order),
-                compute_bitset_offset_loop(log2_page_count, order),
-            );
-        }
+    /// Returns whether the tree contains an address.
+    pub fn contains(&self, addr: usize) -> bool {
+        let tree_addr_lo = self.base_ptr as usize;
+        let tree_addr_hi = tree_addr_lo + (PAGE_SIZE << self.size_class);
+        (tree_addr_lo..tree_addr_hi).contains(&addr)
     }
 }
-*/