Fix UB.

1 files changed, 107 insertions, 132 deletions

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
     }
 }