Learn Zig Series (#134) - Type Checking

What will I learn?

Why a program that parses perfectly can still be complete nonsense, and what a type checker actually is -- the stage that asks "does this mean something?" of every node;
How to model the types of a tiny language as a Zig enum, and why a checker is really just one more pass over the AST we built last episode;
How to walk an expression tree bottom-up to infer the type of every sub-expression, and reject 1 + true before it ever runs;
How a type environment (a scoped symbol table) records which names are defined and what type they hold, so x + 1 can be checked against the let x = ... that introduced x;
How the source span we added last episode lets every type error point at exactly the line and column where it went wrong;
How to test a checker by feeding it programs that must pass and programs that must fail, and asserting on the kind of failure;
How C, Rust, Go and Zig itself structure their real type checkers, and why yours is the same shape scaled down;
Three exercises to push your checker toward a real one before the next episode.

Requirements

A working modern computer running macOS, Windows or Ubuntu;
An installed Zig 0.14+ distribution (download from ziglang.org);
The redesigned AST from episode 133 (the Expr struct with a span and a kind union) and the visitor machinery fresh in mind -- today we write a real pass over exactly that tree;
Tagged unions from episode 6, pointers from episode 8, allocators from episode 7, and hash maps from episode 22;
The ambition to learn Zig programming.

Difficulty

Advanced

Curriculum (of the `Learn Zig Series`):

Learn Zig Series (#134) - Type Checking

Last episode I ended on a promise disguised as a complaint. I wrote let x = 1 + true;, pointed out that it walks perfectly through every pass we had built -- the visitor visits it, the constant folder folds it, the pretty-printer prints it -- and yet it means absolutely nothing. You cannot add an integer to a boolean. The parser was happy. The traversals were happy. And the program was still wrong. I called catching that class of nonsense "the next brick", and said it would walk this exact tree with this exact visitor machinery but ask a sharper question of every node: not "what shape is this?" but "does this actually mean something?". This is that brick.

The stage that asks that question is the type checker, and it is the first place in a compiler where the machine starts to understand your program rather than merely re-arrange it. Up to now every pass has been about form: characters into tokens, tokens into a tree, tree into a cleaner tree. Type checking is the first pass about meaning. It is where 1 + true finally gets told "no", where a name used before it is defined finally gets caught, and -- crucially -- where the error message can say exactly where thanks to the span we bolted onto every node last episode. Let's dive right in!

Solutions to Episode 133 Exercises

As always, the three exercises first, with complete code. All three assume the redesigned Expr from last episode -- a struct with a span and a kind union whose variants include int, float, boolean, ident, unary, binary and call.

Exercise 1 -- A depth-measuring visitor. The task was to report the maximum nesting depth of an expression: 2 for 2 + 3, and 3 for 2 + 3 * 4. As I hinted, the flat visit(node) visitor genuinely cannot do this, because it has no notion of "how deep am I". So we write it as an explicit recursion in the spirit of printTree: each leaf is depth 1, and each internal node is one more than the deepest of its children.

fn maxDepth(e: *const Expr) usize {
    return switch (e.kind) {
        .int, .float, .boolean, .ident => 1,
        .unary => |u| 1 + maxDepth(u.rhs),
        .binary => |b| 1 + @max(maxDepth(b.lhs), maxDepth(b.rhs)),
        .call => |c| blk: {
            var deepest: usize = 0;
            for (c.args) |arg| deepest = @max(deepest, maxDepth(arg));
            break :blk 1 + deepest;
        },
    };
}

The whole trick is that depth is a value that flows up from the leaves, and the generic visitor only ever flows information sideways (into the visitor's own state). Any question whose answer is "computed from the children's answers" wants an explicit recursion that returns something, not a visit that mutates. The test pins both cases down:

test "maxDepth: 2 + 3 is depth 2, 2 + 3 * 4 is depth 3" {
    var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
    defer arena.deinit();
    const a = arena.allocator();
    const sp = Span{ .line = 1, .col = 1 };
    const two = try makeExpr(a, sp, .{ .int = 2 });
    const three = try makeExpr(a, sp, .{ .int = 3 });
    const four = try makeExpr(a, sp, .{ .int = 4 });
    const shallow = try makeExpr(a, sp, .{ .binary = .{ .op = .plus, .lhs = two, .rhs = three } });
    try std.testing.expectEqual(@as(usize, 2), maxDepth(shallow));

    const mul = try makeExpr(a, sp, .{ .binary = .{ .op = .star, .lhs = three, .rhs = four } });
    const deep = try makeExpr(a, sp, .{ .binary = .{ .op = .plus, .lhs = two, .rhs = mul } });
    try std.testing.expectEqual(@as(usize, 3), maxDepth(deep));
}

Exercise 2 -- Extend constant folding to booleans and comparisons. Last episode fold only handled integer arithmetic. The task was to fold comparisons of two integer constants down to a boolean node (3 < 4 becomes bool true), and to fold and/or of two boolean constants, keeping the "return null, fall through to rebuild" discipline for anything that cannot fold. The additions live entirely inside the .binary arm:

fn foldBinary(op: TokenKind, l: *Expr, r: *Expr) ?Expr.Kind {
    if (l.kind == .int and r.kind == .int) {
        const lv = l.kind.int;
        const rv = r.kind.int;
        return switch (op) {
            .plus => .{ .int = lv + rv },
            .minus => .{ .int = lv - rv },
            .star => .{ .int = lv * rv },
            .lt => .{ .boolean = lv < rv },
            .lte => .{ .boolean = lv <= rv },
            .gt => .{ .boolean = lv > rv },
            .gte => .{ .boolean = lv >= rv },
            else => null,
        };
    }
    if (l.kind == .boolean and r.kind == .boolean) {
        const lb = l.kind.boolean;
        const rb = r.kind.boolean;
        return switch (op) {
            .kw_and => .{ .boolean = lb and rb },
            .kw_or => .{ .boolean = lb or rb },
            else => null,
        };
    }
    return null;
}

I pulled the decision out into its own helper that returns an optional Expr.Kind -- null means "cannot fold this, rebuild the node instead". That keeps the recursive fold tidy: it folds both children, asks foldBinary, and either stamps the folded kind onto a fresh node or rebuilds. The chained folding is the satisfying part -- 2 * 3 < 10 folds 2 * 3 to 6 first, then 6 < 10 to true, because the recursion handles the children before the parent:

test "2 * 3 < 10 folds all the way to bool true" {
    var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
    defer arena.deinit();
    const a = arena.allocator();
    const sp = Span{ .line = 1, .col = 1 };
    const two = try makeExpr(a, sp, .{ .int = 2 });
    const three = try makeExpr(a, sp, .{ .int = 3 });
    const ten = try makeExpr(a, sp, .{ .int = 10 });
    const mul = try makeExpr(a, sp, .{ .binary = .{ .op = .star, .lhs = two, .rhs = three } });
    const cmp = try makeExpr(a, sp, .{ .binary = .{ .op = .lt, .lhs = mul, .rhs = ten } });
    const folded = try fold(a, cmp);
    try std.testing.expect(folded.kind == .boolean);
    try std.testing.expectEqual(true, folded.kind.boolean);
}

Exercise 3 -- Finding undefined variables, the first half. The task was a flat set-difference: walk a statement list, gather every name introduced by a let, gather every name used in an expression, and report any used name that was never defined. No scopes, no ordering -- just the naive version, because building the naive version first is the best way to feel why the real one (which we start today) needs more machinery. Two std.StringHashMap(void) sets do it:

fn collectUsed(e: *const Expr, used: *std.StringHashMap(void)) !void {
    switch (e.kind) {
        .ident => |name| try used.put(name, {}),
        .unary => |u| try collectUsed(u.rhs, used),
        .binary => |b| {
            try collectUsed(b.lhs, used);
            try collectUsed(b.rhs, used);
        },
        else => {},
    }
}

fn firstUndefined(alloc: std.mem.Allocator, stmts: []const Stmt) !?[]const u8 {
    var defined = std.StringHashMap(void).init(alloc);
    defer defined.deinit();
    var used = std.StringHashMap(void).init(alloc);
    defer used.deinit();
    for (stmts) |s| switch (s) {
        .let => |l| {
            try collectUsed(l.value, &used);
            try defined.put(l.name, {});
        },
        .expr => |e| try collectUsed(e, &used),
    };
    var it = used.keyIterator();
    while (it.next()) |key| {
        if (!defined.contains(key.*)) return key.*;
    }
    return null;
}

I called this "a genuine miniature" of today's analysis, and you can already smell what is missing. It has no idea that let y = y + 1; uses y before it is defined -- it lumps all definitions and all uses into two flat bags and compares them, blind to order. That blindness is exactly the machinery gap today's real checker closes: it processes statements top to bottom, so a name only counts as defined once its let has actually been seen. On to the real thing.

What a type checker actually is

Strip away the intimidation and a type checker is a startlingly simple idea: walk the AST and give every expression a type, refusing to continue the moment two types that should agree do not. That is the whole game. 2 has type int. true has type boolean. 2 + 3 has type int, because + takes two integers and produces an integer. And 2 + true? That has no type -- there is no rule that says what integer-plus-boolean means -- so the checker stops and reports an error. Every type system you have ever fought with, from C's to Rust's, is this same loop wearing progressively heavier armour.

For our little language, the universe of types is tiny, so we model it as a plain enum. Three types: integers, floats, booleans. An enum is exactly right here -- a closed set of named alternatives (episode 6), cheap to copy and cheap to compare:

const Type = enum {
    int,
    float,
    boolean,

    fn name(self: Type) []const u8 {
        return switch (self) {
            .int => "int",
            .float => "float",
            .boolean => "boolean",
        };
    }
};

That name helper is not decoration -- it is what turns a machine-friendly Type.int into the human-friendly string a diagnostic prints. A real compiler's Type is a far richer thing (it carries struct fields, function signatures, generic parameters, lifetimes), but the role is identical: it is the answer to the question "what is this expression?", and the checker's job is to compute that answer for every node and make sure the answers are consistent.

Inferring the type of an expression

The checker mirrors the visitor's shape from last episode, but with one decisive difference: where walkExpr returned nothing and mutated a visitor, checkExpr returns a Type (or an error). It is a function from a node to its type, defined recursively over the tree. Leaves are trivial -- an integer literal is an int, no thought required. Internal nodes check their children first, then apply the rule for their operator. Here is the skeleton for the literal and unary cases:

const Checker = struct {
    env: std.StringHashMap(Type),
    err_span: Span = .{ .line = 0, .col = 0 },
    err_msg: []const u8 = "",

    const TypeError = error{ UndefinedVariable, TypeMismatch, OutOfMemory };

    fn fail(self: *Checker, span: Span, msg: []const u8) TypeError {
        self.err_span = span;
        self.err_msg = msg;
        return error.TypeMismatch;
    }

    fn checkExpr(self: *Checker, e: *const Expr) TypeError!Type {
        switch (e.kind) {
            .int => return .int,
            .float => return .float,
            .boolean => return .boolean,
            .ident => |nm| return self.env.get(nm) orelse {
                self.err_span = e.span;
                self.err_msg = "use of undefined variable";
                return error.UndefinedVariable;
            },
            .unary => |u| {
                const rhs = try self.checkExpr(u.rhs);
                return switch (u.op) {
                    .minus => if (rhs == .int or rhs == .float) rhs else self.fail(e.span, "unary '-' needs a numeric operand"),
                    .bang => if (rhs == .boolean) .boolean else self.fail(e.span, "unary '!' needs a boolean operand"),
                    else => self.fail(e.span, "unsupported unary operator"),
                };
            },
            .binary => |b| {
                const l = try self.checkExpr(b.lhs);
                const r = try self.checkExpr(b.rhs);
                return self.checkBinary(e.span, b.op, l, r);
            },
        }
    }
};

Look at how the recursion carries the error out. Because checkExpr returns TypeError!Type, a try on the left operand of a binary node means: if checking the left child already failed, we never even look at the right one -- the error propagates straight up to the caller, carrying the span and message we stashed. That is Zig's error unions (episode 4) doing precisely what they are for. The orelse on the identifier case is the same idea in miniature: env.get(nm) returns an optional, and orelse is where "this name is not in scope" becomes a real, positioned error in stead of a silent null.

The rules live in one place

I split the binary operator rules into their own function on purpose. Every interesting type rule in this language is a statement about a binary operator -- arithmetic wants two matching numbers and gives that number back, comparison wants two matching numbers and gives a boolean, and/or want two booleans, equality wants two of anything so long as they match. Putting them in one switch makes the whole type system readable at a glance, and makes adding a new operator a one-line change:

fn checkBinary(self: *Checker, span: Span, op: TokenKind, l: Type, r: Type) Checker.TypeError!Type {
    switch (op) {
        .plus, .minus, .star, .slash => {
            if (l == r and (l == .int or l == .float)) return l;
            return self.fail(span, "arithmetic needs two operands of the same numeric type");
        },
        .lt, .lte, .gt, .gte => {
            if (l == r and (l == .int or l == .float)) return .boolean;
            return self.fail(span, "comparison needs two operands of the same numeric type");
        },
        .eq, .neq => {
            if (l == r) return .boolean;
            return self.fail(span, "'==' needs both sides to be the same type");
        },
        .kw_and, .kw_or => {
            if (l == .boolean and r == .boolean) return .boolean;
            return self.fail(span, "'and'/'or' need boolean operands");
        },
        else => return self.fail(span, "unsupported binary operator"),
    }
}

Notice the design decision hiding in .plus: we demand l == r and numeric. That is a language design choice, not a law of nature -- it means 1 + 1.5 is an error in our language, because we refuse to silently promote an int to a float. C would happily promote; Zig itself refuses implicit widening between runtime i32 and f64 and makes you write the cast, which is exactly the philosophy I am copying here. "No hidden coercions" is a defensible, honest stance, and it makes the checker's rules dead simple to state and to test. If later you want mixed arithmetic, this one function is the only place you would loosen the rule -- the entire policy of the language sits in this single switch, which is the payoff of not scattering it across the tree walk.

Names and scope: the type environment

Expressions do not live in a vacuum -- they mention variables, and a variable's type comes from wherever it was introduced. That "wherever" is the type environment: a map from name to type that we fill in as we process let bindings and consult whenever we hit an identifier. We already leaned on it above (self.env.get(nm)); now let us drive it from statements. Statements are a second small union -- a let that introduces a name, or a bare expression:

const Stmt = union(enum) {
    let: struct { name: []const u8, value: *Expr, span: Span },
    expr: *Expr,
};

fn checkStmt(self: *Checker, s: *const Stmt) Checker.TypeError!void {
    switch (s.*) {
        .let => |l| {
            const t = try self.checkExpr(l.value);
            try self.env.put(l.name, t);
        },
        .expr => |e| {
            _ = try self.checkExpr(e);
        },
    }
}

fn checkProgram(self: *Checker, stmts: []const Stmt) Checker.TypeError!void {
    for (stmts) |*s| try self.checkStmt(s);
}

Here is the subtle, important thing, and it is exactly what exercise 3 was missing. We check the let's value before we put the name into the environment. So in let x = x + 1;, when we check x + 1, x is not in the environment yet, and the checker correctly reports x as undefined -- order matters, and processing statements top-to-bottom while updating the environment as we go is what gives us ordering for free. The flat two-bag version from exercise 3 could never see that, because it gathered all definitions before comparing. This is the machinery gap closing in real time.

The type of x is inferred, by the way -- notice there is no type annotation in let x = 2 + 3;. We compute the value's type (int) and bind the name to it. That is a baby version of the local type inference that Rust, Go, and modern C++ all do: you do not write the type, the compiler works it out from the initialiser. Our version is trivial because expressions have exactly one obvious type, but the shape -- "check the right-hand side, bind its type to the name" -- is the real thing.

Testing the checker: programs that must pass, and programs that must fail

A checker earns trust by rejecting exactly the wrong programs and accepting exactly the right ones, so its tests come in two flavours. First, the happy path: a small well-typed program must type-check cleanly, and the environment must end up holding the types we expect. We build let x = 2 + 3; let ok = x < 10; by hand and assert x is int and ok is boolean:

test "a well-typed program checks and infers the right types" {
    var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
    defer arena.deinit();
    const a = arena.allocator();
    const sp = Span{ .line = 1, .col = 1 };

    const two = try makeExpr(a, sp, .{ .int = 2 });
    const three = try makeExpr(a, sp, .{ .int = 3 });
    const sum = try makeExpr(a, sp, .{ .binary = .{ .op = .plus, .lhs = two, .rhs = three } });
    const x_ref = try makeExpr(a, sp, .{ .ident = "x" });
    const ten = try makeExpr(a, sp, .{ .int = 10 });
    const cmp = try makeExpr(a, sp, .{ .binary = .{ .op = .lt, .lhs = x_ref, .rhs = ten } });

    const stmts = [_]Stmt{
        .{ .let = .{ .name = "x", .value = sum, .span = sp } },
        .{ .let = .{ .name = "ok", .value = cmp, .span = sp } },
    };

    var checker = Checker{ .env = std.StringHashMap(Type).init(a) };
    defer checker.env.deinit();
    try checker.checkProgram(&stmts);
    try std.testing.expectEqual(Type.int, checker.env.get("x").?);
    try std.testing.expectEqual(Type.boolean, checker.env.get("ok").?);
}

Second, and just as important, the sad path: the program that must be rejected, and rejected for the right reason. This is where 1 + true finally gets its comeuppance. We assert not just that checking fails, but that it fails with error.TypeMismatch specifically, and that the recorded message is the arithmetic one -- because a checker that fails with the wrong error is as buggy as one that does not fail at all:

test "1 + true is rejected as a type mismatch" {
    var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
    defer arena.deinit();
    const a = arena.allocator();
    const sp = Span{ .line = 3, .col = 9 };

    const one = try makeExpr(a, sp, .{ .int = 1 });
    const tru = try makeExpr(a, sp, .{ .boolean = true });
    const bad = try makeExpr(a, sp, .{ .binary = .{ .op = .plus, .lhs = one, .rhs = tru } });

    var checker = Checker{ .env = std.StringHashMap(Type).init(a) };
    defer checker.env.deinit();
    try std.testing.expectError(error.TypeMismatch, checker.checkExpr(bad));
    try std.testing.expectEqual(@as(usize, 3), checker.err_span.line);
    try std.testing.expectEqual(@as(usize, 9), checker.err_span.col);
}

That last pair of assertions is the whole reason we suffered through adding a span to every node last episode. The checker did not just say "type error somewhere" -- it handed back line 3, column 9, the exact spot the offending + was parsed from. A real compiler would feed that span, plus err_msg, into a diagnostic renderer that underlines the source line with a little caret. We built the plumbing; the pretty-printing is a formatting detail on top of it.

Where this sits among the grown-ups

None of this is a toy simplification of how real compilers work -- it is the same architecture with more cases bolted on, and the same vocabulary. In C, the type checker (in GCC or Clang's Sema) walks the AST computing a type for every expression node and issuing the diagnostics you know and love ("invalid operands to binary +"), gated by C's promotion rules, which are exactly the coercions we chose not to have. In Rust, rustc's HIR type-checking phase does local inference (that is what lets you skip most type annotations) followed by trait resolution, and its error messages carry Span values that are quite literally the same idea as ours, line and column of the offending code. In Go, the go/types package is a standalone type checker over the go/ast tree that any tool can call, and Go's own compiler tracks a source position on every node so go build can point at the right line -- I mentioned this last episode, and here is where that position finally earns its keep.

And Zig itself? Zig's type checking is unusual because so much of it happens at comptime during semantic analysis, where types are first-class values you can compute with (episode 32's @typeInfo is a peek into that machinery). But the core loop is the same one we wrote today: figure out the type of each expression, and refuse to compile when the types do not line up. The error: incompatible types: 'i64' and 'bool' that Zig throws at you is our error.TypeMismatch in a much bigger coat. What you have built here is the small, honest heart of the very thing every one of these compilers stands on.

Where this is heading

Take stock. We have a Type, a checkExpr that infers the type of any expression by walking the tree bottom-up, a checkBinary that holds the entire type policy of the language in one readable switch, a type environment that gives ordered scope to names, and a checkProgram that ties it together -- rejecting 1 + true and let x = x + 1; alike, each with a span pointing at the crime scene. The program has finally been asked "do you mean something?" and made to answer.

Here is what that unlocks. Once you know the type of every expression, you know things the earlier stages could only guess at: how many bytes each value needs, which machine operation a + should become (integer add versus float add), whether a comparison is signed. A parser cannot tell you that; a type checker can. That knowledge is precisely the input the next stage of this arc needs, because we are about to stop interpreting the tree by walking it and start compiling it down to something lower and faster -- a flat, linear sequence of simple instructions that a small machine can chew through without ever touching the tree again. The tree got us to meaning; the next brick turns meaning into something that runs. Build the checker, feed it a handful of good and bad programs until you trust its verdicts, and you will have the last tree-shaped stage before we go flat. ;-)

Exercises

Add a string type and typed concatenation. Extend the Type enum with a string variant, teach checkExpr that a string-literal node has type string, and add a rule to checkBinary so that + applied to two strings is legal and produces a string -- while "a" + 1 remains a type mismatch. Write one test that a string + string checks to string, and one that string + int is rejected with error.TypeMismatch.
Report the two conflicting types in the message. Right now checkBinary fails with a generic "arithmetic needs two operands of the same numeric type". Change fail (or add a variant of it) so the recorded err_msg names the actual offenders -- something like "cannot add int and boolean" -- by formatting the two Type.name() strings into an allocated string. Test that checking 1 + true produces a message that contains both "int" and "boolean".
Introduce block scoping. Our single flat environment means a name, once defined, is defined forever. Give the checker a stack of environments: pushing a new scope on entering a block, popping it on exit, and looking a name up by searching from the innermost scope outward. Then a name defined inside a block must be invisible after the block ends. You do not have if/block statements in the AST yet, so add a minimal block: []const Stmt statement variant to drive it. Test that a name let-bound inside a block is reported as undefined when used after the block closes.

Bedankt en tot de volgende keer -- one brick further down the road, and the tree finally starts to run! ;-)

@scipio

Learn Zig Series (#134) - Type Checking

Learn Zig Series (#134) - Type Checking

What will I learn?

Requirements

Difficulty

Curriculum (of the Learn Zig Series):

Learn Zig Series (#134) - Type Checking

Solutions to Episode 133 Exercises

What a type checker actually is

Inferring the type of an expression

The rules live in one place

Names and scope: the type environment

Testing the checker: programs that must pass, and programs that must fail

Where this sits among the grown-ups

Where this is heading

Exercises

Curriculum (of the `Learn Zig Series`):