use std::boxed::Box;
use std::cmp::{Ord, Ordering};
use std::iter::Iterator;
use std::ptr::null_mut;
#[derive(Copy, Clone)]
enum Color {
Red,
Black,
}
pub struct RBNode<K: Ord, V> {
key: K,
value: V,
color: Color,
parent: *mut RBNode<K, V>,
left: *mut RBNode<K, V>,
right: *mut RBNode<K, V>,
}
impl<K: Ord, V> RBNode<K, V> {
fn new(key: K, value: V) -> RBNode<K, V> {
RBNode {
key,
value,
color: Color::Red,
parent: null_mut(),
left: null_mut(),
right: null_mut(),
}
}
}
pub struct RBTree<K: Ord, V> {
root: *mut RBNode<K, V>,
}
impl<K: Ord, V> Default for RBTree<K, V> {
fn default() -> Self {
Self::new()
}
}
impl<K: Ord, V> RBTree<K, V> {
pub fn new() -> RBTree<K, V> {
RBTree::<K, V> { root: null_mut() }
}
pub fn find(&self, key: &K) -> Option<&mut V> {
unsafe {
let mut node = self.root;
while !node.is_null() {
node = match (*node).key.cmp(key) {
Ordering::Less => (*node).right,
Ordering::Equal => return Some(&mut (*node).value),
Ordering::Greater => (*node).left,
}
}
}
None
}
pub fn insert(&mut self, key: K, value: V) {
unsafe {
let mut parent = null_mut();
let mut node = self.root;
while !node.is_null() {
parent = node;
node = match (*node).key.cmp(&key) {
Ordering::Less => (*node).right,
Ordering::Equal => {
(*node).value = value;
return;
}
Ordering::Greater => (*node).left,
}
}
node = Box::into_raw(Box::new(RBNode::new(key, value)));
if !parent.is_null() {
if (*node).key < (*parent).key {
(*parent).left = node;
} else {
(*parent).right = node;
}
} else {
self.root = node;
}
(*node).parent = parent;
insert_fixup(self, node);
}
}
pub fn delete(&mut self, key: &K) {
unsafe {
let mut parent = null_mut();
let mut node = self.root;
while !node.is_null() {
node = match (*node).key.cmp(key) {
Ordering::Less => {
parent = node;
(*node).right
}
Ordering::Equal => break,
Ordering::Greater => {
parent = node;
(*node).left
}
};
}
if node.is_null() {
return;
}
/* cl and cr denote left and right child of node, respectively. */
let cl = (*node).left;
let cr = (*node).right;
let mut deleted_color;
if cl.is_null() {
replace_node(self, parent, node, cr);
if cr.is_null() {
/*
* Case 1 - cl and cr are both NULL
* (n could be either color here)
*
* (n) NULL
* / \ -->
* NULL NULL
*/
deleted_color = (*node).color;
} else {
/*
* Case 2 - cl is NULL and cr is not NULL
*
* N Cr
* / \ --> / \
* NULL cr NULL NULL
*/
(*cr).parent = parent;
(*cr).color = Color::Black;
deleted_color = Color::Red;
}
} else if cr.is_null() {
/*
* Case 3 - cl is not NULL and cr is NULL
*
* N Cl
* / \ --> / \
* cl NULL NULL NULL
*/
replace_node(self, parent, node, cl);
(*cl).parent = parent;
(*cl).color = Color::Black;
deleted_color = Color::Red;
} else {
let mut victim = (*node).right;
while !(*victim).left.is_null() {
victim = (*victim).left;
}
if victim == (*node).right {
/* Case 4 - victim is the right child of node
*
* N N n
* / \ / \ / \
* (cl) cr (cl) Cr Cl Cr
*
* N n
* / \ / \
* (cl) Cr Cl Cr
* \ \
* crr crr
*/
replace_node(self, parent, node, victim);
(*victim).parent = parent;
deleted_color = (*victim).color;
(*victim).color = (*node).color;
(*victim).left = cl;
(*cl).parent = victim;
if (*victim).right.is_null() {
parent = victim;
} else {
deleted_color = Color::Red;
(*(*victim).right).color = Color::Black;
}
} else {
/*
* Case 5 - victim is not the right child of node
*/
/* vp and vr denote parent and right child of victim, respectively. */
let vp = (*victim).parent;
let vr = (*victim).right;
(*vp).left = vr;
if vr.is_null() {
deleted_color = (*victim).color;
} else {
deleted_color = Color::Red;
(*vr).parent = vp;
(*vr).color = Color::Black;
}
replace_node(self, parent, node, victim);
(*victim).parent = parent;
(*victim).color = (*node).color;
(*victim).left = cl;
(*victim).right = cr;
(*cl).parent = victim;
(*cr).parent = victim;
parent = vp;
}
}
/* release resource */
Box::from_raw(node);
if matches!(deleted_color, Color::Black) {
delete_fixup(self, parent);
}
}
}
pub fn iter<'a>(&self) -> RBTreeIterator<'a, K, V> {
let mut iterator = RBTreeIterator { stack: Vec::new() };
let mut node = self.root;
unsafe {
while !node.is_null() {
iterator.stack.push(&*node);
node = (*node).left;
}
}
iterator
}
}
#[inline]
unsafe fn insert_fixup<K: Ord, V>(tree: &mut RBTree<K, V>, mut node: *mut RBNode<K, V>) {
let mut parent: *mut RBNode<K, V> = (*node).parent;
let mut gparent: *mut RBNode<K, V>;
let mut tmp: *mut RBNode<K, V>;
loop {
/*
* Loop invariant:
* - node is red
*/
if parent.is_null() {
(*node).color = Color::Black;
break;
}
if matches!((*parent).color, Color::Black) {
break;
}
gparent = (*parent).parent;
tmp = (*gparent).right;
if parent != tmp {
/* parent = (*gparent).left */
if !tmp.is_null() && matches!((*tmp).color, Color::Red) {
/*
* Case 1 - color flips and recurse at g
*
* G g
* / \ / \
* p u --> P U
* / /
* n n
*/
(*parent).color = Color::Black;
(*tmp).color = Color::Black;
(*gparent).color = Color::Red;
node = gparent;
parent = (*node).parent;
continue;
}
tmp = (*parent).right;
if node == tmp {
/* node = (*parent).right */
/*
* Case 2 - left rotate at p (then Case 3)
*
* G G
* / \ / \
* p U --> n U
* \ /
* n p
*/
left_rotate(tree, parent);
parent = node;
}
/*
* Case 3 - right rotate at g
*
* G P
* / \ / \
* p U --> n g
* / \
* n U
*/
(*parent).color = Color::Black;
(*gparent).color = Color::Red;
right_rotate(tree, gparent);
} else {
/* parent = (*gparent).right */
tmp = (*gparent).left;
if !tmp.is_null() && matches!((*tmp).color, Color::Red) {
/*
* Case 1 - color flips and recurse at g
* G g
* / \ / \
* u p --> U P
* \ \
* n n
*/
(*parent).color = Color::Black;
(*tmp).color = Color::Black;
(*gparent).color = Color::Red;
node = gparent;
parent = (*node).parent;
continue;
}
tmp = (*parent).left;
if node == tmp {
/*
* Case 2 - right rotate at p (then Case 3)
*
* G G
* / \ / \
* U p --> U n
* / \
* n p
*/
right_rotate(tree, parent);
parent = node;
}
/*
* Case 3 - left rotate at g
*
* G P
* / \ / \
* U p --> g n
* \ /
* n U
*/
(*parent).color = Color::Black;
(*gparent).color = Color::Red;
left_rotate(tree, gparent);
}
break;
}
}
#[inline]
unsafe fn delete_fixup<K: Ord, V>(tree: &mut RBTree<K, V>, mut parent: *mut RBNode<K, V>) {
let mut node: *mut RBNode<K, V> = null_mut();
let mut sibling: *mut RBNode<K, V>;
/* sl and sr denote left and right child of sibling, respectively. */
let mut sl: *mut RBNode<K, V>;
let mut sr: *mut RBNode<K, V>;
loop {
/*
* Loop invariants:
* - node is black (or null on first iteration)
* - node is not the root (so parent is not null)
* - All leaf paths going through parent and node have a
* black node count that is 1 lower than other leaf paths.
*/
sibling = (*parent).right;
if node != sibling {
/* node = (*parent).left */
if matches!((*sibling).color, Color::Red) {
/*
* Case 1 - left rotate at parent
*
* P S
* / \ / \
* N s --> p Sr
* / \ / \
* Sl Sr N Sl
*/
left_rotate(tree, parent);
(*parent).color = Color::Red;
(*sibling).color = Color::Black;
sibling = (*parent).right;
}
sl = (*sibling).left;
sr = (*sibling).right;
if !sl.is_null() && matches!((*sl).color, Color::Red) {
/*
* Case 2 - right rotate at sibling and then left rotate at parent
* (p and sr could be either color here)
*
* (p) (p) (sl)
* / \ / \ / \
* N S --> N sl --> P S
* / \ \ / \
* sl (sr) S N (sr)
* \
* (sr)
*/
(*sl).color = (*parent).color;
(*parent).color = Color::Black;
right_rotate(tree, sibling);
left_rotate(tree, parent);
} else if !sr.is_null() && matches!((*sr).color, Color::Red) {
/*
* Case 3 - left rotate at parent
* (p could be either color here)
*
* (p) S
* / \ / \
* N S --> (p) (sr)
* / \ / \
* Sl sr N Sl
*/
(*sr).color = (*parent).color;
left_rotate(tree, parent);
} else {
/*
* Case 4 - color clip
* (p could be either color here)
*
* (p) (p)
* / \ / \
* N S --> N s
* / \ / \
* Sl Sr Sl Sr
*/
(*sibling).color = Color::Red;
if matches!((*parent).color, Color::Black) {
node = parent;
parent = (*node).parent;
continue;
}
(*parent).color = Color::Black;
}
} else {
/* node = (*parent).right */
sibling = (*parent).left;
if matches!((*sibling).color, Color::Red) {
/*
* Case 1 - right rotate at parent
*/
right_rotate(tree, parent);
(*parent).color = Color::Red;
(*sibling).color = Color::Black;
sibling = (*parent).right;
}
sl = (*sibling).left;
sr = (*sibling).right;
if !sr.is_null() && matches!((*sr).color, Color::Red) {
/*
* Case 2 - left rotate at sibling and then right rotate at parent
*/
(*sr).color = (*parent).color;
(*parent).color = Color::Black;
left_rotate(tree, sibling);
right_rotate(tree, parent);
} else if !sl.is_null() && matches!((*sl).color, Color::Red) {
/*
* Case 3 - right rotate at parent
*/
(*sl).color = (*parent).color;
right_rotate(tree, parent);
} else {
/*
* Case 4 - color flip
*/
(*sibling).color = Color::Red;
if matches!((*parent).color, Color::Black) {
node = parent;
parent = (*node).parent;
continue;
}
(*parent).color = Color::Black;
}
}
break;
}
}
#[inline]
unsafe fn left_rotate<K: Ord, V>(tree: &mut RBTree<K, V>, x: *mut RBNode<K, V>) {
/*
* Left rotate at x
* (x could also be the left child of p)
*
* p p
* \ \
* x --> y
* / \ / \
* y x
* / \ / \
* c c
*/
let p = (*x).parent;
let y = (*x).right;
let c = (*y).left;
(*y).left = x;
(*x).parent = y;
(*x).right = c;
if !c.is_null() {
(*c).parent = x;
}
if p.is_null() {
tree.root = y;
} else if (*p).left == x {
(*p).left = y;
} else {
(*p).right = y;
}
(*y).parent = p;
}
#[inline]
unsafe fn right_rotate<K: Ord, V>(tree: &mut RBTree<K, V>, x: *mut RBNode<K, V>) {
/*
* Right rotate at x
* (x could also be the left child of p)
*
* p p
* \ \
* x --> y
* / \ / \
* y x
* / \ / \
* c c
*/
let p = (*x).parent;
let y = (*x).left;
let c = (*y).right;
(*y).right = x;
(*x).parent = y;
(*x).left = c;
if !c.is_null() {
(*c).parent = x;
}
if p.is_null() {
tree.root = y;
} else if (*p).left == x {
(*p).left = y;
} else {
(*p).right = y;
}
(*y).parent = p;
}
#[inline]
unsafe fn replace_node<K: Ord, V>(
tree: &mut RBTree<K, V>,
parent: *mut RBNode<K, V>,
node: *mut RBNode<K, V>,
new: *mut RBNode<K, V>,
) {
if parent.is_null() {
tree.root = new;
} else if (*parent).left == node {
(*parent).left = new;
} else {
(*parent).right = new;
}
}
pub struct RBTreeIterator<'a, K: Ord, V> {
stack: Vec<&'a RBNode<K, V>>,
}
impl<'a, K: Ord, V> Iterator for RBTreeIterator<'a, K, V> {
type Item = &'a RBNode<K, V>;
fn next(&mut self) -> Option<Self::Item> {
match self.stack.pop() {
Some(node) => {
let mut next = node.right;
unsafe {
while !next.is_null() {
self.stack.push(&*next);
next = (*next).left;
}
}
Some(node)
}
None => None,
}
}
}
#[cfg(test)]
mod tests {
use super::RBTree;
#[test]
fn find() {
let mut tree = RBTree::<usize, char>::new();
for (k, v) in String::from("hello, world!").chars().enumerate() {
tree.insert(k, v);
}
assert_eq!(*tree.find(&3).unwrap_or(&mut '*'), 'l');
assert_eq!(*tree.find(&6).unwrap_or(&mut '*'), ' ');
assert_eq!(*tree.find(&8).unwrap_or(&mut '*'), 'o');
assert_eq!(*tree.find(&12).unwrap_or(&mut '*'), '!');
}
#[test]
fn insert() {
let mut tree = RBTree::<usize, char>::new();
for (k, v) in String::from("hello, world!").chars().enumerate() {
tree.insert(k, v);
}
let s: String = tree.iter().map(|x| x.value).collect();
assert_eq!(s, "hello, world!");
}
#[test]
fn delete() {
let mut tree = RBTree::<usize, char>::new();
for (k, v) in String::from("hello, world!").chars().enumerate() {
tree.insert(k, v);
}
tree.delete(&1);
tree.delete(&3);
tree.delete(&5);
tree.delete(&7);
tree.delete(&11);
let s: String = tree.iter().map(|x| x.value).collect();
assert_eq!(s, "hlo orl!");
}
}