Immutable AVL tree

immutable-avl.rs

#![feature(let_else)]

use std::cmp::Ordering;
use std::fmt::{Display, Formatter};
use std::sync::Arc;

pub enum Root<K, V> {
    Empty,
    More(Arc<Node<K, V>>)
}

impl <K: Ord + Display, V: Display + Copy> Display for Root<K, V> {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            Root::Empty => write!(f, "Empty"),
            Root::More(node) => write!(f, "{}", node)
        }
    }
}

struct Pair<K, V> {
    key: K,
    value: V
}

#[derive(Clone)]
pub struct Node<K, V> {
    height: u32,
    data: Arc<Pair<K, V>>,
    left_opt: Option<Arc<Node<K, V>>>,
    right_opt: Option<Arc<Node<K, V>>>
}

impl <K: Ord + Display, V: Display> Display for Node<K, V> {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        let left = self.left_opt.as_ref()
            .map(|n| format!(",\"left\": {}", n))
            .unwrap_or("".to_string());
        let right = self.right_opt.as_ref()
            .map(|n| format!(",\"right\": {}", n))
            .unwrap_or("".to_string());
        write!(f, "{{\"key\": \"{}\", \"value\": \"{}\"{}{}}}", self.data.key, self.data.value, left, right)
    }
}

impl <K: Ord, V> Root<K, V> {
    pub(crate) fn new() -> Root<K, V> {
        Root::Empty
    }

    pub(crate) fn lookup(&self, key: K) -> Option<&V> {
        match self {
            Root::Empty => None,
            Root::More(node) => node.search(key)
        }
    }

    pub(crate) fn insert(&self, key: K, value: V) -> Self {
        match self {
            Root::Empty => Root::More(
                Arc::new(
                    Node::new_with_child(
                        Arc::new(Pair { key, value }),
                        None,
                        None
                    )
                )
            ),
            Root::More(node) => {
                Root::More(Arc::new(node.insert(key, value)))
            }
        }
    }

    pub(crate) fn remove(&self, key: K) -> (Self, Option<&V>) {
        match self {
            Root::Empty => (Root::Empty, None),
            Root::More(node) => {
                match node.remove(key) {
                    RemoveResult::End(value) => (Root::Empty, Some(value)),
                    RemoveResult::None => (Root::More(node.clone()), None),
                    RemoveResult::Has(value, new_node) => (Root::More(new_node.clone()), Some(value))
                }
            }
        }
    }
}

enum RemoveResult<V, P> {
    End(V),
    Has(V, P),
    None,
}

impl <K: Ord, V> Node<K, V> {
    fn new(key: K, value: V) -> Self {
        Node {
            height: 1,
            data: Arc::new(Pair { key, value }),
            left_opt: None,
            right_opt: None
        }
    }

    fn new_with_child(pair: Arc<Pair<K, V>>,
                      left_opt: Option<Arc<Self>>,
                      right_opt: Option<Arc<Self>>) -> Self {
        let height = 1 + left_opt.as_ref().map(|n| n.height).unwrap_or(0)
            .max(right_opt.as_ref().map(|n| n.height).unwrap_or(0));
        Node {
            height,
            data: pair,
            left_opt,
            right_opt
        }
    }

    fn search(&self, key: K) -> Option<&V> {
        match &key.cmp(&self.data.key) {
            Ordering::Equal => Some(&self.data.value),
            Ordering::Less => self.left_opt.as_ref().and_then(|n| n.search(key)),
            Ordering::Greater => self.right_opt.as_ref().and_then(|n| n.search(key))
        }
    }

    fn insert(&self, key: K, value: V) -> Self {
        match &key.cmp(&self.data.key) {
            Ordering::Equal => Self::new(key, value),
            Ordering::Less => {
                let new_left = if let Some(x) = self.left_opt.as_ref() {
                    x.insert(key, value)
                } else {
                    Self::new(key, value)
                };

                let new_node = Self::new_with_child(
                    self.data.clone(),
                    Some(Arc::new(new_left)),
                    self.right_opt.clone()
                );

                new_node.left_heavy_rebalance()
            },
            Ordering::Greater => {
                let new_right = if let Some(x) = self.right_opt.as_ref() {
                    x.insert(key, value)
                } else {
                    Self::new(key, value)
                };

                let new_node = Self::new_with_child(
                    self.data.clone(),
                    self.left_opt.clone(),
                    Some(Arc::new(new_right))
                );

                new_node.right_heavy_rebalance()
            }
        }
    }

    fn pick_leftmost(&self) -> (Arc<Pair<K, V>>, Option<Arc<Self>>) {
        match self.left_opt {
            None => (self.data.clone(), self.right_opt.clone()),
            Some(ref left) => {
                let (leftmost, new_left_node) = left.pick_leftmost();
                let new_node = Self::new_with_child(
                    self.data.clone(),
                    new_left_node,
                    self.right_opt.clone()
                );
                (leftmost, Some(Arc::new(new_node.right_heavy_rebalance())))
            }
        }
    }

    fn remove(&self, key: K) -> RemoveResult<&V, Arc<Self>> {
        match &key.cmp(&self.data.key) {
            Ordering::Equal => {
                match (self.left_opt.as_ref(), self.right_opt.as_ref()) {
                    (None, None) => RemoveResult::End(&self.data.as_ref().value),
                    (Some(left), None) => RemoveResult::Has(&self.data.as_ref().value, left.clone()),
                    (None, Some(right)) => RemoveResult::Has(&self.data.as_ref().value, right.clone()),
                    (Some(left), Some(right)) => {
                        let (lift_pair, new_right) = right.pick_leftmost();
                        let new_node = Self::new_with_child(
                            lift_pair.clone(),
                            Some(left.clone()),
                            new_right
                        );

                        let balanced_new_node = new_node.left_heavy_rebalance();
                        RemoveResult::Has(&self.data.as_ref().value, Arc::new(balanced_new_node))
                    },
                }
            },
            Ordering::Less => {
                match self.left_opt.as_ref() {
                    None => RemoveResult::None,
                    Some(left) => {
                        match left.remove(key) {
                            RemoveResult::None => RemoveResult::None,
                            RemoveResult::End(value) => {
                                let new_node = Self::new_with_child(
                                    self.data.clone(),
                                    None,
                                    self.right_opt.clone()
                                );
                                RemoveResult::Has(value, Arc::new(new_node.right_heavy_rebalance()))
                            },
                            RemoveResult::Has(value, new_left) => {
                                let new_node = Self::new_with_child(
                                    self.data.clone(),
                                    Some(new_left.clone()),
                                    self.right_opt.clone()
                                );
                                RemoveResult::Has(value, Arc::new(new_node.right_heavy_rebalance()))
                            },
                        }
                    }
                }
            },
            Ordering::Greater => {
                match self.right_opt.as_ref() {
                    None => RemoveResult::None,
                    Some(right) => {
                        match right.remove(key) {
                            RemoveResult::None => RemoveResult::None,
                            RemoveResult::End(value) => {

                                let new_node = Self::new_with_child(
                                    self.data.clone(),
                                    self.left_opt.clone(),
                                    None
                                );
                                RemoveResult::Has(value, Arc::new(new_node.left_heavy_rebalance()))
                            },
                            RemoveResult::Has(value, new_right) => {
                                let new_node = Self::new_with_child(
                                    self.data.clone(),
                                    self.left_opt.clone(),
                                    Some(new_right.clone())
                                );
                                RemoveResult::Has(value, Arc::new(new_node.left_heavy_rebalance()))
                            },
                        }
                    }
                }
            },
        }
    }

    fn balance_factor(&self) -> i32 {
        let right_height = self.right_opt.as_ref().map(|x| x.height).unwrap_or(0);
        let left_height = self.left_opt.as_ref().map(|x| x.height).unwrap_or(0);
        right_height as i32 - left_height as i32
    }

    fn right_heavy_rebalance(self) -> Self {
        if self.balance_factor() < 2 {
            self
        } else {
            let rbf = self.right_opt.as_ref()
                .map(|n| n.balance_factor())
                .unwrap_or(0);
            if rbf == -1 {
                unsafe { self.rl_rotate() }
            } else {
                unsafe { self.left_rotate() }
            }
        }
    }

    fn left_heavy_rebalance(self) -> Self {
        if self.balance_factor() > -2 {
            self
        } else {
            let lbf = self.left_opt.as_ref()
                .map(|n| n.balance_factor())
                .unwrap_or(0);
            if  lbf == 1 {
                unsafe { self.lr_rotate() }
            } else {
                unsafe { self.right_rotate() }
            }
        }
    }

    unsafe fn right_rotate(&self) -> Self {
        let old_left = self.left_opt.as_ref().unwrap_unchecked();
        let new_right = Self::new_with_child(
            self.data.clone(),
            old_left.right_opt.clone(),
            self.right_opt.clone(),
        );
        let new_root = Self::new_with_child(
            old_left.data.clone(),
            old_left.left_opt.clone(),
            Some(Arc::new(new_right)),
        );
        new_root
    }

    unsafe fn left_rotate(&self) -> Self {
        let old_right = self.right_opt.as_ref().unwrap_unchecked();
        let new_left = Self::new_with_child(
            self.data.clone(),
            self.left_opt.clone(),
            old_right.left_opt.clone()
        );
        let new_root = Self::new_with_child(
            old_right.data.clone(),
            Some(Arc::new(new_left)),
            old_right.right_opt.clone(),
        );
        new_root
    }

    unsafe fn lr_rotate(&self) -> Self {
        let new_left = self.left_opt.as_ref()
            .map(|n| Arc::new(n.left_rotate()));
        let tmp_root = Self::new_with_child(
            self.data.clone(),
            new_left,
            self.right_opt.clone()
        );
        tmp_root.right_rotate()
    }

    unsafe fn rl_rotate(&self) -> Self {
        let new_right = self.right_opt.as_ref()
            .map(|n| Arc::new(n.right_rotate()));
        let tmp_root = Self::new_with_child(
            self.data.clone(),
            self.left_opt.clone(),
            new_right
        );
        tmp_root.left_rotate()
    }
}

fn make_test_tree() -> Root<i32, bool> {
    let mut fake_data = (0..=20).collect::<Vec<i32>>();
    let fake_data2 = (-20..=-1).collect::<Vec<i32>>();
    fake_data.extend(fake_data2);

    let mut root = Root::<i32, bool>::new();
    for i in fake_data {
        root = root.insert(i, true);
    }
    root
}

#[test]
fn seq() {
    let root = make_test_tree();
    println!("{}", root);
}

#[test]
fn search() {
    let root = make_test_tree();
    println!("{}", root.lookup(-1).unwrap());
}

#[test]
fn remove() {
    let mut root = make_test_tree();
    for i in -5..5 {
        let (new_root, value) = root.remove(i);
        println!("{}", new_root);
        root = new_root;
    }
}

#[test]
fn other() {
    let root1 = Root::<i32, bool>::new();
    let root2 = root1.insert(5, true);
    let root3 = root2.insert(10, true);
    let root4 = root3.insert(3, true);
    let root5 = root4.insert(1, true);
    let root6 = root5.insert(2, true);
    let root7 = root6.insert(6, true);
    let root8 = root7.insert(9, true);
    let Root::More(ref n8) = root8 else {
        unreachable!()
    };
    let root9 = root8.insert(7, true);
    let Root::More(ref n9) = root9 else {
        unreachable!()
    };
    let l8 = n8.left_opt.as_ref().unwrap();
    let r8 = n8.right_opt.as_ref().unwrap();
    println!("{:p} {:p} {} {}", l8, r8, l8, r8);
    let l9 = n9.left_opt.as_ref().unwrap();
    let r9 = n9.right_opt.as_ref().unwrap();
    println!("{:p} {:p} {} {}", l9, r9, l9, r9);

    println!("{}", root9);
    unsafe {
        let c = match root9 {
            Root::Empty => unreachable!(),
            Root::More(ref n) => n.as_ref().clone()
        };
        println!("origin: {}", c);
        println!("rotate: {}", c.rl_rotate());
    }
}