Minor cleanup.

1 files changed, 10 insertions, 2 deletions

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)
     }