Rust for C-Programmers

13.8 Creating Iterators for Complex Data Structures

Complex data structures (like trees or graphs) may need custom traversal. Rust’s iterator traits accommodate these scenarios just as well.

13.8.1 An In-Order Binary Tree Iterator

Tree Definition:

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

#[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:

#![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() {
    // Build a simple 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));

    // Traverse with InOrderIter
    let iter = InOrderIter::new(root.clone());
    let traversal: Vec<i32> = iter.collect();
    println!("{:?}", traversal); // [1, 2, 3, 4, 5, 6, 7]
}