Add behaviour to untyped function statements

7 files changed, 28 insertions, 1 deletions

diff --git a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedBlockReturnNode.java b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedBlockReturnNode.java
index 55d6cca0..9afbab90 100644
--- a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedBlockReturnNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedBlockReturnNode.java
@@ -6,4 +6,8 @@ public record UntypedBlockReturnNode(
     UntypedExpressionNode expression,
     Source source
 ) implements UntypedFunctionStatementNode {
+    @Override
+    public UntypedFunctionStatementBehaviour behaviour() {
+        return UntypedFunctionStatementBehaviour.BLOCK_RETURN;
+    }
 }
diff --git a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedExpressionStatementNode.java b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedExpressionStatementNode.java
index dadc4be6..34b275c6 100644
--- a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedExpressionStatementNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedExpressionStatementNode.java
@@ -6,4 +6,8 @@ public record UntypedExpressionStatementNode(
     UntypedExpressionNode expression,
     Source source
 ) implements UntypedFunctionStatementNode {
+    @Override
+    public UntypedFunctionStatementBehaviour behaviour() {
+        return UntypedFunctionStatementBehaviour.NONE;
+    }
 }
diff --git a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionReturnNode.java b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionReturnNode.java
index c3f5663e..94eb7a65 100644
--- a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionReturnNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionReturnNode.java
@@ -6,4 +6,8 @@ public record UntypedFunctionReturnNode(
     UntypedExpressionNode expression,
     Source source
 ) implements UntypedFunctionStatementNode {
+    @Override
+    public UntypedFunctionStatementBehaviour behaviour() {
+        return UntypedFunctionStatementBehaviour.TERMINATE;
+    }
 }
diff --git a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionStatementBehaviour.java b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionStatementBehaviour.java
new file mode 100644
index 00000000..71cfef75
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionStatementBehaviour.java
@@ -0,0 +1,7 @@
+package org.shedlang.compiler.ast.untyped;
+
+public enum UntypedFunctionStatementBehaviour {
+    TERMINATE,
+    BLOCK_RETURN,
+    NONE,
+}
diff --git a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionStatementNode.java b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionStatementNode.java
index c3e0466d..62322251 100644
--- a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionStatementNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionStatementNode.java
@@ -1,4 +1,5 @@
 package org.shedlang.compiler.ast.untyped;
 
 public sealed interface UntypedFunctionStatementNode extends UntypedNode permits UntypedBlockReturnNode, UntypedExpressionStatementNode, UntypedFunctionReturnNode, UntypedValNode {
+    UntypedFunctionStatementBehaviour behaviour();
 }
diff --git a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedValNode.java b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedValNode.java
index 7f9926ea..13826cb7 100644
--- a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedValNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedValNode.java
@@ -7,4 +7,8 @@ public record UntypedValNode(
     UntypedExpressionNode expression,
     Source source
 ) implements UntypedFunctionStatementNode, UntypedNamespaceStatementNode {
+    @Override
+    public UntypedFunctionStatementBehaviour behaviour() {
+        return UntypedFunctionStatementBehaviour.NONE;
+    }
 }
diff --git a/src/main/java/org/shedlang/compiler/parser/Parser.java b/src/main/java/org/shedlang/compiler/parser/Parser.java
index bfb9c822..3243eb17 100644
--- a/src/main/java/org/shedlang/compiler/parser/Parser.java
+++ b/src/main/java/org/shedlang/compiler/parser/Parser.java
@@ -118,7 +118,10 @@ public class Parser {
 
             body.add(statement);
 
-            if (tokens.isNext(TokenType.SYMBOL_BRACE_CLOSE) || statement instanceof UntypedBlockReturnNode) {
+            if (
+                tokens.isNext(TokenType.SYMBOL_BRACE_CLOSE) ||
+                statement.behaviour() == UntypedFunctionStatementBehaviour.BLOCK_RETURN
+            ) {
                 tokens.skip(TokenType.SYMBOL_BRACE_CLOSE);
                 return body;
             }