diff --git a/crates/rbtree/src/lib.rs b/crates/rbtree/src/lib.rs index eccd47b..946f349 100644 --- a/crates/rbtree/src/lib.rs +++ b/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(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 +} diff --git a/crates/rbtree/src/raw_node.rs b/crates/rbtree/src/raw_node.rs index ae0f66f..88277ca 100644 --- a/crates/rbtree/src/raw_node.rs +++ b/crates/rbtree/src/raw_node.rs @@ -54,6 +54,26 @@ pub enum SearchResult { Empty, } +impl SearchResult> { + fn next_extant(self, inclusive: bool) -> Option> { + 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> { + 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 Clone for Handle { impl !Sync for Handle {} impl Handle { - fn from_node(node: Option>) -> Self { + pub fn from_node(node: Option>) -> Self { let Some(node) = node else { return Handle::EmptyRoot; }; @@ -218,8 +238,7 @@ impl Handle { } } - #[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 Handle { Some(side.map(|_| parent)) } + fn remove_from_parent_into_node(self) -> (Option>, Option) { + 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.parent_and_side().map(LeftOrRight::into_inner) } @@ -528,6 +575,18 @@ impl Handle { 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 { @@ -839,7 +898,7 @@ impl RBTree { pub fn pop_min(&mut self) -> Option> { let min = self.root_handle().minimum_of()?; - self.remove_node(min) + self.remove_handle(min) } #[must_use] @@ -852,11 +911,11 @@ impl RBTree { return None; }; - self.remove_node(z) + self.remove_handle(z) } #[must_use] - fn remove_node(&mut self, z: Handle) -> Option> { + pub fn remove_handle(&mut self, z: Handle) -> Option> { // 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 RBTree { range: TreeRange::full_range(self.root_handle()), } } + + pub fn range(&self, range: R) -> TreeIter<'_, N> + where + T: Ord + ?Sized, + N::Key: core::borrow::Borrow + Ord, + R: core::ops::RangeBounds, + { + 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, + _pd: PhantomData<&'a mut RBTree>, +} + +impl<'a, N: UnsafeNode> Drain<'a, N> { + fn next(&mut self) -> Option> { + 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; + + fn next(&mut self) -> Option { + self.next() + } } impl Default for RBTree { @@ -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) -> 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::::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::::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::>()); + } + + #[test] + fn range() { + let mut tree = RBTree::::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]); + } }