diff --git a/crates/rbtree/Cargo.toml b/crates/rbtree/Cargo.toml index 6b15f41..ce35266 100644 --- a/crates/rbtree/Cargo.toml +++ b/crates/rbtree/Cargo.toml @@ -3,4 +3,8 @@ name = "rbtree" version = "0.1.0" edition = "2024" +[features] +default = [] +std = [] + [dependencies] diff --git a/crates/rbtree/src/lib.rs b/crates/rbtree/src/lib.rs index 1252a8a..76079f2 100644 --- a/crates/rbtree/src/lib.rs +++ b/crates/rbtree/src/lib.rs @@ -1,5 +1,6 @@ #![cfg_attr(not(test), no_std)] #![feature(negative_impls)] +#![cfg_attr(test, feature(box_vec_non_null))] mod raw_node; diff --git a/crates/rbtree/src/raw_node.rs b/crates/rbtree/src/raw_node.rs index 2dbf81a..106baea 100644 --- a/crates/rbtree/src/raw_node.rs +++ b/crates/rbtree/src/raw_node.rs @@ -1,6 +1,56 @@ +use core::marker::PhantomData; +use core::mem; use core::ops::Not; use core::ptr::NonNull; +#[cfg(all(test, feature = "std"))] +extern crate std; + +#[cfg(test)] +pub trait TestDebug: core::fmt::Debug {} +#[cfg(test)] +impl TestDebug for T {} + +#[cfg(not(test))] +pub trait TestDebug {} +#[cfg(not(test))] +impl TestDebug for T {} + +pub trait TestDebugInspect: Sized { + fn test_debug_inspect_with(self, f: F) -> Self + where + Self: TestDebug, + { + f(&self); + self + } +} + +impl TestDebugInspect for T {} + +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, + { + f(&self); + self + } +} + +impl DebugInspect for T {} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum SearchResult { + FoundAt(T), + NotFoundAt(T), + Empty, +} + #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum Color { Red, @@ -24,8 +74,8 @@ impl Not for Side { } } -pub trait UnsafeNode { - type Key: Eq + Ord; +pub unsafe trait UnsafeNode: TestDebug { + type Key: Eq + Ord + TestDebug; fn left(&self) -> Option>; fn right(&self) -> Option>; fn parent(&self) -> Option>; @@ -36,14 +86,38 @@ pub trait UnsafeNode { fn set_right(&self, right: Option>); fn set_parent(&self, parent: Option>); fn set_color(&self, color: Color); + + fn copy_meta_from(&self, other: &Self) { + self.set_color(other.color()); + self.set_parent(other.parent()); + self.set_left(other.left()); + self.set_right(other.right()); + } } +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Copy)] pub enum LeftOrRight { Left(T), Right(T), } +impl LeftOrRight> { + fn transpose(self) -> Option> { + match self { + LeftOrRight::Left(Some(t)) => Some(LeftOrRight::Left(t)), + LeftOrRight::Right(Some(t)) => Some(LeftOrRight::Right(t)), + _ => None, + } + } +} + impl LeftOrRight { + 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> { match self { LeftOrRight::Left(t) => LeftOrRight::Left(t), @@ -63,9 +137,17 @@ impl LeftOrRight { LeftOrRight::Right(t) => t, } } + + fn into_inner(self) -> T { + match self { + LeftOrRight::Left(t) => t, + LeftOrRight::Right(t) => t, + } + } } -pub enum Handle { +#[derive(Debug)] +pub enum Handle { EmptyRoot, Root(NonNull), Child { @@ -74,13 +156,103 @@ pub enum Handle { }, } -impl !Sync for Handle {} +impl Eq for Handle {} +impl PartialEq for Handle { + fn eq(&self, other: &Self) -> bool { + match (self, other) { + (Self::Root(l0), Self::Root(r0)) => l0 == r0, + ( + Self::Child { + parent: l_parent, + node: l_node, + }, + Self::Child { + parent: r_parent, + node: r_node, + }, + ) => l_parent == r_parent && l_node == r_node, + _ => core::mem::discriminant(self) == core::mem::discriminant(other), + } + } +} + +impl Clone for Handle { + fn clone(&self) -> Self { + match self { + Handle::EmptyRoot => Handle::EmptyRoot, + Handle::Root(node) => Handle::Root(*node), + Handle::Child { parent, node } => Handle::Child { + parent: *parent, + node: *node, + }, + } + } +} + +impl !Sync for Handle {} impl Handle { + fn from_node(node: Option>) -> Self { + let Some(node) = node else { + return Handle::EmptyRoot; + }; + + let parent = unsafe { node.as_ref().parent() }; + + match parent { + Some(parent) => { + let side = if unsafe { parent.as_ref().left() } == Some(node) { + LeftOrRight::Left(Some(node)) + } else { + LeftOrRight::Right(Some(node)) + }; + + Handle::Child { parent, node: side } + } + None => Handle::Root(node), + } + } + + fn refresh_from_node(&mut self) { + *self = Self::from_node(self.node()); + } + fn is_root(&self) -> bool { matches!(self, Handle::Root(_) | Handle::EmptyRoot) } + fn is_red(&self) -> bool { + self.color() == Color::Red + } + + fn is_black(&self) -> bool { + self.color() == Color::Black + } + + fn set_color(&mut self, color: Color) { + if let Some(node) = self.node() { + unsafe { node.as_ref().set_color(color) }; + } + } + + fn side(&self) -> Option { + match self { + Handle::EmptyRoot => None, + Handle::Root(_) => None, + Handle::Child { node, .. } => match node { + LeftOrRight::Left(_) => Some(Side::Left), + LeftOrRight::Right(_) => Some(Side::Right), + }, + } + } + + fn child(&self, side: Side) -> Option { + match side { + Side::Left => self.left_child(), + Side::Right => self.right_child(), + } + } + fn left_child(&self) -> Option { let &parent = match self { Handle::EmptyRoot => return None, @@ -111,29 +283,74 @@ impl Handle { }) } + fn right_child_extant(&self) -> Option { + self.right_child().filter(|child| !child.is_nil()) + } + + fn left_child_extant(&self) -> Option { + self.left_child().filter(|child| !child.is_nil()) + } + + fn into_has_node(self) -> Option { + self.node().map(|_| self) + } + + fn children(&self) -> (Option, Option) { + (self.left_child(), self.right_child()) + } + + fn non_nil_children(&self) -> (Option, Option) { + ( + self.left_child().and_then(Self::into_has_node), + self.right_child().and_then(Self::into_has_node), + ) + } + + fn sibling(&self) -> Option { + let parent = self.parent()?; + + match self { + Self::Child { node, .. } => match node { + LeftOrRight::Left(_) => parent.right_child(), + LeftOrRight::Right(_) => parent.left_child(), + }, + _ => unreachable!(), + } + } + fn set_left_child(&mut self, child: &mut Self) { - if let Some(parent) = match self { - Handle::EmptyRoot => return, - Handle::Root(parent) => Some(parent), - Handle::Child { node, .. } => node.as_inner_mut().as_mut(), - } { + if let Some(parent) = self.node() { + let mut child_node = child.node(); + unsafe { - parent.as_ref().set_left(child.node()); - child.set_parent(LeftOrRight::Left(Some(*parent))); + parent.as_ref().set_left(child_node); + if let Some(child) = child_node.as_mut().map(|c| c.as_mut()) { + child.set_parent(Some(parent)); + } } + + *child = Handle::Child { + parent, + node: LeftOrRight::Left(child_node), + }; } } fn set_right_child(&mut self, child: &mut Self) { - if let Some(parent) = match self { - Handle::EmptyRoot => return, - Handle::Root(parent) => Some(parent), - Handle::Child { node, .. } => node.as_inner_mut().as_mut(), - } { + if let Some(parent) = self.node() { + let mut child_node = child.node(); + unsafe { - parent.as_ref().set_right(child.node()); - child.set_parent(LeftOrRight::Right(Some(*parent))); + parent.as_ref().set_right(child_node); + if let Some(child) = child_node.as_mut().map(|c| c.as_mut()) { + child.set_parent(Some(parent)); + } } + + *child = Handle::Child { + parent, + node: LeftOrRight::Right(child_node), + }; } } @@ -145,49 +362,62 @@ impl Handle { } } - fn parent(&self) -> Option { - let &parent = match self { + fn parent_and_side(&self) -> Option> { + let (&parent, side) = match self { Handle::EmptyRoot => return None, Handle::Root(_) => return None, - Handle::Child { parent, .. } => parent, + Handle::Child { parent, node } => (parent, node.map(|_| ())), }; let grandparent = unsafe { parent.as_ref().parent() }; - if let Some(grandparent) = grandparent { + let parent = if let Some(grandparent) = grandparent { let node = if unsafe { grandparent.as_ref().left() } == Some(parent) { LeftOrRight::Left(Some(parent)) } else { LeftOrRight::Right(Some(parent)) }; - Some(Handle::Child { + Handle::Child { parent: grandparent, node, - }) + } } else { - Some(Handle::Root(parent)) + Handle::Root(parent) + }; + + Some(side.map(|_| parent)) + } + + fn parent(&self) -> Option { + self.parent_and_side().map(LeftOrRight::into_inner) + } + + fn set_parent_handle(&mut self, parent: Option>>) { + match parent { + Some(LeftOrRight::Left(mut parent)) => { + parent.set_left_child(self); + } + Some(LeftOrRight::Right(mut parent)) => { + parent.set_right_child(self); + } + None => match self.node() { + Some(node) => { + unsafe { + node.as_ref().set_parent(None); + } + *self = Handle::Root(node); + } + None => { + *self = Handle::EmptyRoot; + } + }, } } - fn set_parent(&mut self, parent: LeftOrRight>>) { - use LeftOrRight::*; - let node = self.node(); - - *self = match parent { - Left(Some(parent)) => Self::Child { - parent, - node: Left(node), - }, - Right(Some(parent)) => Self::Child { - parent, - node: Right(node), - }, - Left(None) | Right(None) => match node { - Some(node) => Self::Root(node), - None => Self::EmptyRoot, - }, - }; + fn set_parent(&mut self, parent: Option>>) { + let parent = parent.map(|p| p.map(|p| Handle::from_node(Some(p)))); + self.set_parent_handle(parent); } fn color(&self) -> Color { @@ -200,7 +430,9 @@ impl Handle { } } - fn rotate_left(mut self) -> Result { + /// rotate `self` left, returning the new root of the subtree, or + /// `Err(self)` if `self` has no left child + fn rotate_left(mut self, on_root: F) -> Result { let Some(mut y) = self.right_child() else { return Err(self); }; @@ -209,14 +441,26 @@ impl Handle { return Err(self); }; + let parent = self + .parent_and_side() + .and_then(|p| p.map(|p| p.node()).transpose()); + + y.set_parent(parent); + self.set_right_child(&mut b); y.set_left_child(&mut self); + if matches!(y, Handle::Root(_)) { + on_root(y.clone()); + } + Ok(y) } - fn rotate_right(mut self) -> Result { + /// rotate `self` right, returning the new root of the subtree, or + /// `Err(self)` if `self` has no left child + fn rotate_right(mut self, on_root: F) -> Result { let Some(mut y) = self.left_child() else { return Err(self); }; @@ -225,17 +469,904 @@ impl Handle { return Err(self); }; + let parent = self + .parent_and_side() + .and_then(|p| p.map(|p| p.node()).transpose()); + + y.set_parent(parent); + self.set_left_child(&mut b); y.set_right_child(&mut self); + if matches!(y, Handle::Root(_)) { + on_root(y.clone()); + } + Ok(y) } - pub fn rotate(self, side: Side) -> Result { + /// rotate `self` into the position of its parent. + /// Returns `Ok(self)` in the new position, or `Err(self)` if `self` is the + /// root of the tree. + fn rotate_up(self, on_root: F) -> Result { + let Some(parent) = self.parent() else { + on_root(self.clone()); + return Err(self); + }; + + match self.side() { + Some(Side::Left) => parent.rotate_right(on_root).map_err(|_| self), + Some(Side::Right) => parent.rotate_left(on_root).map_err(|_| self), + None => unreachable!(), + } + } + + /// rotate `self` into the direction of `side`, returning the new root of the subtree, or + /// `Err(self)` if `self` has no left child + pub fn rotate(self, side: Side, on_root: F) -> Result { match side { - Side::Left => self.rotate_left(), - Side::Right => self.rotate_right(), + Side::Left => self.rotate_left(on_root), + Side::Right => self.rotate_right(on_root), + } + } + + /// Returns the least non-nil node in the subtree rooted at `self`, or + /// `None` if the subtree is empty. + pub fn minimum_of(&self) -> Option { + let mut current = self.clone(); + while let Some(left) = current.left_child_extant() { + current = left; + } + Some(current) + } + + /// 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 { + let mut current = self.clone(); + while let Some(right) = current.right_child_extant() { + current = right; + } + Some(current) + } + + /// Returns the next greater non-nil node in the tree, or `None` if `self` + /// is the greatest node. + pub fn next_of(&self) -> Option { + match self.right_child_extant() { + Some(right) => right.minimum_of(), + _ => { + let mut current = self.clone(); + while let Some(parent) = current.parent() { + if current.side() == Some(Side::Left) { + return Some(parent); + } + + current = parent; + } + None + } + } + } + + /// Returns the next smaller non-nil node in the tree, or `None` if `self` + /// is the smallest node. + pub fn next_back_of(&self) -> Option { + match self.left_child_extant() { + Some(left) => left.maximum_of(), + None => { + let mut current = self.clone(); + while let Some(parent) = current.parent() { + if parent.right_child().as_ref() == Some(¤t) { + return Some(parent); + } + current = parent; + } + None + } + } + } + + pub fn is_nil(&self) -> bool { + match self { + Handle::EmptyRoot => true, + Handle::Root(_) => false, + Handle::Child { node, .. } => node.as_inner().is_none(), + } + } + + pub fn make_nil(&mut self) -> Option> { + match self { + Handle::EmptyRoot => None, + &mut Handle::Root(node) => { + *self = Handle::EmptyRoot; + Some(node) + } + &mut Handle::Child { mut node, .. } => unsafe { + let (left, right) = self.non_nil_children(); + + if let Some(mut left) = left { + left.set_parent(None); + } + + if let Some(mut right) = right { + right.set_parent(None); + } + + if let Some(parent) = self + .parent_and_side() + .and_then(|p| p.map(|p| p.node()).transpose()) + { + match parent { + LeftOrRight::Left(parent) => parent.as_ref().set_left(None), + LeftOrRight::Right(parent) => parent.as_ref().set_right(None), + } + } + + node.as_inner_mut().take() + }, + } + } + + #[must_use = "inserting a node may replace an existing node, which must be deallocated"] + pub fn insert(&mut self, new_node: NonNull) -> Option> { + match self { + Self::EmptyRoot => { + *self = Handle::Root(new_node); + None + } + Self::Root(old) => { + let old = *old; + *self = Handle::Root(new_node); + Some(old) + } + Self::Child { node, parent } => { + let old = match node { + LeftOrRight::Left(old) => unsafe { + parent.as_ref().set_left(Some(new_node)); + old.replace(new_node) + }, + LeftOrRight::Right(old) => unsafe { + parent.as_ref().set_right(Some(new_node)); + old.replace(new_node) + }, + }; + + unsafe { + if let Some(old) = old { + new_node.as_ref().copy_meta_from(old.as_ref()); + } else { + new_node.as_ref().set_parent(Some(*parent)); + } + } + + old + } + } + } +} + +struct RBTree { + root: Option>, +} + +impl RBTree { + fn new() -> Self { + Self { root: None } + } + + fn root_handle(&self) -> Handle { + match self.root { + Some(root) => Handle::Root(root), + None => Handle::EmptyRoot, + } + } + + fn set_root_handle(&mut self, mut handle: Handle) { + handle.set_parent(None); + handle.set_color(Color::Black); + self.root = handle.node(); + } + + 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 mut current = self.root_handle(); + + loop { + let node = match ¤t { + Handle::Root(node) => *node, + Handle::Child { node, .. } => match node.into_inner() { + Some(node) => node, + None => return SearchResult::NotFoundAt(current), + }, + _ => return SearchResult::Empty, + }; + + match unsafe { node.as_ref().key().borrow().cmp(key) } { + Less => { + current = current.right_child().expect("current is an occupied node"); + } + Greater => { + current = current.left_child().expect("current is an occupied node"); + } + Equal => { + return SearchResult::FoundAt(current); + } + } + } + } + + pub fn insert_node(&mut self, new_node: NonNull) -> Option> { + let node_ref = unsafe { new_node.as_ref() }; + + let mut entry = match self + .find_by_key(node_ref.key()) + .test_debug_inspect_with(|_res| { + #[cfg(all(test, feature = "std"))] + eprintln!("find_by_key({:?}) = {:?}", node_ref.key(), _res); + }) { + SearchResult::FoundAt(mut entry) => { + return entry.insert(new_node); + } + SearchResult::Empty => { + node_ref.set_color(Color::Black); + node_ref.set_parent(None); + self.set_root_handle(Handle::Root(new_node)); + + return None; + } + SearchResult::NotFoundAt(mut entry) => { + unsafe { new_node.as_ref().set_color(Color::Red) }; + _ = entry.insert(new_node); + + entry + } + }; + + // Fixing + + // we've introduced a new red node, but one of the invariants of the red-black tree is that red nodes cannot have red children. + // In the case that our new node's parent is red, we need to fix the tree: + while let Some(mut parent) = entry.parent() + && parent.is_red() + { + // since the parent is red, it must have a grandparent (and an + // uncle), since the root of the tree is always black. + let mut gp = parent + .parent() + .expect("parent is red, so it must have a grandparent"); + + let mut uncle = parent + .sibling() + .expect("parent is red, so it must have a grandparent, and thus an uncle"); + + if uncle.is_red() { + // Case 1: + // If the uncle is red, we recolour both the parent and uncle black + // and the grandparent red, preserving the black-height of the tree. + // + // In the next loop, we look at the grandparent, which might have + // had a red parent. + + parent.set_color(Color::Black); + uncle.set_color(Color::Black); + gp.set_color(Color::Red); + + entry = gp; + } else { + // Case 2: + // If the uncle is black, we perform one (two) rotations. + // our current subtree (starting from the grandparent) looks like this: + // + // (1) (2) + // GP(B) GP(B) + // / \ / \ + // P(R) U(B) P(R) U(B) + // / \ / \ + // ... E(R) E(R) ... + + if uncle.side() == entry.side() { + // If the uncle has the same sidedness as the entry (1), + // we rotate the entry up into the parent's position and end + // up with a tree with the same shape. + + entry + .rotate_up(|_| panic!("entry cannot become root")) + .unwrap_or_else(|_| panic!("entry is not root, so it must have a parent")); + } + + // We rotate the grandparent towards the uncle to end up with + // the following shape: + // + // P(B) + // / \ + // E(R) GP(R) + // / \ + // ... U(B) + // the grandparent is painted red in order to preserve the black-height of the tree. + gp.set_color(Color::Red); + + let mut parent = gp + .rotate( + uncle + .side() + .expect("uncle is not root, so it must have a side"), + |root| self.set_root_handle(root), + ) + .unwrap_or_else(|_| panic!("grandparent has children")); + + // Whatever node ends up in the grandparent's position (either + // `entry` or `parent`) is painted black, resolving the red-red + // violation and replacing the grandparents black-level within + // the tree. + parent.set_color(Color::Black); + + // We are done fixing the tree, so we can break out of the loop. + break; + } + } + + self.root_handle().set_color(Color::Black); + + None + } + + fn remove(&mut self, key: &Q) -> Option> + where + N::Key: core::borrow::Borrow, + Q: Ord + ?Sized, + { + let SearchResult::FoundAt(z) = self.find_by_key(key) else { + return None; + }; + + // 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 z.non_nil_children() { + (Some(_), Some(_)) => z + .next_of() + .expect("z has a right child, so it must have a successor"), + _ => z.clone(), + }; + + // In either case, X is the only child of Y, if it exists. + let mut x = match y.children() { + (Some(left), Some(right)) => { + if !right.is_nil() { + right + } else { + left + } + } + _ => unreachable!("y is present, so it has children"), + }; + + let parent = y.parent_and_side(); + let color = y.color(); + + // If Y is Z's successor, move Y's data into Z (or move Z's meta into Y). + if y != z { + unsafe { + y.node() + .expect("z is occupied") + .as_ref() + .copy_meta_from(z.node().expect("z is occupied").as_ref()); + } + } + + let Some(parent) = parent else { + // If Y was the root, and the tree is empty. + self.set_root_handle(x); + return z.node(); + }; + + // X is promoted to Y's position, and Y is unlinked from the tree. + x.set_parent_handle(Some(parent)); + + if !x.is_nil() { + // If X is not nil, it must be red, since its parent, Y, must be + // black (otherwise a red node would have a red child), and since Y + // has only one child and nil leafs are implicitly black, Y would + // have been imbalanced if X were black. + // Since X is replacing a black node (Y), the black-height is + // preserved by painting X black. + x.set_color(Color::Black); + + return z.node(); + } + + // Fixing + + if color == Color::Black { + // If Y was black, then we have a black-height violation, since X + // was nil (and thus black) and replaced a black node (Y). + + // x is a NIL leaf and doubly black. + let mut x = x; + + while x != self.root_handle() && x.is_black() { + let mut parent = x.parent().expect("x is not root, so it must have a parent"); + + // W is X's sibling, and must exist because the subtree at X has + // a black-height of 2, and the two subtrees of the parent must + // have the same black-height. + let mut w = x.sibling().expect("w exists because x is black-deficient"); + + // Case 1: W is red -> parent and W's children are black + if w.is_red() { + w.set_color(Color::Black); + parent.set_color(Color::Red); + + parent = w + .rotate_up(|root| self.set_root_handle(root)) + .unwrap_or_else(|_| panic!("w is the child of parent")); + + // X's sibling has changed + w = x.sibling().expect("w exists because x is black-deficient"); + } + + // Case 2: W is black + assert!( + w.is_black(), + "w is black because it was red in the previous case" + ); + + // W's children exist because W exists + match ( + w.left_child().map(|c| c.color()).unwrap_or(Color::Black), + w.right_child().map(|c| c.color()).unwrap_or(Color::Black), + ) { + (Color::Black, Color::Black) => { + // 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. + w.set_color(Color::Red); + x = parent; + continue; + } + (Color::Red, Color::Black) | (Color::Black, Color::Red) + if w.child(x.side().unwrap()).unwrap().is_red() => + { + // 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 + let mut near_cousin = w + .child(x.side().unwrap()) + .expect("near cousin exists because it is red"); + + near_cousin.set_color(Color::Black); + w.set_color(Color::Red); + w = near_cousin + .rotate_up(|_| panic!("this shouldn't be root")) + .unwrap_or_else(|_| panic!("near cousin is the child of w")); + + // 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 X's far-cousin + // black. + + w.set_color(parent.color()); + parent.set_color(Color::Black); + let mut far_cousin = w + .child(x.side().unwrap().not()) + .expect("far cousin exists because it is red"); + far_cousin.set_color(Color::Black); + + w.rotate_up(|root| self.set_root_handle(root)) + .unwrap_or_else(|_| panic!("w is the child of parent")); + } + } + + z.node() + } + + fn iter(&self) -> TreeIter<'_, N> { + TreeIter { + range: TreeRange::full_range(self.root_handle()), + } + } +} + +#[derive(Debug, Clone)] +enum RangeHandle { + Root(Handle), + Node(Handle), +} + +impl RangeHandle { + fn into_inner(self) -> Handle { + match self { + RangeHandle::Root(handle) => handle, + RangeHandle::Node(handle) => handle, + } + } +} + +impl Eq for RangeHandle {} + +impl PartialEq for RangeHandle { + fn eq(&self, other: &Self) -> bool { + match (self, other) { + (Self::Root(l0), Self::Root(r0)) => l0 == r0, + (Self::Node(l0), Self::Node(r0)) => l0 == r0, + _ => false, + } + } +} + +struct TreeRange<'a, N: UnsafeNode> { + start: Option>, + end: Option>, + _pd: PhantomData<&'a ()>, +} + +impl<'a, N: UnsafeNode + 'a> TreeRange<'a, N> { + fn new(start: Handle, end: Handle) -> Self { + Self { + start: Some(RangeHandle::Node(start)), + end: Some(RangeHandle::Node(end)), + _pd: PhantomData, + } + } + fn full_range(root: Handle) -> Self { + Self { + start: Some(RangeHandle::Root(root.clone())), + end: Some(RangeHandle::Root(root)), + _pd: PhantomData, + } + } + + fn init_front(&mut self) -> Option<&mut Handle> { + if let Some(RangeHandle::Root(root)) = &self.start { + self.start = Some(RangeHandle::Node( + root.minimum_of().unwrap_or_else(|| root.clone()), + )); + } + + match &mut self.start { + None => None, + Some(RangeHandle::Node(handle)) => Some(handle), + _ => unreachable!(), + } + } + + fn init_back(&mut self) -> Option<&mut Handle> { + if let Some(RangeHandle::Root(root)) = &self.end { + self.end = Some(RangeHandle::Node( + root.maximum_of().unwrap_or_else(|| root.clone()), + )); + } + + match &mut self.end { + None => None, + Some(RangeHandle::Node(handle)) => Some(handle), + _ => unreachable!(), + } + } + + fn next(&mut self) -> Option> { + let next = self.init_front()?.next_of().map(RangeHandle::Node); + let current = mem::replace(&mut self.start, next); + + match self.end { + None => None, + _ => { + if self.start == self.end { + self.end = None; + } + + current.map(RangeHandle::into_inner) + } + } + } + + fn next_back(&mut self) -> Option> { + let next = self.init_back()?.next_back_of().map(RangeHandle::Node); + let current = mem::replace(&mut self.end, next); + + match self.start { + None => None, + _ => { + if self.start == self.end { + self.start = None; + } + + current.map(RangeHandle::into_inner) + } + } + } +} + +struct TreeIter<'a, N: UnsafeNode> { + range: TreeRange<'a, N>, +} + +impl<'a, N: UnsafeNode + 'a> Iterator for TreeIter<'a, N> { + type Item = Handle; + + fn next(&mut self) -> Option { + self.range.next() + } +} + +impl<'a, N: UnsafeNode + 'a> DoubleEndedIterator for TreeIter<'a, N> { + fn next_back(&mut self) -> Option { + self.range.next_back() + } +} + +#[cfg(test)] +mod tests { + use std::cell::Cell; + + use super::*; + + #[derive(Debug)] + struct TestNode { + key: i32, + left: Cell>>, + right: Cell>>, + parent: Cell>>, + color: Cell, + } + + unsafe impl UnsafeNode for TestNode { + type Key = i32; + + fn left(&self) -> Option> { + self.left.get() + } + + fn right(&self) -> Option> { + self.right.get() + } + + fn parent(&self) -> Option> { + self.parent.get() + } + + fn key(&self) -> &Self::Key { + &self.key + } + + fn color(&self) -> Color { + self.color.get() + } + + fn set_left(&self, left: Option>) { + self.left.set(left); + } + + fn set_right(&self, right: Option>) { + self.right.set(right); + } + + fn set_parent(&self, parent: Option>) { + self.parent.set(parent); + } + + fn set_color(&self, color: Color) { + self.color.set(color); + } + } + + impl TestNode { + fn new(key: i32) -> Self { + Self { + key, + left: Cell::new(None), + right: Cell::new(None), + parent: Cell::new(None), + color: Cell::new(Color::Red), + } + } + } + + #[test] + fn next_of() { + let mut tree = RBTree::::new(); + tree.insert_node(Box::into_non_null(Box::new(TestNode::new(1)))); + + assert_eq!(tree.root_handle().next_of(), None); + assert_eq!(tree.root_handle().next_back_of(), None); + } + + #[test] + fn rotate() { + struct DummyTree { + a: NonNull, + x: NonNull, + b: NonNull, + y: NonNull, + c: NonNull, + } + + impl DummyTree { + fn new() -> Self { + let mut a = Box::into_non_null(Box::new(TestNode::new(1))); + let mut x = Box::into_non_null(Box::new(TestNode::new(2))); + let mut b = Box::into_non_null(Box::new(TestNode::new(3))); + let mut y = Box::into_non_null(Box::new(TestNode::new(4))); + let mut c = Box::into_non_null(Box::new(TestNode::new(5))); + + unsafe { + x.as_ref().set_left(Some(a)); + x.as_ref().set_right(Some(y)); + y.as_ref().set_left(Some(b)); + y.as_ref().set_right(Some(c)); + + a.as_ref().set_parent(Some(x)); + y.as_ref().set_parent(Some(x)); + + b.as_ref().set_parent(Some(y)); + c.as_ref().set_parent(Some(y)); + } + + Self { a, x, b, y, c } + } + } + + impl Drop for DummyTree { + fn drop(&mut self) { + unsafe { + _ = Box::from_raw(self.a.as_ptr()); + _ = Box::from_raw(self.x.as_ptr()); + _ = Box::from_raw(self.b.as_ptr()); + _ = Box::from_raw(self.y.as_ptr()); + _ = Box::from_raw(self.c.as_ptr()); + } + } + } + + struct PanicOnDrop; + impl Drop for PanicOnDrop { + fn drop(&mut self) { + panic!("PanicOnDrop dropped"); + } + } + + let mut panic = Some(PanicOnDrop); + + let tree = DummyTree::new(); + let new_root = Handle::Root(tree.x).rotate_left(|_| mem::forget(panic.take())); + assert_eq!(new_root, Ok(Handle::Root(tree.y))); + + assert_eq!(unsafe { tree.y.as_ref().left() }, Some(tree.x)); + assert_eq!(unsafe { tree.y.as_ref().right() }, Some(tree.c)); + assert_eq!(unsafe { tree.x.as_ref().left() }, Some(tree.a)); + assert_eq!(unsafe { tree.x.as_ref().right() }, Some(tree.b)); + + assert_eq!(unsafe { tree.a.as_ref().parent() }, Some(tree.x)); + assert_eq!(unsafe { tree.b.as_ref().parent() }, Some(tree.x)); + assert_eq!(unsafe { tree.c.as_ref().parent() }, Some(tree.y)); + assert_eq!(unsafe { tree.x.as_ref().parent() }, Some(tree.y)); + assert_eq!(unsafe { tree.y.as_ref().parent() }, None); + + let mut panic = Some(PanicOnDrop); + let new_root = Handle::Root(tree.y).rotate_right(|_| mem::forget(panic.take())); + assert_eq!(new_root, Ok(Handle::Root(tree.x))); + + assert_eq!(unsafe { tree.x.as_ref().left() }, Some(tree.a)); + assert_eq!(unsafe { tree.x.as_ref().right() }, Some(tree.y)); + assert_eq!(unsafe { tree.y.as_ref().left() }, Some(tree.b)); + assert_eq!(unsafe { tree.y.as_ref().right() }, Some(tree.c)); + + assert_eq!(unsafe { tree.a.as_ref().parent() }, Some(tree.x)); + assert_eq!(unsafe { tree.y.as_ref().parent() }, Some(tree.x)); + assert_eq!(unsafe { tree.b.as_ref().parent() }, Some(tree.y)); + assert_eq!(unsafe { tree.c.as_ref().parent() }, Some(tree.y)); + assert_eq!(unsafe { tree.x.as_ref().parent() }, None); + + let mut panic = Some(PanicOnDrop); + let root = Handle::from_node(Some(tree.y)).rotate_up(|_| mem::forget(panic.take())); + assert_eq!(root, Ok(Handle::Root(tree.y))); + + assert_eq!(unsafe { tree.y.as_ref().left() }, Some(tree.x)); + assert_eq!(unsafe { tree.y.as_ref().right() }, Some(tree.c)); + assert_eq!(unsafe { tree.x.as_ref().left() }, Some(tree.a)); + assert_eq!(unsafe { tree.x.as_ref().right() }, Some(tree.b)); + + assert_eq!(unsafe { tree.a.as_ref().parent() }, Some(tree.x)); + assert_eq!(unsafe { tree.b.as_ref().parent() }, Some(tree.x)); + assert_eq!(unsafe { tree.c.as_ref().parent() }, Some(tree.y)); + assert_eq!(unsafe { tree.x.as_ref().parent() }, Some(tree.y)); + assert_eq!(unsafe { tree.y.as_ref().parent() }, None); + } + + #[test] + fn insert() { + let mut tree = RBTree::::new(); + + let nodes: Vec<_> = (0..10) + .map(|i| Box::into_raw(Box::new(TestNode::new(i)))) + .collect(); + + for &node in &nodes { + eprintln!("Inserting node with key: {}", unsafe { (*node).key }); + tree.insert_node(unsafe { NonNull::new_unchecked(node) }); + + eprintln!("Tree after insertion:"); + for n in tree.iter() { + eprintln!( + "\tNode: {:?} => {:?}", + unsafe { n.node().map(|n| n.as_ref().key) }, + n + ); + } + } + + for i in 0..10 { + let result = tree.find_by_key(&i); + eprintln!("{i}: {result:?}"); + assert!(matches!(result, SearchResult::FoundAt(_))); + } + } + + #[test] + fn remove() { + let mut tree = RBTree::::new(); + + let nodes: Vec<_> = (0..10) + .map(|i| Box::into_raw(Box::new(TestNode::new(i)))) + .collect(); + + for &node in &nodes { + tree.insert_node(unsafe { NonNull::new_unchecked(node) }); + } + + for i in 0..10 { + assert!(matches!(tree.find_by_key(&i), SearchResult::FoundAt(_))); + + eprintln!("Removing node with key: {}", i); + let removed_node = tree.remove(&i); + assert!(removed_node.is_some()); + + assert!(matches!( + tree.find_by_key(&i), + SearchResult::NotFoundAt(_) | SearchResult::Empty + )); + + eprintln!("Tree after removal:"); + for n in tree.iter() { + eprintln!( + "\tNode: {:?} => {:?}", + unsafe { n.node().map(|n| n.as_ref().key) }, + n + ); + } + } + + for i in 0..10 { + let result = tree.find_by_key(&i); + eprintln!("{i}: {result:?}"); + assert!(matches!(result, SearchResult::Empty)); } } } diff --git a/flake.nix b/flake.nix index dda3eb3..fcf47cc 100644 --- a/flake.nix +++ b/flake.nix @@ -33,6 +33,7 @@ rust-pkg clang gcc + gdb mold limine-full