From 5cd4cc3e845f7cf2379b0d2bd409b9e9eeaf30f4 Mon Sep 17 00:00:00 2001
From: Sam Vervaeck <samvv@pm.me>
Date: Thu, 11 Jul 2024 21:05:14 +0200
Subject: [PATCH] Improve support for typechecking mutually recursive functions

---
 src/Checker.cc | 25 +++++++++++++++++++------
 1 file changed, 19 insertions(+), 6 deletions(-)

diff --git a/src/Checker.cc b/src/Checker.cc
index 83d96d3f5..a3f2aead0 100644
--- a/src/Checker.cc
+++ b/src/Checker.cc
@@ -394,9 +394,8 @@ ConstraintSet Checker::inferFunctionDeclaration(TypeEnv& Env, FunctionDeclaratio
     Out.push_back(new CTypesEqual(RetTy, BodyTy, Body));
   }
 
-  // Env.add(D->getNameAsString(), Ty, SymbolKind::Var);
-
-  D->setType(Ty);
+  // inferMany() will have set the type of the node to a fresh type variable.
+  Out.push_back(new CTypesEqual { D->getType(), Ty, D });
 
   return Out;
 }
@@ -537,9 +536,20 @@ ConstraintSet Checker::inferMany(TypeEnv& Env, std::vector<Node*>& Elements, Typ
     }
   }
 
-  for (auto Nodes: zen::toposort(G)) {
+  for (auto Mutual: zen::toposort(G)) {
+
     ConstraintSet Out;
-    for (auto N: Nodes) {
+
+    for (auto N: Mutual) {
+      if (isa<FunctionDeclaration>(N)) {
+        auto Func = static_cast<FunctionDeclaration*>(N);
+        Type* Ty = createTVar();
+        Func->setType(Ty);
+        Env.add(Func->getNameAsString(), Ty, SymbolKind::Var);
+      }
+    }
+
+    for (auto N: Mutual) {
       if (isa<FunctionDeclaration>(N)) {
         mergeTo(Out, inferFunctionDeclaration(Env, static_cast<FunctionDeclaration*>(N)));
       } else if (isa<VariableDeclaration>(N)) {
@@ -548,8 +558,10 @@ ConstraintSet Checker::inferMany(TypeEnv& Env, std::vector<Node*>& Elements, Typ
         ZEN_UNREACHABLE
       }
     }
+
     solve(Out);
-    for (auto N: Nodes) {
+
+    for (auto N: Mutual) {
       if (isa<FunctionDeclaration>(N)) {
         auto Func = static_cast<FunctionDeclaration*>(N);
         auto Unbound = getUnbound(Env, Func->getType());
@@ -560,6 +572,7 @@ ConstraintSet Checker::inferMany(TypeEnv& Env, std::vector<Node*>& Elements, Typ
         );
       }
     }
+
   }
 
   ConstraintSet Out;