The Visitor Pattern in Rust With Traits and Enums Compared
In languages without sum types, the visitor pattern exists to add operations to a type hierarchy. Rust's enums and match usually make it unnecessary, with a few exceptions.
The visitor pattern separates an operation from the data structure it walks, so a new operation (print, evaluate, type-check, optimise) can be added without editing the node types. In Java it takes an interface with one visit method per node class and an accept method on every node. In Rust there are two honest ways to get the same result, and the language's enums make the simpler one the default: a plain match over an enum is the visitor, with the compiler checking that every case is handled. This post builds the same small expression tree three ways (a match, a recursive trait visitor, a fold with closures), says which to pick, and shows what each costs. Everything compiled on Rust 1.99, edition 2024.
The data: a small expression tree
Throughout, the tree is an arithmetic expression with numbers, variables, addition, multiplication and negation. A closed set of node kinds is exactly what an enum is for:
use std::collections::HashMap;
#[derive(Debug, Clone)]
enum Expr {
Num(f64),
Var(String),
Add(Box<Expr>, Box<Expr>),
Mul(Box<Expr>, Box<Expr>),
Neg(Box<Expr>),
}
use Expr::*;
fn num(n: f64) -> Box<Expr> {
Box::new(Num(n))
}
fn var(name: &str) -> Box<Expr> {
Box::new(Var(name.to_string()))
}
fn eval(e: &Expr, env: &HashMap<String, f64>) -> Result<f64, String> {
Ok(match e {
Num(n) => *n,
Var(name) => *env.get(name).ok_or_else(|| format!("unbound variable {name}"))?,
Add(a, b) => eval(a, env)? + eval(b, env)?,
Mul(a, b) => eval(a, env)? * eval(b, env)?,
Neg(a) => -eval(a, env)?,
})
}
fn show(e: &Expr) -> String {
match e {
Num(n) => n.to_string(),
Var(name) => name.clone(),
Add(a, b) => format!("({} + {})", show(a), show(b)),
Mul(a, b) => format!("{} * {}", show(a), show(b)),
Neg(a) => format!("-{}", show(a)),
}
}
fn main() {
let expr = Add(Box::new(Mul(num(3.0), var("x"))), Box::new(Neg(var("y"))));
let mut env = HashMap::new();
env.insert("x".to_string(), 4.0);
env.insert("y".to_string(), 2.5);
println!("{}", show(&expr));
println!("{:?}", eval(&expr, &env));
env.remove("y");
println!("{:?}", eval(&expr, &env));
}(3 * x + -y)
Ok(9.5)
Err("unbound variable y")Each operation is one function with one match. There is no accept method, no visitor interface, and adding a third operation is adding a third function. Rust's exhaustiveness check does the job a visitor interface does in Java: add a variant to Expr and every match without a wildcard stops compiling until it handles the new case. That is the double dispatch the pattern exists to provide, done at compile time with no ceremony.
What the compiler checks when you add a variant
Add a Sub node and forget to update show:
enum Expr {
Num(f64),
Add(Box<Expr>, Box<Expr>),
Sub(Box<Expr>, Box<Expr>),
}
fn show(e: &Expr) -> String {
match e {
Expr::Num(n) => n.to_string(),
Expr::Add(a, b) => format!("({} + {})", show(a), show(b)),
}
}
fn main() {}error[E0004]: non-exhaustive patterns: `&Expr::Sub(_, _)` not covered
--> src/main.rs:8:11
|
8 | match e {
| ^ pattern `&Expr::Sub(_, _)` not covered
|
note: `Expr` defined here
--> src/main.rs:1:6
|
1 | enum Expr {
| ^^^^
...
4 | Sub(Box<Expr>, Box<Expr>),
| --- not covered
= note: the matched value is of type `&Expr`
help: ensure that all possible cases are being handled by adding a match arm with a wildcard pattern or an explicit pattern as shown
|
10 ~ Expr::Add(a, b) => format!("({} + {})", show(a), show(b)),
11 ~ &Expr::Sub(_, _) => todo!(),
|That is the property to protect. A _ => arm in an operation silences it, so avoid wildcards in visitors over your own enums, even when some cases are no-ops: write the cases out. The error family is explained in pattern matching.
When you want the trait version: open sets of operations
The enum approach makes operations easy to add and node types hard (each new variant touches every function). The classic visitor trait is the mirror image only if you also make the nodes open, which an enum cannot be. The case where a trait visitor earns its place is when the traversal order or recursion itself is shared, so each visitor implements only the interesting nodes and inherits the walk:
#[derive(Debug)]
enum Expr {
Num(f64),
Var(String),
Add(Box<Expr>, Box<Expr>),
Mul(Box<Expr>, Box<Expr>),
}
trait Visitor {
fn visit_num(&mut self, _n: f64) {}
fn visit_var(&mut self, _name: &str) {}
fn visit_expr(&mut self, e: &Expr) {
walk_expr(self, e);
}
}
fn walk_expr<V: Visitor + ?Sized>(v: &mut V, e: &Expr) {
match e {
Expr::Num(n) => v.visit_num(*n),
Expr::Var(name) => v.visit_var(name),
Expr::Add(a, b) | Expr::Mul(a, b) => {
v.visit_expr(a);
v.visit_expr(b);
}
}
}
#[derive(Default)]
struct VarCollector(Vec<String>);
impl Visitor for VarCollector {
fn visit_var(&mut self, name: &str) {
if !self.0.iter().any(|v| v == name) {
self.0.push(name.to_string());
}
}
}
struct Depth {
now: usize,
max: usize,
}
impl Visitor for Depth {
fn visit_expr(&mut self, e: &Expr) {
self.now += 1;
self.max = self.max.max(self.now);
walk_expr(self, e);
self.now -= 1;
}
}
fn main() {
use Expr::*;
let e = Add(
Box::new(Mul(Box::new(Var("x".into())), Box::new(Num(2.0)))),
Box::new(Add(Box::new(Var("y".into())), Box::new(Var("x".into())))),
);
let mut vars = VarCollector::default();
vars.visit_expr(&e);
println!("{:?}", vars.0);
let mut depth = Depth { now: 0, max: 0 };
depth.visit_expr(&e);
println!("depth {}", depth.max);
}["x", "y"]
depth 3This is the shape of syn::visit, rustc's own AST visitors and the swc and tree-sitter walkers. Each visitor overrides only the hooks it cares about, and walk_expr is the one place that knows how to descend. A visitor that collects variables never mentions addition or multiplication. The cost is a trait with one method per node kind (and one walk_* function per kind), written once for the whole library. For a compiler front end with forty node kinds, that boilerplate is cheaper than forty copies of the recursion. For an expression tree with five kinds it is not.
A third way: fold with closures
If every operation is a bottom-up computation (each node's result depends on its children's), write the traversal once as a fold and pass one function per case. This is a catamorphism, and in Rust it fits as a struct of closures:
enum Expr {
Num(f64),
Add(Box<Expr>, Box<Expr>),
Mul(Box<Expr>, Box<Expr>),
}
fn fold<T>(
e: &Expr,
num: &dyn Fn(f64) -> T,
add: &dyn Fn(T, T) -> T,
mul: &dyn Fn(T, T) -> T,
) -> T {
match e {
Expr::Num(n) => num(*n),
Expr::Add(a, b) => add(fold(a, num, add, mul), fold(b, num, add, mul)),
Expr::Mul(a, b) => mul(fold(a, num, add, mul), fold(b, num, add, mul)),
}
}
fn main() {
use Expr::*;
let e = Add(Box::new(Num(2.0)), Box::new(Mul(Box::new(Num(3.0)), Box::new(Num(4.0)))));
let value = fold(&e, &|n| n, &|a, b| a + b, &|a, b| a * b);
let text = fold(&e, &|n| n.to_string(), &|a, b| format!("({a} + {b})"), &|a, b| format!("{a} * {b}"));
let size = fold(&e, &|_| 1usize, &|a, b| a + b + 1, &|a, b| a + b + 1);
println!("{value} {text} {size}");
}14 (2 + 3 * 4) 5Three operations, three one-liners, one traversal. The limits: it handles only bottom-up computation (no context passed down, like an environment or a depth counter), and a node kind added to the enum adds a parameter to fold and to every call. It shines for small trees with many simple queries.
Rewriting the tree while walking it
A read-only visitor answers questions. The other half of the job is transformation: constant folding, renaming, desugaring. The idiomatic Rust form takes the node by mutable reference and rewrites it in place. Children are visited first, so by the time a parent looks at them they are already simplified:
#[derive(Debug, PartialEq)]
enum Expr {
Num(f64),
Var(String),
Add(Box<Expr>, Box<Expr>),
Mul(Box<Expr>, Box<Expr>),
}
fn fold_constants(e: &mut Expr) {
match e {
Expr::Add(a, b) | Expr::Mul(a, b) => {
fold_constants(a);
fold_constants(b);
}
_ => return,
}
let replacement = match e {
Expr::Add(a, b) => match (&**a, &**b) {
(Expr::Num(x), Expr::Num(y)) => Some(Expr::Num(x + y)),
(Expr::Num(z), other) | (other, Expr::Num(z)) if *z == 0.0 => Some(clone_shallow(other)),
_ => None,
},
Expr::Mul(a, b) => match (&**a, &**b) {
(Expr::Num(x), Expr::Num(y)) => Some(Expr::Num(x * y)),
(Expr::Num(z), _) | (_, Expr::Num(z)) if *z == 0.0 => Some(Expr::Num(0.0)),
_ => None,
},
_ => None,
};
if let Some(r) = replacement {
*e = r;
}
}
fn clone_shallow(e: &Expr) -> Expr {
match e {
Expr::Num(n) => Expr::Num(*n),
Expr::Var(v) => Expr::Var(v.clone()),
Expr::Add(a, b) => Expr::Add(Box::new(clone_shallow(a)), Box::new(clone_shallow(b))),
Expr::Mul(a, b) => Expr::Mul(Box::new(clone_shallow(a)), Box::new(clone_shallow(b))),
}
}
fn main() {
let mut e = Expr::Add(
Box::new(Expr::Mul(Box::new(Expr::Num(2.0)), Box::new(Expr::Num(3.0)))),
Box::new(Expr::Mul(Box::new(Expr::Var("x".into())), Box::new(Expr::Num(0.0)))),
);
fold_constants(&mut e);
println!("{e:?}");
}Num(6.0)That listing is deliberately ugly in one spot, the x + 0 arm, and the ugliness is the lesson. Moving a child out of a Box that the match still borrows is something the borrow checker rejects, so the code clones the survivor with a helper. The real fix is to take the node by value and return a new one, which sidesteps the problem entirely: fn simplify(e: Expr) -> Expr destructures with full ownership and rebuilds only what changed. For a tree that is rewritten once at compile time, that is simpler and no slower; reach for the &mut form only when the nodes are large and rewrites are local. Many Rust compilers and linters (syn's VisitMut, rustc_ast::mut_visit) use the by-reference form because their trees are huge and a typical pass changes a few nodes.
A second consideration for transformations is order. Bottom-up (children first) suits folding, because a parent sees simplified children. Top-down suits desugaring, where the parent rewrites itself into forms the children then still need to process. Choosing the wrong order gives results that are right on small tests and wrong when one rewrite exposes another, so test with nested inputs that need two passes and consider running the pass until nothing changes.
Choosing
| Situation | Use |
|---|---|
| Closed set of node kinds, a handful of operations | A match per operation (default) |
| Many node kinds, many visitors that share the recursion | A visitor trait with a walk_* function |
| Pure bottom-up queries on a small tree | A fold with closures |
| Open set of node types added by plugins | A trait on the node with Box<dyn Node> and a method per operation, see trait objects |
| Mutating the tree while walking it | A VisitMut-style trait taking &mut Expr |
Costs, and what this pattern is not for
A match over an enum compiles to a jump table, and a trait visitor to direct calls (if monomorphized) or vtable calls (if used as &mut dyn Visitor). Neither matters next to the allocation behind each Box. For a tree large enough for that to hurt, store nodes in an arena indexed by integers, as described in arena allocation. The visitor is also the wrong tool for a flat collection: iterate it. Depth-first recursion on a very deep tree can overflow the stack, which is a real failure mode for parsers fed adversarial input (see debugging stack overflows), so cap the depth when the tree comes from outside. The sibling pattern for swapping an algorithm is the strategy pattern.
