Rust AVL Tree

// Завдання: написати програму, яка створює збалансоване бінарне дерево. Написати рекурсивну функцію, що підраховує суму елементів дерева.


use std::{
    cmp::{max, Ordering},
    iter::{self, FromIterator},
    mem,
};

fn main() {
    let tree: Tree = vec![5, 15, 16, 3, 6, 7, 2, 1, 4].into_iter().collect();
    tree.print();
    println!("{}", tree.product());
}

struct Tree {
    root: Option<Node>,
}

impl Tree {
    fn new() -> Self {
        Self { root: None }
    }

    fn insert(&mut self, value: i64) {
        match &mut self.root {
            Some(root) => root.insert(value),
            None => self.root = Some(Node::new(value)),
        }
    }

    fn product(&self) -> i64 {
        let root = match &self.root {
            Some(root) => root,
            None => return 1,
        };

        root.get_all_children()
            .iter()
            .chain(iter::once(&root))
            .map(|node| node.value)
            .product()
    }

    fn print(&self) {
        match &self.root {
            Some(root) => root.print(0),
            None => println!("The tree is empty"),
        }
    }
}

impl FromIterator<i64> for Tree {
    fn from_iter<T: IntoIterator<Item = i64>>(iter: T) -> Self {
        let mut tree = Tree::new();
        iter.into_iter().for_each(|x| tree.insert(x));
        tree
    }
}

enum RotationType {
    NoRotation,
    LL,
    RR,
    LR,
    RL,
}

struct Node {
    value: i64,
    left: Option<Box<Node>>,
    right: Option<Box<Node>>,
}

impl Node {
    fn new(value: i64) -> Self {
        Self {
            value,
            left: None,
            right: None,
        }
    }

    fn insert(&mut self, value: i64) {
        let branch = match value.cmp(&self.value) {
            Ordering::Less => &mut self.left,
            Ordering::Greater => &mut self.right,
            Ordering::Equal => panic!("duplicate value: {}", value),
        };

        match branch {
            Some(branch) => branch.insert(value),
            None => *branch = Some(Box::new(Node::new(value))),
        }

        self.rebalance();
    }

    fn rebalance(&mut self) {
        match self.pick_rotation_type() {
            RotationType::NoRotation => (),
            RotationType::LL => self.rotate_ll(),
            RotationType::RR => self.rotate_rr(),
            RotationType::LR => self.rotate_lr(),
            RotationType::RL => self.rotate_rl(),
        }
    }

    fn pick_rotation_type(&self) -> RotationType {
        match self.balance_factor() {
            -1 | 0 | 1 => RotationType::NoRotation,
            2 => match self.left.as_ref().unwrap().balance_factor() {
                0 | 1 | 2 => RotationType::LL,
                -2 | -1 => RotationType::LR,
                bf => panic!("unexpected balance factor: {}", bf),
            },
            -2 => match self.right.as_ref().unwrap().balance_factor() {
                -2 | -1 | 0 => RotationType::RR,
                1 | 2 => RotationType::RL,
                bf => panic!("unexpected balance factor: {}", bf),
            },
            bf => panic!("unexpected balance factor: {}", bf),
        }
    }

    fn rotate_ll(&mut self) {
        let parent = self;

        let mut left = parent.left.take().unwrap();
        mem::swap(parent, &mut left);
        match &mut parent.right {
            Some(right) => {
                mem::swap(right, &mut left);
                right.left = Some(left);
            }
            None => parent.right = Some(left),
        }
    }

    fn rotate_rr(&mut self) {
        let parent = self;

        let mut right = parent.right.take().unwrap();
        mem::swap(parent, &mut right);
        match &mut parent.left {
            Some(left) => {
                mem::swap(left, &mut right);
                left.right = Some(right);
            }
            None => parent.left = Some(right),
        }
    }

    fn rotate_lr(&mut self) {
        self.left.as_mut().unwrap().rotate_rr();
        self.rotate_ll();
    }

    fn rotate_rl(&mut self) {
        self.right.as_mut().unwrap().rotate_ll();
        self.rotate_rr();
    }

    fn get_all_children(&self) -> Vec<&Node> {
        let mut children = Vec::with_capacity(2);
        if let Some(left) = self.left.as_deref() {
            children.push(left);
            children.extend(left.get_all_children());
        }

        if let Some(right) = self.right.as_deref() {
            children.push(right);
            children.extend(right.get_all_children());
        }

        children
    }

    fn balance_factor(&self) -> i8 {
        (self.left_branch_height() as i64 - self.right_branch_height() as i64) as i8
    }

    fn height(&self) -> usize {
        1 + max(self.left_branch_height(), self.right_branch_height())
    }

    fn left_branch_height(&self) -> usize {
        self.left.as_ref().map_or(0, |node| node.height())
    }

    fn right_branch_height(&self) -> usize {
        self.right.as_ref().map_or(0, |node| node.height())
    }

    fn print(&self, indentation: usize) {
        println!("{} {}", "    ".repeat(indentation), self.value,);
        match (&self.left, &self.right) {
            (Some(left), Some(right)) => {
                left.print(indentation + 1);
                right.print(indentation + 1);
            }
            (Some(left), None) => {
                left.print(indentation + 1);
                println!("{} -", "    ".repeat(indentation + 1));
            }
            (None, Some(right)) => {
                println!("{} -", "    ".repeat(indentation + 1));
                right.print(indentation + 1);
            }
            (None, None) => {}
        }
    }
}