radn-rs/src/mrds/trees/avl.rs

use std::{fmt::Display, marker::PhantomData, rc::Rc};

use crate::flow::{
    binary::{avl::*, *},
    comparator::*,
};
use crate::func::{context::*, *};

#[derive(Clone)]
struct AvlN<A> {
    l: AvlT<A>,
    r: AvlT<A>,
    key: A,
}

impl<A: Display> Display for AvlN<A> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "({} {} {})", self.l, self.key, self.r)
    }
}

#[cfg(test)]
impl<A> AvlN<A> {
    fn balanced(&self) -> bool {
        let (hl, hr) = (self.l.height, self.r.height);
        std::cmp::max(hl, hr) - std::cmp::min(hl, hr) < 2 && self.l.balanced() && self.r.balanced()
    }
}

struct AvlR<A> {
    node: Rc<AvlN<A>>,
    height: u64,
}

impl<A> Clone for AvlR<A> {
    fn clone(&self) -> Self {
        Self {
            node: self.node.clone(),
            height: self.height,
        }
    }
}

impl<A: Display> Display for AvlR<A> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.node)
    }
}

struct AvlT<A> {
    reference: Option<Rc<AvlN<A>>>,
    height: u64,
}

impl<A> Clone for AvlT<A> {
    fn clone(&self) -> Self {
        Self {
            reference: self.reference.clone(),
            height: self.height,
        }
    }
}

impl<A: Display> Display for AvlT<A> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match &self.reference {
            Some(reference) => write!(f, "{}", reference),
            None => write!(f, "-"),
        }
    }
}

#[cfg(test)]
impl<A> AvlT<A> {
    fn balanced(&self) -> bool {
        match &self.reference {
            Some(reference) => reference.balanced(),
            None => true,
        }
    }
}

struct AvlTs<A> {
    _a: PhantomData<A>,
}

impl<A> AvlTs<A> {
    fn new() -> Self {
        Self { _a: PhantomData }
    }
}

impl<A> Clone for AvlTs<A> {
    fn clone(&self) -> Self {
        Self::new()
    }
}

impl<'a, A: 'a> FunctorContext<'a> for AvlTs<A> {
    type T = instances::solo::SoloInstance;
}

impl<'a, A: 'a + PartialOrd + Clone> BinaryTrees<'a> for AvlTs<A> {
    type Node = AvlN<A>;

    type Reference = AvlR<A>;

    type Tree = AvlT<A>;

    type Key = A;

    type Comparator = DefaultComparator;

    type _Tm = Self::T;

    fn comparator(&self) -> &Self::Comparator {
        &DefaultComparator
    }

    fn split(&self, node: &Self::Node) -> Split<'a, Self> {
        (node.l.clone(), node.r.clone(), node.key.clone())
    }

    fn resolve(&self, reference: &Self::Reference) -> BTWrap<'a, Self, Self::Node> {
        reference.node.as_ref().clone()
    }

    fn equal(&self, _rhs: &Self::Reference, _lhs: &Self::Reference) -> bool {
        false
    }

    fn refer(&self, tree: &Self::Tree) -> Option<Self::Reference> {
        Some(AvlR {
            node: tree.reference.clone()?,
            height: tree.height,
        })
    }
}

impl<'a, A: 'a + PartialOrd + Clone> BinaryTreesTreeOf<'a> for AvlTs<A> {
    fn tree_of(&self, node: Self::Node) -> BTWrap<'a, Self, Self::Tree> {
        AvlT {
            height: std::cmp::max(node.l.height, node.r.height)
                .checked_add(1)
                .unwrap(),
            reference: Some(node.into()),
        }
    }
}

impl<'a, A: 'a + PartialOrd + Clone> BinaryTreesEmpty<'a> for AvlTs<A> {
    fn empty(&self) -> Self::Tree {
        AvlT {
            reference: None,
            height: 0,
        }
    }

    fn split_key_empty(
        &self,
        _tree: Self::Tree,
        _key: Self::Key,
    ) -> BTWrap<'a, Self, KeySplit<'a, Self>> {
        (self.empty(), self.empty())
    }
}

impl<'a, A: 'a + PartialOrd + Clone> BinaryTreesMutable<'a> for AvlTs<A> {
    fn join_key(
        self,
        tl: Self::Tree,
        key: Self::Key,
        tr: Self::Tree,
    ) -> BTWrap<'a, Self, Self::Node> {
        self.join_key_balanced(tl, key, tr)
    }
}

impl<'a, A: 'a + PartialOrd + Clone> BinaryTreesHeight<'a> for AvlTs<A> {
    fn height(&self, tree: &Self::Tree) -> u64 {
        tree.height
    }

    fn height_r(&self, reference: &Self::Reference) -> u64 {
        reference.height
    }

    fn leaf_height_error<T: 'a>(&self, height: u64) -> BTWrap<'a, Self, T> {
        panic!("leaf height error: {height}.")
    }
}

impl<'a, A: 'a + PartialOrd + Clone> BinaryTreesTryJoin<'a> for AvlTs<A> {
    fn try_join(
        &self,
        tl: Self::Tree,
        key: Self::Key,
        tr: Self::Tree,
    ) -> BTWrap<'a, Self, Self::Node> {
        AvlN { l: tl, r: tr, key }
    }
}

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

    #[test]
    fn test() {
        let trees = AvlTs::new();
        let mut tree = trees.empty();
        for i in [
            8, 3, 10, 17, 0, 13, 6, 1, 11, 5, 4, 7, 18, 14, 15, 9, 2, 19, 16, 12,
        ] {
            tree = trees.clone().add_tree(tree, i);
            assert!(tree.balanced());
            // println!("{} {}", tree.height, tree);
        }
        for i in [
            2, 9, 4, 7, 8, 10, 17, 1, 13, 15, 18, 12, 5, 0, 3, 6, 16, 19, 14, 11,
        ] {
            tree = trees.clone().remove(tree, i);
            assert!(tree.balanced());
            // println!("{} {}", tree.height, tree);
        }
        // assert!(false);
    }
}