Tagged union pattern in Typescript

To show the pattern, let's implement a simple interpreter.

Start by defining each union variant, or kind, of the type:

type ExprKind =
    | { tag: "Int", value: number }
    | { tag: "Add", left: Expr, right: Expr }

Then define the type itself as a class:

class Expr {
    constructor(
        public line: number,
        public kind: ExprKind,
    ) {}
    // ...
}

Each instance of Expr will have a line (source line number), but only IntExpr will have value, and only AddExpr will have left and right.

To match each variant, a type representing an Expr with a specific variant:

type ExprTag = ExprKind["tag"];

type ExprWithKind<Tag extends ExprTag> = Expr & { kind: { tag: Tag } };

The type ExprKind["Tag"] is a union of each tag value, i.e. in this example "Int" | "Add". Saying Tag extends ExprTag means that Tag (type) is either "Int" or "Add", i.e. for .is("Int"), Tag will be "Int" and for .is("Add"), Tag will be "Add". The narrowed type ExprWithKind is defined by saying that we want an Expr that specifically has the kind.tag value of Tag.

Then define a method that with a type predicate:

class Expr {
    // ...
    is<Tag extends ExprTag>(tag: Tag): this is ExprWithKind<Tag> {
        return this.kind.tag === tag;
    }
    // ...
}

The type predicate tells Typescript that it can expect the Expr to conform to ExprWithKind<Tag> if .is evaluates to true. Using this, access can be gained to the variant fields in a type safe manner.

Now, it's possible to match each variant:

function eval(expr: Expr): number {

    if (expr.is("Int")) {
        return expr.kind.value;
    }
    
    if (expr.is("Add")) {
        return eval(expr.kind.left) + eval(expr.kind.right);
    }

    throw new Error(`'${expr.kind.tag}' not handled`);
}

By having Expr be a class, we can also define functions like this as methods:

class Expr {
    // ...
    eval(): number {
        if (expr.is("Int"))
            return expr.kind.value;
        if (expr.is("Add"))
            return expr.kind.left.eval() + expr.kind.right.eval();
        throw new Error(`'${expr.kind.tag}' not handled`);
    }
    // ...
}

By having the tag invariant, the Typescript compiler can check that you exhaust each variant in a match. To match each variant, I use this pattern ('k-pattern'):

class Expr {
    // ...
    eval(): number {
        const k = this.kind;
        switch (k.tag) { // Typescript language servers can auto-complete all variants
            case "Int":
                return k.value;
            case "Add":
                return k.right.eval() + k.left.eval();
            default:
                k satisfies never; // compile time exhaustiveness check
                throw new Error("unexhausted"): // run time exhaustiveness check
        }
    }
    // ...
}

Sometimes you need to narrow to a variant that you already know is the correct variant. For that purpose, define these helper functions:

class Expr {
    // ...
    assertIs<Tag extends ExprTag>(tag: Tag): asserts this is ExprWithKind<Tag> {
        if (!this.is(tag)) {
            throw new Error();
        }
    }

    as<Tag extends ExprTag>(tag: Tag): ExprWithKind<Tag> {
        this.assertIs(tag);
        return this;
    }
    // ...
}

Using these, narrowing can be done more conveniently:

// this will always be an IntExpr
let expr: Expr = ...;

expr.kind.value; // Property 'value' does not exist ...

expr.as("Int").kind.value; // ok

expr.assertIs("Int"); // will throw if expr isn't an IntExpr
expr.kind.value; // ok

For some applications, it might be more conveniently with a factory-function like this:

type ExprKindData<Tag extends ExprTag> = Omit<ExprWithKind<Tag>, "tag">;

class Expr {
    // ...
    static create<Tag extends ExprTag>(
        tag: Tag,
        kind: ExprKindData<Tag>,
        line: number,
    ): ExprWithKind<Tag> {
        return new Expr(line, {...kind, tag} as ExprKind & { tag: Ty });
    }
    // ...
}

And use it like this:

const myInt = Expr.create("Int", { value: 123 }, line);

For often-used variants, it might be more convenient with a specific type and constructor:

type IntExpr = ExprWithKind<"Int">;
const IntExpr = (line: number, value: number): IntExpr =>
    Expr.create("Int", { value }, line);

That can be used like this:

function exprIdentity(expr: Expr): Expr { return expr; }
function intExprIdentity(expr: IntExpr): IntExpr { return expr; }

const myIntExpr = IntExpr(123, line);
myIntExpr.kind.value;
exprIdentity(myIntExpr);
intExprIdentity(myIntExpr);

const myOtherExpr = Expr.create("Add", ...);
exprIdentity(myIntExpr);
intExprIdentity(myIntExpr); // error

Using this pattern, it's easy to implement a visitor pattern:

interface ExprVisitor {
    visit(expr: Expr): void
}

class Expr {
    // ...
    visit(v: ExprVisitor) {
        v.visit(this);
        const k = this.kind;
        switch (k.tag) {
            case "Int":
                break;
            case "Add":
                k.left.visit(v);
                k.right.visit(v);
                break;
            default:
                k satisfies never;
        }
    }
    // ...
}

The visitor interface can be used like:

const expr: Expr = parse("2 + (3 + 4)");

const intValues: number[] = [];

expr.visit({
    visit(expr) {
        if (expr.is("Int")) {
            intValues.push(expr.kind.value);
        }
    }
});

5.5 KiB Raw Permalink Blame History

Tagged union pattern in Typescript

5.5 KiB

Raw Permalink Blame History