Infer return type of function expressions properly

1 files changed, 7 insertions, 7 deletions

diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 15623e98..17cf42a3 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -627,13 +627,6 @@ public class TypeChecker {
     ) {
         var typedSignature = typeCheckFunctionSignature(untyped.signature(), context);
 
-        var type = new FunctionType(
-            functionSignatureToPositionalParams(typedSignature),
-            functionSignatureToNamedParams(typedSignature),
-            functionSignatureToEffect(typedSignature),
-            functionSignatureToReturnType(typedSignature).orElse(Types.UNIT)
-        );
-
         var typedBodyResult = typeCheckFunctionBody(
             untyped.body(),
             typedSignature,
@@ -642,6 +635,13 @@ public class TypeChecker {
         var typedBody = typedBodyResult.value();
         context = typedBodyResult.context();
 
+        var type = new FunctionType(
+            functionSignatureToPositionalParams(typedSignature),
+            functionSignatureToNamedParams(typedSignature),
+            functionSignatureToEffect(typedSignature),
+            typedBody.type()
+        );
+
         var typed = new TypedFunctionExpressionNode(
             typedSignature,
             typedBody,