Add handled effect to try body effects

6 files changed, 58 insertions, 6 deletions

diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 0725c78e..29376e9f 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -1591,14 +1591,21 @@ public class TypeChecker {
         UntypedTryNode untyped,
         TypeCheckerContext context
     ) {
-        var typedBodyResult = typeCheckBlockNested(
-            untyped.body(),
-            context
+        var typedHandledEffect = typeCheckEffect(untyped.handledEffect(), context);
+
+        var declaredVariables = DeclaredVariables.forFunctionStatements(
+            untyped.body().statements(),
+            P.vector()
         );
+        var bodyEnvironment = EnvironmentBuilder.EMPTY.addUndefined(declaredVariables);
+
+        context = context.enterTryBody(bodyEnvironment.build(), typedHandledEffect.value());
+
+        var typedBodyResult = typeCheckBlock(untyped.body(), bodyEnvironment.names(), context);
         var typedBody = typedBodyResult.value();
         context = typedBodyResult.context();
 
-        var typedHandledEffect = typeCheckEffect(untyped.handledEffect(), context);
+        context = context.exit();
 
         return context.toResult(new TypedTryNode(
             typedBody,
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
index de989863..8209cd16 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
@@ -189,6 +189,18 @@ public class TypeCheckerContext {
         );
     }
 
+    public TypeCheckerContext enterTryBody(
+        PMap<String, Variable> environment,
+        Effect handledEffect
+    ) {
+        return enter(
+            this.namespaceName(),
+            environment,
+            Effects.union(this.effect(), handledEffect),
+            Optional.empty()
+        );
+    }
+
     public TypeCheckerContext enterBlock(PMap<String, Variable> environment) {
         return enter(this.namespaceName(), environment, this.effect(), Optional.empty());
     }
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerSubTyping.java b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerSubTyping.java
index ea86f964..d1857276 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerSubTyping.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerSubTyping.java
@@ -76,6 +76,19 @@ public class TypeCheckerSubTyping {
     }
 
     static boolean isSubEffect(Effect subEffect, Effect superEffect, TypeCheckerContext context) {
-        return subEffect == Effects.NONE || subEffect.equals(superEffect);
+        if (subEffect == Effects.NONE) {
+            return true;
+        }
+
+        if (subEffect.equals(superEffect)) {
+            return true;
+        }
+
+        if (superEffect instanceof EffectUnion superEffectUnion) {
+            return isSubEffect(subEffect, superEffectUnion.left(), context) ||
+                isSubEffect(subEffect, superEffectUnion.right(), context);
+        }
+
+        return false;
     }
 }
diff --git a/src/main/java/org/shedlang/compiler/types/Effect.java b/src/main/java/org/shedlang/compiler/types/Effect.java
index 2a583320..01ba00a3 100644
--- a/src/main/java/org/shedlang/compiler/types/Effect.java
+++ b/src/main/java/org/shedlang/compiler/types/Effect.java
@@ -1,4 +1,4 @@
 package org.shedlang.compiler.types;
 
-public sealed interface Effect extends TypeLevelValue permits IoEffect, NoneEffect, UserDefinedEffect {
+public sealed interface Effect extends TypeLevelValue permits EffectUnion, IoEffect, NoneEffect, UserDefinedEffect {
 }
diff --git a/src/main/java/org/shedlang/compiler/types/EffectUnion.java b/src/main/java/org/shedlang/compiler/types/EffectUnion.java
new file mode 100644
index 00000000..7ab25453
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/types/EffectUnion.java
@@ -0,0 +1,8 @@
+package org.shedlang.compiler.types;
+
+public record EffectUnion(Effect left, Effect right) implements Effect {
+    @Override
+    public String describe() {
+        return left().describe() + " | " + right.describe();
+    }
+}
diff --git a/src/main/java/org/shedlang/compiler/types/Effects.java b/src/main/java/org/shedlang/compiler/types/Effects.java
index 206574d7..f9782f4a 100644
--- a/src/main/java/org/shedlang/compiler/types/Effects.java
+++ b/src/main/java/org/shedlang/compiler/types/Effects.java
@@ -14,4 +14,16 @@ public class Effects {
     ) {
         return new UserDefinedEffect(namespaceName, name);
     }
+
+    public static Effect union(Effect left, Effect right) {
+        if (left.equals(right)) {
+            return left;
+        } else if (left == NONE) {
+            return right;
+        } else if (right == NONE) {
+            return left;
+        } else {
+            return new EffectUnion(left, right);
+        }
+    }
 }