Type check break statements

9 files changed, 136 insertions, 4 deletions

diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedBreakNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedBreakNode.java
new file mode 100644
index 00000000..bc7ea34a
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedBreakNode.java
@@ -0,0 +1,13 @@
+package org.shedlang.compiler.ast.typed;
+
+import org.shedlang.compiler.sources.Source;
+
+public record TypedBreakNode(
+    TypedExpressionNode expression,
+    Source source
+) implements TypedFunctionStatementNode {
+    @Override
+    public TypedFunctionStatementBehaviour behaviour() {
+        return TypedFunctionStatementBehaviour.TERMINATE;
+    }
+}
diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionStatementNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionStatementNode.java
index a79437ec..82c90eea 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionStatementNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionStatementNode.java
@@ -1,5 +1,5 @@
 package org.shedlang.compiler.ast.typed;
 
-public sealed interface TypedFunctionStatementNode extends TypedNode permits TypedAssignNode, TypedBlockReturnNode, TypedExpressionStatementNode, TypedFunctionReturnNode, TypedResumeNode, TypedWhileNode {
+public sealed interface TypedFunctionStatementNode extends TypedNode permits TypedAssignNode, TypedBlockReturnNode, TypedBreakNode, TypedExpressionStatementNode, TypedFunctionReturnNode, TypedResumeNode, TypedWhileNode {
     TypedFunctionStatementBehaviour behaviour();
 }
diff --git a/src/main/java/org/shedlang/compiler/shid/ShidCodeGen.java b/src/main/java/org/shedlang/compiler/shid/ShidCodeGen.java
index 8c930b29..5ed3ffdb 100644
--- a/src/main/java/org/shedlang/compiler/shid/ShidCodeGen.java
+++ b/src/main/java/org/shedlang/compiler/shid/ShidCodeGen.java
@@ -336,6 +336,9 @@ public class ShidCodeGen {
             case TypedBlockReturnNode blockReturn ->
                 compileBlockReturn(blockReturn, context);
 
+            case TypedBreakNode break_ ->
+                throw new UnsupportedOperationException("TODO");
+
             case TypedExpressionStatementNode expressionStatement ->
                 compileExpressionStatement(expressionStatement, context);
 
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index e7052ba3..ae91c3d0 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -280,7 +280,35 @@ public class TypeChecker {
         UntypedBreakNode untyped,
         TypeCheckerContext context
     ) {
-        throw new UnsupportedOperationException("TODO");
+        switch (context.returnTypeState()) {
+            case ReturnTypeState.Explicit explicit -> {
+                throw new CannotBreakOutsideOfEffectHandlerMethodError(untyped.source());
+            }
+            case ReturnTypeState.ImplicitUnknown implicitUnknown -> {
+                throw new CannotBreakOutsideOfEffectHandlerMethodError(untyped.source());
+            }
+            case ReturnTypeState.ImplicitLowerBound implicitLowerBound -> {
+                throw new CannotBreakOutsideOfEffectHandlerMethodError(untyped.source());
+            }
+            case ReturnTypeState.NotInFunction notInFunction -> {
+                throw new CannotBreakOutsideOfEffectHandlerMethodError(untyped.source());
+            }
+            case ReturnTypeState.EffectHandler effectHandler -> {
+                var typedExpressionResult = typeCheckExpressionOptional(
+                    untyped.expression(),
+                    Optional.of(effectHandler.resumeType()),
+                    untyped.source(),
+                    context
+                );
+                var typedExpression = typedExpressionResult.value();
+                context = typedExpressionResult.context();
+
+                return context.toResult(new TypedBreakNode(
+                    typedExpression,
+                    untyped.source()
+                ));
+            }
+        }
     }
 
     private static TypeCheckerResult<TypedExpressionNode> typeCheckCall(
diff --git a/src/main/java/org/shedlang/compiler/typechecker/errors/CannotBreakOutsideOfEffectHandlerMethodError.java b/src/main/java/org/shedlang/compiler/typechecker/errors/CannotBreakOutsideOfEffectHandlerMethodError.java
new file mode 100644
index 00000000..3a32ecaa
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/typechecker/errors/CannotBreakOutsideOfEffectHandlerMethodError.java
@@ -0,0 +1,10 @@
+package org.shedlang.compiler.typechecker.errors;
+
+import org.shedlang.compiler.errors.SourceError;
+import org.shedlang.compiler.sources.Source;
+
+public class CannotBreakOutsideOfEffectHandlerMethodError extends SourceError {
+    public CannotBreakOutsideOfEffectHandlerMethodError(Source source) {
+        super("Cannot break outside of effect handler method", source);
+    }
+}
diff --git a/src/test/java/org/shedlang/compiler/ast/typed/TypedNodeMatchers.java b/src/test/java/org/shedlang/compiler/ast/typed/TypedNodeMatchers.java
index bb787024..0014e4f0 100644
--- a/src/test/java/org/shedlang/compiler/ast/typed/TypedNodeMatchers.java
+++ b/src/test/java/org/shedlang/compiler/ast/typed/TypedNodeMatchers.java
@@ -68,6 +68,15 @@ public class TypedNodeMatchers {
         );
     }
 
+    public static Matcher<TypedFunctionStatementNode> isTypedBreakNode(
+        Matcher<TypedExpressionNode> expression
+    ) {
+        return instanceOf(
+            TypedBreakNode.class,
+            has("expression", x -> x.expression(), expression)
+        );
+    }
+
     public static TypedCallNodeMatcher isTypedCallNode() {
         return TypedCallNodeMatcher.empty();
     }
diff --git a/src/test/java/org/shedlang/compiler/ast/untyped/Untyped.java b/src/test/java/org/shedlang/compiler/ast/untyped/Untyped.java
index 51b6c2b6..8d105dd9 100644
--- a/src/test/java/org/shedlang/compiler/ast/untyped/Untyped.java
+++ b/src/test/java/org/shedlang/compiler/ast/untyped/Untyped.java
@@ -43,6 +43,14 @@ public class Untyped {
         return new UntypedBoolLiteralNode(value, Sources.NULL);
     }
 
+    public static UntypedBreakNode break_() {
+        return new UntypedBreakNode(Optional.empty(), Sources.NULL);
+    }
+
+    public static UntypedBreakNode break_(UntypedExpressionNode expression) {
+        return new UntypedBreakNode(Optional.of(expression), Sources.NULL);
+    }
+
     public static UntypedCallNode call(
         UntypedExpressionNode receiver,
         List<UntypedExpressionNode> positionalArgs
diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerBreakTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerBreakTests.java
new file mode 100644
index 00000000..32f7b127
--- /dev/null
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerBreakTests.java
@@ -0,0 +1,61 @@
+package org.shedlang.compiler.typechecker;
+
+import org.junit.jupiter.api.Test;
+import org.shedlang.compiler.ast.untyped.Untyped;
+import org.shedlang.compiler.typechecker.errors.CannotBreakOutsideOfEffectHandlerMethodError;
+import org.shedlang.compiler.types.Effects;
+import org.shedlang.compiler.types.NamespaceName;
+import org.shedlang.compiler.types.Types;
+import org.shedlang.compiler.util.P;
+
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.shedlang.compiler.ast.typed.TypedNodeMatchers.*;
+import static org.zwobble.precisely.AssertThat.assertThat;
+
+public class TypeCheckerBreakTests {
+    @Test
+    public void expressionIsTypeChecked() {
+        var untyped = Untyped.break_(Untyped.stringLiteral("hello"));
+        var context = TypeCheckerContext.EMPTY
+            .enterFunction(P.map(), Effects.NONE, ReturnTypeState.effectHandler(Types.INT));
+
+        var typed = TypeChecker.typeCheckFunctionStatement(untyped, context);
+
+        assertThat(typed.value(), isTypedBreakNode(isTypedStringLiteralNode("hello")));
+    }
+
+    @Test
+    public void whenExpressionIsAbsentThenUnitIsReturned() {
+        var untyped = Untyped.break_();
+        var context = TypeCheckerContext.EMPTY
+            .enterFunction(P.map(), Effects.NONE, ReturnTypeState.effectHandler(Types.UNIT));
+
+        var typed = TypeChecker.typeCheckFunctionStatement(untyped, context);
+
+        assertThat(typed.value(), isTypedBreakNode(isTypedUnitLiteralNode()));
+    }
+
+    @Test
+    public void whenInNamespaceThenResumeIsForbidden() {
+        var untyped = Untyped.break_(Untyped.stringLiteral("hello"));
+        var context = TypeCheckerContext.EMPTY
+            .enterNamespace(NamespaceName.of("Example"), P.map());
+
+        assertThrows(
+            CannotBreakOutsideOfEffectHandlerMethodError.class,
+            () -> TypeChecker.typeCheckFunctionStatement(untyped, context)
+        );
+    }
+
+    @Test
+    public void whenInFunctionThenResumeIsForbidden() {
+        var untyped = Untyped.break_(Untyped.stringLiteral("hello"));
+        var context = TypeCheckerContext.EMPTY
+            .enterFunction(P.map(), Effects.NONE, ReturnTypeState.explicit(Types.STRING));
+
+        assertThrows(
+            CannotBreakOutsideOfEffectHandlerMethodError.class,
+            () -> TypeChecker.typeCheckFunctionStatement(untyped, context)
+        );
+    }
+}
diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerResumeTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerResumeTests.java
index f0f18839..b5bab880 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerResumeTests.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerResumeTests.java
@@ -43,7 +43,7 @@ public class TypeCheckerResumeTests {
     }
 
     @Test
-    public void whenNotInFunctionThenResumeIsForbidden() {
+    public void whenInNamespaceThenResumeIsForbidden() {
         var untyped = Untyped.resume(Untyped.stringLiteral("hello"));
         var context = TypeCheckerContext.EMPTY
             .enterNamespace(NamespaceName.of("Example"), P.map());
@@ -67,7 +67,7 @@ public class TypeCheckerResumeTests {
     }
 
     @Test
-    public void canResumeWIthSubtypeOfResumeType() {
+    public void canResumeWithSubtypeOfResumeType() {
         var untyped = Untyped.resume(Untyped.stringLiteral("hello"));
         var context = TypeCheckerContext.EMPTY
             .enterFunction(P.map(), Effects.NONE, ReturnTypeState.effectHandler(Types.TOP));