Invert a binary tree in Rust

use std::cell::RefCell;
use std::rc::Rc;

// Definition of a binary tree node
#[derive(Debug)]
struct TreeNode {
    val: i32,
    left: Option<Rc<RefCell<TreeNode>>>,
    right: Option<Rc<RefCell<TreeNode>>>,
}

impl TreeNode {
    fn new(val: i32) -> Self {
        TreeNode {
            val,
            left: None,
            right: None,
        }
    }
}

type Node = Option<Rc<RefCell<TreeNode>>>;

/// Recursively inverts a binary tree by swapping left and right children.
fn invert_tree(root: Node) -> Node {
    if let Some(node) = root.as_ref() {
        let mut node_ref = node.borrow_mut();

        // Swap left and right subtrees
        let left = node_ref.left.take();
        let right = node_ref.right.take();
        node_ref.left = invert_tree(right);
        node_ref.right = invert_tree(left);
    }
    root
}

// ---------- helpers for testing ----------

fn new_node(val: i32, left: Node, right: Node) -> Node {
    Some(Rc::new(RefCell::new(TreeNode {
        val,
        left,
        right,
    })))
}

/// Returns an in-order traversal as a Vec for easy comparison.
fn inorder(root: &Node) -> Vec<i32> {
    let mut result = Vec::new();
    if let Some(node) = root {
        let node_ref = node.borrow();
        result.extend(inorder(&node_ref.left));
        result.push(node_ref.val);
        result.extend(inorder(&node_ref.right));
    }
    result
}

fn main() {
    //        4
    //       / \
    //      2   7
    //     / \ / \
    //    1  3 6  9
    let tree = new_node(
        4,
        new_node(2, new_node(1, None, None), new_node(3, None, None)),
        new_node(7, new_node(6, None, None), new_node(9, None, None)),
    );

    println!("Before: {:?}", inorder(&tree));  // [1, 2, 3, 4, 6, 7, 9]

    let inverted = invert_tree(tree);

    println!("After:  {:?}", inorder(&inverted)); // [9, 7, 6, 4, 3, 2, 1]
    //        4
    //       / \
    //      7   2
    //     / \ / \
    //    9  6 3  1
}

use std::collections::VecDeque;

fn invert_tree_iterative(root: Node) -> Node {
    let mut queue = VecDeque::new();
    if let Some(ref node) = root {
        queue.push_back(Rc::clone(node));
    }

    while let Some(node) = queue.pop_front() {
        let mut n = node.borrow_mut();

        // Swap children
        std::mem::swap(&mut n.left, &mut n.right);

        // Enqueue children for processing
        if let Some(ref left) = n.left {
            queue.push_back(Rc::clone(left));
        }
        if let Some(ref right) = n.right {
            queue.push_back(Rc::clone(right));
        }
    }

    root
}