Introduce TypedFunctionBodyNode

4 files changed, 32 insertions, 17 deletions

diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionBodyNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionBodyNode.java
new file mode 100644
index 00000000..8f617ca8
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionBodyNode.java
@@ -0,0 +1,19 @@
+package org.shedlang.compiler.ast.typed;
+
+import org.shedlang.compiler.sources.Source;
+
+import java.util.List;
+
+public record TypedFunctionBodyNode(
+    FreeVariables freeVariables,
+    TypedBlockNode block,
+    Source source
+) implements TypedNode {
+    public List<String> boundVariableNames() {
+        return block.boundVariableNames();
+    }
+
+    public TypedFunctionStatementBehaviour behaviour() {
+        return block.behaviour();
+    }
+}
diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionDefinitionNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionDefinitionNode.java
index e03018c6..c98a7b1e 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionDefinitionNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionDefinitionNode.java
@@ -10,8 +10,7 @@ import org.shedlang.compiler.types.Type;
 public record TypedFunctionDefinitionNode(
     String name,
     TypedFunctionSignatureNode signature,
-    FreeVariables freeVariables,
-    TypedBlockNode body,
+    TypedFunctionBodyNode body,
     Source source
 ) implements TypedNamespaceStatementNode {
     public static Builder builder() {
@@ -35,8 +34,11 @@ public record TypedFunctionDefinitionNode(
             return new TypedFunctionDefinitionNode(
                 name,
                 signature.build(),
-                freeVariables.build(),
-                new TypedBlockNode(body, boundVariableNames, Sources.NULL),
+                new TypedFunctionBodyNode(
+                    freeVariables.build(),
+                    new TypedBlockNode(body, boundVariableNames, Sources.NULL),
+                    Sources.NULL
+                ),
                 Sources.NULL
             );
         }
diff --git a/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java b/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
index 33d84149..2bd1229a 100644
--- a/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
+++ b/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
@@ -203,7 +203,7 @@ public class StackCompiler {
             ))
             .toList();
 
-        context = compileBlockWithYield(node.body(), context);
+        context = compileBlockWithYield(node.body().block(), context);
 
         switch (node.body().behaviour()) {
             case TypedFunctionStatementBehaviour.BlockReturn blockReturn -> {
@@ -221,7 +221,7 @@ public class StackCompiler {
         context = context.exitBlock();
 
         var freeVariables = FreeVariables.builder();
-        for (var shedFreeVariable : node.freeVariables()) {
+        for (var shedFreeVariable : node.body().freeVariables()) {
             var freeVariableLocalName = context.shedVariableToLocal(shedFreeVariable.name());
             for (var freeVariableSource : shedFreeVariable.referenceSources()) {
                 freeVariables = freeVariables.plus(
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index fe1f41a9..d51674a8 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -464,13 +464,7 @@ public class TypeChecker {
         }
     }
 
-    private record TypedFunctionBody(
-        FreeVariables freeVariables,
-        TypedBlockNode block
-    ) {
-    }
-
-    private static TypeCheckerResult<TypedFunctionBody> typeCheckFunctionBody(
+    private static TypeCheckerResult<TypedFunctionBodyNode> typeCheckFunctionBody(
         UntypedBlockNode untyped,
         TypedFunctionSignatureNode typedSignature,
         TypeCheckerContext context
@@ -532,9 +526,10 @@ public class TypeChecker {
             }
         }
 
-        var typed = new TypedFunctionBody(
+        var typed = new TypedFunctionBodyNode(
             freeVariables,
-            typedBody
+            typedBody,
+            typedBody.source()
         );
 
         return context.toResult(typed);
@@ -584,8 +579,7 @@ public class TypeChecker {
         var typed = new TypedFunctionDefinitionNode(
             untyped.name(),
             typedSignature,
-            typedBody.freeVariables,
-            typedBody.block,
+            typedBody,
             untyped.source()
         );