link_cut_tree.rs

use std::{
    cmp::Ordering,
    fmt,
    ops::{Index, IndexMut},
};

use super::collections::{Label, Pool};

enum Parent<V> {
    Path(Label<Node<V>>),
    BinTree(Label<Node<V>>),
    Root,
}

impl<V> Copy for Parent<V> {}

impl<V> Clone for Parent<V> {
    fn clone(&self) -> Parent<V> {
        *self
    }
}

impl<V> fmt::Debug for Parent<V> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Parent::Path(label) => f
                .debug_struct("Parent::Path")
                .field("label", label)
                .finish(),
            Parent::BinTree(label) => f
                .debug_struct("Parent::BinTree")
                .field("label", label)
                .finish(),
            Parent::Root => f.debug_struct("Parent::Root").finish(),
        }
    }
}

impl<V> PartialEq for Parent<V> {
    fn eq(&self, other: &Self) -> bool {
        match self {
            Parent::Path(label1) => {
                if let Parent::Path(label2) = other {
                    label1 == label2
                } else {
                    false
                }
            }
            Parent::BinTree(label1) => {
                if let Parent::BinTree(label2) = other {
                    label1 == label2
                } else {
                    false
                }
            }
            Parent::Root => matches!(other, Parent::Root),
        }
    }
}

#[derive(Copy, Clone, Debug, PartialEq)]
enum ChildType {
    Left,
    Right,
    Path,
    Root,
}

impl ChildType {
    fn mirror(&self) -> Option<Self> {
        match self {
            ChildType::Left => Some(ChildType::Right),
            ChildType::Right => Some(ChildType::Left),
            _ => None,
        }
    }
}

impl From<bool> for ChildType {
    fn from(var: bool) -> Self {
        if var {
            ChildType::Right
        } else {
            ChildType::Left
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq)]
struct Children<V>(Option<Label<Node<V>>>, Option<Label<Node<V>>>);

impl<V> Index<ChildType> for Children<V> {
    type Output = Option<Label<Node<V>>>;

    fn index(&self, child_type: ChildType) -> &Self::Output {
        match child_type {
            ChildType::Left => &self.0,
            ChildType::Right => &self.1,
            _ => unreachable!(),
        }
    }
}

impl<V> IndexMut<ChildType> for Children<V> {
    fn index_mut(&mut self, child_type: ChildType) -> &mut Self::Output {
        match child_type {
            ChildType::Left => &mut self.0,
            ChildType::Right => &mut self.1,
            _ => unreachable!(),
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub struct Node<V> {
    value: V,
    parent: Parent<V>,
    children: Children<V>,
    d_flip: bool,
    d_depth: usize,
}

impl<V> Pool<Node<V>> {
    pub fn add_node(&mut self, value: V) -> Label<Node<V>> {
        self.insert(Node {
            value,
            parent: Parent::Root,
            children: Children(None, None),
            d_flip: false,
            d_depth: 0,
        })
    }

    pub fn find_root(&mut self, u: Label<Node<V>>) -> Label<Node<V>> {
        self.expose(u);
        let mut x = u;
        let mut flip = self[u].d_flip;
        while let Some(y) = self[x].children[ChildType::from(flip)] {
            x = y;
            flip ^= self[x].d_flip;
        }
        self.expose(x);
        x
    }

    pub fn cut(&mut self, c: Label<Node<V>>) {
        self.expose(c);
        let d_flip = self[c].d_flip;
        if let Some(p) = self[c].children[ChildType::from(d_flip)] {
            self[p].parent = Parent::Root;
            self[c].children[ChildType::from(d_flip)] = None;

            self[p].d_depth = self[c].d_depth - 1;
            self[c].d_depth = 0;
        }
    }

    pub fn link(&mut self, p: Label<Node<V>>, c: Label<Node<V>>) {
        self.expose(c);
        self.expose(p);
        let d_flip = self[c].d_flip;
        self[c].children[ChildType::from(d_flip)] = Some(p);
        self[p].parent = Parent::BinTree(c);

        self[c].d_depth = self[p].d_depth + 1;
        self[p].d_depth = 1;
    }

    pub fn evert(&mut self, u: Label<Node<V>>) {
        self.expose(u);
        self[u].d_flip = !self[u].d_flip;
        self[u].d_depth = 0;
    }

    /*     pub fn first_ancestor(&mut self, u: Label<Node<V>>, v: Label<Node<V>>) -> Label<Node<V>> {
        self.expose(u);
        self.expose(v);

        let mut x = u;
        while let Parent::BinTree(y) = self[x].parent {
            x = y;
        }

        match self[x].parent {
            Parent::BinTree(_) => unreachable!(),
            Parent::Path(y) => y,
            Parent::Root => v,
        }
    } */

    pub fn depth(&self, u: Label<Node<V>>) -> usize {
        let mut label = u;
        let mut depth: isize = 0;
        loop {
            if self.parent(label).is_some() {
                let d_depth = self[label].d_depth as isize;
                let d_flip = self[self.parent(label).unwrap()].d_flip;
                let child_type = self.child_type(label);

                match child_type {
                    ChildType::Left => depth -= d_depth,
                    ChildType::Right => depth += d_depth,
                    ChildType::Path => depth += d_depth,
                    ChildType::Root => unreachable!(),
                }

                if d_flip {
                    depth = -depth;
                }
                label = self.parent(label).unwrap();
            } else {
                let d_depth = self[label].d_depth as isize;
                return (depth + d_depth).abs() as usize;
            }
        }
    }

    // always expose before indexing
    pub fn index_depth(&self, u: Label<Node<V>>, i: usize) -> Label<Node<V>> {
        assert!(i <= self.depth(u));
        let mut label = u;
        let mut flip = self[u].d_flip;
        let mut depth = self[u].d_depth;
        loop {
            let side = match i.cmp(&depth) {
                Ordering::Less => flip,
                Ordering::Equal => return label,
                Ordering::Greater => !flip,
            };

            label = self[label].children[ChildType::from(side)].unwrap();
            if flip ^ side {
                depth += self[label].d_depth;
            } else {
                depth -= self[label].d_depth;
            }
            flip ^= self[label].d_flip;
        }
    }

    pub fn value_mut(&mut self, u: Label<Node<V>>) -> &mut V {
        &mut self[u].value
    }

    pub fn value(&self, u: Label<Node<V>>) -> &V {
        &self[u].value
    }

    fn expose(&mut self, u: Label<Node<V>>) {
        self.splay(u);
        if let Some(c) = self[u].children[ChildType::from(true)] {
            self[c].parent = Parent::Path(u);
            self[u].children[ChildType::from(true)] = None;
        }

        while let Parent::Path(p) = self[u].parent {
            self.splay(p);
            if let Some(x) = self[p].children[ChildType::from(true)] {
                self[x].parent = Parent::Path(p);
            }
            self[p].children[ChildType::from(true)] = Some(u);
            self[u].parent = Parent::BinTree(p);
            self.rotate(u);
        }
    }

    fn splay(&mut self, u: Label<Node<V>>) {
        self.cascade_flipped(u);
        while let Parent::BinTree(p) = self[u].parent {
            if let Parent::BinTree(_) = self[p].parent {
                if self[p].d_flip == (self.child_type(u) == self.child_type(p)) {
                    self.rotate(p);
                    self.rotate(u);
                } else {
                    self.rotate(u);
                    self.rotate(u);
                };
            } else {
                self.rotate(u);
                break;
            }
        }
    }

    fn cascade_flipped(&mut self, u: Label<Node<V>>) {
        if let Parent::BinTree(p) = self[u].parent {
            self.cascade_flipped(p);
        }

        if self[u].d_flip {
            let Children(a_opt, b_opt) = self[u].children;
            self[u].children = Children(b_opt, a_opt);
            self[u].d_flip = false;
            if let Some(a) = a_opt {
                self[a].d_flip = !self[a].d_flip;
            }
            if let Some(b) = b_opt {
                self[b].d_flip = !self[b].d_flip;
            }
        }
    }

    fn rotate(&mut self, u: Label<Node<V>>) {
        let u_child_type = self.child_type(u);
        let x_opt = self[u].children[u_child_type.mirror().unwrap()];
        if let Parent::BinTree(p) = self[u].parent {
            let u_d_depth = self[u].d_depth;
            let p_d_depth = self[p].d_depth;

            self[p].d_depth = u_d_depth;
            self[u].d_depth = if let Parent::BinTree(_) = self[p].parent {
                if self[p].d_flip ^ (self.child_type(u) == self.child_type(p)) {
                    u_d_depth + p_d_depth
                } else {
                    u_d_depth.abs_diff(p_d_depth)
                }
            } else {
                u_d_depth.abs_diff(p_d_depth)
            };

            if let Some(x) = x_opt {
                let x_d_depth = self[x].d_depth;
                self[x].d_depth = x_d_depth.abs_diff(u_d_depth);
            }

            if let Parent::BinTree(g) = self[p].parent {
                let p_child_type = self.child_type(p);
                self[g].children[p_child_type] = Some(u);
            }

            self[p].children[u_child_type] = x_opt;
            self[u].children[u_child_type.mirror().unwrap()] = Some(p);

            if let Some(x) = x_opt {
                self[x].parent = Parent::BinTree(p);
            }

            self[u].parent = self[p].parent;
            self[p].parent = Parent::BinTree(u);
        } else {
            unreachable!();
        };
    }

    fn child_type(&self, u: Label<Node<V>>) -> ChildType {
        match self[u].parent {
            Parent::BinTree(parent) => {
                if let Some(x) = self[parent].children[ChildType::from(false)] {
                    if u == x {
                        return ChildType::Left;
                    }
                }
                if let Some(x) = self[parent].children[ChildType::from(true)] {
                    if u == x {
                        return ChildType::Right;
                    }
                }

                unreachable!();
            }
            Parent::Path(_) => ChildType::Path,
            Parent::Root => ChildType::Root,
        }
    }

    fn parent(&self, u: Label<Node<V>>) -> Option<Label<Node<V>>> {
        match self[u].parent {
            Parent::BinTree(parent) => Some(parent),
            Parent::Path(parent) => Some(parent),
            Parent::Root => None,
        }
    }
}

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

    #[test]
    fn it_works() {
        let mut lct = Pool::with_capacity(1);
        let v = lct.add_node(());
        lct.splay(v);
        lct.expose(v);
    }

    #[test]
    fn rotate_basic() {
        let mut lct = Pool::with_capacity(2);

        // define nodes
        let o = lct.add_node("node");
        let p = lct.add_node("parent");

        // link nodes
        lct[p].children[ChildType::from(false)] = Some(o);
        lct[o].parent = Parent::BinTree(p);

        // rotate
        lct.rotate(o);
        dbg!(&lct);

        // check children
        assert_eq!(lct[o].children, Children(None, Some(p)));
        assert_eq!(lct[p].children, Children(None, None));

        // check parents
        assert_eq!(lct[o].parent, Parent::Root);
        assert_eq!(lct[p].parent, Parent::BinTree(o));
    }

    #[test]
    fn rotate_attachments() {
        let mut lct = Pool::with_capacity(7);

        // define nodes
        let o = lct.add_node("node");
        let p = lct.add_node("parent");
        let g = lct.add_node("grandparent");
        let a = lct.add_node("a");
        let b = lct.add_node("b");
        let c = lct.add_node("c");
        let d = lct.add_node("d");

        // link nodes
        lct[o].children = Children(Some(a), Some(b));
        lct[p].children = Children(Some(o), Some(c));
        lct[g].children = Children(Some(p), Some(d));
        lct[o].parent = Parent::BinTree(p);
        lct[p].parent = Parent::BinTree(g);
        lct[a].parent = Parent::BinTree(o);
        lct[b].parent = Parent::BinTree(o);
        lct[c].parent = Parent::BinTree(p);
        lct[d].parent = Parent::BinTree(g);

        // rotate
        lct.rotate(o);
        dbg!(&lct);

        // check children
        assert_eq!(lct[o].children, Children(Some(a), Some(p)));
        assert_eq!(lct[p].children, Children(Some(b), Some(c)));
        assert_eq!(lct[g].children, Children(Some(o), Some(d)));
        assert_eq!(lct[a].children, Children(None, None));
        assert_eq!(lct[b].children, Children(None, None));
        assert_eq!(lct[c].children, Children(None, None));
        assert_eq!(lct[d].children, Children(None, None));

        // check parents
        assert_eq!(lct[o].parent, Parent::BinTree(g));
        assert_eq!(lct[p].parent, Parent::BinTree(o));
        assert_eq!(lct[g].parent, Parent::Root);
        assert_eq!(lct[a].parent, Parent::BinTree(o));
        assert_eq!(lct[b].parent, Parent::BinTree(p));
        assert_eq!(lct[c].parent, Parent::BinTree(p));
        assert_eq!(lct[d].parent, Parent::BinTree(g));
    }

    #[test]
    fn rotate_flipped() {
        let mut lct = Pool::with_capacity(7);

        // define nodes
        let o = lct.add_node("node");
        let p = lct.add_node("parent");
        let g = lct.add_node("grandparent");
        let a = lct.add_node("a");
        let b = lct.add_node("b");
        let c = lct.add_node("c");
        let d = lct.add_node("d");

        // link nodes
        lct[o].children = Children(Some(a), Some(b));
        lct[p].children = Children(Some(o), Some(c));
        lct[g].children = Children(Some(p), Some(d));
        lct[o].parent = Parent::BinTree(p);
        lct[p].parent = Parent::BinTree(g);
        lct[a].parent = Parent::BinTree(o);
        lct[b].parent = Parent::BinTree(o);
        lct[c].parent = Parent::BinTree(p);
        lct[d].parent = Parent::BinTree(g);

        // reverse origin
        lct[o].d_flip = true;

        // rotate
        lct.cascade_flipped(o);
        lct.rotate(o);
        dbg!(&lct);

        // check children
        assert_eq!(lct[o].children, Children(Some(b), Some(p)));
        assert_eq!(lct[p].children, Children(Some(a), Some(c)));
        assert_eq!(lct[g].children, Children(Some(o), Some(d)));
        assert_eq!(lct[a].children, Children(None, None));
        assert_eq!(lct[b].children, Children(None, None));
        assert_eq!(lct[c].children, Children(None, None));
        assert_eq!(lct[d].children, Children(None, None));

        // check parents
        assert_eq!(lct[o].parent, Parent::BinTree(g));
        assert_eq!(lct[p].parent, Parent::BinTree(o));
        assert_eq!(lct[g].parent, Parent::Root);
        assert_eq!(lct[a].parent, Parent::BinTree(p));
        assert_eq!(lct[b].parent, Parent::BinTree(o));
        assert_eq!(lct[c].parent, Parent::BinTree(p));
        assert_eq!(lct[d].parent, Parent::BinTree(g));

        // check flipped
        assert!(!lct[o].d_flip);
        assert!(!lct[p].d_flip);
        assert!(!lct[g].d_flip);
        assert!(lct[a].d_flip);
        assert!(lct[b].d_flip);
        assert!(!lct[c].d_flip);
        assert!(!lct[d].d_flip);
    }

    #[test]
    fn expose_line() {
        let mut lct = Pool::with_capacity(5);

        // define nodes
        let a = lct.add_node("a");
        let b = lct.add_node("b");
        let c = lct.add_node("c");
        let d = lct.add_node("d");
        let e = lct.add_node("e");

        // link nodes
        lct[b].children = Children(None, Some(c));
        lct[d].children = Children(None, Some(e));
        lct[c].parent = Parent::BinTree(b);
        lct[e].parent = Parent::BinTree(d);
        lct[b].parent = Parent::Path(a);
        lct[d].parent = Parent::Path(c);

        // expose
        lct.expose(d);
        dbg!(&lct);

        // check parents
        assert_eq!(lct[a].parent, Parent::BinTree(d));
        assert_eq!(lct[b].parent, Parent::BinTree(c));
        assert_eq!(lct[c].parent, Parent::BinTree(a));
        assert_eq!(lct[d].parent, Parent::Root);
        assert_eq!(lct[e].parent, Parent::Path(d));
    }

    #[test]
    fn find_root() {
        let mut lct = Pool::with_capacity(7);

        // define nodes
        let o = lct.add_node("node");
        let p = lct.add_node("parent");
        let g = lct.add_node("grandparent");
        let a = lct.add_node("a");
        let b = lct.add_node("b");
        let c = lct.add_node("c");
        let d = lct.add_node("d");

        // link nodes
        lct[o].children = Children(Some(a), Some(b));
        lct[p].children = Children(Some(o), Some(c));
        lct[g].children = Children(Some(p), Some(d));
        lct[o].parent = Parent::BinTree(p);
        lct[p].parent = Parent::BinTree(g);
        lct[a].parent = Parent::BinTree(o);
        lct[b].parent = Parent::BinTree(o);
        lct[c].parent = Parent::BinTree(p);
        lct[d].parent = Parent::BinTree(g);

        // find root
        assert_eq!(lct.find_root(g), a);
    }

    #[test]
    fn link_cut() {
        let mut lct = Pool::with_capacity(2);

        // define nodes
        let a = lct.add_node("a");
        let b = lct.add_node("b");

        // link nodes
        lct.link(a, b);

        // cut nodes
        lct.cut(b);
    }

    #[test]
    fn evert_root() {
        let mut lct = Pool::with_capacity(7);

        // define nodes
        let o = lct.add_node("node");
        let p = lct.add_node("parent");
        let g = lct.add_node("grandparent");
        let a = lct.add_node("a");
        let b = lct.add_node("b");
        let c = lct.add_node("c");
        let d = lct.add_node("d");

        // link nodes
        lct[o].children = Children(Some(a), Some(b));
        lct[p].children = Children(Some(o), Some(c));
        lct[g].children = Children(Some(p), Some(d));
        lct[o].parent = Parent::BinTree(p);
        lct[p].parent = Parent::BinTree(g);
        lct[a].parent = Parent::BinTree(o);
        lct[b].parent = Parent::BinTree(o);
        lct[c].parent = Parent::BinTree(p);
        lct[d].parent = Parent::BinTree(g);

        // evert
        lct.evert(o);
        dbg!(&lct);
        assert_eq!(lct.find_root(g), o);
    }

    #[test]
    fn basic_depth() {
        // generate a basic lct
        let size = 3;
        let mut lct = Pool::with_capacity(size);
        for i in 0..size {
            let _ = lct.add_node(i);
        }
        for i in 0..(size - 1) {
            let parent = Label::from(i);
            let child = Label::from(i + 1);
            lct.link(parent, child);
        }

        // check initial depth
        for i in 0..size {
            let label = Label::from(i);
            assert_eq!(lct.depth(label), i);
        }

        // evert middle node and check new depths
        lct.evert(Label::from(1));
        dbg!(&lct);
        assert_eq!(lct.depth(Label::from(0)), 1);
        assert_eq!(lct.depth(Label::from(1)), 0);
        assert_eq!(lct.depth(Label::from(2)), 1);
    }

    #[test]
    fn random_scan() {
        // generate a basic lct
        let size = 3;
        let mut lct = Pool::with_capacity(size);
        for i in 0..size {
            let _ = lct.add_node(i);
        }
        for i in 0..(size - 1) {
            let parent = Label::from(i);
            let child = Label::from(i + 1);
            lct.link(parent, child);
        }

        // randomly mutate lct and check integrity
        for _ in 0..(size * size * size) {
            let label1 = Label::from(fastrand::usize(0..size));
            let label2 = Label::from(fastrand::usize(0..size));
            if label1 != label2 {
                dbg!(&lct);
                lct.evert(label1);
                dbg!(&lct);
                lct.expose(label2);
                dbg!(&lct);
                dbg!(label1);
                dbg!(label2);
                let depth = lct.depth(label2);
                assert_depth(&lct, label2);
                let index = fastrand::usize(0..depth);
                let cut = lct.index_depth(label2, index + 1);
                lct.cut(cut);
                lct.evert(label2);
                lct.link(label1, label2);
            }
            assert_clean(&lct);
        }
    }

    fn assert_depth<V: fmt::Debug>(lct: &Pool<Node<V>>, label: Label<Node<V>>) {
        let max_depth = lct.depth(label);
        for depth in 0..max_depth {
            let x = lct.index_depth(label, depth);
            assert_eq!(lct.depth(x), depth);
        }
    }

    fn assert_clean<V>(lct: &Pool<Node<V>>) {
        for label in lct.labels() {
            // check parent reciprocrity
            match lct[label].parent {
                Parent::Root => (),
                Parent::Path(_) => (),
                Parent::BinTree(p) => {
                    assert!(lct[p].children.0 == Some(label) || lct[p].children.1 == Some(label))
                }
            }

            // check child reciprocrity
            if let Some(c) = lct[label].children.0 {
                assert_eq!(lct[c].parent, Parent::BinTree(label));
            }
            if let Some(c) = lct[label].children.1 {
                assert_eq!(lct[c].parent, Parent::BinTree(label));
            }
        }
    }
}