Extract TypedBlockNode::behaviour()

3 files changed, 15 insertions, 18 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
index a44caaa3..b5887627 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedBlockNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedBlockNode.java
@@ -1,6 +1,7 @@
 package org.shedlang.compiler.ast.typed;
 
 import org.shedlang.compiler.sources.Source;
+import org.shedlang.compiler.util.Lists;
 
 import java.util.List;
 
@@ -8,4 +9,11 @@ public record TypedBlockNode(
     List<? extends TypedFunctionStatementNode> statements,
     Source source
 ) implements TypedNode {
+    public TypedFunctionStatementBehaviour behaviour() {
+        var lastStatement = Lists.last(statements);
+
+        return lastStatement
+            .map(statement -> statement.behaviour())
+            .orElse(TypedFunctionStatementBehaviour.CONTINUE);
+    }
 }
diff --git a/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java b/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
index b4c8f9f0..70f38cf6 100644
--- a/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
+++ b/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
@@ -4,7 +4,6 @@ import org.pcollections.PVector;
 import org.shedlang.compiler.ShedBuiltins;
 import org.shedlang.compiler.ast.typed.*;
 import org.shedlang.compiler.stackir.instructions.*;
-import org.shedlang.compiler.util.Lists;
 import org.shedlang.compiler.util.P;
 
 import java.util.stream.Collectors;
@@ -101,14 +100,7 @@ public class StackCompiler {
             .flatMap(statement -> compileFunctionStatement(statement).stream())
             .collect(Collectors.toList());
 
-        var lastStatement = Lists.last(node.body().statements());
-        // TODO: remove duplication with type checker (default to none behaviour).
-        // Extract block node?
-        var lastStatementBehaviour = lastStatement
-            .map(statement -> statement.behaviour())
-            .orElse(TypedFunctionStatementBehaviour.CONTINUE);
-
-        switch (lastStatementBehaviour) {
+        switch (node.body().behaviour()) {
             case TypedFunctionStatementBehaviour.BlockReturn blockReturn -> {
                 body.add(new ReturnInstruction(node.source()));
             }
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 63488029..90276186 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -7,7 +7,6 @@ import org.shedlang.compiler.typechecker.errors.*;
 import org.shedlang.compiler.typechecker.typesets.TypeSetTypeLevelValue;
 import org.shedlang.compiler.typechecker.typesets.TypeSets;
 import org.shedlang.compiler.types.*;
-import org.shedlang.compiler.util.Lists;
 import org.shedlang.compiler.util.PCollectors;
 
 import java.util.*;
@@ -206,8 +205,11 @@ public class TypeChecker {
         UntypedBlockNode untyped,
         TypeCheckerContext context
     ) {
-        return context.mapWithContext(untyped.statements(), TypeChecker::typeCheckFunctionStatement)
-            .map(statements -> new TypedBlockNode(statements, untyped.source()));
+        return context.mapWithContext(
+            untyped.statements(),
+            TypeChecker::typeCheckFunctionStatement
+        )
+        .map(statements -> new TypedBlockNode(statements, untyped.source()));
     }
 
     private static TypeCheckerResult<TypedFunctionStatementNode> typeCheckBlockReturn(
@@ -372,12 +374,7 @@ public class TypeChecker {
         var typedBody = bodyResult.value();
         bodyContext = bodyResult.context();
 
-        var lastStatement = Lists.last(typedBody.statements());
-        var lastStatementBehaviour = lastStatement
-            .map(statement -> statement.behaviour())
-            .orElse(TypedFunctionStatementBehaviour.CONTINUE);
-
-        switch (lastStatementBehaviour) {
+        switch (typedBody.behaviour()) {
             case TypedFunctionStatementBehaviour.BlockReturn blockReturn -> {
                 if (!isSubType(blockReturn.type(), typedReturnType.value(), context)) {
                     throw new UnexpectedTypeError(