如何让Rust二叉树遍历函数同时支持&和&mut引用类型？

问题描述 #

我正在实现一个二叉树数据结构，编写了如下的前序遍历高阶函数pre_traverse：

fn pre_traverse<F>(&mut self, mut visit: F)
where
    F: FnMut(&mut T),
{
    let mut stack = vec![self];
    while let Some(tree) = stack.pop() {
        if let Some(node) = &mut tree.root {
            visit(&mut node.data);
            stack.push(&mut node.right);
            stack.push(&mut node.left);
        }
    }
}

该函数要求self为&mut引用，visit闭包接收&mut T类型参数。但当visit仅需只读访问节点（如统计节点数量）时，可变引用的要求并不必要。希望通过指定RefType参数，让函数在调用时可选择使用&或&mut引用类型，而非编写两个重复的函数，请问是否有可行的实现方案？

解决方案 #

在Rust中，我们可以通过宏抽象重复逻辑或者Trait抽象访问行为来实现需求，以下是两种最实用的方案：

方案1：使用宏生成复用函数 #

由于不可变和可变遍历的逻辑几乎完全一致，仅引用类型不同，用宏可以快速生成两个版本的函数，避免代码重复：

struct Node<T> {
    data: T,
    left: Option<Box<Node<T>>>,
    right: Option<Box<Node<T>>>,
}

struct BinaryTree<T> {
    root: Option<Box<Node<T>>>,
}

impl<T> BinaryTree<T> {
    // 定义宏，用于生成不同引用类型的遍历函数
    macro_rules! impl_pre_traverse {
        ($func_name:ident, $self_mut:tt, $node_ref:ty, $data_ref:ty) => {
            pub fn $func_name<F>(&$self_mut self, mut visit: F)
            where
                F: FnMut($data_ref),
            {
                let mut stack = vec![&$self_mut self.root];
                while let Some(node_opt) = stack.pop() {
                    if let Some(node) = node_opt {
                        visit($data_ref node.data);
                        stack.push(&$self_mut node.right);
                        stack.push(&$self_mut node.left);
                    }
                }
            }
        };
    }

    // 生成不可变引用版本的前序遍历
    impl_pre_traverse!(pre_traverse_ref, , &, &T);
    // 生成可变引用版本的前序遍历
    impl_pre_traverse!(pre_traverse_mut, mut, &mut, &mut T);
}

使用时可以根据需求选择调用：

let mut tree = BinaryTree { /* 初始化树 */ };
// 只读访问：统计节点数量
let mut count = 0;
tree.pre_traverse_ref(|_| count += 1);

// 可变访问：修改节点数据
tree.pre_traverse_mut(|data| *data += 1);

方案2：使用Trait抽象访问行为 #

如果希望用更优雅的泛型方式实现，可以定义Trait来抽象节点的访问逻辑，让通用遍历函数适配两种引用类型：

use std::ops::{Deref, DerefMut};

struct Node<T> {
    data: T,
    left: Option<Box<Node<T>>>,
    right: Option<Box<Node<T>>>,
}

struct BinaryTree<T> {
    root: Option<Box<Node<T>>>,
}

// 定义Trait，抽象节点的访问能力
trait NodeAccess<T> {
    type DataRef;
    type NodeRef;

    // 获取节点数据的引用
    fn data(&self) -> Self::DataRef;
    // 获取左子节点的引用
    fn left(&self) -> Option<Self::NodeRef>;
    // 获取右子节点的引用
    fn right(&self) -> Option<Self::NodeRef>;
}

// 为不可变节点引用实现Trait
impl<T> NodeAccess<T> for &Node<T> {
    type DataRef = &T;
    type NodeRef = &Node<T>;

    fn data(&self) -> Self::DataRef {
        &self.data
    }

    fn left(&self) -> Option<Self::NodeRef> {
        self.left.as_ref().map(|boxed| boxed.deref())
    }

    fn right(&self) -> Option<Self::NodeRef> {
        self.right.as_ref().map(|boxed| boxed.deref())
    }
}

// 为可变节点引用实现Trait
impl<T> NodeAccess<T> for &mut Node<T> {
    type DataRef = &mut T;
    type NodeRef = &mut Node<T>;

    fn data(&self) -> Self::DataRef {
        &mut self.data
    }

    fn left(&self) -> Option<Self::NodeRef> {
        self.left.as_mut().map(|boxed| boxed.deref_mut())
    }

    fn right(&self) -> Option<Self::NodeRef> {
        self.right.as_mut().map(|boxed| boxed.deref_mut())
    }
}

impl<T> BinaryTree<T> {
    // 通用前序遍历函数，适配任意实现NodeAccess的引用类型
    fn pre_traverse<F, N>(&self, mut visit: F)
    where
        F: FnMut(N::DataRef),
        for<'a> &'a Node<T>: Into<Option<N>>,
        N: NodeAccess<T>,
    {
        let mut stack = Vec::new();
        let mut current = self.root.as_ref().map(|n| n.deref());
        
        while let Some(node) = current {
            let node: N = node.into().unwrap();
            visit(node.data());
            stack.push(node.right());
            current = node.left();
        }

        while let Some(Some(node)) = stack.pop() {
            visit(node.data());
            stack.push(node.right());
            let mut current = node.left();
            while let Some(node) = current {
                visit(node.data());
                stack.push(node.right());
                current = node.left();
            }
        }
    }
}

这种方式更符合Rust的泛型设计思想，但实现相对复杂，适合需要高度抽象的场景。

总结 #

如果只是避免代码重复，宏方案是最简洁高效的选择；如果需要更灵活的泛型抽象，Trait方案可以满足更复杂的扩展需求。两种方案都不需要编写重复的遍历逻辑，同时支持只读和可变两种访问场景。

内容的提问来源于stack exchange，提问作者Fei Yu