From 3ba16577defcf78c958e83b01f28465843d707e8 Mon Sep 17 00:00:00 2001 From: janis Date: Fri, 31 Jul 2026 19:22:22 +0200 Subject: [PATCH] raw node --- crates/rbtree/src/lib.rs | 3 + crates/rbtree/src/raw_node.rs | 241 ++++++++++++++++++++++++++++++++++ 2 files changed, 244 insertions(+) create mode 100644 crates/rbtree/src/raw_node.rs diff --git a/crates/rbtree/src/lib.rs b/crates/rbtree/src/lib.rs index 5212ec5..1252a8a 100644 --- a/crates/rbtree/src/lib.rs +++ b/crates/rbtree/src/lib.rs @@ -1,4 +1,7 @@ #![cfg_attr(not(test), no_std)] +#![feature(negative_impls)] + +mod raw_node; extern crate alloc; diff --git a/crates/rbtree/src/raw_node.rs b/crates/rbtree/src/raw_node.rs new file mode 100644 index 0000000..2dbf81a --- /dev/null +++ b/crates/rbtree/src/raw_node.rs @@ -0,0 +1,241 @@ +use core::ops::Not; +use core::ptr::NonNull; + +#[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 trait UnsafeNode { + type Key: Eq + Ord; + fn left(&self) -> Option>; + fn right(&self) -> Option>; + fn parent(&self) -> Option>; + fn key(&self) -> &Self::Key; + fn color(&self) -> Color; + + fn set_left(&self, left: Option>); + fn set_right(&self, right: Option>); + fn set_parent(&self, parent: Option>); + fn set_color(&self, color: Color); +} + +pub enum LeftOrRight { + Left(T), + Right(T), +} + +impl LeftOrRight { + 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 { + match self { + LeftOrRight::Left(t) => t, + LeftOrRight::Right(t) => t, + } + } +} + +pub enum Handle { + EmptyRoot, + Root(NonNull), + Child { + parent: NonNull, + node: LeftOrRight>>, + }, +} + +impl !Sync for Handle {} + +impl Handle { + fn is_root(&self) -> bool { + matches!(self, Handle::Root(_) | Handle::EmptyRoot) + } + + fn left_child(&self) -> Option { + 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 { + 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 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(), + } { + unsafe { + parent.as_ref().set_left(child.node()); + child.set_parent(LeftOrRight::Left(Some(*parent))); + } + } + } + + 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(), + } { + unsafe { + parent.as_ref().set_right(child.node()); + child.set_parent(LeftOrRight::Right(Some(*parent))); + } + } + } + + fn node(&self) -> Option> { + match self { + Handle::EmptyRoot => None, + Handle::Root(parent) => Some(*parent), + Handle::Child { node, .. } => node.as_inner().as_ref().copied(), + } + } + + fn parent(&self) -> Option { + let &parent = match self { + Handle::EmptyRoot => return None, + Handle::Root(_) => return None, + Handle::Child { parent, .. } => parent, + }; + + let grandparent = unsafe { parent.as_ref().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 { + parent: grandparent, + node, + }) + } else { + Some(Handle::Root(parent)) + } + } + + 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 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, + }, + } + } + + fn rotate_left(mut self) -> Result { + let Some(mut y) = self.right_child() else { + return Err(self); + }; + + let Some(mut b) = y.left_child() else { + return Err(self); + }; + + self.set_right_child(&mut b); + + y.set_left_child(&mut self); + + Ok(y) + } + + fn rotate_right(mut self) -> Result { + let Some(mut y) = self.left_child() else { + return Err(self); + }; + + let Some(mut b) = y.right_child() else { + return Err(self); + }; + + self.set_left_child(&mut b); + + y.set_right_child(&mut self); + + Ok(y) + } + + pub fn rotate(self, side: Side) -> Result { + match side { + Side::Left => self.rotate_left(), + Side::Right => self.rotate_right(), + } + } +}