raw node

crates/rbtree/src/lib.rs

@@ -1,4 +1,7 @@
 #![cfg_attr(not(test), no_std)]
+#![feature(negative_impls)]
+
+mod raw_node;
 
 extern crate alloc;
 

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<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);
+}
+
+pub enum LeftOrRight<T> {
+    Left(T),
+    Right(T),
+}
+
+impl<T> LeftOrRight<T> {
+    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<N: UnsafeNode> {
+    EmptyRoot,
+    Root(NonNull<N>),
+    Child {
+        parent: NonNull<N>,
+        node: LeftOrRight<Option<NonNull<N>>>,
+    },
+}
+
+impl<N: UnsafeNode> !Sync for Handle<N> {}
+
+impl<N: UnsafeNode> Handle<N> {
+    fn is_root(&self) -> bool {
+        matches!(self, Handle::Root(_) | Handle::EmptyRoot)
+    }
+
+    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 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<NonNull<N>> {
+        match self {
+            Handle::EmptyRoot => None,
+            Handle::Root(parent) => Some(*parent),
+            Handle::Child { node, .. } => node.as_inner().as_ref().copied(),
+        }
+    }
+
+    fn parent(&self) -> Option<Self> {
+        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<Option<NonNull<N>>>) {
+        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<Self, Self> {
+        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<Self, Self> {
+        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<Self, Self> {
+        match side {
+            Side::Left => self.rotate_left(),
+            Side::Right => self.rotate_right(),
+        }
+    }
+}