Add TypedBlockNode

4 files changed, 26 insertions, 11 deletions

diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedBlockNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedBlockNode.java
new file mode 100644
index 00000000..a44caaa3
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedBlockNode.java
@@ -0,0 +1,11 @@
+package org.shedlang.compiler.ast.typed;
+
+import org.shedlang.compiler.sources.Source;
+
+import java.util.List;
+
+public record TypedBlockNode(
+    List<? extends TypedFunctionStatementNode> statements,
+    Source source
+) implements TypedNode {
+}
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 9867c726..3a156671 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionDefinitionNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionDefinitionNode.java
@@ -6,7 +6,6 @@ import org.shedlang.compiler.sources.Source;
 import org.shedlang.compiler.sources.Sources;
 import org.shedlang.compiler.types.*;
 
-import java.util.List;
 import java.util.Optional;
 
 public record TypedFunctionDefinitionNode(
@@ -15,7 +14,7 @@ public record TypedFunctionDefinitionNode(
     Optional<TypedTypeLevelExpressionNode<Effect>> effect,
     TypedTypeLevelExpressionNode<Type> returnType,
     FreeVariables freeVariables,
-    List<? extends TypedFunctionStatementNode> body,
+    TypedBlockNode body,
     FunctionType type,
     Source source
 ) implements TypedNamespaceStatementNode {
@@ -25,7 +24,7 @@ public record TypedFunctionDefinitionNode(
         Optional<TypedTypeLevelExpressionNode<Effect>> effect,
         TypedTypeLevelExpressionNode<Type> returnType,
         FreeVariables freeVariables,
-        List<? extends TypedFunctionStatementNode> body,
+        TypedBlockNode body,
         Source source
     ) {
         this(
@@ -77,7 +76,7 @@ public record TypedFunctionDefinitionNode(
                 effect,
                 returnType,
                 freeVariables.build(),
-                body,
+                new TypedBlockNode(body, 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 c4165b0a..b4c8f9f0 100644
--- a/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
+++ b/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
@@ -97,11 +97,11 @@ public class StackCompiler {
             .map(param -> new DefineFunctionInstruction.Param(param.name(), param.source()))
             .toList();
 
-        var body = node.body().stream()
+        var body = node.body().statements().stream()
             .flatMap(statement -> compileFunctionStatement(statement).stream())
             .collect(Collectors.toList());
 
-        var lastStatement = Lists.last(node.body());
+        var lastStatement = Lists.last(node.body().statements());
         // TODO: remove duplication with type checker (default to none behaviour).
         // Extract block node?
         var lastStatementBehaviour = lastStatement
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 60df3484..63488029 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -202,6 +202,14 @@ public class TypeChecker {
         throw new UnsupportedBinaryOperationError(operator, left, right, source);
     }
 
+    private static TypeCheckerResult<TypedBlockNode> typeCheckBlock(
+        UntypedBlockNode untyped,
+        TypeCheckerContext context
+    ) {
+        return context.mapWithContext(untyped.statements(), TypeChecker::typeCheckFunctionStatement)
+            .map(statements -> new TypedBlockNode(statements, untyped.source()));
+    }
+
     private static TypeCheckerResult<TypedFunctionStatementNode> typeCheckBlockReturn(
         UntypedBlockReturnNode untyped,
         TypeCheckerContext context
@@ -360,14 +368,11 @@ public class TypeChecker {
             typedReturnType.value()
         );
 
-        var bodyResult = bodyContext.mapWithContext(
-            untyped.body().statements(),
-            TypeChecker::typeCheckFunctionStatement
-        );
+        var bodyResult = typeCheckBlock(untyped.body(), bodyContext);
         var typedBody = bodyResult.value();
         bodyContext = bodyResult.context();
 
-        var lastStatement = Lists.last(typedBody);
+        var lastStatement = Lists.last(typedBody.statements());
         var lastStatementBehaviour = lastStatement
             .map(statement -> statement.behaviour())
             .orElse(TypedFunctionStatementBehaviour.CONTINUE);