Fix UB.

4 files changed, 174 insertions, 161 deletions

diff --git a/crates/buddy_allocator/Cargo.toml b/crates/buddy_allocator/Cargo.toml
index e2a85d3..94ab236 100644
--- a/crates/buddy_allocator/Cargo.toml
+++ b/crates/buddy_allocator/Cargo.toml
@@ -9,7 +9,6 @@ allocator-api2 = { version = "0.2.18", default-features = false }
 contracts = { version = "0.6.3", default-features = false }
 static_assertions = { version = "1.1.0", default-features = false }
 vernos_physmem_free_list = { path = "../physmem_free_list" }
-vernos_utils = { path = "../utils" }
 
 [dev-dependencies]
 proptest = { version = "0.9.6", default-features = false, features = ["std"] }
diff --git a/crates/buddy_allocator/src/free_list.rs b/crates/buddy_allocator/src/free_list.rs
index b794b8c..9b470fd 100644
--- a/crates/buddy_allocator/src/free_list.rs
+++ b/crates/buddy_allocator/src/free_list.rs
@@ -1,55 +1,72 @@
 //! A circular doubly-linked list of power-of-two-sized subregions.
 
-use core::{
-    fmt,
-    marker::PhantomPinned,
-    pin::Pin,
-    ptr::{self, NonNull},
-};
-
-use contracts::requires;
+use contracts::{ensures, requires};
+use core::{fmt, ptr::NonNull};
 
 /// The magic number a non-sentinel node should have.
 const FREE_NODE_MAGIC: u64 = u64::from_ne_bytes(*b"FREENODE");
 
+/// The magic number set into the spot for a node's magic number when it is initialized, but not
+/// linked into a list.
+const INITED_NODE_MAGIC: u64 = u64::from_ne_bytes(*b"NODEINIT");
+
 /// The magic number a sentinel node should have.
 const SENTINEL_MAGIC: u64 = u64::from_ne_bytes(*b"SENTINEL");
 
 /// The magic number set into the spot for a node's magic number when it gets unlinked.
 const USED_NODE_MAGIC: u64 = u64::from_ne_bytes(*b"ALLOC'ED");
 
-/// A circular doubly-linked list of power-of-two-sized subregions.
-///
-/// This type is valid when zero-initialized.
-///
-/// # Safety
-///
-/// - This type is valid only if it is zero-initialized or if it contains a valid circular
-///   doubly-linked list of subregions of the correct size, which it owns.
-pub struct FreeList {
-    sentinel: FreeListNode,
-}
+/// A pointer to the sentinel node of a circular doubly-linked list of power-of-two-sized subregions.
+pub struct FreeList(NonNull<FreeListNode>);
 
 impl FreeList {
+    /// Creates a free-list wrapper for a sentinel node's pointer.
+    ///
+    /// # Safety
+    ///
+    /// - `sentinel` must be a pointer to the sentinel of a valid circular doubly-linked list.
+    #[requires(unsafe { (*sentinel.as_ptr()).next }.is_some())]
+    #[requires(unsafe { (*sentinel.as_ptr()).prev }.is_some())]
+    #[requires(unsafe { (*sentinel.as_ptr()).magic } == SENTINEL_MAGIC)]
+    #[requires(unsafe { (*sentinel.as_ptr()).self_addr } == sentinel.as_ptr() as usize)]
+    pub unsafe fn from_sentinel(sentinel: NonNull<FreeListNode>) -> FreeList {
+        FreeList(sentinel)
+    }
+
+    /// Returns whether the free-list is empty.
+    pub fn is_empty(&self) -> bool {
+        let sentinel = self.0.as_ptr();
+        // SAFETY: This is sound as long as the sentinel was valid, which is a precondition of
+        // `FreeList::from_sentinel()`.
+        unsafe {
+            // UNWRAP: This is a precondition of the method.
+            let next = (*sentinel).next.unwrap();
+            (*sentinel).self_addr == next.as_ptr() as usize
+        }
+    }
+
     /// Pops a subregion from the free-list, returning ownership to the caller.
-    pub fn pop(mut self: Pin<&mut FreeList>) -> Option<NonNull<u8>> {
-        assert_ne!(self.sentinel.next, ptr::null_mut());
-        assert_ne!(self.sentinel.prev, ptr::null_mut());
-
-        // Check if the sentinel's next pointer points to the sentinel itself.
-        let sentinel = &self.sentinel as *const FreeListNode as *mut FreeListNode;
-        let next = *self.as_mut().sentinel().next_mut();
-        if sentinel == next {
-            // If so, the list is empty.
+    pub fn pop(&mut self) -> Option<NonNull<u8>> {
+        if self.is_empty() {
+            // If the list is empty, return `None`.
             None
         } else {
-            // Otherwise, unlink that node.
-            //
-            // UNWRAP: The pointer was non-null, by the invariants of FreeList.
-            let mut node = NonNull::new(next).unwrap();
-            // SAFETY: The list was valid, by the invariants of FreeList.
-            unsafe { node.as_mut().unlink() };
-            Some(node.cast())
+            // SAFETY: This is sound as long as the sentinel was valid, which is a precondition of
+            // `FreeList::from_sentinel()`.
+            unsafe {
+                // UNWRAP: In a valid circular doubly-linked list, the pointers are all non-null.
+                let node = (*self.0.as_ptr()).next.unwrap();
+                let prev = (*node.as_ptr()).prev.unwrap();
+                let next = (*node.as_ptr()).next.unwrap();
+
+                (*prev.as_ptr()).next = Some(next);
+                (*next.as_ptr()).prev = Some(prev);
+
+                (*node.as_ptr()).next = None;
+                (*node.as_ptr()).prev = None;
+                (*node.as_ptr()).magic = USED_NODE_MAGIC;
+                Some(node.cast())
+            }
         }
     }
 
@@ -58,132 +75,90 @@ impl FreeList {
     /// # Safety
     ///
     /// - `subregion_ptr` must convey ownership over untyped memory of the correct size.
-    pub unsafe fn push(mut self: Pin<&mut FreeList>, subregion_ptr: NonNull<u8>) {
-        assert_ne!(self.sentinel.next, ptr::null_mut());
-        assert_ne!(self.sentinel.prev, ptr::null_mut());
-
-        // Write the header the node should have to the node.
-        //
-        // SAFETY: This &mut is not used to move out of the node.
-        let sentinel = unsafe { self.as_mut().sentinel().get_unchecked_mut() };
-        let mut node = subregion_ptr.cast();
-        node.write(FreeListNode {
-            next: sentinel.next,
-            prev: sentinel,
-            magic: FREE_NODE_MAGIC,
-            _phantom: PhantomPinned,
-        });
-
-        // Link the node into the list.
-        //
-        // SAFETY: There are no other references to the node live.
-        let node = unsafe { node.as_mut() };
-        (*node.prev).next = node;
-        (*node.next).prev = node;
-    }
+    pub unsafe fn push(&mut self, subregion_ptr: NonNull<u8>) {
+        let node = FreeListNode::init(subregion_ptr);
 
-    /// Performs self-linking. This must be done before the free list is used in any other way.
-    pub fn self_link(self: Pin<&mut FreeList>) {
-        assert_eq!(self.sentinel.next, ptr::null_mut());
-        assert_eq!(self.sentinel.prev, ptr::null_mut());
-
-        let mut sentinel = self.sentinel_unchecked();
-        // SAFETY: We do not use this pointer to move out of the sentinel.
-        let sentinel_ptr = unsafe { sentinel.as_mut().get_unchecked_mut() as *mut FreeListNode };
-        *sentinel.as_mut().next_mut() = sentinel_ptr;
-        *sentinel.as_mut().prev_mut() = sentinel_ptr;
-        *sentinel.as_mut().magic_mut() = SENTINEL_MAGIC;
-    }
+        let prev = self.0;
+        // UNWRAP: In a valid circular doubly-linked list, the pointers are all non-null.
+        let next = (*self.0.as_ptr()).next.unwrap();
 
-    /// Projects the sentinel node out of the list.
-    fn sentinel(self: Pin<&mut FreeList>) -> Pin<&mut FreeListNode> {
-        let mut sentinel = self.sentinel_unchecked();
-        assert_eq!(*sentinel.as_mut().magic_mut(), SENTINEL_MAGIC);
-        sentinel
-    }
+        (*node.as_ptr()).prev = Some(prev);
+        (*node.as_ptr()).next = Some(next);
+        (*node.as_ptr()).magic = FREE_NODE_MAGIC;
 
-    /// Projects the sentinel node out of the list, without checking the magic number.
-    fn sentinel_unchecked(self: Pin<&mut FreeList>) -> Pin<&mut FreeListNode> {
-        // SAFETY: sentinel is structurally pinned.
-        unsafe { self.map_unchecked_mut(|list| &mut list.sentinel) }
+        (*prev.as_ptr()).next = Some(node);
+        (*next.as_ptr()).prev = Some(node);
     }
 }
 
 /// This impl will panic until the sentinel has been self-linked.
 impl fmt::Debug for FreeList {
     fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
-        assert_ne!(self.sentinel.next, ptr::null_mut());
-        assert_ne!(self.sentinel.prev, ptr::null_mut());
+        // SAFETY: This is sound as long as the sentinel was valid, which is a precondition of
+        // `FreeList::from_sentinel()`.
+        let sentinel_addr = unsafe { (*self.0.as_ptr()).self_addr };
 
-        let sentinel = &self.sentinel as *const FreeListNode;
-
-        // SAFETY: The invariants of FreeList guarantee that the list is valid. In a valid list,
-        // the next pointer is always valid and non-null.
-        let mut ptr = unsafe { sentinel.as_ref().unwrap().next as *const FreeListNode };
         let mut fmt = fmt.debug_list();
-        while ptr != sentinel {
-            fmt.entry(&ptr);
+        let mut ptr = self.0;
+        loop {
             // SAFETY: The invariants of FreeList guarantee that the list is valid. In a valid
             // list, the next pointer is always valid and non-null.
-            ptr = unsafe { ptr.as_ref().unwrap().next as *const FreeListNode };
+            ptr = unsafe { (*ptr.as_ptr()).next.unwrap() };
+            if unsafe { (*ptr.as_ptr()).self_addr } == sentinel_addr {
+                break fmt.finish();
+            }
+            fmt.entry(&ptr);
         }
-        fmt.finish()
     }
 }
 
+/// A node in a free list. This should not be moved without careful thought.
+///
+/// This type is valid when zero-initialized.
+///
+/// # Safety
+///
+/// - This type is valid only if it is zero-initialized or if it contains a valid circular
+///   doubly-linked list of subregions of the correct size, which it owns.
 #[derive(Debug)]
-struct FreeListNode {
-    next: *mut FreeListNode,
-    prev: *mut FreeListNode,
+pub struct FreeListNode {
+    next: Option<NonNull<FreeListNode>>,
+    prev: Option<NonNull<FreeListNode>>,
     magic: u64,
-    _phantom: PhantomPinned,
+    self_addr: usize,
 }
 
 impl FreeListNode {
-    /// Projects the `next` pointer out of the node.
-    fn next_mut(self: Pin<&mut FreeListNode>) -> &mut *mut FreeListNode {
-        // SAFETY: next is not structurally pinned.
-        unsafe { &mut self.get_unchecked_mut().next }
-    }
-
-    /// Projects the `prev` pointer out of the node.
-    fn prev_mut(self: Pin<&mut FreeListNode>) -> &mut *mut FreeListNode {
-        // SAFETY: prev is structurally pinned.
-        unsafe { &mut self.get_unchecked_mut().prev }
-    }
-
-    /// Projects the `magic` number out of the node.
-    fn magic_mut(self: Pin<&mut FreeListNode>) -> &mut u64 {
-        // SAFETY: magic is structurally pinned.
-        unsafe { &mut self.get_unchecked_mut().magic }
-    }
-
-    /// Links the node into a circular doubly-linked list.
+    /// Initializes a sentinel node.
     ///
     /// # Safety
     ///
-    /// This must be a node that's not part of any linked list.
-    unsafe fn link(&mut self) {
-        todo!()
+    /// The pointer must convey exclusive access and point to a large-enough allocation.
+    #[ensures(unsafe { (*ptr.as_ptr()).magic } == SENTINEL_MAGIC)]
+    #[ensures(unsafe { (*ptr.as_ptr()).self_addr } == ptr.as_ptr() as usize)]
+    pub unsafe fn init_sentinel(ptr: NonNull<FreeListNode>) {
+        (*ptr.as_ptr()).next = Some(ptr);
+        (*ptr.as_ptr()).prev = Some(ptr);
+        (*ptr.as_ptr()).magic = SENTINEL_MAGIC;
+        (*ptr.as_ptr()).self_addr = ptr.as_ptr() as usize;
     }
 
-    /// Unlinks a node from a circular doubly-linked list.
+    /// Initializes a non-sentinel node.
     ///
     /// # Safety
     ///
-    /// This must be a non-sentinel node in a valid circular doubly-linked list.
-    #[requires(self.magic == FREE_NODE_MAGIC)]
-    #[ensures(self.magic == USED_NODE_MAGIC)]
-    unsafe fn unlink(&mut self) {
-        let prev = self.prev;
-        let next = self.next;
-
-        use core::ptr::addr_of_mut;
-        *addr_of_mut!((*prev).next) = next;
-        *addr_of_mut!((*next).prev) = prev;
-
-        self.next = ptr::null_mut();
-        self.prev = ptr::null_mut();
-        self.magic = USED_NODE_MAGIC;
+    /// The pointer must convey ownership and point to a large-enough allocation.
+    #[ensures(ptr.as_ptr() as usize == ret.as_ptr() as usize)]
+    #[ensures(unsafe { (*ret.as_ptr()).magic } == INITED_NODE_MAGIC)]
+    #[ensures(unsafe { (*ret.as_ptr()).self_addr } == ptr.as_ptr() as usize)]
+    unsafe fn init(ptr: NonNull<u8>) -> NonNull<FreeListNode> {
+        let ptr = ptr.cast();
+        ptr.write(FreeListNode {
+            next: None,
+            prev: None,
+            magic: INITED_NODE_MAGIC,
+            self_addr: ptr.as_ptr() as usize,
+        });
+        ptr
     }
 }
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,
diff --git a/crates/buddy_allocator/tests/hosted_test.rs b/crates/buddy_allocator/tests/hosted_test.rs
index d6d3389..89398ce 100644
--- a/crates/buddy_allocator/tests/hosted_test.rs
+++ b/crates/buddy_allocator/tests/hosted_test.rs
@@ -20,7 +20,7 @@ proptest! {
 #[test]
 fn test_simple_scenario() {
     let scenario = Scenario {
-        range_sizes: vec![16],
+        range_sizes: vec![5],
         actions: vec![
             Action::Debug,
             Action::Alloc {