Rust for C-Programmers

13.8 Creating Iterators for Complex Data Structures

13.8.1 Implementing an Iterator for a Binary Tree

Definition of the Binary Tree:

#![allow(unused)]
fn main() {
use std::rc::Rc;
use std::cell::RefCell;

#[derive(Debug)]
struct TreeNode {
    value: i32,
    left: Option<Rc<RefCell<TreeNode>>>,
    right: Option<Rc<RefCell<TreeNode>>>,
}
impl TreeNode {
    fn new(value: i32) -> Rc<RefCell<Self>> {
        Rc::new(RefCell::new(TreeNode { value, left: None, right: None }))
    }
}
}

In-Order Iterator Implementation:

#![allow(unused)]
fn main() {
struct InOrderIter {
    stack: Vec<Rc<RefCell<TreeNode>>>,
    current: Option<Rc<RefCell<TreeNode>>>,
}
impl InOrderIter {
    fn new(root: Rc<RefCell<TreeNode>>) -> Self {
        InOrderIter { stack: Vec::new(), current: Some(root) }
    }
}
impl Iterator for InOrderIter {
    type Item = i32;
    fn next(&mut self) -> Option<Self::Item> {
        while let Some(node) = self.current.clone() {
            self.stack.push(node.clone());
            self.current = node.borrow().left.clone();
        }
        if let Some(node) = self.stack.pop() {
            let value = node.borrow().value;
            self.current = node.borrow().right.clone();
            Some(value)
        } else {
            None
        }
    }
}
}

Using the Iterator:

use std::rc::Rc;
use std::cell::RefCell;
#[derive(Debug)]
struct TreeNode {
    value: i32,
    left: Option<Rc<RefCell<TreeNode>>>,
    right: Option<Rc<RefCell<TreeNode>>>,
}
impl TreeNode {
    fn new(value: i32) -> Rc<RefCell<Self>> {
        Rc::new(RefCell::new(TreeNode { value, left: None, right: None }))
    }
}
struct InOrderIter {
    stack: Vec<Rc<RefCell<TreeNode>>>,
    current: Option<Rc<RefCell<TreeNode>>>,
}
impl InOrderIter {
    fn new(root: Rc<RefCell<TreeNode>>) -> Self {
        InOrderIter { stack: Vec::new(), current: Some(root) }
    }
}
impl Iterator for InOrderIter {
    type Item = i32;
    fn next(&mut self) -> Option<Self::Item> {
        while let Some(node) = self.current.clone() {
            self.stack.push(node.clone());
            self.current = node.borrow().left.clone();
        }
        if let Some(node) = self.stack.pop() {
            let value = node.borrow().value;
            self.current = node.borrow().right.clone();
            Some(value)
        } else {
            None
        }
    }
}
fn main() {
    // Building the binary tree
    let root = TreeNode::new(4);
    let left = TreeNode::new(2);
    let right = TreeNode::new(6);
    root.borrow_mut().left = Some(left.clone());
    root.borrow_mut().right = Some(right.clone());
    left.borrow_mut().left = Some(TreeNode::new(1));
    left.borrow_mut().right = Some(TreeNode::new(3));
    right.borrow_mut().left = Some(TreeNode::new(5));
    right.borrow_mut().right = Some(TreeNode::new(7));
    // Creating the iterator
    let iter = InOrderIter::new(root.clone());
    // Collecting the in-order traversal
    let traversal: Vec<i32> = iter.collect();
    println!("{:?}", traversal); // Output: [1, 2, 3, 4, 5, 6, 7]
}