Rust for C-Programmers

13.7 Iterating Over Complex Structures: Binary Tree Example

Iterators are not limited to linear sequences like vectors or arrays. They can encapsulate the traversal logic for more complex data structures, such as trees or graphs, providing a standard Iterator interface for consuming code.

Here’s an example of implementing an in-order traversal iterator for a simple binary tree. We use Rc<RefCell<TreeNode<T>>> to handle shared ownership and potential mutation (though mutation isn’t used in this traversal itself), which is common in graph-like structures in Rust where nodes might be reachable via multiple paths.

use std::rc::Rc;
use std::cell::RefCell;
use std::collections::VecDeque; // Using VecDeque as a stack

// Node definition using shared ownership via Rc and interior mutability via RefCell
type TreeNodeLink<T> = Option<Rc<RefCell<TreeNode<T>>>>;

#[derive(Debug)]
struct TreeNode<T> {
    value: T,
    left: TreeNodeLink<T>,
    right: TreeNodeLink<T>,
}

impl<T> TreeNode<T> {
    // Helper to create a new node wrapped in Rc<RefCell<...>>
    fn new(value: T) -> Rc<RefCell<Self>> {
        Rc::new(RefCell::new(TreeNode {
            value,
            left: None,
            right: None,
        }))
    }
}

// Iterator struct for in-order traversal
struct InOrderIter<T: Clone> { // Require T: Clone to yield owned values
    // Stack holds nodes waiting to be visited (after their left subtree is done)
    stack: VecDeque<Rc<RefCell<TreeNode<T>>>>,
    // Current node pointer, used to navigate down left branches
    current: TreeNodeLink<T>,
}

impl<T: Clone> InOrderIter<T> {
    // Creates a new iterator starting traversal from the root
    fn new(root: TreeNodeLink<T>) -> Self {
        let mut iter = InOrderIter {
            stack: VecDeque::new(),
            current: root,
        };
        // Initialize by pushing the left spine onto the stack
        iter.push_left_spine();
        iter
    }

    // Helper: Pushes the current node and all its left children onto the stack.
    // Sets `self.current` to None after finishing.
    fn push_left_spine(&mut self) {
        while let Some(node) = self.current.take() { // Take ownership of current link
            self.stack.push_back(node.clone()); // Push node onto stack
            // Prepare to move left: borrow immutably to get left child link
            let left_link = node.borrow().left.clone();
            self.current = left_link; // Update current to the left child
        }
    }
}

impl<T: Clone> Iterator for InOrderIter<T> {
    type Item = T; // Yield owned copies of node values

    fn next(&mut self) -> Option<Self::Item> {
        // If current is Some, it means we just moved right from a popped node.
        // Push the new current node and its left spine onto the stack.
        if self.current.is_some() {
             self.push_left_spine();
        }

        // Pop the next node from the stack (this is the next in-order node)
        if let Some(node_to_visit) = self.stack.pop_back() {
            // Borrow node to access value and right child
            let node_ref = node_to_visit.borrow();
            let value_to_return = node_ref.value.clone(); // Clone value for return

            // Prepare for the *next* call: move to the right child.
            // The next call to `next()` will handle pushing this right child
            // and its left spine (if it exists) via `push_left_spine`.
            self.current = node_ref.right.clone();

            Some(value_to_return)
        } else {
            // Stack is empty and current is None -> Traversal complete
            None
        }
    }
}

// Add a convenience method to initiate the iteration from a root node
impl<T: Clone> TreeNode<T> {
     // Creates the in-order iterator for a tree rooted at `link`
     fn in_order_iter(link: TreeNodeLink<T>) -> InOrderIter<T> {
         InOrderIter::new(link)
     }
}

fn main() {
    // Build a simple binary search tree:
    //         4
    //        / \
    //       2   6
    //      / \ / \
    //     1  3 5  7
    let root = TreeNode::new(4);
    let node1 = TreeNode::new(1);
    let node3 = TreeNode::new(3);
    let node5 = TreeNode::new(5);
    let node7 = TreeNode::new(7);

    let node2 = TreeNode::new(2);
    node2.borrow_mut().left = Some(node1.clone());
    node2.borrow_mut().right = Some(node3.clone());

    let node6 = TreeNode::new(6);
    node6.borrow_mut().left = Some(node5.clone());
    node6.borrow_mut().right = Some(node7.clone());

    root.borrow_mut().left = Some(node2.clone());
    root.borrow_mut().right = Some(node6.clone());

    // Use the iterator and collect the results
    println!("Tree nodes (in-order traversal):");
    let traversal: Vec<i32> = TreeNode::in_order_iter(Some(root)).collect();
    println!("{:?}", traversal); // Expected: [1, 2, 3, 4, 5, 6, 7]
    assert_eq!(traversal, vec![1, 2, 3, 4, 5, 6, 7]);

    // Example of using the iterator step-by-step with a single node tree
    let root_single = TreeNode::new(10);
    let mut iter_manual = TreeNode::in_order_iter(Some(root_single));
    assert_eq!(iter_manual.next(), Some(10));
    assert_eq!(iter_manual.next(), None);
    assert_eq!(iter_manual.next(), None); // Fused behavior
}

This example demonstrates how the Iterator trait can encapsulate complex stateful traversal logic (managing a stack and current node pointer for tree traversal), exposing it through the simple, standard next() interface familiar to users of standard collection iterators. The T: Clone bound is necessary here because the iterator only has shared references (Rc<RefCell<...>>) to the nodes but needs to yield owned T values. An alternative design could yield references or require T: Copy.