diff --git a/crates/rbtree/src/lib.rs b/crates/rbtree/src/lib.rs index 76079f2..9ab3676 100644 --- a/crates/rbtree/src/lib.rs +++ b/crates/rbtree/src/lib.rs @@ -6,595 +6,4 @@ mod raw_node; extern crate alloc; -pub trait Node { - type Id: Eq + Copy; - type Key: Ord; - - fn parent(&self) -> Option; - fn set_parent(&mut self, parent: Option); - - fn left(&self) -> Option; - fn set_left(&mut self, left: Option); - - fn right(&self) -> Option; - fn set_right(&mut self, right: Option); - - fn child(&self, left: bool) -> Option { - if left { self.left() } else { self.right() } - } - - fn children(&self) -> (Option, Option) { - (self.left(), self.right()) - } - - fn key(&self) -> &Self::Key; - fn color(&self) -> bool; - fn set_color(&mut self, color: bool); -} - -#[allow(clippy::missing_safety_doc)] -pub trait NodeStore { - fn get(&self, id: N::Id) -> Option<&N>; - unsafe fn get_unchecked(&self, id: N::Id) -> &N { - self.get(id).unwrap() - } - fn get_mut(&mut self, id: N::Id) -> Option<&mut N>; - unsafe fn get_mut_unchecked(&mut self, id: N::Id) -> &mut N { - self.get_mut(id).unwrap() - } - fn insert(&mut self, node: N) -> N::Id; - fn remove(&mut self, id: N::Id) -> Option; -} - -pub struct RBTree> { - root: Option, - store: S, -} - -impl> RBTree { - pub fn new(store: S) -> Self { - Self { root: None, store } - } - - fn try_get(&self, id: N::Id) -> Option<&N> { - self.store.get(id) - } - - fn try_get_mut(&mut self, id: N::Id) -> Option<&mut N> { - self.store.get_mut(id) - } - - pub fn insert(&mut self, node: N) -> N::Id { - self.store.insert(node) - } - - pub fn minimum_of(&self, mut id: N::Id) -> N::Id { - while let Some(left_id) = self.store.get(id).and_then(|node| node.left()) { - id = left_id; - } - id - } - - pub fn maximum_of(&self, mut id: N::Id) -> N::Id { - while let Some(right_id) = self.store.get(id).and_then(|node| node.right()) { - id = right_id; - } - id - } - - pub fn next_of(&self, mut id: N::Id) -> Option { - let node = self.store.get(id)?; - match node.right() { - Some(r) => Some(self.minimum_of(r)), - None => { - let mut p = node.parent()?; - while id == self.store.get(p)?.right()? { - id = p; - p = self.store.get(p)?.parent()?; - } - Some(p) - } - } - } - - pub fn next_back_of(&self, mut id: N::Id) -> Option { - let node = self.store.get(id)?; - match node.left() { - Some(l) => Some(self.maximum_of(l)), - None => { - let mut p = node.parent()?; - while id == self.store.get(p)?.left()? { - id = p; - p = self.store.get(p)?.parent()?; - } - Some(p) - } - } - } - - // x y - // / \ / \ - // a y => x c - // / \ / \ - // b c a b - pub fn rotate_left(&mut self, x: N::Id) -> Option<()> { - let y = self.right_child_of(x)?; - let b = self.left_child_of(y); - - let x_parent = self.parent_of(x); - - if let Some(node) = self.store.get_mut(x) { - node.set_right(b) - } - - if let Some(b_mut) = b.and_then(|id| self.store.get_mut(id)) { - b_mut.set_parent(Some(x)); - } - - if let Some(node) = self.store.get_mut(y) { - node.set_left(Some(x)); - - if let Some(parent) = x_parent { - node.set_parent(Some(parent)); - if self.store.get(parent).unwrap().left() == Some(x) { - self.store.get_mut(parent).unwrap().set_left(Some(y)); - } else { - self.store.get_mut(parent).unwrap().set_right(Some(y)); - } - } else { - self.root = Some(y); - self.store.get_mut(y).unwrap().set_parent(None); - } - } - - Some(()) - } - - // x y - // / \ / \ - // y c => a x - // / \ / \ - // a b b c - pub fn rotate_right(&mut self, x: N::Id) -> Option<()> { - let y = self.left_child_of(x)?; - let b = self.right_child_of(y); - - let x_parent = self.parent_of(x); - - if let Some(node) = self.store.get_mut(x) { - node.set_left(b) - } - - if let Some(b_mut) = b.and_then(|id| self.store.get_mut(id)) { - b_mut.set_parent(Some(x)); - } - - if let Some(node) = self.store.get_mut(y) { - node.set_right(Some(x)); - - if let Some(parent) = x_parent { - node.set_parent(Some(parent)); - if self.store.get(parent).unwrap().left() == Some(x) { - self.store.get_mut(parent).unwrap().set_left(Some(y)); - } else { - self.store.get_mut(parent).unwrap().set_right(Some(y)); - } - } else { - self.root = Some(y); - self.store.get_mut(y).unwrap().set_parent(None); - } - } - - Some(()) - } - - pub fn rotate(&mut self, x: N::Id, left: bool) -> Option<()> { - if left { - self.rotate_left(x) - } else { - self.rotate_right(x) - } - } - - pub fn find_by_key(&self, key: &Q) -> SearchResult - where - N::Key: core::borrow::Borrow, - Q: Ord + ?Sized, - { - use core::borrow::Borrow; - use core::cmp::Ordering::*; - - let Some(mut current) = self.root else { - return SearchResult::Empty; - }; - - loop { - let node = self.store.get(current).unwrap(); - - match key.cmp(node.key().borrow()) { - Less => { - if let Some(left) = node.left() { - current = left; - } else { - return SearchResult::NotFoundLeftOf(current); - } - } - Greater => { - if let Some(right) = node.right() { - current = right; - } else { - return SearchResult::NotFoundRightOf(current); - } - } - Equal => return SearchResult::FoundAt(current), - } - } - } - - pub fn insert_id(&mut self, id: N::Id) { - let node = self.store.get(id).unwrap(); - - match self.find_by_key(node.key()) { - SearchResult::FoundAt(_) => { - panic!("Duplicate key insertion is not allowed"); - } - SearchResult::NotFoundRightOf(parent) => { - self.store.get_mut(parent).unwrap().set_right(Some(id)); - self.store.get_mut(id).unwrap().set_parent(Some(parent)); - self.store.get_mut(id).unwrap().set_color(true); // new node is always red - } - SearchResult::NotFoundLeftOf(parent) => { - self.store.get_mut(parent).unwrap().set_left(Some(id)); - self.store.get_mut(id).unwrap().set_parent(Some(parent)); - self.store.get_mut(id).unwrap().set_color(true); // new node is always red - } - SearchResult::Empty => { - let node = self.store.get_mut(id).unwrap(); - node.set_parent(None); - node.set_color(false); // root is always black - - self.root = Some(id); - } - } - - self.fix_insert(id); - } - - fn fix_insert(&mut self, mut id: N::Id) { - while let Some(parent) = self - .store - .get(id) - .and_then(|node| node.parent()) - .filter(|&p| self.store.get(p).unwrap().color()) - { - // gp is guaranteed to exist because parent is red and the root is black - let grandparent = self - .store - .get(parent) - .and_then(|node| node.parent()) - .unwrap(); - let (uncle, is_left) = if self.store.get(grandparent).unwrap().left() == Some(parent) { - (self.store.get(grandparent).unwrap().right(), true) - } else { - (self.store.get(grandparent).unwrap().left(), false) - }; - - if let Some(uncle_id) = uncle - && self.store.get(uncle_id).unwrap().color() - { - // Case 1: Uncle is red - self.store.get_mut(parent).unwrap().set_color(false); - self.store.get_mut(uncle_id).unwrap().set_color(false); - self.store.get_mut(grandparent).unwrap().set_color(true); - id = grandparent; - } else { - // Case 2: Uncle is black - if self.store.get(parent).unwrap().child(!is_left) == Some(id) { - // Case 2a: id is on the same side as uncle - self.rotate(parent, is_left); - id = parent; - } - - // Case 2b: id is on the opposite side as uncle - self.store.get_mut(parent).unwrap().set_color(false); - self.store.get_mut(grandparent).unwrap().set_color(true); - self.rotate(grandparent, !is_left); - } - } - - if let Some(root_id) = self.root { - self.store.get_mut(root_id).unwrap().set_color(false); - } - } - - fn color_of(&self, id: Option) -> bool { - id.map(|id| self.store.get(id).unwrap().color()) - .unwrap_or(false) - } - - fn set_color_of(&mut self, id: Option, color: bool) { - if let Some(id) = id { - self.store.get_mut(id).unwrap().set_color(color); - } - } - - fn parent_of(&self, id: N::Id) -> Option { - self.store.get(id).unwrap().parent() - } - - fn left_child_of(&self, id: N::Id) -> Option { - self.store.get(id).unwrap().left() - } - - fn right_child_of(&self, id: N::Id) -> Option { - self.store.get(id).unwrap().right() - } - - fn child_of(&self, id: N::Id, left: bool) -> Option { - self.store.get(id).unwrap().child(left) - } - - fn children_of(&self, id: N::Id) -> (Option, Option) { - self.store.get(id).unwrap().children() - } - - // When removing an node from an RB tree, we have to potentially fix the - // invariants of the tree starting from the node X that replaced a deleted - // node Y. - - // X is either root, or X is None and the child of a parent node P which is - // guaranteed to have a non-nil sibling W, since X is doubly-black, and so - // the path through W must hold at least 2 black nodes, including implicit - // black nil-leafs. - pub fn remove(&mut self, z: N::Id) { - // 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). - - let y = match self.children_of(z) { - (None, None) | (Some(_), None) | (None, Some(_)) => z, - // z has a successor, since it has two children. - _ => self.next_of(z).unwrap(), - }; - - // In either case, X is the only child of Y, if it exists. - let x = match self.children_of(y) { - (Some(left), _) => Some(left), - (None, Some(right)) => Some(right), - (None, None) => None, - }; - // Therefore, X must be red or None, and if X is red Y must be black. - - let parent = self.store.get(y).unwrap().parent(); - - // If Y's parent was None, then Y was the root, and the tree is empty. - let Some(parent) = parent else { - self.root = x; - - if let Some(x) = x { - self.store.get_mut(x).unwrap().set_parent(None); - self.store.get_mut(x).unwrap().set_color(false); - } - - return; - }; - - // X is promoted to Y's position, and Y is unlinked from the tree. - if self.store.get(parent).unwrap().left() == Some(y) { - self.store.get_mut(parent).unwrap().set_left(x); - } else { - self.store.get_mut(parent).unwrap().set_right(x); - } - - let color = self.store.get(y).unwrap().color(); - // If Y is Z's successor, move Y's data into Z (or move Z's meta into Y). - if y != z { - self.copy_meta_to(z, y); - } - - if let Some(x_id) = x { - // If X was red, color it black. Since it replaces a black node, the - // black-height of the subtree is preserved. - let x = self.store.get_mut(x_id).unwrap(); - x.set_parent(Some(parent)); - x.set_color(false); - - return; - } - - // If X is None and Y was black, then a black node was removed, and the tree - // needs to be rebalanced. - if !color { - // x is the NIL leaf child of parent, and is doubly black. - let mut x = None; - let mut parent = Some(parent); - while x != self.root && !self.color_of(x) { - // safe because x is not root. - let parent_id = parent.unwrap(); - - // w exists because it is the sibling of x; the subtree at x has - // a black-height of 2, the subtree at w must equally have a - // black-height of 2. - let (mut w, is_left) = if self.store.get(parent_id).unwrap().left() == x { - (self.store.get(parent_id).unwrap().right().unwrap(), true) - } else { - (self.store.get(parent_id).unwrap().left().unwrap(), false) - }; - - // Case 1: w is red -> parent and w's children are black - if self.color_of(Some(w)) { - // set w to black and parent to red - self.set_color_of(Some(w), false); - self.set_color_of(Some(parent_id), true); - - // rotate around parent such that w becomes the parent of parent - self.rotate(parent_id, is_left); - - // x hasn't change, but now the sibling is the child of w, which is black - w = if is_left { - self.store.get(parent_id).unwrap().right().unwrap() - } else { - self.store.get(parent_id).unwrap().left().unwrap() - }; - } - - // Case 2: w is black - - match ( - self.color_of(self.left_child_of(w)), - self.color_of(self.right_child_of(w)), - ) { - (false, false) => { - // Case 2a: w's children are both black - // x carries a phantom black, and its sibling w is black - // we can remove one black from x and w by colouring w - // red and giving x's phantom black to the - // parent. - self.set_color_of(Some(w), true); - x = Some(parent_id); - parent = self.store.get(parent_id).unwrap().parent(); - continue; - } - (near_cousin @ true, false) | (near_cousin @ false, true) => { - // Case 2b: x's near-cousin is red and its far-cousin is black - // colour the near-cousin black and w red, rotate so - // that the near cousin becomes the sibling of x - if near_cousin == is_left { - self.set_color_of(self.child_of(w, !is_left), false); - self.set_color_of(Some(w), true); - self.rotate(w, !is_left); - } - - // fall through to case 2c - } - _ => {} - } - - // Case 2c: x's far-cousin is red - // we can now rotate around the parent to balance the - // subtree at parent without increasing the - // black-height. - // However, in the case that the far-cousin is red and - // the parent is black, this will result, after - // rotating, in the far-cousin's path having fewer black - // nodes than the parent's path: we have taken a black - // node from above the far-cousin and moved it into its - // sibling branch. - // This is fixed by colouring w the colour of the - // parent, and colouring the parent and w's far-cousin - // black. - self.set_color_of(Some(w), self.color_of(Some(parent_id))); - self.set_color_of(Some(parent_id), false); - self.set_color_of(self.child_of(w, !is_left), false); - self.rotate(parent_id, is_left); - - // After the rotation, the inbalance has been resolved. - break; - } - } - } - - fn copy_meta_to(&mut self, from: N::Id, to: N::Id) { - let (parent, color, left, right) = { - let from_node = self.store.get(from).unwrap(); - ( - from_node.parent(), - from_node.color(), - from_node.left(), - from_node.right(), - ) - }; - - let to_node = self.store.get_mut(to).unwrap(); - - to_node.set_parent(parent); - to_node.set_color(color); - to_node.set_left(left); - to_node.set_right(right); - - if let Some(parent) = parent { - let parent_node = self.store.get_mut(parent).unwrap(); - if parent_node.left() == Some(from) { - parent_node.set_left(Some(to)); - } else { - parent_node.set_right(Some(to)); - } - } else { - self.root = Some(to); - } - - if let Some(left) = left { - self.store.get_mut(left).unwrap().set_parent(Some(to)); - } - if let Some(right) = right { - self.store.get_mut(right).unwrap().set_parent(Some(to)); - } - } -} - -pub enum SearchResult { - FoundAt(T), - NotFoundRightOf(T), - NotFoundLeftOf(T), - Empty, -} - -mod default_node { - use crate::Node; - - type DefaultNodeId = u64; - - struct DefaultNodeStore { - nodes: alloc::collections::BTreeMap, - next_id: DefaultNodeId, - } - - struct DefaultNode { - parent: Option, - left: Option, - right: Option, - key: u64, - color: bool, - } - - impl Node for DefaultNode { - type Id = DefaultNodeId; - type Key = u64; - - fn parent(&self) -> Option { - self.parent - } - - fn set_parent(&mut self, parent: Option) { - self.parent = parent; - } - - fn left(&self) -> Option { - self.left - } - - fn set_left(&mut self, left: Option) { - self.left = left; - } - - fn right(&self) -> Option { - self.right - } - - fn set_right(&mut self, right: Option) { - self.right = right; - } - - fn key(&self) -> &Self::Key { - &self.key - } - - fn color(&self) -> bool { - self.color - } - - fn set_color(&mut self, color: bool) { - self.color = color; - } - } -} +pub use raw_node::{RBTree, TreeIter, TreeNodeIter}; diff --git a/crates/rbtree/src/raw_node.rs b/crates/rbtree/src/raw_node.rs index 106baea..7d82702 100644 --- a/crates/rbtree/src/raw_node.rs +++ b/crates/rbtree/src/raw_node.rs @@ -28,11 +28,13 @@ pub trait TestDebugInspect: Sized { impl TestDebugInspect for T {} +#[allow(dead_code)] pub trait DebugInspect: Sized { // fn debug_inspect(self) -> Self { // dbg!(&self); // self // } + fn debug_inspect_with(self, f: F) -> Self where Self: core::fmt::Debug, @@ -102,7 +104,7 @@ pub enum LeftOrRight { } impl LeftOrRight> { - fn transpose(self) -> Option> { + pub fn transpose(self) -> Option> { match self { LeftOrRight::Left(Some(t)) => Some(LeftOrRight::Left(t)), LeftOrRight::Right(Some(t)) => Some(LeftOrRight::Right(t)), @@ -112,33 +114,35 @@ impl LeftOrRight> { } impl LeftOrRight { - fn map U>(self, f: F) -> LeftOrRight { + pub fn map U>(self, f: F) -> LeftOrRight { match self { LeftOrRight::Left(t) => LeftOrRight::Left(f(t)), LeftOrRight::Right(t) => LeftOrRight::Right(f(t)), } } - fn as_ref(&self) -> LeftOrRight<&T> { + + pub fn as_ref(&self) -> LeftOrRight<&T> { match self { LeftOrRight::Left(t) => LeftOrRight::Left(t), LeftOrRight::Right(t) => LeftOrRight::Right(t), } } - fn as_inner(&self) -> &T { - match self { - LeftOrRight::Left(t) => t, - LeftOrRight::Right(t) => t, - } - } - fn as_inner_mut(&mut self) -> &mut T { + pub fn as_inner(&self) -> &T { match self { LeftOrRight::Left(t) => t, LeftOrRight::Right(t) => t, } } - fn into_inner(self) -> T { + pub fn as_inner_mut(&mut self) -> &mut T { + match self { + LeftOrRight::Left(t) => t, + LeftOrRight::Right(t) => t, + } + } + + pub fn into_inner(self) -> T { match self { LeftOrRight::Left(t) => t, LeftOrRight::Right(t) => t, @@ -213,10 +217,12 @@ impl Handle { } } + #[expect(dead_code)] fn refresh_from_node(&mut self) { *self = Self::from_node(self.node()); } + #[expect(dead_code)] fn is_root(&self) -> bool { matches!(self, Handle::Root(_) | Handle::EmptyRoot) } @@ -647,12 +653,12 @@ impl Handle { } } -struct RBTree { +pub struct RBTree { root: Option>, } impl RBTree { - fn new() -> Self { + pub fn new() -> Self { Self { root: None } } @@ -669,7 +675,7 @@ impl RBTree { self.root = handle.node(); } - fn find_by_key(&self, key: &Q) -> SearchResult> + pub fn find_by_key(&self, key: &Q) -> SearchResult> where N::Key: core::borrow::Borrow, Q: Ord + ?Sized, @@ -818,7 +824,7 @@ impl RBTree { None } - fn remove(&mut self, key: &Q) -> Option> + pub fn remove(&mut self, key: &Q) -> Option> where N::Key: core::borrow::Borrow, Q: Ord + ?Sized, @@ -987,13 +993,19 @@ impl RBTree { z.node() } - fn iter(&self) -> TreeIter<'_, N> { - TreeIter { + pub fn iter(&self) -> TreeNodeIter<'_, N> { + TreeNodeIter { range: TreeRange::full_range(self.root_handle()), } } } +impl Default for RBTree { + fn default() -> Self { + Self::new() + } +} + #[derive(Debug, Clone)] enum RangeHandle { Root(Handle), @@ -1028,6 +1040,7 @@ struct TreeRange<'a, N: UnsafeNode> { } impl<'a, N: UnsafeNode + 'a> TreeRange<'a, N> { + #[expect(dead_code)] fn new(start: Handle, end: Handle) -> Self { Self { start: Some(RangeHandle::Node(start)), @@ -1104,11 +1117,11 @@ impl<'a, N: UnsafeNode + 'a> TreeRange<'a, N> { } } -struct TreeIter<'a, N: UnsafeNode> { +pub struct TreeNodeIter<'a, N: UnsafeNode> { range: TreeRange<'a, N>, } -impl<'a, N: UnsafeNode + 'a> Iterator for TreeIter<'a, N> { +impl<'a, N: UnsafeNode + 'a> Iterator for TreeNodeIter<'a, N> { type Item = Handle; fn next(&mut self) -> Option { @@ -1116,12 +1129,34 @@ impl<'a, N: UnsafeNode + 'a> Iterator for TreeIter<'a, N> { } } -impl<'a, N: UnsafeNode + 'a> DoubleEndedIterator for TreeIter<'a, N> { +impl<'a, N: UnsafeNode + 'a> DoubleEndedIterator for TreeNodeIter<'a, N> { fn next_back(&mut self) -> Option { self.range.next_back() } } +pub struct TreeIter<'a, N: UnsafeNode> { + range: TreeRange<'a, N>, +} + +impl<'a, N: UnsafeNode + 'a> Iterator for TreeIter<'a, N> { + type Item = &'a N::Key; + + fn next(&mut self) -> Option { + self.range + .next() + .map(|n| unsafe { n.node().unwrap().as_ref().key() }) + } +} + +impl<'a, N: UnsafeNode + 'a> DoubleEndedIterator for TreeIter<'a, N> { + fn next_back(&mut self) -> Option { + self.range + .next_back() + .map(|n| unsafe { n.node().unwrap().as_ref().key() }) + } +} + #[cfg(test)] mod tests { use std::cell::Cell;