From 02412c2a0940717e1ce1a183c4f5585624796548 Mon Sep 17 00:00:00 2001
From: Nilstrieb <48135649+Nilstrieb@users.noreply.github.com>
Date: Mon, 24 Jul 2023 20:30:06 +0200
Subject: [PATCH] Improve type inference

Fixes the union-find inference by properly unioning the roots
and also resolves all variables after typeck.
---
 src/index.ts  |  2 +-
 src/typeck.ts | 34 +++++++++++++++++++++++-----------
 2 files changed, 24 insertions(+), 12 deletions(-)

diff --git a/src/index.ts b/src/index.ts
index 09cd0c9..f4edec4 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -7,7 +7,7 @@ import { typeck } from "./typeck";
 
 const input = `
 function main() = (
-  let true = false in print(true)
+  let a = 0 in 0;
 );
 `;
 
diff --git a/src/typeck.ts b/src/typeck.ts
index 437dcb4..64b2fb2 100644
--- a/src/typeck.ts
+++ b/src/typeck.ts
@@ -252,19 +252,18 @@ export function checkBody(
    * before calling this.
    */
   function constrainVar(variable: number, ty: Ty) {
+    let root = variable;
+    let nextVar;
+    while ((nextVar = tyVars[root]).kind === "unified") {
+      root = nextVar.index;
+    }
+
     if (ty.kind === "var") {
       // Point the lhs to the rhs.
-      tyVars[variable] = { kind: "unified", index: ty.index };
+      tyVars[root] = { kind: "unified", index: ty.index };
+    } else {
+      tyVars[root] = { kind: "final", ty };
     }
-
-    let idx = variable;
-    let nextVar;
-    while ((nextVar = tyVars[idx]).kind === "unified") {
-      idx = nextVar.index;
-    }
-
-    const root = idx;
-    tyVars[root] = { kind: "final", ty };
   }
 
   function resolveIfPossible(ty: Ty): Ty {
@@ -472,7 +471,20 @@ export function checkBody(
 
   assign(fnTy.returnTy, checked.ty!, body.span);
 
-  return checked;
+  const resolver: Folder = {
+    ...DEFAULT_FOLDER,
+    expr(expr) {
+      const ty = resolveIfPossible(expr.ty!);
+      if (!ty) {
+        throw new CompilerError("cannot infer type", expr.span);
+      }
+      return { ...expr, ty };
+    },
+  };
+
+  const resolved = resolver.expr(checked);
+
+  return resolved;
 }
 
 function checkBinary(expr: Expr & ExprBinary): Expr {