2 + 3 * 4 into a single 14 node -- as the worked example;Expr/Stmt tree from episode 132 fresh in mind -- today we redesign that tree and learn to walk it;Learn Zig Series):Last episode we wrote a recursive-descent parser that turned a flat token stream into a tree. It worked -- 2 + 3 * 4 came out with the multiplication correctly nested under the addition, and malformed input became a typed error.UnexpectedToken in stead of a crash. But I was honest at the time that the tree itself was deliberately bare: a tagged union with just enough shape to prove the parser did its job, and one ad-hoc printExpr hacked together so we could eyeball the result. I called the proper design of that tree, and the patterns for walking it, "the very next brick in this arc." This is that brick.
Here is the thing that bare tree cannot do. The moment you want to report a type error, you need to say where -- "line 4, column 12" -- and our nodes do not remember where they came from. The moment you want to analyse a program (count its nodes, collect every variable it mentions, check that no name is used before it is defined), you need a disciplined way to visit every node exactly once, and copy-pasting a switch into every pass is how bugs breed. And the moment you want to transform a program -- fold constants, desugar syntax, optimise -- you need a pass that reads one tree and produces another. Design, traversal, transformation. Those three verbs are what an AST is for, and they are what today is about. Let's dive right in!
As always, the three exercises from last episode first, with complete code. All three build on the Parser, Expr union and tokenize helper exactly as episode 132 left them.
Exercise 1 -- Modulo and a right-associative power operator. The task was to add a ^ operator that binds tighter than * and is right-associative, so 2 ^ 3 ^ 2 becomes 2 ^ (3 ^ 2). Three changes. Add a .caret token kind and a lexer case for ^ (a one-line addition to the operator switch, omitted here for brevity). Give it precedence 7, above *///% at 6. And -- the real trick -- teach parseExpr to recurse with prec in stead of prec + 1 for right-associative operators, because demanding the right-hand side be built from equally tight operators is exactly what makes it nest to the right:
fn precedence(kind: TokenKind) u8 {
return switch (kind) {
.kw_or => 1,
.kw_and => 2,
.eq, .neq => 3,
.lt, .lte, .gt, .gte => 4,
.plus, .minus => 5,
.star, .slash, .percent => 6,
.caret => 7,
else => 0,
};
}
fn isRightAssoc(kind: TokenKind) bool {
return kind == .caret;
}
fn parseExpr(self: *Parser, min_prec: u8) ParseError!*Expr {
var lhs = try self.parseUnary();
while (true) {
const op = self.peek().kind;
const prec = precedence(op);
if (prec == 0 or prec < min_prec) break;
_ = self.advance();
const next_min = if (isRightAssoc(op)) prec else prec + 1;
const rhs = try self.parseExpr(next_min);
lhs = try self.makeExpr(.{ .binary = .{ .op = op, .lhs = lhs, .rhs = rhs } });
}
return lhs;
}
The whole difference between left- and right-associativity is that next_min -- prec keeps the door open for another operator of the same rank on the right (so it folds rightward), prec + 1 slams it shut (so it folds leftward). One line of policy. The test pins the nesting down: the top node is a caret, and its right child is another caret, which is only true if the parse nested to the right:
test "power is right-associative: 2 ^ 3 ^ 2 => 2 ^ (3 ^ 2)" {
var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
defer arena.deinit();
const a = arena.allocator();
const tokens = try tokenize(a, "2 ^ 3 ^ 2;");
var p = Parser{ .tokens = tokens, .arena = a };
const e = try p.parseExpr(1);
try std.testing.expectEqual(TokenKind.caret, e.binary.op);
try std.testing.expect(e.binary.lhs.* == .int); // left child is the bare 2
try std.testing.expect(e.binary.rhs.* == .binary); // right child nests
try std.testing.expectEqual(TokenKind.caret, e.binary.rhs.binary.op);
}
Exercise 2 -- Better error messages. error.UnexpectedToken says that something went wrong but not what. The task was to record what the parser expected and what it actually saw, then print a real diagnostic. Two nullable fields on the Parser, filled in the one place that raises the error -- expect -- plus a small reporter that leans on the line and column the lexer already stamped onto every token:
// added to the Parser struct's fields:
// expected: ?TokenKind = null,
// found: ?Token = null,
fn expect(self: *Parser, kind: TokenKind) ParseError!Token {
if (self.check(kind)) return self.advance();
self.expected = kind;
self.found = self.peek();
return error.UnexpectedToken;
}
fn reportError(self: *Parser) void {
const f = self.found.?;
std.debug.print("expected '{s}' but found '{s}' at line {d}, column {d}\n", .{
@tagName(self.expected.?),
@tagName(f.kind),
f.line,
f.col,
});
}
Because expect is the single choke point every required token passes through, we do not have to sprinkle bookkeeping across twenty parsing functions -- we instrument the one function they all call. Feeding let z = 1 (no semicolon), the failing expect(.semicolon) records .semicolon as expected and the eof token as found, and the reported position points at exactly where the ; should have been:
test "expect records what it wanted and what it saw" {
var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
defer arena.deinit();
const a = arena.allocator();
const tokens = try tokenize(a, "let z = 1");
var p = Parser{ .tokens = tokens, .arena = a };
try std.testing.expectError(error.UnexpectedToken, p.parseProgram());
try std.testing.expectEqual(TokenKind.semicolon, p.expected.?);
try std.testing.expectEqual(TokenKind.eof, p.found.?.kind);
}
Exercise 3 -- A parenthesis-balance pre-check. The task was to scan the token slice once, before parsing, and verify every (, ), {, } is balanced -- catching the pathological "ten thousand open parens" case cheaply and pointing at the first unmatched bracket. A stack of open-bracket tokens does it: push on every opener, and on every closer make sure the top matches. Whatever is left on the stack at the end is unclosed. We hand the offending token back through an out-parameter so the caller can print its position:
const BalanceError = error{ UnmatchedOpen, UnmatchedClose };
fn checkBalance(
alloc: std.mem.Allocator,
tokens: []const Token,
offending: *?Token,
) (BalanceError || std.mem.Allocator.Error)!void {
var stack: std.ArrayList(Token) = .empty;
defer stack.deinit(alloc);
for (tokens) |t| {
switch (t.kind) {
.lparen, .lbrace => try stack.append(alloc, t),
.rparen, .rbrace => {
const want: TokenKind = if (t.kind == .rparen) .lparen else .lbrace;
if (stack.items.len == 0 or stack.items[stack.items.len - 1].kind != want) {
offending.* = t;
return error.UnmatchedClose;
}
_ = stack.pop();
},
else => {},
}
}
if (stack.items.len != 0) {
offending.* = stack.items[0];
return error.UnmatchedOpen;
}
}
Feeding foo((1 + 2); -- two openers, one closer -- leaves one ( stranded on the stack, so the check reports UnmatchedOpen and hands back the first unclosed paren, all before the parser is even started:
test "balance check flags an unmatched open bracket" {
var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
defer arena.deinit();
const a = arena.allocator();
const tokens = try tokenize(a, "foo((1 + 2);");
var offending: ?Token = null;
try std.testing.expectError(error.UnmatchedOpen, checkBalance(a, tokens, &offending));
try std.testing.expect(offending != null);
try std.testing.expectEqual(TokenKind.lparen, offending.?.kind);
}
That is the kind of cheap, fail-fast guard production parsers put in front of untrusted input. On to today's real subject.
The bare Expr from last episode was a tagged union and nothing more. The single most important thing we can add is a source span -- the line and column the node was parsed from -- because every diagnostic any later pass ever prints will want it. The clean way to do that in Zig is to stop making Expr a raw union and make it a struct: a span plus a kind union. The variants live inside the struct as a nested type:
const Span = struct { line: usize, col: usize };
const Expr = struct {
span: Span,
kind: Kind,
const Kind = union(enum) {
int: i64,
float: f64,
boolean: bool,
ident: []const u8,
unary: struct { op: TokenKind, rhs: *Expr },
binary: struct { op: TokenKind, lhs: *Expr, rhs: *Expr },
call: struct { callee: []const u8, args: []const *Expr },
};
};
The shape is the same tree it always was -- recursive variants still box their children as *Expr on the arena, because a union still cannot contain itself by value. What changed is that every node now carries span, uniformly, whatever its kind. That uniformity is the point: a traversal that wants to report an error on any node can read node.span without caring what kind of node it is. The variant payloads move one level down, so where last episode we wrote e.binary.op, we now write e.kind.binary.op -- a small tax for a big gain in what the tree can tell us. A single constructor stamps the span so the parser never forgets to set it:
fn makeExpr(arena: std.mem.Allocator, span: Span, kind: Expr.Kind) !*Expr {
const node = try arena.create(Expr);
node.* = .{ .span = span, .kind = kind };
return node;
}
Every place the parser used to call its old makeExpr now passes the current token's span too -- and since parsePrimary already holds the token it is consuming, the span is right there for free. Statements would get the same treatment (a Span field beside their kind), and I will leave that mechanical change to you; the interesting machinery is all in how we now walk this tree.
Almost every question you can ask about a program has the same skeleton: visit every node, do something small at each. Count the nodes. Collect the identifiers. Find the deepest nesting. Check no return sits outside a function. Writing the recursive switch from scratch for each of those is how you end up with five subtly different traversals, four of which forget the call arguments. The fix is to write the traversal once, generically, and let each pass supply only the "do something small" part. That is the visitor pattern.
In Zig we do it with anytype and comptime duck typing (episode 14): walkExpr takes any visitor that has a visit method, calls it on the current node, then recurses into the children. The visitor decides what visit means:
fn walkExpr(e: *const Expr, visitor: anytype) void {
visitor.visit(e);
switch (e.kind) {
.unary => |u| walkExpr(u.rhs, visitor),
.binary => |b| {
walkExpr(b.lhs, visitor);
walkExpr(b.rhs, visitor);
},
.call => |c| for (c.args) |arg| walkExpr(arg, visitor),
.int, .float, .boolean, .ident => {}, // leaves: nothing to descend into
}
}
Notice the leaves are listed explicitly rather than swept under an else. That is deliberate: when a future episode adds a new expression variant, this switch stops compiling until you decide whether the walker should descend into it. Zig's exhaustive switches turn "I forgot to handle the new node" from a silent runtime bug into a compile error, which for tree-walking code is precisely the safety net you want. Now a concrete pass is just a struct with a visit method and whatever state it accumulates. Counting every node in a tree:
const NodeCounter = struct {
count: usize = 0,
fn visit(self: *NodeCounter, e: *const Expr) void {
_ = e;
self.count += 1;
}
};
And a pass that gathers every variable name mentioned anywhere in an expression -- the exact groundwork a later "is this name defined?" check will stand on -- differs only in what visit does. It ignores every node except identifiers, and appends those:
const IdentCollector = struct {
alloc: std.mem.Allocator,
names: std.ArrayList([]const u8),
fn visit(self: *IdentCollector, e: *const Expr) void {
switch (e.kind) {
.ident => |name| self.names.append(self.alloc, name) catch {},
else => {},
}
}
};
Two completely different analyses, zero duplicated traversal logic. That is the whole appeal: the shape of the tree is encoded in walkExpr once, and every pass you ever write borrows it. Add a node kind later, fix walkExpr once, and all your passes keep working.
The generic visitor is perfect for "do something at each node independently." It falls short the instant a pass needs to know whether it is entering or leaving a node -- and pretty-printing with indentation is the canonical example, because the indent has to grow on the way down and shrink on the way back up. A single visit(node) call has no "before children / after children" structure to hang that on. So for those passes we drop back to an explicit recursion that threads a depth and brackets its children:
fn printTree(e: *const Expr, depth: usize) void {
var i: usize = 0;
while (i < depth) : (i += 1) std.debug.print(" ", .{});
switch (e.kind) {
.int => |v| std.debug.print("int {d}\n", .{v}),
.float => |v| std.debug.print("float {d}\n", .{v}),
.boolean => |b| std.debug.print("bool {}\n", .{b}),
.ident => |s| std.debug.print("ident {s}\n", .{s}),
.unary => |u| {
std.debug.print("unary {s}\n", .{@tagName(u.op)});
printTree(u.rhs, depth + 1);
},
.binary => |b| {
std.debug.print("binary {s}\n", .{@tagName(b.op)});
printTree(b.lhs, depth + 1);
printTree(b.rhs, depth + 1);
},
.call => |c| {
std.debug.print("call {s}\n", .{c.callee});
for (c.args) |arg| printTree(arg, depth + 1);
},
}
}
Run it on 2 + 3 * 4 and you get an indented outline -- binary plus at depth 0, its int 2 and the nested binary star at depth 1, and the star's int 3 and int 4 at depth 2. The rule of thumb is worth stating plainly: reach for the generic visitor when each node is handled in isolation, and reach for explicit recursion when the structure of the descent matters (indentation, scopes, "am I inside a loop right now"). Real compilers use both, and knowing which to reach for is half the craft.
Reading a tree is one thing; rewriting it is the next. A transformation pass takes an AST and produces a new AST -- desugaring, optimising, normalising. The friendliest first example is constant folding: any sub-expression made entirely of constants can be computed at compile time and replaced by its result. 2 + 3 * 4 has no variables in it at all, so it can collapse to a single int 14 node before we ever run the program.
The pass recurses children-first (fold the operands, then look at the operator), and only when both folded operands are integer literals does it actually compute. Otherwise it rebuilds the node with its folded children and moves on. Every new node comes from the arena through the same makeExpr, so the transformed tree has the same clean single-deinit lifetime as the original:
fn fold(arena: std.mem.Allocator, e: *Expr) !*Expr {
switch (e.kind) {
.binary => |b| {
const l = try fold(arena, b.lhs);
const r = try fold(arena, b.rhs);
if (l.kind == .int and r.kind == .int) {
const lv = l.kind.int;
const rv = r.kind.int;
const folded: ?i64 = switch (b.op) {
.plus => lv + rv,
.minus => lv - rv,
.star => lv * rv,
.slash => if (rv == 0) null else @divTrunc(lv, rv),
else => null,
};
if (folded) |v| return makeExpr(arena, e.span, .{ .int = v });
}
return makeExpr(arena, e.span, .{ .binary = .{ .op = b.op, .lhs = l, .rhs = r } });
},
.unary => |u| {
const rhs = try fold(arena, u.rhs);
if (u.op == .minus and rhs.kind == .int)
return makeExpr(arena, e.span, .{ .int = -rhs.kind.int });
return makeExpr(arena, e.span, .{ .unary = .{ .op = u.op, .rhs = rhs } });
},
else => return e, // leaves fold to themselves
}
}
Two design decisions worth calling out. First, the folded value is an optional ?i64: null means "cannot fold this" (a division by zero, or an operator like a comparison that does not produce an integer), and we fall through to rebuilding the node rather than inventing a wrong answer -- Zig's optionals modelling "maybe no result" exactly, again. Second, notice we carry e.span onto every rebuilt node, so a folded constant still remembers the source position it came from; a later error message about that 14 can still point at the original 2 + 3 * 4. In a production compiler you would also guard against integer overflow here (folding huge + huge must not itself panic), but for our small language the shape is what matters. The pass is honest: it only ever replaces an expression with one that computes the identical value.
A traversal you cannot test is a traversal you cannot trust, so we pin both down. We build 2 + 3 * 4 by hand with makeExpr (five nodes: the three integer leaves, the star, and the plus), then assert the counter sees exactly five, and that folding collapses the whole thing to a single int 14:
test "walker visits every node exactly once" {
var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
defer arena.deinit();
const a = arena.allocator();
const span = Span{ .line = 1, .col = 1 };
const n2 = try makeExpr(a, span, .{ .int = 2 });
const n3 = try makeExpr(a, span, .{ .int = 3 });
const n4 = try makeExpr(a, span, .{ .int = 4 });
const mul = try makeExpr(a, span, .{ .binary = .{ .op = .star, .lhs = n3, .rhs = n4 } });
const add = try makeExpr(a, span, .{ .binary = .{ .op = .plus, .lhs = n2, .rhs = mul } });
var counter = NodeCounter{};
walkExpr(add, &counter);
try std.testing.expectEqual(@as(usize, 5), counter.count);
const folded = try fold(a, add);
try std.testing.expect(folded.kind == .int);
try std.testing.expectEqual(@as(i64, 14), folded.kind.int);
}
Both halves of today's machinery, verified against one tiny tree: the visitor touches every node once, and the transform preserves the value while shrinking the tree. Building the node by hand in the test, in stead of routing through the parser, keeps the test focused on exactly the thing under test -- traversal, not parsing.
Everything so far boxes each node on the arena and points at its children with *Expr. That is the natural, readable representation, and for a hobby language it is the right call. But there is a second design that every serious compiler eventually reaches for, and it is worth understanding now because the choice shapes everything downstream. In stead of pointers, store all nodes in one flat array and refer to children by their index into that array:
const NodeIndex = u32;
const FlatNode = union(enum) {
int: i64,
ident: []const u8,
unary: struct { op: TokenKind, rhs: NodeIndex },
binary: struct { op: TokenKind, lhs: NodeIndex, rhs: NodeIndex },
};
const FlatAst = struct {
nodes: std.ArrayList(FlatNode),
fn add(self: *FlatAst, alloc: std.mem.Allocator, node: FlatNode) !NodeIndex {
const idx: NodeIndex = @intCast(self.nodes.items.len);
try self.nodes.append(alloc, node);
return idx;
}
};
You build it bottom-up: add the leaves, get their indices back, then add the parent referring to those indices. 1 + 2 becomes three add calls, and the root is simply the last index handed out:
var ast = FlatAst{ .nodes = .empty };
defer ast.nodes.deinit(alloc);
const one = try ast.add(alloc, .{ .int = 1 });
const two = try ast.add(alloc, .{ .int = 2 });
const root = try ast.add(alloc, .{ .binary = .{ .op = .plus, .lhs = one, .rhs = two } });
Why bother? Three real wins. Memory: a u32 index is half the size of a 64-bit pointer, so nodes are smaller and you fit more of the tree in cache -- and for a compiler that walks the tree many times, cache locality is the whole ballgame. Locality: all nodes live in one contiguous block, so walking them is a linear scan over memory the prefetcher loves, in stead of chasing pointers all over the heap. Serialization: indices survive being written to disk and read back; raw pointers do not, so an index-based AST is trivially cacheable between compiler runs. The cost is readability -- ast.nodes.items[b.lhs] is clumsier than b.lhs.* -- and the loss of Zig's null-safety on the "reference" (an index can be garbage in a way a typed pointer cannot). Zig's own compiler stores its AST exactly this way, and pushes it further with std.MultiArrayList, which splits the struct into one array per field (a "struct of arrays") so that a pass touching only node tags never loads the payloads it does not need. For us, the pointer tree stays; but now you know the door the big compilers walk through, and why.
None of this is a teaching simplification -- it is how production front-ends are built, right down to the vocabulary. In C, the reference here is Clang, whose RecursiveASTVisitor is a giant generated visitor that calls a VisitFoo hook for every node kind: exactly our walkExpr-plus-visitor split, scaled to a few hundred node types, with the traversal written once and each analysis supplying only its Visit methods. The GCC internals use tree "walk" callbacks in the same spirit.
In Rust, the compiler's HIR and AST both come with Visitor traits (rustc_hir::intravisit, rustc_ast::visit) whose default methods recurse into children and which you override only for the nodes you care about -- the identical "borrow the traversal, supply the leaf logic" design, expressed with traits in stead of comptime duck typing. The syn crate, which every procedural macro in the ecosystem leans on, ships both a Visit trait (read-only walking) and a Fold trait (tree-to-tree transformation) -- and Fold is our constant-folding pass generalised: take a node, return a possibly-different node. Rust even splits, as we discussed, into an Ident-interning arena so equal names share storage.
In Go, the standard library hands you tree-walking directly: go/ast exposes an ast.Visitor interface and ast.Walk/ast.Inspect that traverse a parsed Go file, and the entire go vet and gopls toolchain is built on visitors over that tree. Go's own compiler front-end tracks a source position on every node just like our span, for exactly the same reason: so the error messages can point at real code. Across all three languages -- and Zig itself -- the pattern is stable: a tree of typed nodes, each remembering its source position, walked by a traversal written once that individual passes borrow. You have now built the small honest version of the very thing clang, rustc, go build and zig all stand on.
Take stock of what the tree can do now. Every node remembers where it was parsed from. A single generic walker powers any number of analysis passes without a line of duplicated traversal. An explicit recursion handles the passes where entering and leaving matter. And a transformation pass can read one tree and produce a cleaner one, all inside the same arena we free in a single line. Design, traversal, transformation -- the three verbs, all in hand.
But a tree that parses is still not a program that makes sense. let x = 1 + true; walks perfectly through every pass we wrote today, and means absolutely nothing -- you cannot add an integer to a boolean. Our parser is happy with it; our traversals visit it without complaint; and yet it is wrong. Catching that class of nonsense -- a name used before it is defined, an operator applied to types it does not accept, a function called with the wrong number of arguments -- is a whole stage of its own, and it is the next brick. It walks this exact tree, using this exact visitor machinery, but it asks a sharper question of every node: not "what shape is this?" but "does this actually mean something?" The span we added today is what lets it answer "no, and here is precisely where." Build the redesigned AST, wire up the visitor, fold a few constant expressions until the tree shrinks the way you expect, and you will have the backbone every remaining stage of this arc hangs off. ;-)
A depth-measuring visitor. Write a pass that reports the maximum nesting depth of an expression -- 2 for 2 + 3 (the plus over its leaves), 3 for 2 + 3 * 4. You will find the plain visit(node) visitor cannot do this alone (it has no notion of "how deep am I"), so write it as an explicit recursion in the style of printTree, returning the depth of each subtree and taking the max of the children plus one. Test it against both expressions above.
Extend constant folding to booleans and comparisons. Our fold only folds integer arithmetic. Extend it so that comparisons of two integer constants fold to a boolean node (3 < 4 becomes bool true), and so that and/or of two boolean constants fold too. Keep the "return null, fall through to rebuild" discipline for anything it cannot fold. Write a test proving 2 * 3 < 10 folds all the way down to a single bool true.
A visitor that finds undefined variables -- the first half. Write an IdentCollector-style pass over a whole statement list that walks every let binding to gather the set of defined names, then walks every expression to gather the set of used names, and returns any used name that was never defined. Do not worry about scopes or ordering yet (that is a real can of worms) -- just the flat set difference. This is a genuine miniature of the analysis the next stage does properly, and building the naive version first is the best way to feel why the real one needs more machinery.
Thanks for reading -- de groeten, and I'll see you at the next brick! ;-)
@scipio