Fix a bitset indexing bug, change a panic to be more informative.

3 files changed, 58 insertions, 37 deletions

diff --git a/crates/alloc_buddy/src/bitvec.rs b/crates/alloc_buddy/src/bitset.rs
index 5cabc5e..b5f12d2 100644
--- a/crates/alloc_buddy/src/bitvec.rs
+++ b/crates/alloc_buddy/src/bitset.rs
@@ -1,22 +1,12 @@
 use contracts::requires;
-use core::{fmt, mem::transmute};
+use core::fmt;
 
 /// A fixed-length vector of bits used for determining whether a subregion is in the free lists or
 /// not.
 pub struct Bitset([usize]);
 
 impl Bitset {
-    fn from_mut(words: &mut [usize]) -> &mut Bitset {
-        // SAFETY: The types have a newtype relationship.
-        unsafe { transmute(words) }
-    }
-
-    fn from_ref(words: &[usize]) -> &Bitset {
-        // SAFETY: The types have a newtype relationship.
-        unsafe { transmute(words) }
-    }
-
-    /// Retrieves the value of a bit from the Bitset.
+    /// Retrieves the value of a bit from the bitset.
     #[requires(i < self.len())]
     pub fn get(&self, i: usize) -> SubregionStatus {
         let word_index = i / usize::BITS as usize;
@@ -29,17 +19,28 @@ impl Bitset {
         }
     }
 
-    /// Returns whether the Bitset is empty.
-    pub fn is_empty(&self) -> bool {
-        self.0.is_empty()
-    }
-
-    /// Returns the number of bits in the Bitset.
+    /// Returns the number of bits in the bitset.
     pub fn len(&self) -> usize {
         self.0.len() * usize::BITS as usize
     }
 
-    /// Sets the value of a bit in the Bitset.
+    /// Replaces a bit in the bitset.
+    #[requires(i < self.len())]
+    pub fn replace(
+        &mut self,
+        i: usize,
+        status_before: SubregionStatus,
+        status_after: SubregionStatus,
+    ) -> Result<(), UnexpectedBitsetState> {
+        if self.get(i) == status_before {
+            self.set(i, status_after);
+            Ok(())
+        } else {
+            Err(UnexpectedBitsetState)
+        }
+    }
+
+    /// Sets the value of a bit in the bitset.
     #[requires(i < self.len())]
     pub fn set(&mut self, i: usize, status: SubregionStatus) {
         let word_index = i / usize::BITS as usize;
@@ -59,13 +60,23 @@ impl Bitset {
 impl fmt::Debug for Bitset {
     fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
         write!(fmt, "Bitset(")?;
-        for word in &self.0 {
-            write!(fmt, "{word:064b}")?;
+        for i in 0..self.len() {
+            write!(
+                fmt,
+                "{}",
+                match self.get(i) {
+                    SubregionStatus::NotInFreeList => '0',
+                    SubregionStatus::InFreeList => '1',
+                }
+            )?;
         }
         write!(fmt, ")")
     }
 }
 
+/// An error returned when the `replace` method does not find the expected bit.
+pub struct UnexpectedBitsetState;
+
 /// The status of a subregion, as tracked by `Bitset`.
 #[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd)]
 pub enum SubregionStatus {
diff --git a/crates/alloc_buddy/src/lib.rs b/crates/alloc_buddy/src/lib.rs
index d0d4818..fc5d17e 100644
--- a/crates/alloc_buddy/src/lib.rs
+++ b/crates/alloc_buddy/src/lib.rs
@@ -44,12 +44,12 @@
 
 extern crate std; // TODO: Remove me
 
-mod bitvec;
+mod bitset;
 mod free_list;
 mod tree;
 
 use crate::{
-    bitvec::{Bitset, SubregionStatus},
+    bitset::{Bitset, SubregionStatus},
     free_list::{FreeList, FreeListNode},
     tree::Tree,
 };
@@ -230,7 +230,9 @@ impl<
                 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);
+                let Ok(()) = tree.bitset_mark_as_present(bitset, tree.size_class, 0) else {
+                    panic!("the first bit was already set in {tree:#?}\nbitset = {bitset:?}")
+                };
             }
         }
 
@@ -253,7 +255,9 @@ impl<
             let tree = &mut self.trees[tree_index];
 
             // Mark the pointer as no longer being in the tree.
-            tree.bitset_mark_as_absent(&mut self.bitset, size_class, offset);
+            let Ok(()) = tree.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)
@@ -289,7 +293,9 @@ impl<
         }
 
         // Mark the pointer as being in the tree.
-        tree.bitset_mark_as_present(&mut self.bitset, size_class, offset);
+        let Ok(()) = tree.bitset_mark_as_present(self.bitset, size_class, offset) else {
+            panic!("the bit for tree {tree_index}, offset {offset} was already set in {self:#?}")
+        };
 
         // Return the pointer to the appropriate free list.
         self.free_list(size_class).push(ptr);
diff --git a/crates/alloc_buddy/src/tree.rs b/crates/alloc_buddy/src/tree.rs
index 72ee466..b713285 100644
--- a/crates/alloc_buddy/src/tree.rs
+++ b/crates/alloc_buddy/src/tree.rs
@@ -1,4 +1,4 @@
-use crate::bitvec::{Bitset, SubregionStatus};
+use crate::bitset::{Bitset, SubregionStatus, UnexpectedBitsetState};
 use contracts::requires;
 use core::ptr::NonNull;
 
@@ -56,8 +56,8 @@ impl<'buddy, const PAGE_SIZE: usize, const PAGE_SIZE_BITS: usize>
         // 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
+        // The sum of those two with our offset is simply our index.
+        self.bitset_offset + bits_skipped_for_size_class + bits_skipped_for_index
     }
 
     /// Changes a bit in the bitset from `InFreeList` to `NotInFreeList`.
@@ -69,10 +69,12 @@ impl<'buddy, const PAGE_SIZE: usize, const PAGE_SIZE_BITS: usize>
         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)
+    ) -> Result<(), UnexpectedBitsetState> {
+        bitset.replace(
+            self.bitset_index(size_class, offset_bytes),
+            SubregionStatus::InFreeList,
+            SubregionStatus::NotInFreeList,
+        )
     }
 
     /// Changes a bit in the bitset from `NotInFreeList` to `InFreeList`.
@@ -84,10 +86,12 @@ impl<'buddy, const PAGE_SIZE: usize, const PAGE_SIZE_BITS: usize>
         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)
+    ) -> Result<(), UnexpectedBitsetState> {
+        bitset.replace(
+            self.bitset_index(size_class, offset_bytes),
+            SubregionStatus::NotInFreeList,
+            SubregionStatus::InFreeList,
+        )
     }
 
     /// Returns whether the tree contains an address.