foundation: linked list + DormantMutRef

crates/foundation/src/collections/linked_list.rs

@@ -1,19 +1,104 @@
+use core::alloc::{Allocator, Layout};
 use core::ptr::NonNull;
 
+use crate::mem::DormantMutRef;
+
 pub unsafe trait LinkedListNode {
     fn next(this: NonNull<Self>) -> Option<NonNull<Self>>;
     fn set_next(this: NonNull<Self>, next: Option<NonNull<Self>>);
 }
 
+#[derive(Debug)]
 pub struct LinkedList<T: LinkedListNode> {
     head: Option<NonNull<T>>,
 }
 
+pub struct NextOf<'a, T: LinkedListNode> {
+    parent: NonNull<T>,
+    node: NonNull<T>,
+    _marker: core::marker::PhantomData<&'a mut T>,
+}
+pub struct Head<'a, T: LinkedListNode> {
+    list: DormantMutRef<'a, LinkedList<T>>,
+    node: NonNull<T>,
+}
+
+pub enum LinkedListEntry<'a, T: LinkedListNode> {
+    NextOf(NextOf<'a, T>),
+    Head(Head<'a, T>),
+}
+
+impl<'a, T: LinkedListNode> LinkedListEntry<'a, T> {
+    unsafe fn new_head(list: &'a mut LinkedList<T>, head: NonNull<T>) -> Self {
+        Self::Head(Head {
+            list: DormantMutRef::new(list).1,
+            node: head,
+        })
+    }
+    pub fn get(&self) -> &T {
+        match self {
+            LinkedListEntry::NextOf(NextOf { node, .. }) => unsafe { node.as_ref() },
+            LinkedListEntry::Head(Head { node, .. }) => unsafe { node.as_ref() },
+        }
+    }
+
+    pub fn get_mut(&mut self) -> &mut T {
+        match self {
+            LinkedListEntry::NextOf(NextOf { node, .. }) => unsafe { node.as_mut() },
+            LinkedListEntry::Head(Head { node, .. }) => unsafe { node.as_mut() },
+        }
+    }
+
+    pub fn remove(self) -> NonNull<T> {
+        match self {
+            LinkedListEntry::NextOf(NextOf { parent, node, .. }) => {
+                let next = T::next(node);
+                T::set_next(parent, next);
+                node
+            }
+            LinkedListEntry::Head(Head { list, node }) => {
+                let next = T::next(node);
+                unsafe { list.awaken() }.head = next;
+                node
+            }
+        }
+    }
+
+    pub fn inner(&self) -> NonNull<T> {
+        match self {
+            LinkedListEntry::NextOf(NextOf { node, .. }) => *node,
+            LinkedListEntry::Head(Head { node, .. }) => *node,
+        }
+    }
+    pub fn next(&self) -> Option<NonNull<T>> {
+        match self {
+            LinkedListEntry::NextOf(NextOf { node, .. }) => T::next(*node),
+            LinkedListEntry::Head(Head { node, .. }) => T::next(*node),
+        }
+    }
+}
+
 impl<T: LinkedListNode> LinkedList<T> {
     pub const fn new() -> Self {
         Self { head: None }
     }
 
+    pub fn get_or_insert_front_with<F>(&mut self, f: F) -> LinkedListEntry<'_, T>
+    where
+        F: FnOnce() -> NonNull<T>,
+    {
+        let head = if let Some(head) = self.head {
+            head
+        } else {
+            let node = f();
+            self.push_front(node);
+            node
+        };
+
+        // SAFETY: head is the head of the list
+        unsafe { LinkedListEntry::new_head(self, head) }
+    }
+
     pub fn push_front(&mut self, node: NonNull<T>) {
         T::set_next(node, self.head);
         self.head = Some(node);
@@ -25,29 +110,31 @@ impl<T: LinkedListNode> LinkedList<T> {
         Some(node)
     }
 
-    pub fn remove_if(&mut self, mut predicate: impl FnMut(NonNull<T>) -> bool) {
-        let mut current = self.head;
-        let mut prev: Option<NonNull<T>> = None;
-
-        while let Some(node) = current {
-            if predicate(node) {
-                let next = T::next(node);
-                if let Some(prev_node) = prev {
-                    T::set_next(prev_node, next);
-                } else {
-                    self.head = next;
-                }
-                current = next;
-            } else {
-                prev = current;
-                current = T::next(node);
-            }
-        }
+    pub fn remove_if<F>(&mut self, predicate: F) -> RemoveIf<'_, T, F>
+    where
+        F: FnMut(NonNull<T>) -> bool,
+    {
+        RemoveIf::new(self, predicate)
     }
 
     pub fn iter(&self) -> LinkedListIter<T> {
         LinkedListIter { current: self.head }
     }
+    pub fn into_iter(self) -> LinkedListIter<T> {
+        LinkedListIter { current: self.head }
+    }
+    pub fn entries(&mut self) -> LinkedListEntries<'_, T> {
+        LinkedListEntries::new_head(self)
+    }
+
+    pub unsafe fn drop_in<A: Allocator>(mut self, alloc: &A) {
+        let iter = LinkedListIter {
+            current: self.head.take(),
+        };
+        for node in iter {
+            unsafe { alloc.deallocate(node.cast(), Layout::new::<T>()) }
+        }
+    }
 }
 
 impl<T: LinkedListNode> Default for LinkedList<T> {
@@ -56,6 +143,65 @@ impl<T: LinkedListNode> Default for LinkedList<T> {
     }
 }
 
+pub struct RemoveIf<'a, T: LinkedListNode, F> {
+    entries: LinkedListEntries<'a, T>,
+    pred: F,
+}
+
+impl<'a, T: LinkedListNode, F> RemoveIf<'a, T, F> {
+    pub fn new(list: &'a mut LinkedList<T>, pred: F) -> Self {
+        Self {
+            entries: LinkedListEntries::new_head(list),
+            pred,
+        }
+    }
+}
+
+impl<'a, T: LinkedListNode, F> Iterator for RemoveIf<'a, T, F>
+where
+    F: FnMut(NonNull<T>) -> bool,
+{
+    type Item = NonNull<T>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        while let Some(entry) = self.entries.next() {
+            let node = entry.inner();
+            if (self.pred)(node) {
+                return Some(entry.remove());
+            }
+        }
+        None
+    }
+}
+
+pub struct LinkedListEntries<'a, T: LinkedListNode> {
+    current: Option<LinkedListEntry<'a, T>>,
+}
+
+impl<'a, T: LinkedListNode> LinkedListEntries<'a, T> {
+    fn new_head(list: &'a mut LinkedList<T>) -> Self {
+        let current = list
+            .head
+            .map(|head| unsafe { LinkedListEntry::new_head(list, head) });
+        Self { current }
+    }
+    pub fn next(&mut self) -> Option<LinkedListEntry<'_, T>> {
+        if let Some(current) = self.current.take() {
+            self.current = match current.next() {
+                Some(next_node) => Some(LinkedListEntry::NextOf(NextOf {
+                    parent: current.inner(),
+                    node: next_node,
+                    _marker: core::marker::PhantomData,
+                })),
+                None => None,
+            };
+            Some(current)
+        } else {
+            None
+        }
+    }
+}
+
 pub struct LinkedListIter<T: LinkedListNode> {
     current: Option<NonNull<T>>,
 }

crates/foundation/src/lib.rs

@@ -54,6 +54,79 @@ pub mod mem {
             }
         }
     }
+
+    use core::marker::PhantomData;
+    use core::ptr::NonNull;
+
+    /// Models a reborrow of some unique reference, when you know that the reborrow
+    /// and all its descendants (i.e., all pointers and references derived from it)
+    /// will not be used any more at some point, after which you want to use the
+    /// original unique reference again.
+    ///
+    /// The borrow checker usually handles this stacking of borrows for you, but
+    /// some control flows that accomplish this stacking are too complicated for
+    /// the compiler to follow. A `DormantMutRef` allows you to check borrowing
+    /// yourself, while still expressing its stacked nature, and encapsulating
+    /// the raw pointer code needed to do this without undefined behavior.
+    pub(super) struct DormantMutRef<'a, T> {
+        ptr: NonNull<T>,
+        _marker: PhantomData<&'a mut T>,
+    }
+
+    unsafe impl<'a, T> Sync for DormantMutRef<'a, T> where &'a mut T: Sync {}
+    unsafe impl<'a, T> Send for DormantMutRef<'a, T> where &'a mut T: Send {}
+
+    impl<'a, T> DormantMutRef<'a, T> {
+        /// Capture a unique borrow, and immediately reborrow it. For the compiler,
+        /// the lifetime of the new reference is the same as the lifetime of the
+        /// original reference, but you promise to use it for a shorter period.
+        pub(super) fn new(t: &'a mut T) -> (&'a mut T, Self) {
+            let ptr = NonNull::from(t);
+            // SAFETY: we hold the borrow throughout 'a via `_marker`, and we expose
+            // only this reference, so it is unique.
+            let new_ref = unsafe { &mut *ptr.as_ptr() };
+            (
+                new_ref,
+                Self {
+                    ptr,
+                    _marker: PhantomData,
+                },
+            )
+        }
+
+        /// Revert to the unique borrow initially captured.
+        ///
+        /// # Safety
+        ///
+        /// The reborrow must have ended, i.e., the reference returned by `new` and
+        /// all pointers and references derived from it, must not be used anymore.
+        pub(super) unsafe fn awaken(self) -> &'a mut T {
+            // SAFETY: our own safety conditions imply this reference is again unique.
+            unsafe { &mut *self.ptr.as_ptr() }
+        }
+
+        /// Borrows a new mutable reference from the unique borrow initially captured.
+        ///
+        /// # Safety
+        ///
+        /// The reborrow must have ended, i.e., the reference returned by `new` and
+        /// all pointers and references derived from it, must not be used anymore.
+        pub(super) unsafe fn reborrow(&mut self) -> &'a mut T {
+            // SAFETY: our own safety conditions imply this reference is again unique.
+            unsafe { &mut *self.ptr.as_ptr() }
+        }
+
+        /// Borrows a new shared reference from the unique borrow initially captured.
+        ///
+        /// # Safety
+        ///
+        /// The reborrow must have ended, i.e., the reference returned by `new` and
+        /// all pointers and references derived from it, must not be used anymore.
+        pub(super) unsafe fn reborrow_shared(&self) -> &'a T {
+            // SAFETY: our own safety conditions imply this reference is again unique.
+            unsafe { &*self.ptr.as_ptr() }
+        }
+    }
 }
 
 pub struct DropGuard<F: FnOnce()>(::core::mem::ManuallyDrop<F>);