Fix UB.

1 files changed, 66 insertions, 27 deletions

diff --git a/crates/buddy_allocator/src/lib.rs b/crates/buddy_allocator/src/lib.rs
index 4292444..f0c8daa 100644
--- a/crates/buddy_allocator/src/lib.rs
+++ b/crates/buddy_allocator/src/lib.rs
@@ -52,25 +52,23 @@ mod tree;
 
 use crate::{
     bitvec::{Bitset, SubregionStatus},
-    free_list::FreeList,
+    free_list::{FreeList, FreeListNode},
     tree::Tree,
 };
 use allocator_api2::alloc::{AllocError, Layout, LayoutError};
 use contracts::{ensures, requires};
-use core::{mem, pin::Pin, ptr::NonNull};
+use core::{fmt, mem, ptr::NonNull};
 use vernos_physmem_free_list::FreeListAllocator;
-use vernos_utils::pin::{pin_project_slice_mut, pin_project_slice_mut_iter};
 
 /// A buddy allocator.
-#[derive(Debug)]
 pub struct BuddyAllocator<
     'allocator,
     const PAGE_SIZE: usize,
     const PAGE_SIZE_BITS: usize,
     const SIZE_CLASS_COUNT: usize,
 > {
-    /// The free lists.
-    free_lists: Pin<&'allocator mut [FreeList; SIZE_CLASS_COUNT]>,
+    /// The free list sentinels (`SIZE_CLASS_COUNT` of them).
+    free_list_sentinels: NonNull<FreeListNode>,
 
     /// The metadata associated with each tree.
     trees: &'allocator mut [Tree<'allocator, PAGE_SIZE, PAGE_SIZE_BITS>],
@@ -140,11 +138,10 @@ impl<
         //
         // SAFETY: The layout is computed to make all of this sound if the data is valid. All of
         // these types are valid when zero-initialized.
-        let free_lists: &'allocator mut [FreeList; SIZE_CLASS_COUNT] = unsafe {
+        let free_list_sentinels: NonNull<FreeListNode> = unsafe {
             metadata
-                .byte_add(metadata_layout.free_lists_offset)
+                .byte_add(metadata_layout.free_list_sentinels_offset)
                 .cast()
-                .as_mut()
         };
         let trees: &'allocator mut [Tree<PAGE_SIZE, PAGE_SIZE_BITS>] = unsafe {
             NonNull::slice_from_raw_parts(
@@ -165,13 +162,11 @@ impl<
             )
         };
 
-        // Pin and self-link each free list. Although they're valid in a UB sense when zeroed,
-        // they're not yet ready to use until this is done.
-        //
-        // SAFETY: We never provide the ability to move a free list.
-        let mut free_lists = unsafe { Pin::new_unchecked(free_lists) };
-        for free_list in pin_project_slice_mut_iter(free_lists.as_mut()) {
-            free_list.self_link();
+        // Self-link each free list. Although they're valid in a UB sense when zeroed, they're not
+        // yet ready to use until this is done.
+        for size_class in 0..SIZE_CLASS_COUNT {
+            // SAFETY: This is the number of nodes we allocate, so we'll stay in-bounds.
+            unsafe { FreeListNode::init_sentinel(free_list_sentinels.add(size_class)) };
         }
 
         // Initialize the trees, adding subregions into the allocator as we do. We don't actually
@@ -228,7 +223,13 @@ impl<
                 // UNWRAP: The pointer must have been non-null when it was in the free-list
                 // allocator.
                 let ptr = NonNull::new(tree.base_ptr as *mut u8).unwrap();
-                let free_list = pin_project_slice_mut(free_lists.as_mut(), tree.size_class);
+
+                // SAFETY: There are as many sentinels as size classes, so this will be in-bounds.
+                let sentinel = unsafe { free_list_sentinels.add(tree.size_class) };
+
+                // SAFETY: The free list is kept valid.
+                let mut free_list = unsafe { FreeList::from_sentinel(sentinel) };
+
                 // SAFETY: This region is not owned or borrowed by anything else (the only thing it
                 // could be would be the metadata, but that's not in this tree), is of the correct
                 // size, and is composed of untyped memory.
@@ -241,7 +242,7 @@ impl<
 
         // Make the buddy allocator's struct and return it.
         Ok(BuddyAllocator {
-            free_lists,
+            free_list_sentinels,
             trees,
             bitset,
         })
@@ -250,7 +251,7 @@ impl<
     /// Tries to allocate a subregion of a particular size class from the allocator.
     #[requires(size_class < SIZE_CLASS_COUNT)]
     pub fn alloc(&mut self, size_class: usize) -> Result<NonNull<u8>, AllocError> {
-        if let Some(ptr) = self.free_list_mut(size_class).pop() {
+        if let Some(ptr) = self.free_list(size_class).pop() {
             // Fast-path: try allocating from the right size class's free list.
 
             // Find the corresponding tree and the offset of the subregion in the tree.
@@ -295,9 +296,14 @@ impl<
         todo!("{tree:#?}")
     }
 
-    /// Returns a (pinned) mutable reference to the free list with the given size class.
-    fn free_list_mut(&mut self, i: usize) -> Pin<&mut FreeList> {
-        pin_project_slice_mut(self.free_lists.as_mut(), i)
+    /// Returns the free list with the given size class.
+    #[requires(size_class < SIZE_CLASS_COUNT)]
+    fn free_list(&mut self, size_class: usize) -> FreeList {
+        // SAFETY: The bounds-check is a precondition.
+        let sentinel = unsafe { self.free_list_sentinels.add(size_class) };
+
+        // SAFETY: The free list is kept valid.
+        unsafe { FreeList::from_sentinel(sentinel) }
     }
 
     /// Returns the index of the tree this pointer was allocated in.
@@ -320,6 +326,39 @@ impl<
     }
 }
 
+impl<
+        'allocator,
+        const PAGE_SIZE: usize,
+        const PAGE_SIZE_BITS: usize,
+        const SIZE_CLASS_COUNT: usize,
+    > fmt::Debug for BuddyAllocator<'allocator, PAGE_SIZE, PAGE_SIZE_BITS, SIZE_CLASS_COUNT>
+{
+    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
+        struct FreeLists<const SIZE_CLASS_COUNT: usize>(NonNull<FreeListNode>);
+
+        impl<const SIZE_CLASS_COUNT: usize> fmt::Debug for FreeLists<SIZE_CLASS_COUNT> {
+            fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
+                fmt.debug_list()
+                    .entries((0..SIZE_CLASS_COUNT).map(|size_class| {
+                        // SAFETY: The free lists are kept valid, and the range of size classes is
+                        // necessarily in-bounds.
+                        unsafe { FreeList::from_sentinel(self.0.add(size_class)) }
+                    }))
+                    .finish()
+            }
+        }
+
+        fmt.debug_struct("BuddyAllocator")
+            .field(
+                "free_lists",
+                &FreeLists::<SIZE_CLASS_COUNT>(self.free_list_sentinels),
+            )
+            .field("trees", &self.trees)
+            .field("bitset", &self.bitset)
+            .finish()
+    }
+}
+
 #[derive(Debug)]
 struct MetadataLayout<
     const PAGE_SIZE: usize,
@@ -327,7 +366,7 @@ struct MetadataLayout<
     const SIZE_CLASS_COUNT: usize,
 > {
     layout: Layout,
-    free_lists_offset: usize,
+    free_list_sentinels_offset: usize,
     trees_offset: usize,
     trees_len: usize,
     bitset_offset: usize,
@@ -350,17 +389,17 @@ impl<const PAGE_SIZE: usize, const PAGE_SIZE_BITS: usize, const SIZE_CLASS_COUNT
         }
         let bitset_len = (bitset_len_bits + usize::BITS as usize - 1) / usize::BITS as usize;
 
-        let free_lists_layout = Layout::new::<[FreeList; SIZE_CLASS_COUNT]>();
+        let free_list_sentinels_layout = Layout::new::<[FreeListNode; SIZE_CLASS_COUNT]>();
         let trees_layout = Layout::array::<Tree<PAGE_SIZE, PAGE_SIZE_BITS>>(trees_len)?;
         let bitset_layout = Layout::array::<usize>(bitset_len)?;
 
-        let free_lists_offset = 0;
-        let (layout, trees_offset) = free_lists_layout.extend(trees_layout)?;
+        let free_list_sentinels_offset = 0;
+        let (layout, trees_offset) = free_list_sentinels_layout.extend(trees_layout)?;
         let (layout, bitset_offset) = layout.extend(bitset_layout)?;
 
         Ok(MetadataLayout {
             layout,
-            free_lists_offset,
+            free_list_sentinels_offset,
             trees_offset,
             trees_len,
             bitset_offset,