From 5f191c72cc2c1a3147455602bfd68d14aa61ca28 Mon Sep 17 00:00:00 2001
From: Nilstrieb <48135649+Nilstrieb@users.noreply.github.com>
Date: Mon, 31 Jul 2023 16:17:56 +0200
Subject: [PATCH] Strongly typed AST

---
 src/ast.ts     | 275 ++++++++++++++++++++++++++++---------------------
 src/index.ts   |   5 +-
 src/lower.ts   |  42 +++++---
 src/parser.ts  | 111 +++++++++++---------
 src/printer.ts |  23 +++--
 src/resolve.ts |  47 ++++++---
 src/typeck.ts  |  93 +++++++++++------
 7 files changed, 355 insertions(+), 241 deletions(-)

diff --git a/src/ast.ts b/src/ast.ts
index e918574..500ffbe 100644
--- a/src/ast.ts
+++ b/src/ast.ts
@@ -1,88 +1,115 @@
 import { Span } from "./error";
 import { LitIntType } from "./lexer";
 
-export type Ast = {
-  rootItems: Item[];
+export type Ast<P extends Phase> = {
+  rootItems: Item<P>[];
   typeckResults?: TypeckResults;
-  itemsById: Map<ItemId, Item>;
+  itemsById: Map<ItemId, Item<P>>;
   packageName: string;
 };
 
-export type Identifier = {
+export type Phase = {
+  res: unknown;
+};
+
+type NoRes = object;
+type HasRes = { res: Resolution };
+
+export type Parsed = {
+  res: NoRes;
+};
+export type Built = {
+  res: NoRes;
+};
+export type Resolved = {
+  res: HasRes;
+};
+export type Typecked = {
+  res: HasRes;
+};
+export type AnyPhase = {
+  res: NoRes | HasRes;
+};
+
+export type Ident = {
   name: string;
   span: Span;
-  res?: Resolution;
 };
 
+export type IdentWithRes<P extends Phase> = {
+  name: string;
+  span: Span;
+} & P["res"];
+
 export type ItemId = number;
 
-export type ItemKind =
+export type ItemKind<P extends Phase> =
   | {
       kind: "function";
-      node: FunctionDef;
+      node: FunctionDef<P>;
     }
   | {
       kind: "type";
-      node: TypeDef;
+      node: TypeDef<P>;
     }
   | {
       kind: "import";
-      node: ImportDef;
+      node: ImportDef<P>;
     }
   | {
       kind: "mod";
-      node: ModItem;
+      node: ModItem<P>;
     };
 
-export type Item = ItemKind & {
+export type Item<P extends Phase> = ItemKind<P> & {
   span: Span;
   id: ItemId;
   defPath?: string[];
 };
 
-export type FunctionDef = {
+export type FunctionDef<P extends Phase> = {
   name: string;
-  params: FunctionArg[];
-  body: Expr;
-  returnType?: Type;
+  params: FunctionArg<P>[];
+  body: Expr<P>;
+  returnType?: Type<P>;
   ty?: TyFn;
 };
 
-export type FunctionArg = {
+export type FunctionArg<P extends Phase> = {
   name: string;
-  type: Type;
+  type: Type<P>;
   span: Span;
 };
 
-export type TypeDef = {
+export type TypeDef<P extends Phase> = {
   name: string;
-  fields: FieldDef[];
+  fields: FieldDef<P>[];
   ty?: TyStruct;
 };
 
-export type FieldDef = {
-  name: Identifier;
-  type: Type;
+export type FieldDef<P extends Phase> = {
+  name: Ident;
+  type: Type<P>;
 };
 
-export type ImportDef = {
+export type ImportDef<P extends Phase> = {
   module: StringLiteral;
   func: StringLiteral;
   name: string;
-  params: FunctionArg[];
-  returnType?: Type;
+  params: FunctionArg<P>[];
+  returnType?: Type<P>;
   ty?: TyFn;
 };
 
-export type ModItem = {
+export type ModItem<P extends Phase> = {
   name: string;
-  modKind: ModItemKind;
+  modKind: ModItemKind<P>;
 };
 
-export type ModItemKind =
+export type ModItemKind<P extends Phase> =
   | {
       kind: "inline";
-      contents: Item[];
+      contents: Item<P>[];
     }
   | {
       kind: "extern";
@@ -90,11 +117,11 @@ export type ModItemKind =
 
 export type ExprEmpty = { kind: "empty" };
 
-export type ExprLet = {
+export type ExprLet<P extends Phase> = {
   kind: "let";
-  name: Identifier;
-  type?: Type;
-  rhs: Expr;
+  name: Ident;
+  type?: Type<P>;
+  rhs: Expr<P>;
   // IMPORTANT: This is (sadly) shared with ExprBlock.
   // TODO: Stop this sharing and just store the stack of blocks in typeck.
   local?: LocalInfo;
@@ -102,15 +129,15 @@ export type ExprLet = {
 
 // A bit like ExprBinary except there are restrictions
 // on the LHS and precedence is unrestricted.
-export type ExprAssign = {
+export type ExprAssign<P extends Phase> = {
   kind: "assign";
-  lhs: Expr;
-  rhs: Expr;
+  lhs: Expr<P>;
+  rhs: Expr<P>;
 };
 
-export type ExprBlock = {
+export type ExprBlock<P extends Phase> = {
   kind: "block";
-  exprs: Expr[];
+  exprs: Expr<P>[];
   // IMPORTANT: This is (sadly) shared with ExprLet.
   locals?: LocalInfo[];
 };
@@ -120,9 +147,9 @@ export type ExprLiteral = {
   value: Literal;
 };
 
-export type ExprIdent = {
+export type ExprIdent<P extends Phase> = {
   kind: "ident";
-  value: Identifier;
+  value: IdentWithRes<P>;
 };
 
 /**
@@ -139,28 +166,28 @@ export type ExprPath = {
   res: Resolution;
 };
 
-export type ExprBinary = {
+export type ExprBinary<P extends Phase> = {
   kind: "binary";
   binaryKind: BinaryKind;
-  lhs: Expr;
-  rhs: Expr;
+  lhs: Expr<P>;
+  rhs: Expr<P>;
 };
 
-export type ExprUnary = {
+export type ExprUnary<P extends Phase> = {
   kind: "unary";
   unaryKind: UnaryKind;
-  rhs: Expr;
+  rhs: Expr<P>;
 };
 
-export type ExprCall = {
+export type ExprCall<P extends Phase> = {
   kind: "call";
-  lhs: Expr;
-  args: Expr[];
+  lhs: Expr<P>;
+  args: Expr<P>[];
 };
 
-export type ExprFieldAccess = {
+export type ExprFieldAccess<P extends Phase> = {
   kind: "fieldAccess";
-  lhs: Expr;
+  lhs: Expr<P>;
   field: {
     value: string | number;
     span: Span;
@@ -168,18 +195,18 @@ export type ExprFieldAccess = {
   };
 };
 
-export type ExprIf = {
+export type ExprIf<P extends Phase> = {
   kind: "if";
-  cond: Expr;
-  then: Expr;
-  else?: Expr;
+  cond: Expr<P>;
+  then: Expr<P>;
+  else?: Expr<P>;
 };
 
 export type LoopId = number;
 
-export type ExprLoop = {
+export type ExprLoop<P extends Phase> = {
   kind: "loop";
-  body: Expr;
+  body: Expr<P>;
   loopId: LoopId;
 };
 
@@ -188,36 +215,36 @@ export type ExprBreak = {
   target?: LoopId;
 };
 
-export type ExprStructLiteral = {
+export type ExprStructLiteral<P extends Phase> = {
   kind: "structLiteral";
-  name: Identifier;
-  fields: [Identifier, Expr][];
+  name: IdentWithRes<P>;
+  fields: [Ident, Expr<P>][];
 };
 
-export type TupleLiteral = {
+export type TupleLiteral<P extends Phase> = {
   kind: "tupleLiteral";
-  fields: Expr[];
+  fields: Expr<P>[];
 };
 
-export type ExprKind =
+export type ExprKind<P extends Phase> =
   | ExprEmpty
-  | ExprLet
-  | ExprAssign
-  | ExprBlock
+  | ExprLet<P>
+  | ExprAssign<P>
+  | ExprBlock<P>
   | ExprLiteral
-  | ExprIdent
+  | ExprIdent<P>
   | ExprPath
-  | ExprBinary
-  | ExprUnary
-  | ExprCall
-  | ExprFieldAccess
-  | ExprIf
-  | ExprLoop
+  | ExprBinary<P>
+  | ExprUnary<P>
+  | ExprCall<P>
+  | ExprFieldAccess<P>
+  | ExprIf<P>
+  | ExprLoop<P>
   | ExprBreak
-  | ExprStructLiteral
-  | TupleLiteral;
+  | ExprStructLiteral<P>
+  | TupleLiteral<P>;
 
-export type Expr = ExprKind & {
+export type Expr<P extends Phase> = ExprKind<P> & {
   span: Span;
   ty?: Ty;
 };
@@ -291,22 +318,22 @@ export function binaryExprPrecedenceClass(k: BinaryKind): number {
 export type UnaryKind = "!" | "-";
 export const UNARY_KINDS: UnaryKind[] = ["!", "-"];
 
-export type TypeKind =
+export type TypeKind<P extends Phase> =
   | {
       kind: "ident";
-      value: Identifier;
+      value: IdentWithRes<P>;
     }
   | {
       kind: "list";
-      elem: Type;
+      elem: Type<P>;
     }
   | {
       kind: "tuple";
-      elems: Type[];
+      elems: Type<P>[];
     }
   | { kind: "never" };
 
-export type Type = TypeKind & {
+export type Type<P extends Phase> = TypeKind<P> & {
   span: Span;
   ty?: Ty;
 };
@@ -446,60 +473,68 @@ export type TypeckResults = {
 
 // folders
 
-export type FoldFn<T> = (value: T) => T;
+export type FoldFn<From, To> = (value: From) => To;
 
-export type Folder = {
-  ast: () => Ast;
+export type Folder<From extends Phase, To extends Phase> = {
+  newItemsById: Map<ItemId, Item<To>>;
   /**
    * This should not be overridden.
    */
-  item: FoldFn<Item>;
-  itemInner: FoldFn<Item>;
-  expr: FoldFn<Expr>;
-  ident: FoldFn<Identifier>;
-  type: FoldFn<Type>;
+  item: FoldFn<Item<From>, Item<To>>;
+  itemInner: FoldFn<Item<From>, Item<To>>;
+  expr: FoldFn<Expr<From>, Expr<To>>;
+  ident: FoldFn<IdentWithRes<From>, IdentWithRes<To>>;
+  type: FoldFn<Type<From>, Type<To>>;
+};
+
+type ItemFolder<From extends Phase, To extends Phase> = {
+  newItemsById: Map<ItemId, Item<To>>;
+  item: FoldFn<Item<From>, Item<To>>;
+  itemInner: FoldFn<Item<From>, Item<To>>;
 };
 
 const ITEM_DEFAULT = Symbol("item must not be overriden");
 
-export const DEFAULT_FOLDER: Folder = {
-  ast() {
-    throw new Error("folders need to implement `ast`");
-  },
-  item(item) {
-    const newItem = this.itemInner(item);    
-    this.ast().itemsById.set(item.id, newItem);
-    return newItem;
-  },
-  itemInner(item) {
-    return superFoldItem(item, this);
-  },
-  expr(expr) {
-    return superFoldExpr(expr, this);
-  },
-  ident(ident) {
-    return ident;
-  },
-  type(type) {
-    return superFoldType(type, this);
-  },
-};
-(DEFAULT_FOLDER.item as any)[ITEM_DEFAULT] = ITEM_DEFAULT;
+export function mkDefaultFolder<
+  From extends Phase,
+  To extends Phase
+>(): ItemFolder<From, To> {
+  const folder: ItemFolder<From, To> = {
+    newItemsById: new Map(),
+    item(item) {
+      const newItem = this.itemInner(item);
+      this.newItemsById.set(item.id, newItem);
+      return newItem;
+    },
+    itemInner(_item) {
+      throw new Error("unimplemented");
+    },
+  };
+  (folder.item as any)[ITEM_DEFAULT] = ITEM_DEFAULT;
 
-export function foldAst(ast: Ast, folder: Folder): Ast {
+  return folder;
+}
+
+export function foldAst<From extends Phase, To extends Phase>(
+  ast: Ast<From>,
+  folder: Folder<From, To>
+): Ast<To> {
   if ((folder.item as any)[ITEM_DEFAULT] !== ITEM_DEFAULT) {
     throw new Error("must not override `item` on folders");
   }
 
   return {
     rootItems: ast.rootItems.map((item) => folder.item(item)),
-    itemsById: ast.itemsById,
+    itemsById: folder.newItemsById,
     typeckResults: ast.typeckResults,
     packageName: ast.packageName,
   };
 }
 
-export function superFoldItem(item: Item, folder: Folder): Item {
+export function superFoldItem<From extends Phase, To extends Phase>(
+  item: Item<From>,
+  folder: Folder<From, To>
+): Item<To> {
   switch (item.kind) {
     case "function": {
       const args = item.node.params.map(({ name, type, span }) => ({
@@ -550,7 +585,7 @@ export function superFoldItem(item: Item, folder: Folder): Item {
       };
     }
     case "mod": {
-      let kind: ModItemKind;
+      let kind: ModItemKind<To>;
       const { modKind: itemKind } = item.node;
       switch (itemKind.kind) {
         case "inline":
@@ -576,7 +611,10 @@ export function superFoldItem(item: Item, folder: Folder): Item {
   }
 }
 
-export function superFoldExpr(expr: Expr, folder: Folder): Expr {
+export function superFoldExpr<From extends Phase, To extends Phase>(
+  expr: Expr<From>,
+  folder: Folder<From, To>
+): Expr<To> {
   const span = expr.span;
   switch (expr.kind) {
     case "empty": {
@@ -687,7 +725,10 @@ export function superFoldExpr(expr: Expr, folder: Folder): Expr {
   }
 }
 
-export function superFoldType(type: Type, folder: Folder): Type {
+export function superFoldType<From extends Phase, To extends Phase>(
+  type: Type<From>,
+  folder: Folder<From, To>
+): Type<To> {
   const span = type.span;
   switch (type.kind) {
     case "ident": {
diff --git a/src/index.ts b/src/index.ts
index f1c80f1..c64164a 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -9,6 +9,7 @@ import { writeModuleWatToString } from "./wasm/wat";
 import fs from "fs";
 import path from "path";
 import { exec } from "child_process";
+import { Ast, Built, Typecked } from "./ast";
 
 const INPUT = `
 function main() = (
@@ -55,7 +56,7 @@ function main() {
     console.log("-----TOKENS------------");
     console.log(tokens);
 
-    const ast = parse(packageName, tokens);
+    const ast: Ast<Built> = parse(packageName, tokens);
     console.log("-----AST---------------");
 
     console.dir(ast.rootItems, { depth: 50 });
@@ -70,7 +71,7 @@ function main() {
     console.log(resolvedPrinted);
 
     console.log("-----AST typecked------");
-    const typecked = typeck(resolved);
+    const typecked: Ast<Typecked> = typeck(resolved);
     const typeckPrinted = printAst(typecked);
     console.log(typeckPrinted);
 
diff --git a/src/lower.ts b/src/lower.ts
index d16fde2..f207c6e 100644
--- a/src/lower.ts
+++ b/src/lower.ts
@@ -10,6 +10,7 @@ import {
   Ty,
   TyFn,
   TyTuple,
+  Typecked,
   varUnreachable,
 } from "./ast";
 import { ComplexMap, encodeUtf8, unwrap } from "./utils";
@@ -38,7 +39,7 @@ export type Context = {
   funcTypes: ComplexMap<wasm.FuncType, wasm.TypeIdx>;
   reservedHeapMemoryStart: number;
   funcIndices: ComplexMap<Resolution, FuncOrImport>;
-  ast: Ast;
+  ast: Ast<Typecked>;
   relocations: Relocation[];
 };
 
@@ -88,7 +89,7 @@ function appendData(cx: Context, newData: Uint8Array): number {
   }
 }
 
-export function lower(ast: Ast): wasm.Module {
+export function lower(ast: Ast<Typecked>): wasm.Module {
   const mod: wasm.Module = {
     types: [],
     funcs: [],
@@ -122,7 +123,7 @@ export function lower(ast: Ast): wasm.Module {
     relocations: [],
   };
 
-  function lowerMod(items: Item[]) {
+  function lowerMod(items: Item<Typecked>[]) {
     items.forEach((item) => {
       switch (item.kind) {
         case "function": {
@@ -172,7 +173,11 @@ export function lower(ast: Ast): wasm.Module {
   return mod;
 }
 
-function lowerImport(cx: Context, item: Item, def: ImportDef) {
+function lowerImport(
+  cx: Context,
+  item: Item<Typecked>,
+  def: ImportDef<Typecked>
+) {
   const existing = cx.mod.imports.findIndex(
     (imp) => imp.module === def.module.value && imp.name === def.func.value
   );
@@ -201,8 +206,8 @@ function lowerImport(cx: Context, item: Item, def: ImportDef) {
 
 type FuncContext = {
   cx: Context;
-  item: Item;
-  func: FunctionDef;
+  item: Item<Typecked>;
+  func: FunctionDef<Typecked>;
   wasmType: wasm.FuncType;
   wasm: wasm.Func;
   varLocations: VarLocation[];
@@ -216,7 +221,11 @@ type ArgRetAbi = wasm.ValType[];
 
 type VarLocation = { localIdx: number; types: wasm.ValType[] };
 
-function lowerFunc(cx: Context, item: Item, func: FunctionDef) {
+function lowerFunc(
+  cx: Context,
+  item: Item<Typecked>,
+  func: FunctionDef<Typecked>
+) {
   const abi = computeAbi(func.ty!);
   const { type: wasmType, paramLocations } = wasmTypeForAbi(abi);
   const type = internFuncType(cx, wasmType);
@@ -255,7 +264,11 @@ Expression lowering.
 - the result of an expression evaluation is stored on the top of the stack
 */
 
-function lowerExpr(fcx: FuncContext, instrs: wasm.Instr[], expr: Expr) {
+function lowerExpr(
+  fcx: FuncContext,
+  instrs: wasm.Instr[],
+  expr: Expr<Typecked>
+) {
   const ty = expr.ty!;
 
   exprKind: switch (expr.kind) {
@@ -284,7 +297,7 @@ function lowerExpr(fcx: FuncContext, instrs: wasm.Instr[], expr: Expr) {
       const { lhs } = expr;
       switch (lhs.kind) {
         case "ident": {
-          const res = lhs.value.res!;
+          const res = lhs.value.res;
 
           switch (res.kind) {
             case "local": {
@@ -355,7 +368,7 @@ function lowerExpr(fcx: FuncContext, instrs: wasm.Instr[], expr: Expr) {
     }
     case "path":
     case "ident": {
-      const res = expr.kind === "ident" ? expr.value.res! : expr.res;
+      const res = expr.kind === "ident" ? expr.value.res : expr.res;
 
       switch (res.kind) {
         case "local": {
@@ -502,8 +515,7 @@ function lowerExpr(fcx: FuncContext, instrs: wasm.Instr[], expr: Expr) {
         todo("non constant calls");
       }
 
-      const res =
-        expr.lhs.kind === "ident" ? expr.lhs.value.res! : expr.lhs.res;
+      const res = expr.lhs.kind === "ident" ? expr.lhs.value.res : expr.lhs.res;
 
       if (res.kind === "builtin") {
         switch (res.name) {
@@ -575,7 +587,7 @@ function lowerExpr(fcx: FuncContext, instrs: wasm.Instr[], expr: Expr) {
 
       // TODO: Actually do this instead of being naive.
 
-      const _isPlace = (expr: Expr) =>
+      const _isPlace = (expr: Expr<Typecked>) =>
         expr.kind === "ident" || expr.kind === "fieldAccess";
 
       lowerExpr(fcx, instrs, expr.lhs);
@@ -711,7 +723,7 @@ function lowerExpr(fcx: FuncContext, instrs: wasm.Instr[], expr: Expr) {
 
 function lowerExprBlockBody(
   fcx: FuncContext,
-  expr: ExprBlock & Expr
+  expr: ExprBlock<Typecked> & Expr<Typecked>
 ): wasm.Instr[] {
   fcx.currentBlockDepth++;
   const innerInstrs: wasm.Instr[] = [];
@@ -853,7 +865,7 @@ function todo(msg: string): never {
 }
 
 // Make the program runnable using wasi-preview-1
-function addRt(cx: Context, ast: Ast) {
+function addRt(cx: Context, ast: Ast<Typecked>) {
   const { mod } = cx;
 
   const mainCall: wasm.Instr = { kind: "call", func: 9999999 };
diff --git a/src/parser.ts b/src/parser.ts
index be8ca1b..a646b75 100644
--- a/src/parser.ts
+++ b/src/parser.ts
@@ -4,7 +4,7 @@ import {
   Ast,
   BinaryKind,
   COMPARISON_KINDS,
-  DEFAULT_FOLDER,
+  mkDefaultFolder,
   Expr,
   ExprLoop,
   ExprStructLiteral,
@@ -12,7 +12,7 @@ import {
   Folder,
   FunctionArg,
   FunctionDef,
-  Identifier,
+  Ident,
   ImportDef,
   Item,
   LOGICAL_KINDS,
@@ -25,6 +25,8 @@ import {
   foldAst,
   superFoldExpr,
   superFoldItem,
+  Built,
+  Parsed,
 } from "./ast";
 import { CompilerError, Span, spanMerge } from "./error";
 import { BaseToken, Token, TokenIdent, TokenLitString } from "./lexer";
@@ -32,8 +34,8 @@ import { Ids } from "./utils";
 
 type Parser<T> = (t: Token[]) => [Token[], T];
 
-export function parse(packageName: string, t: Token[]): Ast {
-  const items: Item[] = [];
+export function parse(packageName: string, t: Token[]): Ast<Built> {
+  const items: Item<Parsed>[] = [];
 
   while (t.length > 0) {
     let item;
@@ -48,7 +50,7 @@ export function parse(packageName: string, t: Token[]): Ast {
   return ast;
 }
 
-function parseItem(t: Token[]): [Token[], Item] {
+function parseItem(t: Token[]): [Token[], Item<Parsed>] {
   let tok;
   [t, tok] = next(t);
   if (tok.kind === "function") {
@@ -62,7 +64,7 @@ function parseItem(t: Token[]): [Token[], Item] {
 
     [t] = expectNext(t, ";");
 
-    const def: FunctionDef = {
+    const def: FunctionDef<Parsed> = {
       ...sig,
       body,
     };
@@ -84,7 +86,7 @@ function parseItem(t: Token[]): [Token[], Item] {
     [t] = expectNext(t, "{");
 
     let fields;
-    [t, fields] = parseCommaSeparatedList<FieldDef>(t, "}", (t) => {
+    [t, fields] = parseCommaSeparatedList<FieldDef<Parsed>>(t, "}", (t) => {
       let name;
       [t, name] = expectNext<TokenIdent>(t, "identifier");
       [t] = expectNext(t, ":");
@@ -104,7 +106,7 @@ function parseItem(t: Token[]): [Token[], Item] {
 
     [t] = expectNext(t, ";");
 
-    const def: TypeDef = {
+    const def: TypeDef<Parsed> = {
       name: name.ident,
       fields,
     };
@@ -123,7 +125,7 @@ function parseItem(t: Token[]): [Token[], Item] {
 
     [t] = expectNext(t, ";");
 
-    const def: ImportDef = {
+    const def: ImportDef<Parsed> = {
       module: { kind: "str", value: module.value, span: module.span },
       func: { kind: "str", value: func.value, span: func.span },
       ...sig,
@@ -135,7 +137,7 @@ function parseItem(t: Token[]): [Token[], Item] {
     let name;
     [t, name] = expectNext<TokenIdent>(t, "identifier");
 
-    const node: ModItem = {
+    const node: ModItem<Parsed> = {
       name: name.ident,
       modKind: { kind: "extern" },
     };
@@ -149,7 +151,7 @@ function parseItem(t: Token[]): [Token[], Item] {
 
     [t] = expectNext(t, "(");
 
-    const contents: Item[] = [];
+    const contents: Item<Parsed>[] = [];
 
     while (next(t)[1].kind !== ")") {
       let item;
@@ -161,7 +163,7 @@ function parseItem(t: Token[]): [Token[], Item] {
     [t] = expectNext(t, ")");
     [t] = expectNext(t, ";");
 
-    const node: ModItem = {
+    const node: ModItem<Parsed> = {
       name: name.ident,
       modKind: { kind: "inline", contents },
     };
@@ -174,8 +176,8 @@ function parseItem(t: Token[]): [Token[], Item] {
 
 type FunctionSig = {
   name: string;
-  params: FunctionArg[];
-  returnType?: Type;
+  params: FunctionArg<Parsed>[];
+  returnType?: Type<Parsed>;
 };
 
 function parseFunctionSig(t: Token[]): [Token[], FunctionSig] {
@@ -184,7 +186,7 @@ function parseFunctionSig(t: Token[]): [Token[], FunctionSig] {
 
   [t] = expectNext(t, "(");
 
-  let params: FunctionArg[];
+  let params: FunctionArg<Parsed>[];
   [t, params] = parseCommaSeparatedList(t, ")", (t) => {
     let name;
     [t, name] = expectNext<TokenIdent & { span: Span }>(t, "identifier");
@@ -205,7 +207,7 @@ function parseFunctionSig(t: Token[]): [Token[], FunctionSig] {
   return [t, { name: name.ident, params, returnType }];
 }
 
-function parseExpr(t: Token[]): [Token[], Expr] {
+function parseExpr(t: Token[]): [Token[], Expr<Parsed>] {
   /*
   EXPR = ASSIGNMENT
 
@@ -237,16 +239,21 @@ function parseExpr(t: Token[]): [Token[], Expr] {
   return parseExprAssignment(t);
 }
 
-function mkBinaryExpr(lhs: Expr, rhs: Expr, span: Span, kind: string): Expr {
+function mkBinaryExpr(
+  lhs: Expr<Parsed>,
+  rhs: Expr<Parsed>,
+  span: Span,
+  kind: string
+): Expr<Parsed> {
   return { kind: "binary", binaryKind: kind as BinaryKind, lhs, rhs, span };
 }
 
 function mkParserExprBinary(
-  lower: Parser<Expr>,
+  lower: Parser<Expr<Parsed>>,
   kinds: string[],
   mkExpr = mkBinaryExpr
-): Parser<Expr> {
-  function parser(t: Token[]): [Token[], Expr] {
+): Parser<Expr<Parsed>> {
+  function parser(t: Token[]): [Token[], Expr<Parsed>] {
     let lhs;
     [t, lhs] = lower(t);
 
@@ -288,7 +295,7 @@ const parseExprAssignment = mkParserExprBinary(
   (lhs, rhs, span) => ({ kind: "assign", lhs, rhs, span })
 );
 
-function parseExprUnary(t: Token[]): [Token[], Expr] {
+function parseExprUnary(t: Token[]): [Token[], Expr<Parsed>] {
   const [, peak] = next(t);
   if (peak.kind in UNARY_KINDS) {
     let rhs;
@@ -306,8 +313,8 @@ function parseExprUnary(t: Token[]): [Token[], Expr] {
 
   return parseExprCall(t);
 }
-function parseExprCall(t: Token[]): [Token[], Expr] {
-  let lhs: Expr;
+function parseExprCall(t: Token[]): [Token[], Expr<Parsed>] {
+  let lhs: Expr<Parsed>;
   [t, lhs] = parseExprAtom(t);
 
   while (next(t)[1].kind === "(" || next(t)[1].kind === ".") {
@@ -343,13 +350,13 @@ function parseExprCall(t: Token[]): [Token[], Expr] {
   return [t, lhs];
 }
 
-function parseExprAtom(startT: Token[]): [Token[], Expr] {
+function parseExprAtom(startT: Token[]): [Token[], Expr<Parsed>] {
   // eslint-disable-next-line prefer-const
   let [t, tok] = next(startT);
   const span = tok.span;
 
   if (tok.kind === "(") {
-    let expr: Expr;
+    let expr: Expr<Parsed>;
     [t, expr] = parseExpr(t);
 
     // This could be a block or a tuple literal. We can only know after
@@ -447,7 +454,7 @@ function parseExprAtom(startT: Token[]): [Token[], Expr] {
     let rhs;
     [t, rhs] = parseExpr(t);
 
-    const nameIdent: Identifier = { name: name.ident, span: name.span };
+    const nameIdent: Ident = { name: name.ident, span: name.span };
 
     return [
       t,
@@ -493,11 +500,13 @@ function parseExprAtom(startT: Token[]): [Token[], Expr] {
   return [startT, { kind: "empty", span }];
 }
 
-function parseStructInit(t: Token[]): [Token[], ExprStructLiteral["fields"]] {
+function parseStructInit(
+  t: Token[]
+): [Token[], ExprStructLiteral<Parsed>["fields"]] {
   [t] = expectNext(t, "{");
 
   let fields;
-  [t, fields] = parseCommaSeparatedList<[Identifier, Expr]>(t, "}", (t) => {
+  [t, fields] = parseCommaSeparatedList<[Ident, Expr<Parsed>]>(t, "}", (t) => {
     let name;
     [t, name] = expectNext<TokenIdent>(t, "identifier");
     [t] = expectNext(t, ":");
@@ -510,7 +519,7 @@ function parseStructInit(t: Token[]): [Token[], ExprStructLiteral["fields"]] {
   return [t, fields];
 }
 
-function parseType(t: Token[]): [Token[], Type] {
+function parseType(t: Token[]): [Token[], Type<Parsed>] {
   let tok;
   [t, tok] = next(t);
   const span = tok.span;
@@ -649,22 +658,19 @@ function unexpectedToken(token: Token, expected: string): never {
   throw new CompilerError(`unexpected token, expected ${expected}`, token.span);
 }
 
-function validateAst(ast: Ast) {
+function validateAst(ast: Ast<Built>) {
   const seenItemIds = new Set();
 
-  const validator: Folder = {
-    ...DEFAULT_FOLDER,
-    ast() {
-      return ast;
-    },
-    itemInner(item: Item): Item {
+  const validator: Folder<Built, Built> = {
+    ...mkDefaultFolder(),
+    itemInner(item: Item<Built>): Item<Built> {
       if (seenItemIds.has(item.id)) {
         throw new Error(`duplicate item id: ${item.id} for ${item.node.name}`);
       }
       seenItemIds.add(item.id);
       return superFoldItem(item, this);
     },
-    expr(expr: Expr): Expr {
+    expr(expr: Expr<Built>): Expr<Built> {
       if (expr.kind === "block") {
         expr.exprs.forEach((inner) => {
           if (inner.kind === "let") {
@@ -680,7 +686,7 @@ function validateAst(ast: Ast) {
       } else if (expr.kind === "let") {
         throw new CompilerError("let is only allowed in blocks", expr.span);
       } else if (expr.kind === "binary") {
-        const checkPrecedence = (inner: Expr, side: string) => {
+        const checkPrecedence = (inner: Expr<Built>, side: string) => {
           if (inner.kind === "binary") {
             const ourClass = binaryExprPrecedenceClass(expr.binaryKind);
             const innerClass = binaryExprPrecedenceClass(inner.binaryKind);
@@ -702,40 +708,49 @@ function validateAst(ast: Ast) {
         return superFoldExpr(expr, this);
       }
     },
+    ident(ident) {
+      return ident;
+    },
+    type(type) {
+      return type;
+    },
   };
 
   foldAst(ast, validator);
 }
 
-function buildAst(packageName: string, rootItems: Item[]): Ast {
+function buildAst(packageName: string, rootItems: Item<Parsed>[]): Ast<Built> {
   const itemId = new Ids();
   const loopId = new Ids();
 
-  const ast: Ast = {
+  const ast: Ast<Built> = {
     rootItems,
     itemsById: new Map(),
     packageName,
   };
 
-  const assigner: Folder = {
-    ...DEFAULT_FOLDER,
-    ast() {
-      return ast;
-    },
-    itemInner(item: Item): Item {
+  const assigner: Folder<Parsed, Built> = {
+    ...mkDefaultFolder(),
+    itemInner(item: Item<Parsed>): Item<Built> {
       const id = itemId.next();
       ast.itemsById.set(id, item);
       return { ...superFoldItem(item, this), id };
     },
-    expr(expr: Expr): Expr {
+    expr(expr: Expr<Parsed>): Expr<Built> {
       if (expr.kind === "loop") {
         return {
-          ...(superFoldExpr(expr, this) as ExprLoop & Expr),
+          ...(superFoldExpr(expr, this) as ExprLoop<Built> & Expr<Built>),
           loopId: loopId.next(),
         };
       }
       return superFoldExpr(expr, this);
     },
+    ident(ident) {
+      return ident;
+    },
+    type(type) {
+      return type;
+    },
   };
 
   return foldAst(ast, assigner);
diff --git a/src/printer.ts b/src/printer.ts
index 321e3cd..bc7b531 100644
--- a/src/printer.ts
+++ b/src/printer.ts
@@ -1,8 +1,9 @@
 import {
+  AnyPhase,
   Ast,
   Expr,
   FunctionDef,
-  Identifier,
+  IdentWithRes,
   ImportDef,
   Item,
   ModItem,
@@ -14,7 +15,7 @@ import {
   tyIsUnit,
 } from "./ast";
 
-export function printAst(ast: Ast): string {
+export function printAst(ast: Ast<AnyPhase>): string {
   return ast.rootItems.map(printItem).join("\n");
 }
 
@@ -22,7 +23,7 @@ function printStringLiteral(lit: StringLiteral): string {
   return `"${lit.value.replace("\n", "\\n")}"`;
 }
 
-function printItem(item: Item): string {
+function printItem(item: Item<AnyPhase>): string {
   const id = `/*${item.id}*/ `;
 
   switch (item.kind) {
@@ -41,7 +42,7 @@ function printItem(item: Item): string {
   }
 }
 
-function printFunction(func: FunctionDef): string {
+function printFunction(func: FunctionDef<AnyPhase>): string {
   const args = func.params
     .map(({ name, type }) => `${name}: ${printType(type)}`)
     .join(", ");
@@ -49,7 +50,7 @@ function printFunction(func: FunctionDef): string {
   return `function ${func.name}(${args})${ret} = ${printExpr(func.body, 0)};`;
 }
 
-function printTypeDef(type: TypeDef): string {
+function printTypeDef(type: TypeDef<AnyPhase>): string {
   const fields = type.fields.map(
     ({ name, type }) => `${ind(1)}${name.name}: ${printType(type)},`
   );
@@ -60,7 +61,7 @@ function printTypeDef(type: TypeDef): string {
   return `type ${type.name} = ${fieldPart};`;
 }
 
-function printImportDef(def: ImportDef): string {
+function printImportDef(def: ImportDef<AnyPhase>): string {
   const args = def.params
     .map(({ name, type }) => `${name}: ${printType(type)}`)
     .join(", ");
@@ -71,7 +72,7 @@ function printImportDef(def: ImportDef): string {
   )}(${args})${ret};`;
 }
 
-function printMod(mod: ModItem): string {
+function printMod(mod: ModItem<AnyPhase>): string {
   switch (mod.modKind.kind) {
     case "inline":
       return `mod ${mod.name} (\n${mod.modKind.contents
@@ -82,7 +83,7 @@ function printMod(mod: ModItem): string {
   }
 }
 
-function printExpr(expr: Expr, indent: number): string {
+function printExpr(expr: Expr<AnyPhase>, indent: number): string {
   switch (expr.kind) {
     case "empty": {
       return "";
@@ -192,7 +193,7 @@ function printExpr(expr: Expr, indent: number): string {
   }
 }
 
-function printType(type: Type): string {
+function printType(type: Type<AnyPhase>): string {
   switch (type.kind) {
     case "ident":
       return printIdent(type.value);
@@ -217,8 +218,8 @@ function printRes(res: Resolution): string {
   }
 }
 
-function printIdent(ident: Identifier): string {
-  const res = ident.res ? printRes(ident.res) : "";
+function printIdent(ident: IdentWithRes<AnyPhase>): string {
+  const res = "res" in ident ? printRes(ident.res) : "";
   return `${ident.name}${res}`;
 }
 
diff --git a/src/resolve.ts b/src/resolve.ts
index 5a38348..7289d04 100644
--- a/src/resolve.ts
+++ b/src/resolve.ts
@@ -1,18 +1,21 @@
 import {
   Ast,
   BUILTINS,
+  Built,
   BuiltinName,
-  DEFAULT_FOLDER,
   Expr,
   Folder,
-  Identifier,
+  Ident,
   Item,
   ItemId,
   LocalInfo,
   ModItem,
   Resolution,
+  Resolved,
+  mkDefaultFolder,
   superFoldExpr,
   superFoldItem,
+  superFoldType,
 } from "./ast";
 import { CompilerError, spanMerge, todo } from "./error";
 import { unwrap } from "./utils";
@@ -20,13 +23,14 @@ import { unwrap } from "./utils";
 const BUILTIN_SET = new Set<string>(BUILTINS);
 
 type Context = {
-  ast: Ast;
+  ast: Ast<Built>;
   modContentsCache: Map<ItemId, Map<string, ItemId>>;
+  newItemsById: Map<ItemId, Item<Resolved>>;
 };
 
 function resolveModItem(
   cx: Context,
-  mod: ModItem,
+  mod: ModItem<Built>,
   modId: ItemId,
   name: string
 ): ItemId | undefined {
@@ -49,18 +53,26 @@ function resolveModItem(
   }
 }
 
-export function resolve(ast: Ast): Ast {
-  const cx: Context = { ast, modContentsCache: new Map() };
+export function resolve(ast: Ast<Built>): Ast<Resolved> {
+  const cx: Context = {
+    ast,
+    modContentsCache: new Map(),
+    newItemsById: new Map(),
+  };
 
   const rootItems = resolveModule(cx, [ast.packageName], ast.rootItems);
-  return { ...ast, rootItems };
+  return {
+    itemsById: cx.newItemsById,
+    rootItems,
+    packageName: ast.packageName,
+  };
 }
 
 function resolveModule(
   cx: Context,
   modName: string[],
-  contents: Item[]
-): Item[] {
+  contents: Item<Built>[]
+): Item<Resolved>[] {
   const items = new Map<string, number>();
 
   contents.forEach((item) => {
@@ -85,7 +97,7 @@ function resolveModule(
     }
   };
 
-  const resolveIdent = (ident: Identifier): Resolution => {
+  const resolveIdent = (ident: Ident): Resolution => {
     const lastIdx = scopes.length - 1;
     for (let i = lastIdx; i >= 0; i--) {
       const candidate = scopes[i];
@@ -115,11 +127,8 @@ function resolveModule(
 
   const blockLocals: LocalInfo[][] = [];
 
-  const resolver: Folder = {
-    ...DEFAULT_FOLDER,
-    ast() {
-      return cx.ast;
-    },
+  const resolver: Folder<Built, Resolved> = {
+    ...mkDefaultFolder(),
     itemInner(item) {
       const defPath = [...modName, item.node.name];
 
@@ -178,7 +187,9 @@ function resolveModule(
           const prevScopeLength = scopes.length;
           blockLocals.push([]);
 
-          const exprs = expr.exprs.map<Expr>((inner) => this.expr(inner));
+          const exprs = expr.exprs.map<Expr<Resolved>>((inner) =>
+            this.expr(inner)
+          );
 
           scopes.length = prevScopeLength;
           const locals = blockLocals.pop();
@@ -263,6 +274,10 @@ function resolveModule(
       const res = resolveIdent(ident);
       return { name: ident.name, span: ident.span, res };
     },
+    type(type) {
+      return superFoldType(type, this);
+    },
+    newItemsById: cx.newItemsById,
   };
 
   return contents.map((item) => resolver.item(item));
diff --git a/src/typeck.ts b/src/typeck.ts
index 9f9855f..7970f1f 100644
--- a/src/typeck.ts
+++ b/src/typeck.ts
@@ -2,19 +2,21 @@ import {
   Ast,
   BuiltinName,
   COMPARISON_KINDS,
-  DEFAULT_FOLDER,
+  mkDefaultFolder,
   EQUALITY_KINDS,
   Expr,
   ExprBinary,
   ExprUnary,
   foldAst,
   Folder,
-  Identifier,
+  Ident,
+  IdentWithRes,
   ItemId,
   LOGICAL_KINDS,
   LoopId,
   ModItemKind,
   Resolution,
+  Resolved,
   Ty,
   TY_BOOL,
   TY_I32,
@@ -25,7 +27,9 @@ import {
   TyFn,
   tyIsUnit,
   Type,
+  Typecked,
   TyStruct,
+  Item,
 } from "./ast";
 import { CompilerError, Span } from "./error";
 import { printTy } from "./printer";
@@ -84,8 +88,8 @@ function typeOfBuiltinValue(name: BuiltinName, span: Span): Ty {
 
 // TODO: Cleanup, maybe get the ident switch into this function because typeOfItem is unused.
 function lowerAstTyBase(
-  type: Type,
-  lowerIdentTy: (ident: Identifier) => Ty,
+  type: Type<Resolved>,
+  lowerIdentTy: (ident: IdentWithRes<Resolved>) => Ty,
   typeOfItem: (index: number, cause: Span) => Ty
 ): Ty {
   switch (type.kind) {
@@ -112,7 +116,7 @@ function lowerAstTyBase(
   }
 }
 
-export function typeck(ast: Ast): Ast {
+export function typeck(ast: Ast<Resolved>): Ast<Typecked> {
   const itemTys = new Map<number, Ty | null>();
   function typeOfItem(index: ItemId, cause: Span): Ty {
     const item = unwrap(ast.itemsById.get(index));
@@ -165,11 +169,11 @@ export function typeck(ast: Ast): Ast {
     }
   }
 
-  function lowerAstTy(type: Type): Ty {
+  function lowerAstTy(type: Type<Resolved>): Ty {
     return lowerAstTyBase(
       type,
       (ident) => {
-        const res = ident.res!;
+        const res = ident.res;
         switch (res.kind) {
           case "local": {
             throw new Error("Item type cannot refer to local variable");
@@ -186,12 +190,9 @@ export function typeck(ast: Ast): Ast {
     );
   }
 
-  const checker: Folder = {
-    ...DEFAULT_FOLDER,
-    ast() {
-      return ast;
-    },
-    itemInner(item) {
+  const checker: Folder<Resolved, Typecked> = {
+    ...mkDefaultFolder(),
+    itemInner(item: Item<Resolved>): Item<Typecked> {
       switch (item.kind) {
         case "function": {
           const fnTy = typeOfItem(item.id, item.span) as TyFn;
@@ -298,7 +299,7 @@ export function typeck(ast: Ast): Ast {
         case "mod": {
           switch (item.node.modKind.kind) {
             case "inline": {
-              const modKind: ModItemKind = {
+              const modKind: ModItemKind<Typecked> = {
                 kind: "inline",
                 contents: item.node.modKind.contents.map((item) =>
                   this.item(item)
@@ -323,6 +324,15 @@ export function typeck(ast: Ast): Ast {
         }
       }
     },
+    expr(_expr) {
+      throw new Error("expressions need to be handled in checkBody");
+    },
+    ident(ident) {
+      return ident;
+    },
+    type(_type) {
+      throw new Error("all types should be typechecked manually");
+    },
   };
 
   const typecked = foldAst(ast, checker);
@@ -518,10 +528,10 @@ export class InferContext {
 }
 
 export function checkBody(
-  body: Expr,
+  body: Expr<Resolved>,
   fnTy: TyFn,
   typeOfItem: (index: number, cause: Span) => Ty
-): Expr {
+): Expr<Typecked> {
   const localTys = [...fnTy.params];
   const loopState: { hasBreak: boolean; loopId: LoopId }[] = [];
 
@@ -541,11 +551,11 @@ export function checkBody(
     }
   }
 
-  function lowerAstTy(type: Type): Ty {
+  function lowerAstTy(type: Type<Resolved>): Ty {
     return lowerAstTyBase(
       type,
       (ident) => {
-        const res = ident.res!;
+        const res = ident.res;
         switch (res.kind) {
           case "local": {
             const idx = localTys.length - 1 - res.index;
@@ -562,8 +572,8 @@ export function checkBody(
     );
   }
 
-  const checker: Folder = {
-    ...DEFAULT_FOLDER,
+  const checker: Folder<Resolved, Typecked> = {
+    ...mkDefaultFolder(),
     expr(expr) {
       switch (expr.kind) {
         case "empty": {
@@ -584,7 +594,7 @@ export function checkBody(
 
           expr.local!.ty = bindingTy;
 
-          const type: Type | undefined = loweredBindingTy && {
+          const type: Type<Typecked> | undefined = loweredBindingTy && {
             ...expr.type!,
             ty: loweredBindingTy,
           };
@@ -606,7 +616,7 @@ export function checkBody(
 
           switch (lhs.kind) {
             case "ident":
-              if (lhs.value.res!.kind !== "local") {
+              if (lhs.value.res.kind !== "local") {
                 throw new CompilerError("cannot assign to items", expr.span);
               }
               break;
@@ -664,7 +674,7 @@ export function checkBody(
           return { ...expr, ty };
         }
         case "ident": {
-          const ty = typeOf(expr.value.res!, expr.value.span);
+          const ty = typeOf(expr.value.res, expr.value.span);
 
           return { ...expr, ty };
         }
@@ -842,12 +852,11 @@ export function checkBody(
           };
         }
         case "structLiteral": {
-          const fields = expr.fields.map<[Identifier, Expr]>(([name, expr]) => [
-            name,
-            this.expr(expr),
-          ]);
+          const fields = expr.fields.map<[Ident, Expr<Typecked>]>(
+            ([name, expr]) => [name, this.expr(expr)]
+          );
 
-          const structTy = typeOf(expr.name.res!, expr.name.span);
+          const structTy = typeOf(expr.name.res, expr.name.span);
 
           if (structTy.kind !== "struct") {
             throw new CompilerError(
@@ -897,6 +906,15 @@ export function checkBody(
         }
       }
     },
+    itemInner(_item) {
+      throw new Error("cannot deal with items inside body");
+    },
+    ident(ident) {
+      return ident;
+    },
+    type(_type) {
+      throw new Error("all types in the body should be handled elsewhere");
+    },
   };
 
   const checked = checker.expr(body);
@@ -912,8 +930,8 @@ export function checkBody(
     return resTy;
   };
 
-  const resolver: Folder = {
-    ...DEFAULT_FOLDER,
+  const resolver: Folder<Typecked, Typecked> = {
+    ...mkDefaultFolder(),
     expr(expr) {
       const ty = resolveTy(expr.ty!, expr.span);
 
@@ -925,6 +943,13 @@ export function checkBody(
 
       return { ...expr, ty };
     },
+    type(type) {
+      const ty = resolveTy(type.ty!, type.span);
+      return { ...type, ty };
+    },
+    ident(ident) {
+      return ident;
+    },
   };
 
   const resolved = resolver.expr(checked);
@@ -932,7 +957,9 @@ export function checkBody(
   return resolved;
 }
 
-function checkBinary(expr: Expr & ExprBinary): Expr {
+function checkBinary(
+  expr: Expr<Typecked> & ExprBinary<Typecked>
+): Expr<Typecked> {
   const lhsTy = expr.lhs.ty!;
   const rhsTy = expr.rhs.ty!;
 
@@ -973,7 +1000,9 @@ function checkBinary(expr: Expr & ExprBinary): Expr {
   );
 }
 
-function checkUnary(expr: Expr & ExprUnary): Expr {
+function checkUnary(
+  expr: Expr<Typecked> & ExprUnary<Typecked>
+): Expr<Typecked> {
   const rhsTy = expr.rhs.ty!;
 
   if (