rbtree: range and drain

crates/rbtree/src/lib.rs

@@ -6,4 +6,22 @@ mod raw_node;
 
 extern crate alloc;
 
-pub use raw_node::{Color, RBTree, Side, TreeIter, TreeNodeIter, UnsafeNode};
+pub use raw_node::{
+    Color, Drain, Handle, LeftOrRight, RBTree, SearchResult, Side, TreeIter, TreeNodeIter,
+    UnsafeNode,
+};
+
+fn replace<T, R>(v: &mut T, f: impl FnOnce(T) -> (T, R)) -> R {
+    struct Guard;
+    impl Drop for Guard {
+        fn drop(&mut self) {
+            panic!("replace() panicked");
+        }
+    }
+    let guard = Guard;
+    let value = unsafe { core::ptr::read(v) };
+    let (new_value, ret) = f(value);
+    unsafe { core::ptr::write(v, new_value) };
+    core::mem::forget(guard);
+    ret
+}

crates/rbtree/src/raw_node.rs

@@ -54,6 +54,26 @@ pub enum SearchResult<T> {
     Empty,
 }
 
+impl<N: UnsafeNode> SearchResult<Handle<N>> {
+    fn next_extant(self, inclusive: bool) -> Option<Handle<N>> {
+        match self {
+            SearchResult::FoundAt(handle) if inclusive => Some(handle),
+            SearchResult::FoundAt(handle) | SearchResult::NotFoundAt(handle) => handle.next_of(),
+            SearchResult::Empty => None,
+        }
+    }
+
+    fn next_back_extant(self, inclusive: bool) -> Option<Handle<N>> {
+        match self {
+            SearchResult::FoundAt(handle) if inclusive => Some(handle),
+            SearchResult::FoundAt(handle) | SearchResult::NotFoundAt(handle) => {
+                handle.next_back_of()
+            }
+            SearchResult::Empty => None,
+        }
+    }
+}
+
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub enum Color {
     Red,
@@ -197,7 +217,7 @@ impl<N> Clone for Handle<N> {
 impl<N> !Sync for Handle<N> {}
 
 impl<N: UnsafeNode> Handle<N> {
-    fn from_node(node: Option<NonNull<N>>) -> Self {
+    pub fn from_node(node: Option<NonNull<N>>) -> Self {
         let Some(node) = node else {
             return Handle::EmptyRoot;
         };
@@ -218,8 +238,7 @@ impl<N: UnsafeNode> Handle<N> {
         }
     }
 
-    #[expect(dead_code)]
-    fn refresh_from_node(&mut self) {
+    pub fn refresh_from_node(&mut self) {
         *self = Self::from_node(self.node());
     }
 
@@ -396,6 +415,34 @@ impl<N: UnsafeNode> Handle<N> {
         Some(side.map(|_| parent))
     }
 
+    fn remove_from_parent_into_node(self) -> (Option<NonNull<N>>, Option<Self>) {
+        match self {
+            Handle::EmptyRoot => (None, None),
+            Handle::Root(node) => (Some(node), None),
+            Handle::Child {
+                parent,
+                node: LeftOrRight::Left(Some(node)),
+            } => {
+                unsafe { parent.as_ref().set_left(None) };
+                (Some(node), Some(Handle::from_node(Some(parent))))
+            }
+            Handle::Child {
+                parent,
+                node: LeftOrRight::Right(Some(node)),
+            } => {
+                unsafe { parent.as_ref().set_right(None) };
+                (Some(node), Some(Handle::from_node(Some(parent))))
+            }
+            Handle::Child {
+                parent,
+                node: LeftOrRight::Right(None) | LeftOrRight::Left(None),
+            } => {
+                unsafe { parent.as_ref().set_right(None) };
+                (None, Some(Handle::from_node(Some(parent))))
+            }
+        }
+    }
+
     fn parent(&self) -> Option<Self> {
         self.parent_and_side().map(LeftOrRight::into_inner)
     }
@@ -528,6 +575,18 @@ impl<N: UnsafeNode> Handle<N> {
         Some(current)
     }
 
+    /// moves to the least non-nil node in the tree, or `Err(self)`.
+    pub fn into_least_leaf(mut self) -> Self {
+        while let Some(child) = self
+            .left_child_extant()
+            .or_else(|| self.right_child_extant())
+        {
+            self = child;
+        }
+
+        self
+    }
+
     /// Returns the greatest non-nil node in the subtree rooted at `self`, or
     /// `None` if the subtree is empty.
     pub fn maximum_of(&self) -> Option<Self> {
@@ -839,7 +898,7 @@ impl<N: UnsafeNode> RBTree<N> {
     pub fn pop_min(&mut self) -> Option<NonNull<N>> {
         let min = self.root_handle().minimum_of()?;
 
-        self.remove_node(min)
+        self.remove_handle(min)
     }
 
     #[must_use]
@@ -852,11 +911,11 @@ impl<N: UnsafeNode> RBTree<N> {
             return None;
         };
 
-        self.remove_node(z)
+        self.remove_handle(z)
     }
 
     #[must_use]
-    fn remove_node(&mut self, z: Handle<N>) -> Option<NonNull<N>> {
+    pub fn remove_handle(&mut self, z: Handle<N>) -> Option<NonNull<N>> {
         // Y is either Z, the removed node in the case that Z has at most
         // one child, or Y is Z's successor which is guaranteed to have at most one
         // child (the right child).
@@ -1028,6 +1087,74 @@ impl<N: UnsafeNode> RBTree<N> {
             range: TreeRange::full_range(self.root_handle()),
         }
     }
+
+    pub fn range<T, R>(&self, range: R) -> TreeIter<'_, N>
+    where
+        T: Ord + ?Sized,
+        N::Key: core::borrow::Borrow<T> + Ord,
+        R: core::ops::RangeBounds<T>,
+    {
+        use core::ops::Bound;
+        let front = match range.start_bound() {
+            Bound::Included(key) => self.find_by_key(key).next_extant(true),
+            Bound::Excluded(key) => self.find_by_key(key).next_extant(false),
+            Bound::Unbounded => None,
+        }
+        .or_else(|| self.root_handle().minimum_of());
+
+        let back = match range.end_bound() {
+            Bound::Included(key) => self.find_by_key(key).next_back_extant(true),
+            Bound::Excluded(key) => self.find_by_key(key).next_back_extant(false),
+            Bound::Unbounded => None,
+        }
+        .or_else(|| self.root_handle().maximum_of());
+
+        TreeIter {
+            range: TreeRange {
+                start: front.map(RangeHandle::Node),
+                end: back.map(RangeHandle::Node),
+                _pd: PhantomData,
+            },
+        }
+    }
+
+    pub fn drain(&mut self) -> Drain<'_, N> {
+        let node = self.root_handle();
+        self.root = None;
+
+        Drain {
+            node,
+            _pd: PhantomData,
+        }
+    }
+}
+
+pub struct Drain<'a, N: UnsafeNode> {
+    node: Handle<N>,
+    _pd: PhantomData<&'a mut RBTree<N>>,
+}
+
+impl<'a, N: UnsafeNode> Drain<'a, N> {
+    fn next(&mut self) -> Option<NonNull<N>> {
+        if self.node.is_nil() {
+            return None;
+        }
+
+        crate::replace(&mut self.node, |node| {
+            let leaf = node.into_least_leaf();
+            let (node, parent) = leaf.remove_from_parent_into_node();
+
+            (parent.unwrap_or(Handle::EmptyRoot), node)
+        })
+    }
+}
+
+impl<'a, N: UnsafeNode> Iterator for Drain<'a, N> {
+    type Item = NonNull<N>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        self.next()
+    }
 }
 
 impl<N: UnsafeNode> Default for RBTree<N> {
@@ -1078,6 +1205,15 @@ impl<'a, N: UnsafeNode + 'a> TreeRange<'a, N> {
             _pd: PhantomData,
         }
     }
+
+    fn empty() -> Self {
+        Self {
+            start: None,
+            end: None,
+            _pd: PhantomData,
+        }
+    }
+
     fn full_range(root: Handle<N>) -> Self {
         Self {
             start: Some(RangeHandle::Root(root.clone())),
@@ -1121,8 +1257,8 @@ impl<'a, N: UnsafeNode + 'a> TreeRange<'a, N> {
         match self.end {
             None => None,
             _ => {
-                if self.start == self.end {
-                    self.end = None;
+                if current == self.end {
+                    *self = TreeRange::empty();
                 }
 
                 current.map(RangeHandle::into_inner)
@@ -1137,8 +1273,8 @@ impl<'a, N: UnsafeNode + 'a> TreeRange<'a, N> {
         match self.start {
             None => None,
             _ => {
-                if self.start == self.end {
-                    self.start = None;
+                if current == self.start {
+                    *self = TreeRange::empty();
                 }
 
                 current.map(RangeHandle::into_inner)
@@ -1195,6 +1331,8 @@ impl<'a, N: UnsafeNode + 'a> DoubleEndedIterator for TreeIter<'a, N> {
 mod tests {
     use std::cell::Cell;
 
+    use core::sync::atomic::{AtomicI32, Ordering};
+
     use super::*;
 
     #[derive(Debug)]
@@ -1438,4 +1576,69 @@ mod tests {
             assert!(matches!(result, SearchResult::Empty));
         }
     }
+
+    #[test]
+    fn drain() {
+        let mut tree = RBTree::<TestNode>::new();
+
+        let count = AtomicI32::new(0);
+
+        for node in (1..=10).map(|i| {
+            count.fetch_add(i, Ordering::SeqCst);
+            Box::into_raw(Box::new(TestNode::new(i)))
+        }) {
+            let None = tree.insert_node(unsafe { NonNull::new_unchecked(node) }) else {
+                panic!("duplicate node")
+            };
+        }
+
+        let drain = tree.drain();
+
+        drain.for_each(|node| {
+            let node = unsafe { Box::from_raw(node.as_ptr()) };
+            count.fetch_sub(node.key, Ordering::SeqCst);
+        });
+
+        assert_eq!(
+            count.load(Ordering::SeqCst),
+            0,
+            "all nodes should have been drained and dropped"
+        );
+    }
+
+    #[test]
+    fn drain_collect() {
+        let mut tree = RBTree::<TestNode>::new();
+
+        for node in (1..=10).map(|i| Box::into_raw(Box::new(TestNode::new(i)))) {
+            let None = tree.insert_node(unsafe { NonNull::new_unchecked(node) }) else {
+                panic!("duplicate node")
+            };
+        }
+
+        let mut nodes: Vec<_> = tree
+            .drain()
+            .map(|node| unsafe { Box::from_raw(node.as_ptr()).key })
+            .collect();
+
+        nodes.sort();
+
+        assert_eq!(nodes, (1..=10).collect::<Vec<_>>());
+    }
+
+    #[test]
+    fn range() {
+        let mut tree = RBTree::<TestNode>::new();
+
+        for node in (1..=10).map(|i| Box::into_raw(Box::new(TestNode::new(i)))) {
+            let None = tree.insert_node(unsafe { NonNull::new_unchecked(node) }) else {
+                panic!("duplicate node")
+            };
+        }
+
+        let range = tree.range(3..=7);
+
+        let keys: Vec<_> = range.cloned().collect();
+        assert_eq!(keys, vec![3, 4, 5, 6, 7]);
+    }
 }