Binary tree with indexing

binary_tree.rs

// Binary tree with index, copied from https://sachanganesh.com/programming/graph-tree-traversals-in-rust/

pub type Idx = usize;

pub struct TreeNode {
    pub value: usize,
    pub left: Option<Idx>,
    pub right: Option<Idx>,
}

impl TreeNode {
    pub fn new(value: usize, left: Option<Idx>, right: Option<Idx>) -> Self {
        TreeNode { value, left, right }
    }
}
pub struct Tree {
    arena: Vec<Option<TreeNode>>,
    root: Option<Idx>,
}

impl Tree {
    pub fn new() -> Self {
        Tree {
            arena: Vec::new(),
            root: None,
        }
    }

    pub fn iter(&self) -> PreorderIter {
        PreorderIter::new(self.root)
    }

    pub fn set_root(&mut self, root: Option<Idx>) {
        self.root = root;
    }

    pub fn add_node(&mut self, node: TreeNode) -> Idx {
        let index = self.arena.len();
        self.arena.push(Some(node));
        return index;
    }

    pub fn remove_node_at(&mut self, index: Idx) -> Option<TreeNode> {
        if let Some(node) = self.arena.get_mut(index) {
            node.take()
        } else {
            None
        }
    }

    pub fn node_at(&self, index: Idx) -> Option<&TreeNode> {
        return if let Some(node) = self.arena.get(index) {
            node.as_ref()
        } else {
            None
        };
    }

    pub fn node_at_mut(&mut self, index: Idx) -> Option<&mut TreeNode> {
        return if let Some(node) = self.arena.get_mut(index) {
            node.as_mut()
        } else {
            None
        };
    }
}
pub struct PreorderIter {
    stack: Vec<Idx>,
}

impl PreorderIter {
    pub fn new(root: Option<Idx>) -> Self {
        if let Some(index) = root {
            PreorderIter { stack: vec![index] }
        } else {
            PreorderIter { stack: vec![] }
        }
    }

    pub fn next(&mut self, tree: &Tree) -> Option<Idx> {
        while let Some(node_index) = self.stack.pop() {
            if let Some(node) = tree.node_at(node_index) {
                if let Some(right) = node.right {
                    self.stack.push(right)
                }

                if let Some(left) = node.left {
                    self.stack.push(left)
                }

                return Some(node_index);
            }
        }

        return None;
    } // immutable borrow &Tree ends here
}
fn main() {
    let mut tree = Tree::new();

    let a = tree.add_node(TreeNode::new(4, None, None));
    let b = tree.add_node(TreeNode::new(5, None, None));
    let c = tree.add_node(TreeNode::new(2, Some(a), Some(b)));

    let d = tree.add_node(TreeNode::new(3, None, None));
    let e = tree.add_node(TreeNode::new(1, Some(c), Some(d)));

    tree.set_root(Some(e));

    let mut preorder = tree.iter();
    while let Some(i) = preorder.next(&tree) {
        let node = tree.node_at_mut(i).expect("node to exist at given index"); // mutable borrow starts here
        node.value *= 10;
    } // mutable borrow ends here

    let mut preorder = tree.iter();
    while let Some(i) = preorder.next(&tree) {
        let node = tree.node_at(i).expect("node to exist at given index");
        println!("{}", node.value)
    }
}