Handle behaviour of if expressions properly

5 files changed, 54 insertions, 14 deletions

diff --git a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedBranchingExpressionNode.java b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedBranchingExpressionNode.java
index eab2f680..1e7f18ff 100644
--- a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedBranchingExpressionNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedBranchingExpressionNode.java
@@ -1,7 +1,5 @@
 package org.shedlang.compiler.ast.untyped;
 
-import java.util.stream.Stream;
-
 public interface UntypedBranchingExpressionNode {
-    Stream<UntypedBlockNode> branchBodies();
+    Iterable<UntypedBlockNode> exhaustiveBranchBodies();
 }
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 c258257c..da9a04e1 100644
--- a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedIfNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedIfNode.java
@@ -1,9 +1,9 @@
 package org.shedlang.compiler.ast.untyped;
 
 import org.shedlang.compiler.sources.Source;
+import org.shedlang.compiler.util.Lists;
 
 import java.util.List;
-import java.util.stream.Stream;
 
 public record UntypedIfNode(
     List<UntypedConditionalBranchNode> conditionalBranches,
@@ -11,10 +11,10 @@ public record UntypedIfNode(
     Source source
 ) implements UntypedBranchingExpressionNode, UntypedExpressionNode {
     @Override
-    public Stream<UntypedBlockNode> branchBodies() {
-        return Stream.concat(
-            conditionalBranches.stream().map(branch -> branch.body()),
-            Stream.of(elseBody)
+    public Iterable<UntypedBlockNode> exhaustiveBranchBodies() {
+        return Lists.concat(
+            conditionalBranches.stream().map(branch -> branch.body()).toList(),
+            List.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 46a175f0..35efa5ab 100644
--- a/src/main/java/org/shedlang/compiler/parser/Parser.java
+++ b/src/main/java/org/shedlang/compiler/parser/Parser.java
@@ -1,15 +1,13 @@
 package org.shedlang.compiler.parser;
 
 import org.shedlang.compiler.ast.untyped.*;
-import org.shedlang.compiler.parser.errors.IntLiteralCannotHaveLeadingZeroesError;
-import org.shedlang.compiler.parser.errors.PositionalArgAfterNamedArgError;
-import org.shedlang.compiler.parser.errors.PositionalParamAfterNamedParamError;
-import org.shedlang.compiler.parser.errors.UnrecognisedEscapeSequenceError;
+import org.shedlang.compiler.parser.errors.*;
 import org.shedlang.compiler.sources.Source;
 import org.shedlang.compiler.tokeniser.TokenIterator;
 import org.shedlang.compiler.tokeniser.UnexpectedTokenException;
 import org.shedlang.compiler.types.NamespaceName;
 import org.shedlang.compiler.types.Visibility;
+import org.shedlang.compiler.util.Iterables;
 
 import java.math.BigInteger;
 import java.util.ArrayList;
@@ -308,8 +306,7 @@ public class Parser {
         var expression = parseTopLevelExpression(tokens);
 
         var isReturn = expression instanceof UntypedBranchingExpressionNode branchingExpression
-            // TODO: Check for consistency
-            ? branchingExpression.branchBodies().findFirst().orElseThrow().behaviour() == UntypedFunctionStatementBehaviour.BLOCK_RETURN
+            ? expressionStatementBehaviour(branchingExpression.exhaustiveBranchBodies()) == UntypedFunctionStatementBehaviour.BLOCK_RETURN
             : !tokens.trySkip(TokenType.SYMBOL_SEMICOLON);
 
         var end = tokens.position();
@@ -320,6 +317,26 @@ public class Parser {
             : new UntypedExpressionStatementNode(expression, source);
     }
 
+    static UntypedFunctionStatementBehaviour expressionStatementBehaviour(
+        Iterable<UntypedBlockNode> branchBodies
+    ) {
+        return Iterables.foldLeft(
+            UntypedFunctionStatementBehaviour.TERMINATE,
+            branchBodies,
+            (expressionStatementBehaviour, branchBody) -> {
+                if (expressionStatementBehaviour == UntypedFunctionStatementBehaviour.TERMINATE) {
+                    return branchBody.behaviour();
+                } else if (branchBody.behaviour() == UntypedFunctionStatementBehaviour.TERMINATE) {
+                    return expressionStatementBehaviour;
+                } else if (expressionStatementBehaviour == branchBody.behaviour()) {
+                    return expressionStatementBehaviour;
+                } else {
+                    throw new InconsistentBlockBehaviourError(branchBody.source());
+                }
+            }
+        );
+    }
+
     private static class FieldAccessParser implements OperationParser {
         @Override
         public OperatorPrecedence precedence() {
diff --git a/src/main/java/org/shedlang/compiler/parser/errors/InconsistentBlockBehaviourError.java b/src/main/java/org/shedlang/compiler/parser/errors/InconsistentBlockBehaviourError.java
new file mode 100644
index 00000000..5e44de63
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/parser/errors/InconsistentBlockBehaviourError.java
@@ -0,0 +1,10 @@
+package org.shedlang.compiler.parser.errors;
+
+import org.shedlang.compiler.errors.SourceError;
+import org.shedlang.compiler.sources.Source;
+
+public class InconsistentBlockBehaviourError extends SourceError {
+    public InconsistentBlockBehaviourError(Source source) {
+        super("Inconsistent block behaviour: blocks in the same expression must either all yield a value, or all yield no value", source);
+    }
+}
diff --git a/src/main/java/org/shedlang/compiler/util/Iterables.java b/src/main/java/org/shedlang/compiler/util/Iterables.java
index 27fa3e8d..799c9c4f 100644
--- a/src/main/java/org/shedlang/compiler/util/Iterables.java
+++ b/src/main/java/org/shedlang/compiler/util/Iterables.java
@@ -4,6 +4,7 @@ import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Comparator;
 import java.util.List;
+import java.util.function.BiFunction;
 import java.util.function.Function;
 
 public class Iterables {
@@ -34,6 +35,20 @@ public class Iterables {
         return iterable.iterator().next();
     }
 
+    public static <T, R> R foldLeft(
+        R initial,
+        Iterable<T> iterable,
+        BiFunction<R, T, R> combine
+    ) {
+        R result = initial;
+
+        for (var element : iterable) {
+            result = combine.apply(result, element);
+        }
+
+        return result;
+    }
+
     public static <T, E extends Exception> void forEachInterspersed(
         Iterable<T> values,
         OnElement<T, E> onElement,