From 7024b8790de889424ed14ec1628b617b6f1fa31e Mon Sep 17 00:00:00 2001
From: Sam Vervaeck <samvv@pm.me>
Date: Tue, 23 May 2023 21:51:44 +0200
Subject: [PATCH] Fix return-statements not being type-checked

---
 src/Checker.cc | 17 ++++++++---------
 1 file changed, 8 insertions(+), 9 deletions(-)

diff --git a/src/Checker.cc b/src/Checker.cc
index b34b64b6d..86a6c8995 100644
--- a/src/Checker.cc
+++ b/src/Checker.cc
@@ -393,8 +393,9 @@ namespace bolt {
       case NodeKind::LetDeclaration:
       {
         auto Decl = static_cast<LetDeclaration*>(N);
+        bool HasContext = !Decl->Params.empty();
 
-        if (!Decl->Params.empty()) {
+        if (HasContext) {
           Contexts.push_back(Decl->Ctx);
         }
 
@@ -419,7 +420,8 @@ namespace bolt {
             case NodeKind::LetBlockBody:
             {
               auto Block = static_cast<LetBlockBody*>(Decl->Body);
-              RetType = createTypeVar();
+              ZEN_ASSERT(HasContext);
+              RetType = Decl->Ctx->ReturnType;
               for (auto Element: Block->Elements) {
                 infer(Element);
               }
@@ -432,16 +434,13 @@ namespace bolt {
           RetType = createTypeVar();
         }
 
-        if (ParamTypes.empty()) {
-          // Declaration is a plain (typed) variable
-          addConstraint(new CEqual { Decl->Ty, RetType, N });
-        } else {
+        if (HasContext) {
           // Declaration is a function
           addConstraint(new CEqual { Decl->Ty, new TArrow(ParamTypes, RetType), N });
-        }
-
-        if (!Decl->Params.empty()) {
           Contexts.pop_back();
+        } else {
+          // Declaration is a plain (typed) variable
+          addConstraint(new CEqual { Decl->Ty, RetType, N });
         }
 
         break;