Extract typeCheckFunctionBody()

1 files changed, 59 insertions, 31 deletions

diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 6bf20edb..1c763bba 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -464,38 +464,17 @@ public class TypeChecker {
         }
     }
 
-    private static TypeCheckSteps<TypedFunctionDefinitionNode> typeCheckStepsFunctionDefinition(
-        UntypedFunctionDefinitionNode untyped
+    private record TypedFunctionBody(
+        FreeVariables freeVariables,
+        TypedBlockNode block
     ) {
-        return new TypeCheckSteps<TypedFunctionDefinitionNode>() {
-            @Override
-            public TypeCheckerContext declareFunctions(TypeCheckerContext context) {
-                return declareFunctionDefinition(untyped, context);
-            }
-
-            @Override
-            public TypeCheckerResult<TypedFunctionDefinitionNode> finish(TypeCheckerContext context) {
-                return defineFunctionDefinition(untyped, context);
-            }
-        };
     }
 
-    private static TypeCheckerContext declareFunctionDefinition(
-        UntypedFunctionDefinitionNode untyped,
-        TypeCheckerContext context
-    ) {
-        var typedSignature = typeCheckFunctionSignature(untyped.signature(), context);
-
-        return context
-            .define(untyped.name(), typedSignature.type(), untyped.source())
-            .addLexicalState(untyped.name(), typedSignature);
-    }
-
-    private static TypeCheckerResult<TypedFunctionDefinitionNode> defineFunctionDefinition(
-        UntypedFunctionDefinitionNode untyped,
+    private static TypeCheckerResult<TypedFunctionBody> typeCheckFunctionBody(
+        UntypedBlockNode untyped,
+        TypedFunctionSignatureNode typedSignature,
         TypeCheckerContext context
     ) {
-        var typedSignature = context.getLexicalState(TypedFunctionSignatureNode.class, untyped.name());
         var functionType = typedSignature.type();
 
         var bodyEnvironment = EnvironmentBuilder.EMPTY;
@@ -510,7 +489,7 @@ public class TypeChecker {
 
         // TODO: pass interface to collect variables?
         var declaredVariables = DeclaredVariables.forFunctionStatements(
-            untyped.body().statements(),
+            untyped.statements(),
             P.vector()
         );
         bodyEnvironment = bodyEnvironment.addUndefined(declaredVariables);
@@ -521,7 +500,7 @@ public class TypeChecker {
             functionType.returnType()
         );
 
-        var bodyResult = typeCheckBlock(untyped.body(), bodyEnvironment.names(), context);
+        var bodyResult = typeCheckBlock(untyped, bodyEnvironment.names(), context);
         var typedBody = bodyResult.value();
         context = bodyResult.context();
 
@@ -553,11 +532,60 @@ public class TypeChecker {
             }
         }
 
+        var typed = new TypedFunctionBody(
+            freeVariables,
+            typedBody
+        );
+
+        return context.toResult(typed);
+    }
+
+    private static TypeCheckSteps<TypedFunctionDefinitionNode> typeCheckStepsFunctionDefinition(
+        UntypedFunctionDefinitionNode untyped
+    ) {
+        return new TypeCheckSteps<TypedFunctionDefinitionNode>() {
+            @Override
+            public TypeCheckerContext declareFunctions(TypeCheckerContext context) {
+                return declareFunctionDefinition(untyped, context);
+            }
+
+            @Override
+            public TypeCheckerResult<TypedFunctionDefinitionNode> finish(TypeCheckerContext context) {
+                return defineFunctionDefinition(untyped, context);
+            }
+        };
+    }
+
+    private static TypeCheckerContext declareFunctionDefinition(
+        UntypedFunctionDefinitionNode untyped,
+        TypeCheckerContext context
+    ) {
+        var typedSignature = typeCheckFunctionSignature(untyped.signature(), context);
+
+        return context
+            .define(untyped.name(), typedSignature.type(), untyped.source())
+            .addLexicalState(untyped.name(), typedSignature);
+    }
+
+    private static TypeCheckerResult<TypedFunctionDefinitionNode> defineFunctionDefinition(
+        UntypedFunctionDefinitionNode untyped,
+        TypeCheckerContext context
+    ) {
+        var typedSignature = context.getLexicalState(TypedFunctionSignatureNode.class, untyped.name());
+
+        var typedBodyResult = typeCheckFunctionBody(
+            untyped.body(),
+            typedSignature,
+            context
+        );
+        var typedBody = typedBodyResult.value();
+        context = typedBodyResult.context();
+
         var typed = new TypedFunctionDefinitionNode(
             untyped.name(),
             typedSignature,
-            freeVariables,
-            typedBody,
+            typedBody.freeVariables,
+            typedBody.block,
             untyped.source()
         );