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

1 files changed, 15 insertions, 11 deletions

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.