More buddy work, but we've got UB again... This time relating to stacked borrows...

1 files changed, 70 insertions, 10 deletions

diff --git a/crates/buddy_allocator/src/free_list.rs b/crates/buddy_allocator/src/free_list.rs
index 2e3afe2..b794b8c 100644
--- a/crates/buddy_allocator/src/free_list.rs
+++ b/crates/buddy_allocator/src/free_list.rs
@@ -1,4 +1,4 @@
-//! A doubly-linked list of power-of-two-sized subregions.
+//! A circular doubly-linked list of power-of-two-sized subregions.
 
 use core::{
     fmt,
@@ -7,25 +7,52 @@ use core::{
     ptr::{self, NonNull},
 };
 
+use contracts::requires;
+
 /// The magic number a non-sentinel node should have.
 const FREE_NODE_MAGIC: u64 = u64::from_ne_bytes(*b"FREENODE");
 
 /// The magic number a sentinel node should have.
 const SENTINEL_MAGIC: u64 = u64::from_ne_bytes(*b"SENTINEL");
 
-/// A doubly-linked list of power-of-two-sized subregions.
+/// 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 doubly-linked
-///   list of subregions of the correct size, which it owns.
+/// - 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,
 }
 
 impl FreeList {
+    /// 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.
+            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())
+        }
+    }
+
     /// Pushes a subregion into the free-list.
     ///
     /// # Safety
@@ -36,19 +63,23 @@ impl FreeList {
         assert_ne!(self.sentinel.prev, ptr::null_mut());
 
         // Write the header the node should have to the node.
-        let sentinel = &self.sentinel as *const FreeListNode as *mut FreeListNode;
-        let next = *self.as_mut().sentinel().next_mut();
-        let node = subregion_ptr.cast();
+        //
+        // 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,
+            next: sentinel.next,
             prev: sentinel,
             magic: FREE_NODE_MAGIC,
             _phantom: PhantomPinned,
         });
 
         // Link the node into the list.
-        *self.sentinel().next_mut() = node.as_ptr();
-        (*next).prev = node.as_ptr();
+        //
+        // SAFETY: There are no other references to the node live.
+        let node = unsafe { node.as_mut() };
+        (*node.prev).next = node;
+        (*node.next).prev = node;
     }
 
     /// Performs self-linking. This must be done before the free list is used in any other way.
@@ -126,4 +157,33 @@ impl FreeListNode {
         // SAFETY: magic is structurally pinned.
         unsafe { &mut self.get_unchecked_mut().magic }
     }
+
+    /// Links the node into a circular doubly-linked list.
+    ///
+    /// # Safety
+    ///
+    /// This must be a node that's not part of any linked list.
+    unsafe fn link(&mut self) {
+        todo!()
+    }
+
+    /// Unlinks a node from a circular doubly-linked list.
+    ///
+    /// # 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;
+    }
 }