Handle if expression statement

4 files changed, 27 insertions, 2 deletions

diff --git a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedBlockNode.java b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedBlockNode.java
index 2d686147..e69883a0 100644
--- a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedBlockNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedBlockNode.java
@@ -1,6 +1,7 @@
 package org.shedlang.compiler.ast.untyped;
 
 import org.shedlang.compiler.sources.Source;
+import org.shedlang.compiler.util.Lists;
 
 import java.util.List;
 
@@ -8,4 +9,10 @@ public record UntypedBlockNode(
     List<UntypedFunctionStatementNode> statements,
     Source source
 ) implements UntypedNode {
+    public UntypedFunctionStatementBehaviour behaviour() {
+        var lastStatement = Lists.last(statements);
+        return lastStatement
+            .map(statement -> statement.behaviour())
+            .orElse(UntypedFunctionStatementBehaviour.CONTINUE);
+    }
 }
diff --git a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedBranchingExpressionNode.java b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedBranchingExpressionNode.java
new file mode 100644
index 00000000..eab2f680
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedBranchingExpressionNode.java
@@ -0,0 +1,7 @@
+package org.shedlang.compiler.ast.untyped;
+
+import java.util.stream.Stream;
+
+public interface UntypedBranchingExpressionNode {
+    Stream<UntypedBlockNode> branchBodies();
+}
diff --git a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedIfNode.java b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedIfNode.java
index 96f6172a..c258257c 100644
--- a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedIfNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedIfNode.java
@@ -3,10 +3,18 @@ package org.shedlang.compiler.ast.untyped;
 import org.shedlang.compiler.sources.Source;
 
 import java.util.List;
+import java.util.stream.Stream;
 
 public record UntypedIfNode(
     List<UntypedConditionalBranchNode> conditionalBranches,
     UntypedBlockNode elseBody,
     Source source
-) implements UntypedExpressionNode {
+) implements UntypedBranchingExpressionNode, UntypedExpressionNode {
+    @Override
+    public Stream<UntypedBlockNode> branchBodies() {
+        return Stream.concat(
+            conditionalBranches.stream().map(branch -> branch.body()),
+            Stream.of(elseBody)
+        );
+    }
 }
diff --git a/src/main/java/org/shedlang/compiler/parser/Parser.java b/src/main/java/org/shedlang/compiler/parser/Parser.java
index 184c3fdd..bcd02807 100644
--- a/src/main/java/org/shedlang/compiler/parser/Parser.java
+++ b/src/main/java/org/shedlang/compiler/parser/Parser.java
@@ -205,7 +205,10 @@ public class Parser {
 
         var expression = parseTopLevelExpression(tokens);
 
-        var isReturn = !tokens.trySkip(TokenType.SYMBOL_SEMICOLON);
+        var isReturn = expression instanceof UntypedBranchingExpressionNode branchingExpression
+            // TODO: Check for consistency
+            ? branchingExpression.branchBodies().findFirst().orElseThrow().behaviour() == UntypedFunctionStatementBehaviour.BLOCK_RETURN
+            : !tokens.trySkip(TokenType.SYMBOL_SEMICOLON);
 
         var end = tokens.position();
         var source = start.to(end);