rbtree.rs

use std::{cmp::Ordering, ops::Not};

pub struct RBTree<K>
where
    K: Ord,
{
    root: Option<Box<Node<K>>>,
}

impl<K> RBTree<K>
where
    K: Ord,
{
    pub fn new() -> Self {
        Self { root: None }
    }

    pub fn search(&self, key: &K) -> Option<&K> {
        let mut node = &self.root;
        while let Some(n) = node {
            match key.cmp(&n.key) {
                Ordering::Less => node = &n.left,
                Ordering::Greater => node = &n.right,
                Ordering::Equal => return Some(&n.key),
            }
        }
        None
    }

    pub fn insert(&mut self, key: K) {
        let mut root = insert_node(self.root.take(), key);
        root.color = Color::Black;
        self.root = Some(root);
    }
}

fn insert_node<K: Ord>(node: Option<Box<Node<K>>>, key: K) -> Box<Node<K>> {
    if node.is_none() {
        return Box::new(Node::new(key));
    }

    let mut node = node.unwrap();
    if is_red(&node.left) && is_red(&node.right) {
        flip_colors(&mut node);
    }

    match key.cmp(&node.key) {
        Ordering::Less => {
            node.left = Some(insert_node(node.left.take(), key));
        }
        Ordering::Greater => {
            node.right = Some(insert_node(node.right.take(), key));
        }
        Ordering::Equal => {}
    }

    if is_red(&node.right) && !is_red(&node.left) {
        node = rotate_left(node);
    }
    if is_red(&node.left) && is_red(&node.left.as_ref().unwrap().left) {
        node = rotate_right(node);
    }

    node
}

fn is_red<K: Ord>(node: &Option<Box<Node<K>>>) -> bool {
    if let Some(n) = node {
        n.color == Color::Red
    } else {
        false
    }
}

fn rotate_left<K: Ord>(mut h: Box<Node<K>>) -> Box<Node<K>> {
    let mut x = h.right.take().unwrap();
    h.right = x.left.take();
    x.left = Some(h);
    let h = x.left.as_mut().unwrap();
    x.color = h.color;
    h.color = Color::Red;
    x
}

fn rotate_right<K: Ord>(mut h: Box<Node<K>>) -> Box<Node<K>> {
    let mut x = h.left.take().unwrap();
    h.left = x.right.take();
    x.right = Some(h);
    let h = x.right.as_mut().unwrap();
    x.color = h.color;
    h.color = Color::Red;
    x
}

fn flip_colors<K: Ord>(node: &mut Node<K>) {
    node.color = !node.color;
    if let Some(left) = &mut node.left {
        left.color = !left.color;
    }
    if let Some(right) = &mut node.right {
        right.color = !right.color;
    }
}

#[derive(Debug, Clone, Copy, PartialEq)]
enum Color {
    Red,
    Black,
}

impl Not for Color {
    type Output = Self;

    fn not(self) -> Self::Output {
        match self {
            Self::Red => Self::Black,
            Self::Black => Self::Red,
        }
    }
}

struct Node<K>
where
    K: Ord,
{
    key: K,
    color: Color,
    left: Option<Box<Node<K>>>,
    right: Option<Box<Node<K>>>,
}

impl<K> Node<K>
where
    K: Ord,
{
    fn new(key: K) -> Self {
        Self {
            key,
            color: Color::Red,
            left: None,
            right: None,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use proptest::prelude::*;

    fn is_ordered<K: Ord>(node: &Option<Box<Node<K>>>) -> bool {
        if let Some(n) = node {
            if let Some(left) = &n.left {
                if left.key > n.key {
                    return false;
                }
            }
            if let Some(right) = &n.right {
                if right.key < n.key {
                    return false;
                }
            }
            is_ordered(&n.left) && is_ordered(&n.right)
        } else {
            true
        }
    }

    fn satisfy_third_property<K: Ord>(node: &Option<Box<Node<K>>>) -> bool {
        if let Some(n) = node {
            if n.color == Color::Red {
                if is_red(&n.left) || is_red(&n.right) {
                    return false;
                }
            }
            satisfy_third_property(&n.left) && satisfy_third_property(&n.right)
        } else {
            true
        }
    }

    fn satisfy_fourth_property<K: Ord>(node: &Option<Box<Node<K>>>) -> bool {
        if node.is_some() {
            let mut black_count = 0;
            let mut current = node;
            while let Some(c) = current {
                if c.color == Color::Black {
                    black_count += 1;
                }
                current = &c.left;
            }
            satisfy_fourth_property_helper(node, black_count, 0)
        } else {
            true
        }
    }

    fn satisfy_fourth_property_helper<K: Ord>(
        node: &Option<Box<Node<K>>>,
        black_count: usize,
        current_count: usize,
    ) -> bool {
        if let Some(n) = node {
            let mut count = current_count;
            if n.color == Color::Black {
                count += 1;
            }
            satisfy_fourth_property_helper(&n.left, black_count, count)
                && satisfy_fourth_property_helper(&n.right, black_count, count)
        } else {
            current_count == black_count
        }
    }

    proptest! {
        #[test]
        fn test_is_binary_search_tree(v in prop::collection::vec(0..100, 0..100)) {
            let mut rbtree = RBTree::new();
            for i in v {
                rbtree.insert(i);
            }

            assert!(is_ordered(&rbtree.root));
        }

        #[test]
        fn test_satisfy_red_black_properties(v in prop::collection::vec(0..100, 0..100)) {
            let mut rbtree = RBTree::new();
            for i in v {
                rbtree.insert(i);
            }

            if let Some(root) = &rbtree.root {
                assert_eq!(root.color, Color::Black);
            }
            assert!(satisfy_third_property(&rbtree.root));
            assert!(satisfy_fourth_property(&rbtree.root));
        }
    }
}