1 files changed, 100 insertions, 0 deletions
diff --git a/crates/alloc_buddy/src/tree.rs b/crates/alloc_buddy/src/tree.rs
new file mode 100644
index 0000000..72ee466
--- /dev/null
+++ b/crates/alloc_buddy/src/tree.rs
@@ -0,0 +1,100 @@
+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.
+///
+/// This type is valid when zero-initialized.
+#[derive(Debug)]
+pub struct Tree<'buddy, const PAGE_SIZE: usize, const PAGE_SIZE_BITS: usize> {
+    /// The base address of the tree.
+    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,
+
+    /// 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>
+{
+    /// 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))]
+    #[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 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
+    }
+
+    /// 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)
+    }
+
+    /// 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)
+    }
+
+    /// 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_addr();
+        let tree_addr_hi = tree_addr_lo + (PAGE_SIZE << self.size_class);
+        (tree_addr_lo..tree_addr_hi).contains(&addr)
+    }
+}