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

1 files changed, 88 insertions, 0 deletions

diff --git a/crates/alloc_buddy/src/bitset.rs b/crates/alloc_buddy/src/bitset.rs
new file mode 100644
index 0000000..b5f12d2
--- /dev/null
+++ b/crates/alloc_buddy/src/bitset.rs
@@ -0,0 +1,88 @@
+use contracts::requires;
+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 {
+    /// 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;
+        let subword_index = i % usize::BITS as usize;
+        let one_hot = 1 << subword_index;
+        if (self.0[word_index] & one_hot) == 0 {
+            SubregionStatus::NotInFreeList
+        } else {
+            SubregionStatus::InFreeList
+        }
+    }
+
+    /// Returns the number of bits in the bitset.
+    pub fn len(&self) -> usize {
+        self.0.len() * usize::BITS as usize
+    }
+
+    /// 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;
+        let subword_index = i % usize::BITS as usize;
+        let one_hot = 1 << subword_index;
+        match status {
+            SubregionStatus::NotInFreeList => {
+                self.0[word_index] &= !one_hot;
+            }
+            SubregionStatus::InFreeList => {
+                self.0[word_index] |= one_hot;
+            }
+        }
+    }
+}
+
+impl fmt::Debug for Bitset {
+    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
+        write!(fmt, "Bitset(")?;
+        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 {
+    /// The region is not in the free list.
+    NotInFreeList = 0,
+
+    /// The region is in the free list.
+    InFreeList = 1,
+}