Use break statements to determine type of try expression

8 files changed, 76 insertions, 28 deletions

diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionBodyNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionBodyNode.java
index cae9cf43..fbdf2773 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionBodyNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionBodyNode.java
@@ -8,7 +8,10 @@ import java.util.List;
 public record TypedFunctionBodyNode(
     FreeVariables freeVariables,
     TypedBlockNode block,
+    // TODO: better handle this body being used for both functions and effect
+    // handler methods.
     Type type,
+    Type breakType,
     Source source
 ) implements TypedNode {
     public List<BoundVariable> boundVariables() {
diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionDefinitionNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionDefinitionNode.java
index c1a28516..cae0c93c 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionDefinitionNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionDefinitionNode.java
@@ -42,6 +42,7 @@ public record TypedFunctionDefinitionNode(
                     freeVariables.build(),
                     new TypedBlockNode(body, boundVariables, Sources.NULL),
                     Types.UNIT,
+                    Types.BOTTOM,
                     Sources.NULL
                 ),
                 type,
diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionExpressionNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionExpressionNode.java
index fd53a26d..40168641 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionExpressionNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionExpressionNode.java
@@ -38,6 +38,7 @@ public record TypedFunctionExpressionNode(
                     freeVariables.build(),
                     new TypedBlockNode(body, boundVariables, Sources.NULL),
                     Types.UNIT,
+                    Types.BOTTOM,
                     Sources.NULL
                 ),
                 type,
diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedTryHandlerMethodNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedTryHandlerMethodNode.java
index e9cbd4d8..e1f1351d 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedTryHandlerMethodNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedTryHandlerMethodNode.java
@@ -13,6 +13,8 @@ public record TypedTryHandlerMethodNode(
     TypedParamsNode params,
     TypedFunctionBodyNode body,
     FunctionType type,
+    // TODO: combine function type with break type?
+    Type breakType,
     Source source
 ) implements TypedNode {
     public static Builder builder() {
@@ -22,7 +24,8 @@ public record TypedTryHandlerMethodNode(
             FreeVariables.builder(),
             Empty.vector(),
             Empty.vector(),
-            new FunctionType(List.of(), List.of(), Effects.NONE, Types.UNIT)
+            new FunctionType(List.of(), List.of(), Effects.NONE, Types.UNIT),
+            Types.BOTTOM
         );
     }
 
@@ -32,7 +35,8 @@ public record TypedTryHandlerMethodNode(
         FreeVariables.Builder freeVariables,
         PVector<BoundVariable> boundVariables,
         PVector<TypedFunctionStatementNode> body,
-        FunctionType type
+        FunctionType type,
+        Type breakType
     ) {
         public TypedTryHandlerMethodNode build() {
             return new TypedTryHandlerMethodNode(
@@ -42,9 +46,11 @@ public record TypedTryHandlerMethodNode(
                     freeVariables.build(),
                     new TypedBlockNode(body, boundVariables, Sources.NULL),
                     Types.UNIT,
+                    Types.BOTTOM,
                     Sources.NULL
                 ),
                 type,
+                breakType,
                 Sources.NULL
             );
         }
@@ -56,7 +62,8 @@ public record TypedTryHandlerMethodNode(
                 this.freeVariables,
                 this.boundVariables,
                 this.body,
-                this.type
+                this.type,
+                this.breakType
             );
         }
 
@@ -67,7 +74,8 @@ public record TypedTryHandlerMethodNode(
                 this.freeVariables,
                 this.boundVariables,
                 this.body,
-                this.type
+                this.type,
+                this.breakType
             ).addBoundVariable(BoundVariable.immutable(name));
         }
 
@@ -78,7 +86,8 @@ public record TypedTryHandlerMethodNode(
                 this.freeVariables,
                 this.boundVariables,
                 this.body,
-                this.type
+                this.type,
+                this.breakType
             ).addBoundVariable(BoundVariable.immutable(name));
         }
 
@@ -89,7 +98,8 @@ public record TypedTryHandlerMethodNode(
                 this.freeVariables.plus(freeVariable, Sources.NULL),
                 this.boundVariables,
                 this.body,
-                this.type
+                this.type,
+                this.breakType
             );
         }
 
@@ -100,7 +110,8 @@ public record TypedTryHandlerMethodNode(
                 this.freeVariables,
                 this.boundVariables.plus(boundVariable),
                 this.body,
-                this.type
+                this.type,
+                this.breakType
             );
         }
 
@@ -111,7 +122,8 @@ public record TypedTryHandlerMethodNode(
                 this.freeVariables,
                 this.boundVariables,
                 this.body.plus(statement),
-                this.type
+                this.type,
+                this.breakType
             );
         }
     }
diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedTryNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedTryNode.java
index c3c5d32d..5ce4f205 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedTryNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedTryNode.java
@@ -3,10 +3,7 @@ package org.shedlang.compiler.ast.typed;
 import org.pcollections.PVector;
 import org.shedlang.compiler.sources.Source;
 import org.shedlang.compiler.sources.Sources;
-import org.shedlang.compiler.types.Effects;
-import org.shedlang.compiler.types.NamespaceName;
-import org.shedlang.compiler.types.Type;
-import org.shedlang.compiler.types.UserDefinedEffect;
+import org.shedlang.compiler.types.*;
 import org.shedlang.compiler.util.P;
 
 import java.util.List;
@@ -15,31 +12,30 @@ public record TypedTryNode(
     TypedBlockNode body,
     TypedTypeLevelExpressionNode<UserDefinedEffect> handledEffect,
     List<TypedTryHandlerMethodNode> handlerMethods,
+    Type type,
     Source source
 ) implements TypedExpressionNode {
-    @Override
-    public Type type() {
-        return body.type();
-    }
-
     public static Builder builder() {
         return new Builder(
             P.vector(),
             new TypedTypeLevelReferenceNode<>("Effect", Effects.userDefined(NamespaceName.of("Example"), "Effect"), Sources.NULL),
-            P.vector()
+            P.vector(),
+            Types.TOP
         );
     }
 
     public record Builder(
         PVector<TypedFunctionStatementNode> body,
         TypedTypeLevelExpressionNode<UserDefinedEffect> handledEffect,
-        PVector<TypedTryHandlerMethodNode> handlerMethods
+        PVector<TypedTryHandlerMethodNode> handlerMethods,
+        Type type
     ) {
         public TypedTryNode build() {
             return new TypedTryNode(
                 new TypedBlockNode(body, P.vector(), Sources.NULL),
                 handledEffect,
                 handlerMethods,
+                type,
                 Sources.NULL
             );
         }
@@ -48,7 +44,8 @@ public record TypedTryNode(
             return new Builder(
                 this.body.plus(statement),
                 this.handledEffect,
-                this.handlerMethods
+                this.handlerMethods,
+                this.type
             );
         }
 
@@ -56,7 +53,8 @@ public record TypedTryNode(
             return new Builder(
                 this.body,
                 handledEffect,
-                this.handlerMethods
+                this.handlerMethods,
+                this.type
             );
         }
 
@@ -64,7 +62,8 @@ public record TypedTryNode(
             return new Builder(
                 this.body,
                 this.handledEffect,
-                this.handlerMethods.plus(method)
+                this.handlerMethods.plus(method),
+                this.type
             );
         }
     }
diff --git a/src/main/java/org/shedlang/compiler/typechecker/ReturnTypeState.java b/src/main/java/org/shedlang/compiler/typechecker/ReturnTypeState.java
index b43b17ea..201489d5 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/ReturnTypeState.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/ReturnTypeState.java
@@ -1,6 +1,8 @@
 package org.shedlang.compiler.typechecker;
 
+import org.pcollections.PVector;
 import org.shedlang.compiler.types.Type;
+import org.shedlang.compiler.util.P;
 
 public sealed interface ReturnTypeState {
     ReturnTypeState NOT_IN_FUNCTION = new NotInFunction();
@@ -15,12 +17,17 @@ public sealed interface ReturnTypeState {
     }
 
     static ReturnTypeState effectHandler(Type resumeType) {
-        return new EffectHandler(resumeType);
+        return new EffectHandler(resumeType, P.vector());
     }
 
     record NotInFunction() implements ReturnTypeState {}
     record Explicit(Type type) implements ReturnTypeState {}
     record ImplicitUnknown() implements ReturnTypeState {}
     record ImplicitLowerBound(Type type) implements ReturnTypeState {}
-    record EffectHandler(Type resumeType) implements ReturnTypeState {}
+
+    record EffectHandler(Type resumeType, PVector<Type> breakTypes) implements ReturnTypeState {
+        public ReturnTypeState addBreakType(Type breakType) {
+            return new EffectHandler(resumeType, breakTypes.plus(breakType));
+        }
+    }
 }
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index e484f35c..94ba7ea8 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -304,6 +304,10 @@ public class TypeChecker {
                 var typedExpression = typedExpressionResult.value();
                 context = typedExpressionResult.context();
 
+                context = context.updateReturnTypeState(
+                    effectHandler.addBreakType(typedExpression.type())
+                );
+
                 return context.toResult(new TypedBreakNode(
                     typedExpression,
                     untyped.source()
@@ -645,6 +649,7 @@ public class TypeChecker {
         }
 
         Type returnType;
+        Type breakType = Types.BOTTOM;
 
         switch (context.returnTypeState()) {
             case ReturnTypeState.Explicit explicit -> {
@@ -665,6 +670,7 @@ public class TypeChecker {
             }
             case ReturnTypeState.EffectHandler effectHandler -> {
                 returnType = effectHandler.resumeType();
+                breakType = Types.commonSupertype(effectHandler.breakTypes());
             }
         }
 
@@ -692,6 +698,7 @@ public class TypeChecker {
             freeVariables,
             typedBody,
             returnType,
+            breakType,
             typedBody.source()
         );
 
@@ -1745,7 +1752,11 @@ public class TypeChecker {
                     untypedHandlerMethod.source()
                 ));
 
-            var typedHandlerMethodResult = typeCheckTryHandlerMethod(untypedHandlerMethod, effectFieldType, context);
+            var typedHandlerMethodResult = typeCheckTryHandlerMethod(
+                untypedHandlerMethod,
+                effectFieldType,
+                context
+            );
             var typedHandlerMethod = typedHandlerMethodResult.value();
             context = typedHandlerMethodResult.context();
 
@@ -1765,10 +1776,18 @@ public class TypeChecker {
             }
         }
 
+        var type = Types.commonSupertype(Lists.concat(
+            List.of(typedBody.type()),
+            typedHandlerMethods.stream()
+                .map(method -> method.breakType())
+                .toList()
+        ));
+
         return context.toResult(new TypedTryNode(
             typedBody,
             typedHandledEffect,
             typedHandlerMethods,
+            type,
             untyped.source()
         ));
     }
@@ -1845,6 +1864,7 @@ public class TypeChecker {
             typedParams,
             typedBody,
             type,
+            typedBody.breakType(),
             untyped.source()
         );
 
diff --git a/src/main/java/org/shedlang/compiler/types/Types.java b/src/main/java/org/shedlang/compiler/types/Types.java
index 71c931b3..efc9f000 100644
--- a/src/main/java/org/shedlang/compiler/types/Types.java
+++ b/src/main/java/org/shedlang/compiler/types/Types.java
@@ -1,7 +1,7 @@
 package org.shedlang.compiler.types;
 
-import java.util.HashSet;
 import java.util.List;
+import java.util.stream.Collectors;
 
 public class Types {
     private Types() {
@@ -39,8 +39,13 @@ public class Types {
     }
 
     public static Type commonSupertype(List<Type> subtypes) {
-        var uniqueSubtypes = new HashSet<>(subtypes);
-        if (uniqueSubtypes.size() == 1) {
+        var uniqueSubtypes = subtypes.stream()
+            .filter(subtype -> subtype != Types.BOTTOM)
+            .collect(Collectors.toSet());
+
+        if (uniqueSubtypes.isEmpty()) {
+            return BOTTOM;
+        } else if (uniqueSubtypes.size() == 1) {
             return subtypes.get(0);
         } else {
             return TOP;