ast-annotate.ts

interface Lit<K> {
    kind: 'lit';
    value: K;
}

interface BinOp<T> {
    kind: 'binop';
    l: T;
    r: T;
}

type Expr<T, K> = Lit<K> | BinOp<TaggedExpr<T, K>>;
type TaggedExpr<T, K> = Expr<T, K> & T;

type Basic<K> = TaggedExpr<unknown, K>;
type WithDepth<K> = TaggedExpr<{ depth: number }, K>;

interface Annotator<T, K> {
    lit: (k: K) => T;
    binop: (l: T, r: T) => T;
}

function annotateExpr<T, U, K>(annotate: Annotator<U, K>, node: TaggedExpr<T, K>): TaggedExpr<T & U, K> {
    switch (node.kind) {
        case 'lit':
            return { ...node, ...annotate.lit(node.value) };
        case 'binop':
            const l = annotateExpr(annotate, node.l);
            const r = annotateExpr(annotate, node.r);
            return { ...node, l, r, ...annotate.binop(l, r) };
    }
}

const addDepth = <T, K>(node: TaggedExpr<T, K>) => annotateExpr<T, { depth: number }, K>(
    {
        lit: () => ({ depth: 0 }),
        binop: (l: { depth: number }, r: { depth: number }) => ({ depth: Math.max(l.depth, r.depth) + 1 })
    },
    node
);

const addType = <T, K>(node: TaggedExpr<T, K>) => annotateExpr<T, { type: string }, K>(
    {
        lit: (k: K) => ({ type: typeof k }),
        binop: (l: { type: string }, r: { type: string }) => ({ type: l.type === r.type ? l.type : "any" })
    },
    node
);

const node: Basic<number | string | boolean> = { kind: 'binop', l: { kind: 'lit', value: 3 }, r: { kind: 'lit', value: 10 } };
const transformed: TaggedExpr<{ depth: number, type: string }, number | string | boolean> = addType(addDepth(node));