View source for AVL tree/Rust

<syntaxhighlight lang="rust">/// This implementation uses an addressable vector as the tree's store.
/// It is possible to construct a mutable tree using Rc<RefCell<>>,
/// but it adds some complexity.
///
/// "Pointers" to nodes are indices into the vector store, and have
/// trait Copy.
///
/// The index of a node in the vector store should not be confused with its key.
extern crate rand;
extern crate term_painter;

use std::cmp::Ordering;
use std::fmt::{Debug, Display, Formatter, Result};

use rand::distributions::Uniform;
use rand::Rng;
use term_painter::Color::*;
use term_painter::ToStyle;

pub type NodePtr = Option<usize>;

#[derive(Debug, PartialEq, Clone, Copy)]
pub enum Side {
    Left,
    Right,
    Up,
    Root,
}

#[derive(Debug, PartialEq, Clone, Copy)]
enum DisplayElement {
    TrunkSpace,
    SpaceLeft,
    SpaceRight,
    SpaceSpace,
    Root,
}

impl DisplayElement {
    fn string(&self) -> String {
        match *self {
            DisplayElement::TrunkSpace => "    │   ".to_string(),
            DisplayElement::SpaceRight => "    ┌───".to_string(),
            DisplayElement::SpaceLeft => "    └───".to_string(),
            DisplayElement::SpaceSpace => "        ".to_string(),
            DisplayElement::Root => "├──".to_string(),
        }
    }
}

/// Handedness of balanced insert and delete operations differs only by values encapsulated here.
struct BalanceConstants {
    bal_incr: i8,
    this_side: Side,
    that_side: Side,
    key_order: Ordering, // Ins only
    // These are used in the +1/-1 & -1/+1 deletion cases
    gcm1_child_adj: i8, // Del only, balance adjustment to child for b = -1 grandchild
    gcm1_parent_adj: i8, // Del only, balance adjustment to parent for b = -1 grandchild
    gcp1_child_adj: i8, // Del only, balance adjustment to child for b = 1 grandchild
    gcp1_parent_adj: i8, // Del only, balance adjustment to parent for b = 1 grandchild
}

const BALANCE_CONSTANTS_A: BalanceConstants = BalanceConstants {
    bal_incr: -1,
    this_side: Side::Left,
    that_side: Side::Right,
    key_order: Ordering::Greater,
    gcm1_child_adj: 0,
    gcm1_parent_adj: 1,
    gcp1_child_adj: -1,
    gcp1_parent_adj: 0,
};

const BALANCE_CONSTANTS_B: BalanceConstants = BalanceConstants {
    bal_incr: 1,
    this_side: Side::Right,
    that_side: Side::Left,
    key_order: Ordering::Less,
    gcm1_child_adj: 1,
    gcm1_parent_adj: 0,
    gcp1_child_adj: 0,
    gcp1_parent_adj: -1,
};

#[derive(Debug, Clone, Copy)]
pub struct Node<K, V> {
    key: K,
    value: V,
    balance: i8,
    left: NodePtr,
    right: NodePtr,
    up: NodePtr,
}

#[derive(Debug)]
pub struct AVLTree<K, V> {
    root: NodePtr,
    store: Vec<Node<K, V>>,
}

impl<K: Ord + Copy + Debug + Display, V: Debug + Copy + Display> Default for AVLTree<K, V> {
    fn default() -> Self {
        AVLTree::new()
    }
}

impl<K: Ord + Copy + Debug + Display, V: Debug + Copy + Display> AVLTree<K, V> {
    pub fn get_node(&self, np: NodePtr) -> Node<K, V> {
        assert!(np.is_some());
        self.store[np.unwrap()]
    }

    pub fn get_balance(&self, np: NodePtr) -> i8 {
        assert!(np.is_some());
        self.store[np.unwrap()].balance
    }

    pub fn get_key(&self, np: NodePtr) -> K {
        assert!(np.is_some());
        self.store[np.unwrap()].key
    }

    pub fn get_value(&self, np: NodePtr) -> V {
        assert!(np.is_some());
        self.store[np.unwrap()].value
    }

    pub fn get_pointer(&self, np: NodePtr, side: Side) -> NodePtr {
        assert!(np.is_some());
        self.store[np.unwrap()].get_ptr(side)
    }

    pub fn set_balance(&mut self, np: NodePtr, bal: i8) {
        assert!(np.is_some());
        self.store[np.unwrap()].balance = bal;
    }

    pub fn set_key(&mut self, np: NodePtr, to: K) {
        assert!(np.is_some());
        self.store[np.unwrap()].key = to;
    }

    pub fn set_value(&mut self, np: NodePtr, to: V) {
        assert!(np.is_some());
        self.store[np.unwrap()].value = to;
    }

    pub fn set_pointer(&mut self, np: NodePtr, side: Side, to: NodePtr) {
        assert!(np.is_some());
        self.store[np.unwrap()].set_ptr(side, to);
    }

    pub fn increment_balance(&mut self, np: NodePtr, delta: i8) -> i8 {
        assert!(np.is_some());
        self.store[np.unwrap()].balance += delta;
        self.store[np.unwrap()].balance
    }

    pub fn new() -> Self {
        AVLTree {
            root: None,
            store: Vec::<Node<K, V>>::with_capacity(20_000),
        }
    }

    /// Insert key-value
    pub fn insert(&mut self, k: K, v: V) -> Option<Node<K, V>> {
        let (n, _) = self.insert_node(Node::new(k, v));
        n
    }

    /// Insert Node struct
    pub fn insert_node(&mut self, mut n: Node<K, V>) -> (Option<Node<K, V>>, Side) {
        if self.root.is_none() {
            assert!(self.store.is_empty());
            self.store.push(n);
            self.root = Some(0);
            return (Some(n), Side::Root);
        }

        let mut p = self.root; // Possibly None
        let mut prev = p;
        let mut side = Side::Left;
        while p.is_some() {
            prev = p;
            match n.key.cmp(&self.get_key(p)) {
                Ordering::Less => {
                    side = Side::Left;
                    p = self.get_pointer(p, side);
                }
                Ordering::Greater => {
                    side = Side::Right;
                    p = self.get_pointer(p, side);
                }
                Ordering::Equal => {
                    println!("Key exists");
                    return (None, side);
                }
            }
        }
        // Set child's pointer
        n.up = prev;
        // Stow the node
        self.store.push(n);
        // Set parent's pointer
        let ptr = Some(self.store.len() - 1);
        self.set_pointer(prev, side, ptr);
        (Some(n), side)
    }

    /// Insert key-value and rebalance
    pub fn insert_bal(&mut self, k: K, v: V) -> Option<Node<K, V>> {
        self.insert_node_bal(Node::new(k, v))
    }

    /// Insert Node struct and rebalance
    pub fn insert_node_bal(&mut self, n: Node<K, V>) -> Option<Node<K, V>> {
        let (nins, side) = self.insert_node(n);
        if nins.is_none() || side == Side::Root {
            return nins;
        }

        let mut p = nins.unwrap().up;
        let mut is_left = side == Side::Left;

        while p.is_some() {
            let i_c = get_insertion_constants(is_left);

            let b = self.increment_balance(p, i_c.bal_incr);
            if b == 0 {
                break; // No further adjustments necessary
            } else if b.abs() > 1 {
                let child_p = self.get_pointer(p, i_c.this_side);
                match self.get_balance(child_p) * b {
                    2 => {
                        // -2/-1 & +2/+1 patterns
                        self.single_rotation(i_c.this_side, p, child_p);
                        self.set_balance(p, 0);
                        self.set_balance(child_p, 0);
                        break;
                    }
                    -2 => {
                        // -2/+1 & +2/-1 patterns
                        let grand_p = self.get_pointer(child_p, i_c.that_side);
                        self.double_rotation(i_c.this_side, p, child_p, grand_p);
                        if self.get_pointer(child_p, i_c.this_side).is_none() {
                            // Degenerate case, no subtrees
                            self.set_balance(child_p, 0);
                            self.set_balance(p, 0);
                        } else if n.key.cmp(&self.get_key(grand_p)) == i_c.key_order {
                            self.set_balance(child_p, i_c.bal_incr);
                            self.set_balance(p, 0);
                        } else {
                            self.set_balance(child_p, 0);
                            self.set_balance(p, -i_c.bal_incr);
                        }
                        self.set_balance(grand_p, 0);
                        break;
                    }
                    _ => unreachable!(),
                }
            }

            let child_p = p;
            p = self.get_pointer(p, Side::Up);
            if p.is_some() {
                let left_p = self.get_pointer(p, Side::Left);
                is_left = left_p.is_some() && left_p == child_p;
            }
        }

        nins
    }

    /// Remove the node at index from the store and patch the hole in the vector,
    /// modifying pointers in the moved node's parents and children.
    fn remove_carefully(&mut self, p: NodePtr) {
        assert!(p.is_some());
        let index = p.unwrap();
        let old_index = self.store.len() - 1;
        self.store.swap_remove(index);

        if index == old_index {
            // Nothing moved
            return;
        }

        // Element -1 has moved into the spot _index_. The in-pointers that need modifying
        // belong to that element's parent and children.

        // Fix child pointer in parent:
        let parent_p = self.get_pointer(p, Side::Up);
        if parent_p.is_some() {
            let l = self.get_pointer(parent_p, Side::Left);
            if l == Some(old_index) {
                self.set_pointer(parent_p, Side::Left, Some(index));
            } else {
                self.set_pointer(parent_p, Side::Right, Some(index));
            }
        }

        // Fix parent pointers in children:
        let l = self.get_pointer(p, Side::Left);
        let r = self.get_pointer(p, Side::Right);
        if l.is_some() {
            self.set_pointer(l, Side::Up, Some(index));
        }
        if r.is_some() {
            self.set_pointer(r, Side::Up, Some(index));
        }

        // Fix root if necessary
        if self.root == Some(old_index) {
            self.root = Some(index);
        }
    }

    /// Uses delete-by-copy procedure if node with key k has two children.
    /// Returns (parent, side) tuple.
    pub fn delete(&mut self, k: K) -> (NodePtr, Side) {
        let mut p = self.root;
        let mut prev = None;
        let mut res = None;
        let mut side = Side::Root;
        while p.is_some() {
            match k.cmp(&self.get_key(p)) {
                Ordering::Equal => {
                    break;
                }
                Ordering::Less => {
                    prev = p;
                    side = Side::Left;
                    p = self.get_pointer(p, side);
                }
                Ordering::Greater => {
                    prev = p;
                    side = Side::Right;
                    p = self.get_pointer(p, side);
                }
            }
        }

        if p.is_none() {
            println!("Key {:?} not found", k);
            return (res, side);
        }

        let n = self.get_node(p);
        // Is this a leaf?
        if n.is_leaf() {
            if n.key.cmp(&self.get_key(self.root)) == Ordering::Equal {
                self.root = None;
                assert_eq!(self.store.len(), 1);
            } else {
                self.set_pointer(prev, side, None);
            }
            self.remove_carefully(p);
            // The prev pointer is now stale
            res = if prev == Some(self.store.len()) {
                p
            } else {
                prev
            };

        // Is this a one-child node?
        } else if n.left.is_none() || n.right.is_none() {
            let ch = if n.left.is_some() { n.left } else { n.right };
            if n.key.cmp(&self.get_key(self.root)) == Ordering::Equal {
                self.set_pointer(ch, Side::Up, None);
                self.root = ch;
            } else {
                self.set_pointer(prev, side, ch);
                self.set_pointer(ch, Side::Up, prev);
            }
            self.remove_carefully(p);
            // The prev pointer is now stale
            res = if prev == Some(self.store.len()) {
                p
            } else {
                prev
            };

        // Complicated case:  two children, do delete-by-copy. Replace n with its first
        // predecessor (the mirror image using the first successor would work as well).
        } else {
            let mut tmp = n.left;
            let mut last = tmp;
            prev = self.get_pointer(tmp, Side::Up);
            while tmp.is_some() && self.get_pointer(last, Side::Right).is_some() {
                prev = self.get_pointer(tmp, Side::Up);
                last = tmp;
                tmp = self.get_pointer(tmp, Side::Right);
            }
            tmp = last;
            // Copy ...
            let the_key = self.get_key(tmp);
            let the_value = self.get_value(tmp);
            self.set_key(p, the_key);
            self.set_value(p, the_value);

            let left_ptr = self.get_pointer(tmp, Side::Left);
            if prev == p {
                self.set_pointer(p, Side::Left, left_ptr);
                if left_ptr.is_some() {
                    self.set_pointer(left_ptr, Side::Up, p);
                }
                side = Side::Left;
            } else {
                self.set_pointer(prev, Side::Right, left_ptr);
                if left_ptr.is_some() {
                    self.set_pointer(left_ptr, Side::Up, prev);
                }
                side = Side::Right;
            }

            self.remove_carefully(tmp);
            // The prev pointer is now stale
            res = if prev.unwrap() == self.store.len() {
                tmp
            } else {
                prev
            };
        }

        (res, side)
    }

    /// Rebalance on delete
    pub fn delete_bal(&mut self, k: K) -> Option<Node<K, V>> {
        // slug: (pointer to parent of deleted node, side of deleted node)
        let slug = self.delete(k);
        let (pdel, side) = slug;
        if pdel.is_none() {
            return None;
        };
        let ndel = self.get_node(pdel);

        let mut p = pdel;
        let mut is_left = side == Side::Left;

        // Rebalance and update balance factors. There are two different rotation sequences that
        // are the same within handedness,
        // and the +1/-1 / -1/+1 sequence has three possible balance adjustments
        // depending on the grandchild.
        while p.is_some() {
            let d_c = get_deletion_constants(is_left);

            let b = self.increment_balance(p, d_c.bal_incr);
            if b.abs() == 1 {
                break; // No further adjustments necessary
            } else if b.abs() > 1 {
                let child_p = self.get_pointer(p, d_c.this_side);
                match self.get_balance(child_p) * b {
                    2 => {
                        // +1/+1 & -1/-1 patterns
                        self.single_rotation(d_c.this_side, p, child_p);
                        self.set_balance(p, 0);
                        p = self.get_pointer(p, Side::Up);
                        self.set_balance(p, 0);
                    }
                    0 => {
                        // +1/0 & -1/0 patterns
                        self.single_rotation(d_c.this_side, p, child_p);
                        self.set_balance(p, d_c.bal_incr);
                        p = self.get_pointer(p, Side::Up);
                        self.set_balance(p, -d_c.bal_incr);
                        break; // No height change
                    }
                    -2 => {
                        // +1/-1/x & -1/+1/x patterns
                        let grand_p = self.get_pointer(child_p, d_c.that_side);
                        self.double_rotation(d_c.this_side, p, child_p, grand_p);
                        // p is now one child, grand_p is the other, child_p is their parent
                        match self.get_balance(grand_p) {
                            -1 => {
                                self.set_balance(p, d_c.gcm1_parent_adj);
                                self.set_balance(child_p, d_c.gcm1_child_adj);
                            }
                            0 => {
                                self.set_balance(p, 0);
                                self.set_balance(child_p, 0);
                            }
                            1 => {
                                self.set_balance(p, d_c.gcp1_parent_adj);
                                self.set_balance(child_p, d_c.gcp1_child_adj);
                            }
                            _ => unreachable!(),
                        }
                        self.set_balance(grand_p, 0);
                        p = self.get_pointer(p, Side::Up);
                    }
                    _ => unreachable!(),
                }
            }

            let child_p = p;
            p = self.get_pointer(p, Side::Up);
            if p.is_some() {
                let left_p = self.get_pointer(p, Side::Left);
                is_left = left_p.is_some() && left_p == child_p;
            }
        }

        Some(ndel)
    }

    /// Returns node value
    pub fn lookup(&self, k: K) -> Option<V> {
        self.search(k).map(|n| n.value)
    }

    /// Returns node (not pointer)
    pub fn search(&self, k: K) -> Option<Node<K, V>> {
        let mut p = self.root;
        let mut res = None;

        while p.is_some() {
            match k.cmp(&self.get_key(p)) {
                Ordering::Less => {
                    p = self.get_pointer(p, Side::Left);
                }
                Ordering::Greater => {
                    p = self.get_pointer(p, Side::Right);
                }
                Ordering::Equal => {
                    res = Some(self.get_node(p));
                    break;
                }
            }
        }
        res
    }

    /// Do an in-order traversal, where a "visit" prints the row with that node in it.
    fn display(&self, p: NodePtr, side: Side, e: &[DisplayElement], f: &mut Formatter) {
        if p.is_none() {
            return;
        }

        let mut elems = e.to_vec();
        let node = self.get_node(p);
        let mut tail = DisplayElement::SpaceSpace;
        if node.up != self.root {
            // Direction switching, need trunk element to be printed for lines before that node
            // is visited.
            let new_elem = if side == Side::Left && node.right.is_some() {
                DisplayElement::TrunkSpace
            } else {
                DisplayElement::SpaceSpace
            };
            elems.push(new_elem);
        }
        let hindex = elems.len() - 1;
        self.display(node.right, Side::Right, &elems, f);

        if p == self.root {
            elems[hindex] = DisplayElement::Root;
            tail = DisplayElement::TrunkSpace;
        } else if side == Side::Right {
            // Right subtree finished
            elems[hindex] = DisplayElement::SpaceRight;
            // Prepare trunk element in case there is a left subtree
            tail = DisplayElement::TrunkSpace;
        } else
        // if side == Side::Left
        {
            elems[hindex] = DisplayElement::SpaceLeft;
            let parent_p = self.get_pointer(p, Side::Up);
            let gp_p = self.get_pointer(parent_p, Side::Up);
            if gp_p.is_some() && self.get_pointer(gp_p, Side::Right) == parent_p {
                // Direction switched, need trunk element starting with this node/line
                elems[hindex - 1] = DisplayElement::TrunkSpace;
            }
        }

        // Visit node => print accumulated elements. Each node gets a line.
        {
            for e in elems.clone() {
                let _ = write!(f, "{}", e.string());
            }
            let _ = write!(
                f,
                "{key:>width$} ",
                key = Green.bold().paint(node.key),
                width = 2
            );
            let _ = write!(
                f,
                "{value:>width$} ",
                value = Blue.bold().paint(format!("{:.*}", 2, node.value)),
                width = 4
            );
            let _ = write!(
                f,
                "{bal:<-width$}\n",
                bal = Red.bold().paint(node.balance),
                width = 2
            );

            elems[hindex] = tail;
        }

        self.display(node.left, Side::Left, &elems, f);
    }

    pub fn gather_balances(&self) -> (Vec<K>, Vec<i8>) {
        let mut keys = Vec::<K>::new();
        let mut bals = Vec::<i8>::new();

        self.gather_balances_impl(self.root, &mut keys, &mut bals);
        (keys, bals)
    }

    fn gather_balances_impl(&self, p: NodePtr, k: &mut Vec<K>, b: &mut Vec<i8>) {
        if p.is_none() {
            return;
        }
        let r = self.get_pointer(p, Side::Right);
        self.gather_balances_impl(r, k, b);
        k.push(self.get_key(p));
        b.push(self.get_balance(p));
        let l = self.get_pointer(p, Side::Left);
        self.gather_balances_impl(l, k, b)
    }

    pub fn compute_balances(&mut self, p: NodePtr) -> i8 {
        self.compute_balances_impl(p, 0)
    }

    fn compute_balances_impl(&mut self, p: NodePtr, level: i8) -> i8 {
        if p.is_none() {
            return level - 1;
        }
        let r = self.get_pointer(p, Side::Right);
        let l = self.get_pointer(p, Side::Left);
        let rb = self.compute_balances_impl(r, level + 1);
        let lb = self.compute_balances_impl(l, level + 1);
        self.set_balance(p, rb - lb);
        std::cmp::max(rb, lb)
    }

    ///     P                Q
    ///   /   \     =>     /   \
    ///  h     Q          P     h'
    fn rotate_left(&mut self, p: NodePtr, q: NodePtr) {
        assert!(p.is_some());
        assert!(q.is_some());
        let p_parent = self.get_pointer(p, Side::Up);
        // Take care of parent pointers
        self.set_pointer(q, Side::Up, p_parent);
        self.set_pointer(p, Side::Up, q);
        let ql = self.get_pointer(q, Side::Left);
        if ql.is_some() {
            self.set_pointer(ql, Side::Up, p);
        }

        // Take care of child pointers
        self.set_pointer(q, Side::Left, p);
        self.set_pointer(p, Side::Right, ql);
        if p_parent.is_some() {
            let side = if self.get_pointer(p_parent, Side::Right) == p {
                Side::Right
            } else {
                Side::Left
            };
            self.set_pointer(p_parent, side, q);
        } else {
            self.root = q;
        }
    }

    ///     P                Q
    ///   /   \     =>     /   \
    ///  Q     h          h'    P
    fn rotate_right(&mut self, p: NodePtr, q: NodePtr) {
        assert!(p.is_some());
        assert!(q.is_some());
        let p_parent = self.get_pointer(p, Side::Up);
        // Take care of parent pointers
        self.set_pointer(q, Side::Up, p_parent);
        self.set_pointer(p, Side::Up, q);
        let qr = self.get_pointer(q, Side::Right);
        if qr.is_some() {
            self.set_pointer(qr, Side::Up, p);
        }

        // Take care of child pointers
        self.set_pointer(q, Side::Right, p);
        self.set_pointer(p, Side::Left, qr);
        if p_parent.is_some() {
            let side = if self.get_pointer(p_parent, Side::Right) == p {
                Side::Right
            } else {
                Side::Left
            };
            self.set_pointer(p_parent, side, q);
        } else {
            self.root = q;
        }
    }

    fn single_rotation(&mut self, side: Side, p: NodePtr, q: NodePtr) {
        if side == Side::Left {
            self.rotate_right(p, q);
        } else {
            self.rotate_left(p, q);
        }
    }

    fn double_rotation(&mut self, side: Side, p: NodePtr, child_p: NodePtr, grand_p: NodePtr) {
        if side == Side::Left {
            self.rotate_left(child_p, grand_p);
            self.rotate_right(p, grand_p);
        } else {
            self.rotate_right(child_p, grand_p);
            self.rotate_left(p, grand_p);
        }
    }
}

impl<K: Ord + Copy, V: Copy> Node<K, V> {
    pub fn new(k: K, v: V) -> Node<K, V> {
        Node {
            key: k,
            value: v,
            balance: 0,
            left: None,
            right: None,
            up: None,
        }
    }

    pub fn is_leaf(&self) -> bool {
        self.left.is_none() && self.right.is_none()
    }

    pub fn set_ptr(&mut self, side: Side, to: NodePtr) {
        let field = match side {
            Side::Up => &mut self.up,
            Side::Left => &mut self.left,
            _ => &mut self.right,
        };
        *field = to;
    }

    pub fn get_ptr(&self, side: Side) -> NodePtr {
        match side {
            Side::Up => self.up,
            Side::Left => self.left,
            _ => self.right,
        }
    }
}

impl<K: Ord + Copy + Debug + Display, V: Debug + Copy + Display> Display for AVLTree<K, V> {
    fn fmt(&self, f: &mut Formatter) -> Result {
        if self.root.is_none() {
            write!(f, "[empty]")
        } else {
            let v: Vec<DisplayElement> = Vec::new();
            self.display(self.root, Side::Up, &v, f);
            Ok(())
        }
    }
}

fn get_insertion_constants(is_left: bool) -> &'static BalanceConstants {
    if is_left {
        &BALANCE_CONSTANTS_A
    } else {
        &BALANCE_CONSTANTS_B
    }
}

fn get_deletion_constants(is_left: bool) -> &'static BalanceConstants {
    get_insertion_constants(!is_left)
}

pub fn random_bal_tree(n: u32) -> AVLTree<i32, f32> {
    let mut tree: AVLTree<i32, f32> = AVLTree::new();
    let mut rng = rand::thread_rng();
    // `Uniform` rather than `gen_range`'s `Uniform::sample_single` for speed
    let key_range = Uniform::new(-(n as i32) / 2, (n as i32) / 2);
    let value_range = Uniform::new(-1.0, 1.0);
    tree.insert_bal(0, rng.sample(value_range));
    for _ in 0..n {
        tree.insert_bal(rng.sample(key_range), rng.sample(value_range));
    }
    tree
}

#[cfg(test)]
mod tests {
    use rand::{seq, thread_rng};

    use super::AVLTree;
    use random_bal_tree;

    #[test]
    fn test_insert() {
        let mut tree: AVLTree<i32, f32> = AVLTree::new();
        tree.insert(0, 0.0);
        tree.insert(8, 8.8);
        tree.insert(-8, -8.8);
        assert!(tree.insert(4, 4.4).is_some());
        tree.insert(12, 12.12);

        assert_eq!(tree.lookup(4), Some(4.4));
        assert_eq!(tree.lookup(5), None);
        assert_eq!(tree.lookup(-8), Some(-8.8));

        let s = &tree.store;
        assert_eq!(
            s[s[s[tree.root.unwrap()].right.unwrap()].right.unwrap()].value,
            12.12
        );
        assert_eq!(
            s[s[s[tree.root.unwrap()].right.unwrap()].right.unwrap()].left,
            None
        );
    }

    #[test]
    fn test_delete() {
        let mut tree: AVLTree<i32, f32> = AVLTree::new();
        tree.insert(0, 0.0);
        tree.insert(8, 8.8);
        tree.insert(-8, -8.8);
        assert!(tree.insert(4, 4.4).is_some());
        tree.insert(12, 12.12);

        // delete leaf
        tree.delete(12);
        assert_eq!(tree.lookup(12), None);
        let mut n = tree.search(8).unwrap();
        assert_eq!(n.right, None);

        // delete one-child node
        tree.delete(4);
        assert_eq!(tree.lookup(4), None);
        n = tree.search(0).unwrap();
        assert_eq!(tree.store[n.right.unwrap()].key, 8);

        // delete two-child node
        tree.insert(6, 6.6);
        tree.insert(10, 10.10);
        tree.insert(7, 7.7);
        tree.delete(8);
        n = tree.search(7).unwrap();
        assert_eq!(tree.store[n.left.unwrap()].key, 6);
        assert_eq!(tree.store[n.right.unwrap()].key, 10);
        assert_eq!(tree.store[n.up.unwrap()].key, 0);

        // delete two-child root
        tree.delete(0);
        assert_eq!(tree.store[tree.root.unwrap()].key, -8);

        // delete one-child root
        tree.delete(-8);
        assert_eq!(tree.store[tree.root.unwrap()].key, 7);

        // delete no-child root
        tree.delete(6);
        tree.delete(7);
        tree.delete(10);
        assert!(tree.root.is_none());
        assert_eq!(tree.store.len(), 0);
    }

    #[test]
    fn test_rotate_left() {
        let mut tree: AVLTree<i32, f32> = AVLTree::new();
        tree.insert(0, 0.0);
        tree.insert(8, 8.8);
        tree.insert(4, 4.4);
        tree.insert(-8, -8.8);

        let mut r = tree.root;
        let mut right = tree.store[r.unwrap()].right;
        tree.rotate_left(r, right);
        r = tree.root;
        right = tree.store[r.unwrap()].right;
        let left = tree.store[r.unwrap()].left;
        let left_left = tree.store[left.unwrap()].left;
        let left_right = tree.store[left.unwrap()].right;
        assert_eq!(right, None);
        assert_eq!(tree.store[left.unwrap()].key, 0);
        assert_eq!(tree.store[left_left.unwrap()].key, -8);
        assert_eq!(tree.store[left_right.unwrap()].key, 4);
        assert_eq!(tree.store[r.unwrap()].key, 8);
    }

    #[test]
    fn test_rotate_right() {
        let mut tree: AVLTree<i32, f32> = AVLTree::new();
        tree.insert(0, 0.0);
        tree.insert(8, 8.8);
        tree.insert(-8, -8.8);
        tree.insert(-4, 4.4);

        let mut r = tree.root;
        let mut left = tree.store[r.unwrap()].left;
        tree.rotate_right(r, left);
        r = tree.root;
        left = tree.store[r.unwrap()].left;
        let right = tree.store[r.unwrap()].right;
        let right_right = tree.store[right.unwrap()].right;
        let right_left = tree.store[right.unwrap()].left;
        assert_eq!(left, None);
        assert_eq!(tree.store[right.unwrap()].key, 0);
        assert_eq!(tree.store[right_right.unwrap()].key, 8);
        assert_eq!(tree.store[right_left.unwrap()].key, -4);
        assert_eq!(tree.store[r.unwrap()].key, -8);
    }

    #[test]
    // This tree tests all four insertion types
    fn test_balanced_inserts() {
        let mut tree: AVLTree<i32, f32> = AVLTree::new();
        tree.insert_bal(0, 0.0);
        tree.insert_bal(8, 8.8);
        tree.insert_bal(-8, -8.8);
        tree.insert_bal(12, 12.12);
        tree.insert_bal(16, 16.16);
        tree.insert_bal(11, 11.11);
        tree.insert_bal(4, 4.4);
        tree.insert_bal(-10, -8.8);
        tree.insert_bal(-12, -8.8);
        tree.insert_bal(-9, -8.8);

        let mut res = tree.gather_balances();
        let (_, bals) = res;
        assert!(bals.iter().max().unwrap() < &2);
        assert!(bals.iter().min().unwrap() > &-2);

        for _ in 0..10 {
            tree = random_bal_tree(1000);
            res = tree.gather_balances();
            let (_, bals) = res;
            assert!(bals.iter().max().unwrap() < &2);
            assert!(bals.iter().min().unwrap() > &-2);
        }
    }

    #[test]
    /// This sequence hits all five rotation possibilities on each side.
    fn test_balanced_deletes() {
        let mut tree: AVLTree<i32, f32> = AVLTree::new();
        tree.insert_bal(0, 0.0);
        tree.insert_bal(-32, 0.0);
        tree.insert_bal(32, 0.0);
        tree.insert_bal(-64, 0.0);
        tree.insert_bal(64, 0.0);
        tree.delete_bal(64);
        tree.delete_bal(32);
        tree.delete_bal(-32);
        tree.delete_bal(-64);
        tree.delete_bal(0);
        assert_eq!(tree.root, None);
        assert_eq!(tree.store.len(), 0);

        tree.insert_bal(0, 0.0);
        tree.insert_bal(-32, 0.0);
        tree.insert_bal(32, 0.0);
        tree.insert_bal(-64, 0.0);
        tree.insert_bal(64, 0.0);
        tree.insert_bal(-16, 0.0);
        tree.insert_bal(16, 0.0);
        tree.insert_bal(-8, 0.0);
        tree.insert_bal(8, 0.0);
        tree.insert_bal(-12, 0.0);
        tree.insert_bal(-7, 0.0);
        tree.insert_bal(-6, 0.0);
        tree.insert_bal(-11, 0.0);

        tree.delete_bal(-64);
        tree.delete_bal(-32);
        tree.delete_bal(-7);
        tree.delete_bal(-6);
        tree.delete_bal(-16);
        tree.delete_bal(-11);
        tree.delete_bal(-12);
        tree.delete_bal(8);
        tree.delete_bal(-8);
        tree.delete_bal(0);
        tree.insert_bal(24, 0.0);
        tree.insert_bal(8, 0.0);
        tree.insert_bal(4, 0.0);
        tree.insert_bal(128, 0.0);
        tree.insert_bal(48, 0.0);
        tree.delete_bal(32);
        tree.delete_bal(48);

        tree.insert_bal(-24, 0.0);
        tree.insert_bal(-8, 0.0);
        tree.insert_bal(-128, 0.0);
        tree.insert_bal(-48, 0.0);
        tree.insert_bal(-20, 0.0);
        tree.insert_bal(-30, 0.0);
        tree.insert_bal(-22, 0.0);
        tree.insert_bal(-21, 0.0);
        tree.delete_bal(24);
        tree.delete_bal(64);
        tree.delete_bal(-30);
        tree.delete_bal(-22);
        tree.delete_bal(-21);
        tree.delete_bal(-128);
        tree.delete_bal(128);
        tree.delete_bal(-8);
        tree.insert_bal(-96, 0.0);
        tree.insert_bal(-95, 0.0);
        tree.insert_bal(-10, 0.0);
        tree.insert_bal(6, 0.0);
        tree.delete_bal(-24);

        let mut res = tree.gather_balances();
        let (_, bals) = res;
        assert!(bals.iter().max().unwrap() < &2);
        assert!(bals.iter().min().unwrap() > &-2);

        let mut p = tree.root;
        while p.is_some() {
            let key = tree.store[p.unwrap()].key;
            tree.delete_bal(key);
            p = tree.root;
        }
        assert_eq!(tree.root, None);
        assert_eq!(tree.store.len(), 0);

        // */*/+1 patterns
        tree.insert(6, 0.0);
        tree.insert(-1, 0.0);
        tree.insert(9, 0.0);
        tree.insert(7, 0.0);
        tree.insert(3, 0.0);
        tree.insert(-9, 0.0);
        tree.insert(4, 0.0);
        p = tree.root;
        tree.compute_balances(p);
        tree.delete_bal(-9);
        res = tree.gather_balances();
        let (_, bals) = res;
        tree.compute_balances(p);
        res = tree.gather_balances();
        let (_, bals_after) = res;
        assert_eq!(bals, bals_after);

        tree.insert(6, 0.0);
        tree.insert(-1, 0.0);
        tree.insert(3, 0.0);
        tree.insert(9, 0.0);
        tree.insert(7, 0.0);
        tree.insert(11, 0.0);
        tree.insert(8, 0.0);
        p = tree.root;
        tree.compute_balances(p);
        tree.delete_bal(-1);
        res = tree.gather_balances();
        let (_, bals) = res;
        tree.compute_balances(p);
        res = tree.gather_balances();
        let (_, bals_after) = res;
        assert_eq!(bals, bals_after);

        let mut rng = thread_rng();
        for _ in 0..100 {
            tree = random_bal_tree(100);
            let sample = seq::sample_iter(&mut rng, -50..50, 80).unwrap();
            for i in sample {
                tree.delete_bal(i);
            }
        }

        res = tree.gather_balances();
        let (_, bals) = res;

        if bals.len() > 0 {
            assert!(*bals.iter().max().unwrap() < 2);
            assert!(*bals.iter().min().unwrap() > -2);
        }

        return;
    }
}</syntaxhighlight>