curiOS/crates/rbtree/src/raw_node.rs

1416 lines
43 KiB
Rust

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<T: core::fmt::Debug> TestDebug for T {}
#[cfg(not(test))]
pub trait TestDebug {}
#[cfg(not(test))]
impl<T> TestDebug for T {}
pub trait TestDebugInspect: Sized {
fn test_debug_inspect_with<F: FnOnce(&Self)>(self, f: F) -> Self
where
Self: TestDebug,
{
f(&self);
self
}
}
impl<T: Sized> TestDebugInspect for T {}
#[allow(dead_code)]
pub trait DebugInspect: Sized {
// fn debug_inspect(self) -> Self {
// dbg!(&self);
// self
// }
fn debug_inspect_with<F: FnOnce(&Self)>(self, f: F) -> Self
where
Self: core::fmt::Debug,
{
f(&self);
self
}
}
impl<T: Sized> DebugInspect for T {}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SearchResult<T> {
FoundAt(T),
NotFoundAt(T),
Empty,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Color {
Red,
Black,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Side {
Left,
Right,
}
impl Not for Side {
type Output = Self;
fn not(self) -> Self::Output {
match self {
Side::Left => Side::Right,
Side::Right => Side::Left,
}
}
}
pub unsafe trait UnsafeNode: TestDebug {
type Key: Eq + Ord + TestDebug;
fn left(&self) -> Option<NonNull<Self>>;
fn right(&self) -> Option<NonNull<Self>>;
fn parent(&self) -> Option<NonNull<Self>>;
fn key(&self) -> &Self::Key;
fn color(&self) -> Color;
fn set_left(&self, left: Option<NonNull<Self>>);
fn set_right(&self, right: Option<NonNull<Self>>);
fn set_parent(&self, parent: Option<NonNull<Self>>);
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<T> {
Left(T),
Right(T),
}
impl<T> LeftOrRight<Option<T>> {
pub fn transpose(self) -> Option<LeftOrRight<T>> {
match self {
LeftOrRight::Left(Some(t)) => Some(LeftOrRight::Left(t)),
LeftOrRight::Right(Some(t)) => Some(LeftOrRight::Right(t)),
_ => None,
}
}
}
impl<T> LeftOrRight<T> {
pub fn map<U, F: FnOnce(T) -> U>(self, f: F) -> LeftOrRight<U> {
match self {
LeftOrRight::Left(t) => LeftOrRight::Left(f(t)),
LeftOrRight::Right(t) => LeftOrRight::Right(f(t)),
}
}
pub fn as_ref(&self) -> LeftOrRight<&T> {
match self {
LeftOrRight::Left(t) => LeftOrRight::Left(t),
LeftOrRight::Right(t) => LeftOrRight::Right(t),
}
}
pub fn as_inner(&self) -> &T {
match self {
LeftOrRight::Left(t) => t,
LeftOrRight::Right(t) => 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,
}
}
}
#[derive(Debug)]
pub enum Handle<N> {
EmptyRoot,
Root(NonNull<N>),
Child {
parent: NonNull<N>,
node: LeftOrRight<Option<NonNull<N>>>,
},
}
impl<N: UnsafeNode> Eq for Handle<N> {}
impl<N: UnsafeNode> PartialEq for Handle<N> {
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<N> Clone for Handle<N> {
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<N> !Sync for Handle<N> {}
impl<N: UnsafeNode> Handle<N> {
fn from_node(node: Option<NonNull<N>>) -> 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),
}
}
#[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)
}
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<Side> {
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<Self> {
match side {
Side::Left => self.left_child(),
Side::Right => self.right_child(),
}
}
fn left_child(&self) -> Option<Self> {
let &parent = match self {
Handle::EmptyRoot => return None,
Handle::Root(parent) => parent,
Handle::Child { node, .. } => node.as_inner().as_ref()?,
};
let child = unsafe { parent.as_ref().left() };
Some(Handle::Child {
parent,
node: LeftOrRight::Left(child),
})
}
fn right_child(&self) -> Option<Self> {
let &parent = match self {
Handle::EmptyRoot => return None,
Handle::Root(parent) => parent,
Handle::Child { node, .. } => node.as_inner().as_ref()?,
};
let child = unsafe { parent.as_ref().right() };
Some(Handle::Child {
parent,
node: LeftOrRight::Right(child),
})
}
fn right_child_extant(&self) -> Option<Self> {
self.right_child().filter(|child| !child.is_nil())
}
fn left_child_extant(&self) -> Option<Self> {
self.left_child().filter(|child| !child.is_nil())
}
fn into_has_node(self) -> Option<Self> {
self.node().map(|_| self)
}
fn children(&self) -> (Option<Self>, Option<Self>) {
(self.left_child(), self.right_child())
}
fn non_nil_children(&self) -> (Option<Self>, Option<Self>) {
(
self.left_child().and_then(Self::into_has_node),
self.right_child().and_then(Self::into_has_node),
)
}
fn sibling(&self) -> Option<Self> {
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) = self.node() {
let mut child_node = child.node();
unsafe {
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) = self.node() {
let mut child_node = child.node();
unsafe {
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),
};
}
}
fn node(&self) -> Option<NonNull<N>> {
match self {
Handle::EmptyRoot => None,
Handle::Root(parent) => Some(*parent),
Handle::Child { node, .. } => node.as_inner().as_ref().copied(),
}
}
fn parent_and_side(&self) -> Option<LeftOrRight<Self>> {
let (&parent, side) = match self {
Handle::EmptyRoot => return None,
Handle::Root(_) => return None,
Handle::Child { parent, node } => (parent, node.map(|_| ())),
};
let grandparent = unsafe { parent.as_ref().parent() };
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))
};
Handle::Child {
parent: grandparent,
node,
}
} else {
Handle::Root(parent)
};
Some(side.map(|_| parent))
}
fn parent(&self) -> Option<Self> {
self.parent_and_side().map(LeftOrRight::into_inner)
}
fn set_parent_handle(&mut self, parent: Option<LeftOrRight<Handle<N>>>) {
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: Option<LeftOrRight<NonNull<N>>>) {
let parent = parent.map(|p| p.map(|p| Handle::from_node(Some(p))));
self.set_parent_handle(parent);
}
fn color(&self) -> Color {
match self {
Handle::Root(_) | Handle::EmptyRoot => Color::Black,
Handle::Child { node, .. } => match node.as_inner() {
Some(node) => unsafe { node.as_ref().color() },
None => Color::Black,
},
}
}
/// rotate `self` left, returning the new root of the subtree, or
/// `Err(self)` if `self` has no left child
fn rotate_left<F: FnOnce(Self)>(mut self, on_root: F) -> Result<Self, Self> {
let Some(mut y) = self.right_child() else {
return Err(self);
};
let Some(mut b) = y.left_child() else {
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)
}
/// rotate `self` right, returning the new root of the subtree, or
/// `Err(self)` if `self` has no left child
fn rotate_right<F: FnOnce(Self)>(mut self, on_root: F) -> Result<Self, Self> {
let Some(mut y) = self.left_child() else {
return Err(self);
};
let Some(mut b) = y.right_child() else {
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)
}
/// 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<F: FnOnce(Self)>(self, on_root: F) -> Result<Self, Self> {
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<F: FnOnce(Self)>(self, side: Side, on_root: F) -> Result<Self, Self> {
match side {
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<Self> {
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<Self> {
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<Self> {
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<Self> {
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(&current) {
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<NonNull<N>> {
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<N>) -> Option<NonNull<N>> {
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
}
}
}
}
pub struct RBTree<N: UnsafeNode> {
root: Option<NonNull<N>>,
}
impl<N: UnsafeNode> RBTree<N> {
pub fn new() -> Self {
Self { root: None }
}
fn root_handle(&self) -> Handle<N> {
match self.root {
Some(root) => Handle::Root(root),
None => Handle::EmptyRoot,
}
}
fn set_root_handle(&mut self, mut handle: Handle<N>) {
handle.set_parent(None);
handle.set_color(Color::Black);
self.root = handle.node();
}
pub fn find_by_key<Q>(&self, key: &Q) -> SearchResult<Handle<N>>
where
N::Key: core::borrow::Borrow<Q>,
Q: Ord + ?Sized,
{
use core::borrow::Borrow;
use core::cmp::Ordering::*;
let mut current = self.root_handle();
loop {
let node = match &current {
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);
}
}
}
}
#[must_use]
pub fn insert_node(&mut self, new_node: NonNull<N>) -> Option<NonNull<N>> {
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) <or> 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
}
#[must_use]
pub fn remove<Q>(&mut self, key: &Q) -> Option<NonNull<N>>
where
N::Key: core::borrow::Borrow<Q>,
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()
}
pub fn iter_nodes(&self) -> TreeNodeIter<'_, N> {
TreeNodeIter {
range: TreeRange::full_range(self.root_handle()),
}
}
pub fn iter(&self) -> TreeIter<'_, N> {
TreeIter {
range: TreeRange::full_range(self.root_handle()),
}
}
}
impl<N: UnsafeNode> Default for RBTree<N> {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone)]
enum RangeHandle<N> {
Root(Handle<N>),
Node(Handle<N>),
}
impl<N> RangeHandle<N> {
fn into_inner(self) -> Handle<N> {
match self {
RangeHandle::Root(handle) => handle,
RangeHandle::Node(handle) => handle,
}
}
}
impl<N: UnsafeNode + Eq> Eq for RangeHandle<N> {}
impl<N: UnsafeNode> PartialEq for RangeHandle<N> {
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<RangeHandle<N>>,
end: Option<RangeHandle<N>>,
_pd: PhantomData<&'a ()>,
}
impl<'a, N: UnsafeNode + 'a> TreeRange<'a, N> {
#[expect(dead_code)]
fn new(start: Handle<N>, end: Handle<N>) -> Self {
Self {
start: Some(RangeHandle::Node(start)),
end: Some(RangeHandle::Node(end)),
_pd: PhantomData,
}
}
fn full_range(root: Handle<N>) -> Self {
Self {
start: Some(RangeHandle::Root(root.clone())),
end: Some(RangeHandle::Root(root)),
_pd: PhantomData,
}
}
fn init_front(&mut self) -> Option<&mut Handle<N>> {
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<N>> {
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<Handle<N>> {
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<Handle<N>> {
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)
}
}
}
}
pub struct TreeNodeIter<'a, N: UnsafeNode> {
range: TreeRange<'a, N>,
}
impl<'a, N: UnsafeNode + 'a> Iterator for TreeNodeIter<'a, N> {
type Item = Handle<N>;
fn next(&mut self) -> Option<Self::Item> {
self.range.next()
}
}
impl<'a, N: UnsafeNode + 'a> DoubleEndedIterator for TreeNodeIter<'a, N> {
fn next_back(&mut self) -> Option<Self::Item> {
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::Item> {
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::Item> {
self.range
.next_back()
.map(|n| unsafe { n.node().unwrap().as_ref().key() })
}
}
#[cfg(test)]
mod tests {
use std::cell::Cell;
use super::*;
#[derive(Debug)]
struct TestNode {
key: i32,
left: Cell<Option<NonNull<TestNode>>>,
right: Cell<Option<NonNull<TestNode>>>,
parent: Cell<Option<NonNull<TestNode>>>,
color: Cell<Color>,
}
unsafe impl UnsafeNode for TestNode {
type Key = i32;
fn left(&self) -> Option<NonNull<Self>> {
self.left.get()
}
fn right(&self) -> Option<NonNull<Self>> {
self.right.get()
}
fn parent(&self) -> Option<NonNull<Self>> {
self.parent.get()
}
fn key(&self) -> &Self::Key {
&self.key
}
fn color(&self) -> Color {
self.color.get()
}
fn set_left(&self, left: Option<NonNull<Self>>) {
self.left.set(left);
}
fn set_right(&self, right: Option<NonNull<Self>>) {
self.right.set(right);
}
fn set_parent(&self, parent: Option<NonNull<Self>>) {
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::<TestNode>::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<TestNode>,
x: NonNull<TestNode>,
b: NonNull<TestNode>,
y: NonNull<TestNode>,
c: NonNull<TestNode>,
}
impl DummyTree {
fn new() -> Self {
let a = Box::into_non_null(Box::new(TestNode::new(1)));
let x = Box::into_non_null(Box::new(TestNode::new(2)));
let b = Box::into_non_null(Box::new(TestNode::new(3)));
let y = Box::into_non_null(Box::new(TestNode::new(4)));
let 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::<TestNode>::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_nodes() {
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::<TestNode>::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_nodes() {
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));
}
}
}