Return context when type checking type level expressions

2 files changed, 319 insertions, 138 deletions

diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index b8a05746..49072117 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -392,20 +392,24 @@ public class TypeChecker {
         var typedReceiver = receiverResult.value();
         context = receiverResult.context();
 
-        var contextForReceiver = context;
-        var receiverType = untyped.typeLevelArgs()
-            .map(untypedTypeLevelArgNodes -> {
+        var typeCheckReceiverResult = TypeCheckerResult.mapOr(
+            untyped.typeLevelArgs(),
+            context,
+            (untypedTypeLevelArgNodes, receiverContext) -> {
                 Type specificReceiverType;
 
                 if (
                     typedReceiver.type() instanceof TypeLevelValueType initialReceiverType &&
                     initialReceiverType.value() instanceof TypeConstructor<?> typeConstructor
                 ) {
-                    var typeLevelArgs = typeCheckTypeLevelArgs(
+                    var typeCheckTypeLevelArgsResult = typeCheckTypeLevelArgs(
                         untypedTypeLevelArgNodes,
                         typeConstructor.params(),
-                        contextForReceiver
+                        receiverContext
                     );
+                    var typeLevelArgs = typeCheckTypeLevelArgsResult.value();
+                    receiverContext = typeCheckTypeLevelArgsResult.context();
+
                     specificReceiverType = Types.metaType(
                         Types.constructedType(typeConstructor, typeLevelArgs)
                     );
@@ -413,23 +417,31 @@ public class TypeChecker {
                     typedReceiver.type() instanceof FunctionType initialReceiverType &&
                     initialReceiverType.typeLevelParams().isPresent()
                 ) {
-                    var typeLevelArgs = typeCheckTypeLevelArgs(
+                    var typeCheckTypeLevelArgsResult = typeCheckTypeLevelArgs(
                         untypedTypeLevelArgNodes,
                         initialReceiverType.typeLevelParams().get(),
-                        contextForReceiver
+                        receiverContext
                     );
+                    var typeLevelArgs = typeCheckTypeLevelArgsResult.value();
+                    receiverContext = typeCheckTypeLevelArgsResult.context();
+
                     specificReceiverType = initialReceiverType.applyTypeLevelArgs(typeLevelArgs);
                 } else {
                     throw new IllegalArgumentException("TODO");
                 }
 
-                return typeCheckCallReceiverType(
+                return receiverContext.toResult(typeCheckCallReceiverType(
                     specificReceiverType,
                     typedReceiver.source(),
-                    contextForReceiver
-                );
-            })
-            .orElseGet(() -> typeCheckCallReceiverType(typedReceiver, contextForReceiver));
+                    receiverContext
+                ));
+            },
+            receiverContext -> receiverContext.toResult(
+                typeCheckCallReceiverType(typedReceiver, receiverContext)
+            )
+        );
+        var receiverType = typeCheckReceiverResult.value();
+        context = typeCheckReceiverResult.context();
 
         var argsResult = typeCheckArgs(untyped.args(), context, receiverType);
         var typedArgs = argsResult.value().typedArgsNode;
@@ -569,14 +581,14 @@ public class TypeChecker {
         ));
     }
 
-    private static TypedTypeLevelExpressionNode<Effect> typeCheckEffect(
+    private static TypeCheckerResult<TypedTypeLevelExpressionNode<Effect>> typeCheckEffect(
         UntypedTypeLevelExpressionNode untyped,
         TypeCheckerContext context
     ) {
         return typeCheckTypeLevelExpression(TypeSets.EFFECT, untyped, context);
     }
 
-    private static TypedTypeLevelExpressionNode<UserDefinedEffect> typeCheckUserDefinedEffect(
+    private static TypeCheckerResult<TypedTypeLevelExpressionNode<UserDefinedEffect>> typeCheckUserDefinedEffect(
         UntypedTypeLevelExpressionNode untyped,
         TypeCheckerContext context
     ) {
@@ -610,8 +622,11 @@ public class TypeChecker {
                 var fieldsInfoBuilder = FieldsInfo.<FunctionType>builder();
 
                 for (var methodNode : untyped.methods()) {
-                    var typedFunctionTypeNode = typeCheckFunctionType(methodNode.type(), context);
-                    fieldsInfoBuilder = fieldsInfoBuilder.addField(methodNode.name(), typedFunctionTypeNode.value());
+                    var typeCheckTypeResult = typeCheckFunctionType(methodNode.type(), context);
+                    var type = typeCheckTypeResult.value().value();
+                    context = typeCheckTypeResult.context();
+
+                    fieldsInfoBuilder = fieldsInfoBuilder.addField(methodNode.name(), type);
                 }
                 var effectInfo = new UserDefinedEffectInfo(fieldsInfoBuilder.build());
 
@@ -885,11 +900,14 @@ public class TypeChecker {
     ) {
         var namespaceName = context.namespaceName()
             .orElseThrow(() -> new CannotDefineFunctionOutsideOfNamespaceError(source));
-        var typedSignature = typeCheckFunctionSignature(
+
+        var typeCheckSignatureResult = typeCheckFunctionSignature(
             signature,
             Optional.of(new Scope.Function(namespaceName, name)),
             context
         );
+        var typedSignature = typeCheckSignatureResult.value();
+        context = typeCheckSignatureResult.context();
 
         Function<TypedParamNode, PartialFunctionSignature.Param> typedParamNodeToParam = param ->
             new PartialFunctionSignature.Param(
@@ -982,11 +1000,13 @@ public class TypeChecker {
         Optional<Type> hint,
         TypeCheckerContext context
     ) {
-        var typedSignature = typeCheckFunctionSignature(
+        var typeCheckSignatureResult = typeCheckFunctionSignature(
             untyped.signature(),
             Optional.empty(),
             context
         );
+        var typedSignature = typeCheckSignatureResult.value();
+        context = typeCheckSignatureResult.context();
 
         var partialSignature = inferFunctionPartialSignature(typedSignature, hint);
 
@@ -1084,7 +1104,7 @@ public class TypeChecker {
         return context;
     }
 
-    static TypedFunctionSignatureNode typeCheckFunctionSignature(
+    static TypeCheckerResult<TypedFunctionSignatureNode> typeCheckFunctionSignature(
         UntypedFunctionSignatureNode untyped,
         Optional<Scope> scope,
         TypeCheckerContext context
@@ -1106,26 +1126,35 @@ public class TypeChecker {
                 .collect(PCollectors.toMap())
         );
 
-        var functionSignatureContext = context;
+        var typeCheckParamsResult = typeCheckParams(untyped.params(), context);
+        var typedParams = typeCheckParamsResult.value();
+        context = typeCheckParamsResult.context();
 
-        var typedParams = typeCheckParams(untyped.params(), context);
+        var typeCheckEffectResult = TypeCheckerResult.map(
+            untyped.effect(),
+            context,
+            TypeChecker::typeCheckEffect
+        );
+        var typedEffect = typeCheckEffectResult.value();
+        context = typeCheckEffectResult.context();
 
-        var typedEffect = untyped.effect()
-            .map(untypedEffect -> typeCheckEffect(untypedEffect, functionSignatureContext));
+        var typeCheckReturnTypeResult = TypeCheckerResult.map(
+            untyped.returnType(),
+            context,
+            TypeChecker::typeCheckTypeExpression
+        );
+        var typedReturnType = typeCheckReturnTypeResult.value();
+        context = typeCheckReturnTypeResult.context();
 
-        var typedReturnType = untyped.returnType()
-            .map(untypedReturnType -> typeCheckTypeExpression(
-                untypedReturnType,
-                functionSignatureContext
-            ));
+        context = context.exit();
 
-        return new TypedFunctionSignatureNode(
+        return context.toResult(new TypedFunctionSignatureNode(
             typeLevelParams,
             typedParams,
             typedEffect,
             typedReturnType,
             untyped.source()
-        );
+        ));
     }
 
     private static Variable typeLevelParamToVariable(TypeLevelParam typeLevelParam) {
@@ -1245,20 +1274,27 @@ public class TypeChecker {
         };
     }
 
-    private static TypedFunctionTypeNode typeCheckFunctionType(
+    private static TypeCheckerResult<TypedFunctionTypeNode> typeCheckFunctionType(
         UntypedFunctionTypeNode untyped,
         TypeCheckerContext context
     ) {
-        var positionalParams = untyped.params().positional().stream()
-            .map(param -> typeCheckTypeExpression(
-                param,
-                context
-            ).value())
-            .toList();
+        var typeCheckPositionalParamsResult = TypeCheckerResult.map(
+            untyped.params().positional(),
+            context,
+            (untypedParamNode, paramContext) ->
+                typeCheckTypeExpression(untypedParamNode, paramContext)
+                    .map(typedParamNode -> typedParamNode.value())
+        );
+        var positionalParams = typeCheckPositionalParamsResult.value();
+        context = typeCheckPositionalParamsResult.context();
 
-        var namedParams = untyped.params().named().stream()
-            .map(param -> typeCheckFunctionTypeNamedParam(param, context))
-            .toList();
+        var typeCheckNamedParamsResult = TypeCheckerResult.map(
+            untyped.params().named(),
+            context,
+            TypeChecker::typeCheckFunctionTypeNamedParam
+        );
+        var namedParams = typeCheckNamedParamsResult.value();
+        context = typeCheckNamedParamsResult.context();
 
         var seenParamNames = new HashSet<String>();
         for (var untypedParamNode : untyped.params().named()) {
@@ -1270,16 +1306,25 @@ public class TypeChecker {
             }
         }
 
-        var effect = untyped.effect()
-            .map(untypedEffect -> typeCheckEffect(untypedEffect, context).value())
-            .orElse(Effects.NONE);
+        var typeCheckEffectResult = TypeCheckerResult.map(
+            untyped.effect(),
+            context,
+            (untypedEffectNode, effectContext) ->
+                typeCheckEffect(untypedEffectNode, effectContext)
+                    .map(typedEffectNode -> typedEffectNode.value())
+        );
+        var effect = typeCheckEffectResult.value().orElse(Effects.NONE);
+        context = typeCheckEffectResult.context();
 
-        var returnType = typeCheckTypeExpression(
+
+        var typeCheckReturnTypeResult = typeCheckTypeExpression(
             untyped.returnType(),
             context
-        ).value();
+        );
+        var returnType = typeCheckReturnTypeResult.value().value();
+        context = typeCheckEffectResult.context();
 
-        return new TypedFunctionTypeNode(
+        return context.toResult(new TypedFunctionTypeNode(
             new FunctionType(
                 Optional.empty(),
                 positionalParams,
@@ -1288,17 +1333,21 @@ public class TypeChecker {
                 returnType
             ),
             untyped.source()
-        );
+        ));
     }
 
-    private static NamedParam typeCheckFunctionTypeNamedParam(
+    private static TypeCheckerResult<NamedParam> typeCheckFunctionTypeNamedParam(
         UntypedFunctionTypeNamedParamNode untyped,
         TypeCheckerContext context
     ) {
-        return new NamedParam(
+        var typeCheckTypeResult = typeCheckTypeExpression(untyped.type(), context);
+        var typedTypeNode = typeCheckTypeResult.value();
+        context = typeCheckTypeResult.context();
+
+        return context.toResult(new NamedParam(
             untyped.name(),
-            typeCheckTypeExpression(untyped.type(), context).value()
-        );
+            typedTypeNode.value()
+        ));
     }
 
 
@@ -1625,15 +1674,17 @@ public class TypeChecker {
         return new TypeCheckSteps<TypedNativeFunctionDeclarationNode>() {
             @Override
             public TypeCheckerContext declareFunctions(TypeCheckerContext context) {
-                var typedNode = typeCheckFunctionType(untyped.signature(), context);
+                var typeCheckSignatureResult = typeCheckFunctionType(untyped.signature(), context);
+                var typedSignatureNode = typeCheckSignatureResult.value();
+                context = typeCheckSignatureResult.context();
 
                 context = context.declareVariableType(
                     untyped.name(),
-                    typedNode.value(),
+                    typedSignatureNode.value(),
                     untyped.source()
                 );
 
-                return context.addLexicalState(untyped.name(), typedNode);
+                return context.addLexicalState(untyped.name(), typedSignatureNode);
             }
 
             @Override
@@ -1691,30 +1742,52 @@ public class TypeChecker {
         };
     }
 
-    private static TypedParamNode typeCheckParam(
+    private static TypeCheckerResult<TypedParamNode> typeCheckParam(
         UntypedParamNode untyped,
         TypeCheckerContext context
     ) {
-        return new TypedParamNode(
+        var typeCheckTypeResult = TypeCheckerResult.map(
+            untyped.type(),
+            context,
+            TypeChecker::typeCheckTypeExpression
+        );
+        var type = typeCheckTypeResult.value();
+        context = typeCheckTypeResult.context();
+
+        var typedNode = new TypedParamNode(
             untyped.name(),
-            untyped.type().map(p -> typeCheckTypeExpression(p, context)),
+            type,
             untyped.source()
         );
+
+        return context.toResult(typedNode);
     }
 
-    private static TypedParamsNode typeCheckParams(
+    private static TypeCheckerResult<TypedParamsNode> typeCheckParams(
         UntypedParamsNode untyped,
         TypeCheckerContext context
     ) {
-        var typedPositional = untyped.positional().stream()
-            .map(param -> typeCheckParam(param, context))
-            .toList();
+        var typeCheckPositionalResult = TypeCheckerResult.map(
+            untyped.positional(),
+            context,
+            TypeChecker::typeCheckParam
+        );
+        var typedPositional = typeCheckPositionalResult.value();
+        context = typeCheckPositionalResult.context();
 
-        var typedNamed = untyped.named().stream()
-            .map(param -> typeCheckParam(param, context))
-            .toList();
+        var typeCheckNamedResult = TypeCheckerResult.map(
+            untyped.named(),
+            context,
+            TypeChecker::typeCheckParam
+        );
+        var typedNamed = typeCheckNamedResult.value();
+        context = typeCheckNamedResult.context();
 
-        return new TypedParamsNode(typedPositional, typedNamed, untyped.source());
+        return context.toResult(new TypedParamsNode(
+            typedPositional,
+            typedNamed,
+            untyped.source()
+        ));
     }
 
     private static TypeCheckerResult<TypedReferenceNode> typeCheckReference(
@@ -1862,9 +1935,13 @@ public class TypeChecker {
                 var fieldInfosBuilder = FieldsInfo.builder();
 
                 for (var untypedField : untyped.fields()) {
+                    var typeCheckFieldTypeResult = typeCheckTypeExpression(untypedField.type(), context);
+                    var typedFieldTypeNode = typeCheckFieldTypeResult.value();
+                    context = typeCheckFieldTypeResult.context();
+
                     var typedFieldNode = new TypedStructFieldDefinitionNode(
                         untypedField.name(),
-                        typeCheckTypeExpression(untypedField.type(), context),
+                        typedFieldTypeNode,
                         untypedField.source()
                     );
                     typedFieldNodes.add(typedFieldNode);
@@ -2030,34 +2107,46 @@ public class TypeChecker {
 
                 context = context.enterBlock(environment.build());
 
-                var variants = typeCheckVariants(context);
+                var typeCheckVariantsResult = typeCheckVariants(context);
+                var variants = typeCheckVariantsResult.value();
+                context = typeCheckVariantsResult.context();
 
                 context = context.exit();
 
                 return context.addSumTypeInfo(typeOrConstructor.toType(), new SumTypeInfo(variants));
             }
 
-            private OrderedPSet<StructType> typeCheckVariants(TypeCheckerContext context) {
+            private TypeCheckerResult<OrderedPSet<StructType>> typeCheckVariants(TypeCheckerContext context) {
                 // TODO: prevent constructed types with same underlying type from appearing multiple times
                 // (e.g. Option[String], Option[Int])
 
-                return untyped.variants().stream()
-                    .map(untypedVariant -> {
-                        var variantType = typeCheckTypeExpression(
-                            untypedVariant,
-                            context
-                        ).value();
-
-                        if (variantType instanceof StructType variantStructType) {
-                            return variantStructType;
-                        } else {
-                            throw new VariantTypeIsNotStructError(
-                                variantType,
-                                untypedVariant.source()
-                            );
-                        }
-                    })
-                    .collect(PCollectors.toOrderedSet());
+                return TypeCheckerResult.map(
+                    untyped.variants(),
+                    context,
+                    (untypedVariantNode, variantContext) ->
+                        typeCheckVariant(untypedVariantNode, variantContext)
+                ).map(variants -> variants.stream().collect(PCollectors.toOrderedSet()));
+            }
+
+            private TypeCheckerResult<StructType> typeCheckVariant(
+                UntypedTypeLevelExpressionNode untypedVariantNode,
+                TypeCheckerContext context
+            ) {
+                var typeCheckVariantResult = typeCheckTypeExpression(
+                    untypedVariantNode,
+                    context
+                );
+                var variantType = typeCheckVariantResult.value().value();
+                context = typeCheckVariantResult.context();
+
+                if (variantType instanceof StructType variantStructType) {
+                    return context.toResult(variantStructType);
+                } else {
+                    throw new VariantTypeIsNotStructError(
+                        variantType,
+                        untypedVariantNode.source()
+                    );
+                }
             }
 
             @Override
@@ -2075,14 +2164,15 @@ public class TypeChecker {
         };
     }
 
-    static TypeConstructor<?> typeCheckTypeConstructor(
+    static TypeCheckerResult<TypeConstructor<?>> typeCheckTypeConstructor(
         UntypedTypeLevelExpressionNode untyped,
         TypeCheckerContext context
     ) {
-        return typeCheckTypeLevelExpression(TypeSets.TYPE_CONSTRUCTOR, untyped, context).value();
+        return typeCheckTypeLevelExpression(TypeSets.TYPE_CONSTRUCTOR, untyped, context)
+            .map(x -> x.value());
     }
 
-    static TypedTypeLevelExpressionNode<Type> typeCheckTypeExpression(
+    static TypeCheckerResult<TypedTypeLevelExpressionNode<Type>> typeCheckTypeExpression(
         UntypedTypeLevelExpressionNode untyped,
         TypeCheckerContext context
     ) {
@@ -2090,16 +2180,18 @@ public class TypeChecker {
     }
 
     @SuppressWarnings("unchecked")
-    static <T extends TypeLevelValue> TypedTypeLevelExpressionNode<T> typeCheckTypeLevelExpression(
+    static <T extends TypeLevelValue> TypeCheckerResult<TypedTypeLevelExpressionNode<T>> typeCheckTypeLevelExpression(
         TypeSetTypeLevelValue<T> valueType,
         UntypedTypeLevelExpressionNode untyped,
         TypeCheckerContext context
     ) {
-        var typedNode = typeCheckTypeLevelExpression(untyped, context);
+        var result = typeCheckTypeLevelExpression(untyped, context);
+        var typedNode = result.value();
+        context = result.context();
 
         // TODO: tidy this up
         if (valueType.isInstance(typedNode.value())) {
-            return (TypedTypeLevelExpressionNode<T>) typedNode;
+            return context.toResult((TypedTypeLevelExpressionNode<T>) typedNode);
         } else {
             throw new UnexpectedTypeError(
                 valueType,
@@ -2109,23 +2201,33 @@ public class TypeChecker {
         }
     }
 
-    static TypedTypeLevelExpressionNode<?> typeCheckTypeLevelExpression(
+    static TypeCheckerResult<? extends TypedTypeLevelExpressionNode<?>> typeCheckTypeLevelExpression(
         UntypedTypeLevelExpressionNode untyped,
         TypeCheckerContext context
     ) {
         return switch (untyped) {
             case UntypedConstructedTypeNode constructedTypeNode -> {
                 // TODO: check number of args
-                var constructor = typeCheckTypeConstructor(constructedTypeNode.receiver(), context);
-                var args = constructedTypeNode.args().stream()
-                    .map(arg -> typeCheckTypeExpression(arg, context).value())
-                    .toList();
+                var typeCheckConstructorResult = typeCheckTypeConstructor(constructedTypeNode.receiver(), context);
+                var constructor = typeCheckConstructorResult.value();
+                context = typeCheckConstructorResult.context();
+
+                var typeCheckArgsResult = TypeCheckerResult.map(
+                    constructedTypeNode.args(),
+                    context,
+                    (untypedArgNode, argContext) ->
+                        typeCheckTypeExpression(untypedArgNode, argContext)
+                            .map(typedArgNode -> typedArgNode.value())
+                );
+                var args = typeCheckArgsResult.value();
+                context = typeCheckArgsResult.context();
+
                 var constructedType = Types.constructedType(constructor, args);
 
-                yield new TypedConstructedTypeNode(
+                yield context.toResult(new TypedConstructedTypeNode(
                     constructedType,
                     constructedTypeNode.source()
-                );
+                ));
             }
 
             case UntypedFunctionTypeNode functionType ->
@@ -2163,23 +2265,29 @@ public class TypeChecker {
         return context.toResult(new TypedTupleNode(typedElements, type, untyped.source()));
     }
 
-    private static TypedTupleTypeNode typeCheckTupleType(
+    private static TypeCheckerResult<TypedTupleTypeNode> typeCheckTupleType(
         UntypedTupleTypeNode untyped,
         TypeCheckerContext context
     ) {
-        var elementTypes = untyped.elements().stream()
-            .map(elementType -> typeCheckTypeExpression(
-                elementType,
-                context
-            ).value())
-            .toList();
+        var typeCheckElementTypesResult = TypeCheckerResult.map(
+            untyped.elements(),
+            context,
+            (untypedElementTypeNode, elementContext) ->
+                typeCheckTypeExpression(untypedElementTypeNode, elementContext)
+                    .map(typedElementTypeNode -> typedElementTypeNode.value())
+        );
+        var elementTypes = typeCheckElementTypesResult.value();
+        context = typeCheckElementTypesResult.context();
 
         var tupleType = new TupleType(elementTypes);
 
-        return new TypedTupleTypeNode(tupleType, untyped.source());
+        return context.toResult(new TypedTupleTypeNode(
+            tupleType,
+            untyped.source()
+        ));
     }
 
-    private static ArrayList<TypeLevelValue> typeCheckTypeLevelArgs(
+    private static TypeCheckerResult<ArrayList<TypeLevelValue>> typeCheckTypeLevelArgs(
         UntypedTypeLevelArgsNode untyped,
         List<TypeLevelParam> typeLevelParams,
         TypeCheckerContext context
@@ -2200,19 +2308,22 @@ public class TypeChecker {
             var typeLevelParam = typeLevelParams.get(paramIndex);
             var argNode = untyped.args().get(paramIndex);
 
-            var typeLevelArg = switch (typeLevelParam) {
-                case EffectParam effectParam ->
-                    typeCheckEffect(argNode, context).value();
-                case TypeParam typeParam ->
-                    typeCheckTypeExpression(argNode, context).value();
+            var typeCheckTypeLevelArgResult = switch (typeLevelParam) {
+                case EffectParam _ ->
+                    typeCheckEffect(argNode, context);
+                case TypeParam _ ->
+                    typeCheckTypeExpression(argNode, context);
             };
+            var typeLevelArg = typeCheckTypeLevelArgResult.value().value();
+            context = typeCheckTypeLevelArgResult.context();
+
             typeLevelArgs.add(typeLevelArg);
         }
 
-        return typeLevelArgs;
+        return context.toResult(typeLevelArgs);
     }
 
-    private static TypedTypeLevelExpressionNode<TypeLevelValue> typeCheckTypeLevelReference(
+    private static TypeCheckerResult<TypedTypeLevelExpressionNode<TypeLevelValue>> typeCheckTypeLevelReference(
         UntypedTypeLevelReferenceNode untyped,
         TypeCheckerContext context
     ) {
@@ -2223,11 +2334,11 @@ public class TypeChecker {
             throw new UnexpectedTypeError(TypeSets.META_TYPE, type, untyped.source());
         }
 
-        return new TypedTypeLevelReferenceNode<>(
+        return context.toResult(new TypedTypeLevelReferenceNode<>(
             untyped.name(),
             typeLevelValueType.value(),
             untyped.source()
-        );
+        ));
     }
 
     private static TypeCheckerResult<TypedTargetNode> typeCheckTarget(
@@ -2326,7 +2437,10 @@ public class TypeChecker {
         UntypedTryNode untyped,
         TypeCheckerContext context
     ) {
-        var typedHandledEffect = typeCheckUserDefinedEffect(untyped.handledEffect(), context);
+        var typeCheckHandledEffectResult = typeCheckUserDefinedEffect(untyped.handledEffect(), context);
+        var typedHandledEffect = typeCheckHandledEffectResult.value();
+        context = typeCheckHandledEffectResult.context();
+
         var handledEffect = typedHandledEffect.value();
 
         var declaredVariables = DeclaredVariables.forFunctionStatements(
@@ -2335,7 +2449,7 @@ public class TypeChecker {
         );
         var bodyEnvironment = EnvironmentBuilder.EMPTY.addUndefined(declaredVariables);
 
-        context = context.enterTryBody(bodyEnvironment.build(), typedHandledEffect.value());
+        context = context.enterTryBody(bodyEnvironment.build(), handledEffect);
 
         var typedBodyResult = typeCheckBlock(
             untyped.body(),
@@ -2413,41 +2527,61 @@ public class TypeChecker {
         FunctionType effectFieldType,
         TypeCheckerContext context
     ) {
-        var typedParams = typeCheckParams(untyped.params(), context);
+        var typeCheckParamsResult = typeCheckParams(untyped.params(), context);
+        var typedParams = typeCheckParamsResult.value();
+        context = typeCheckParamsResult.context();
+
         var inferredParamTypes = inferParamTypes(typedParams, Optional.of(effectFieldType));
 
-        var initialContext = context;
-        var expectedEffect = initialContext.effect();
-        var effect = untyped.effect()
-            .map(untypedEffectNode -> {
-                var explicitEffect = typeCheckEffect(untypedEffectNode, initialContext).value();
-                if (!isSubEffect(explicitEffect, expectedEffect, initialContext.toTypeInfos())) {
+        var expectedEffect = context.effect();
+        var typeCheckEffectResult = TypeCheckerResult.mapOr(
+            untyped.effect(),
+            context,
+            (untypedEffectNode, effectContext) -> {
+                var typeCheckResult = typeCheckEffect(untypedEffectNode, effectContext);
+                var explicitEffect = typeCheckResult.value().value();
+                effectContext = typeCheckResult.context();
+
+                if (!isSubEffect(explicitEffect, expectedEffect, effectContext.toTypeInfos())) {
                     throw new UnexpectedEffectError(
                         expectedEffect,
                         explicitEffect,
                         untypedEffectNode.source()
                     );
                 }
-                return explicitEffect;
-            })
-            .orElse(expectedEffect);
+
+                return effectContext.toResult(explicitEffect);
+            },
+            effectContext -> effectContext.toResult(expectedEffect)
+        );
+        var effect = typeCheckEffectResult.value();
+        context = typeCheckEffectResult.context();
 
         var expectedResumeType = effectFieldType.returnType();
-        var resumeType = untyped.resumeType()
-            .map(untypedResumeTypeNode -> {
-                var explicitResumeType = typeCheckTypeExpression(
+        var typeCheckResumeTypeResult = TypeCheckerResult.mapOr(
+            untyped.resumeType(),
+            context,
+            (untypedResumeTypeNode, resumeTypeContext) -> {
+                var typeCheckResult = typeCheckTypeExpression(
                     untypedResumeTypeNode,
-                    initialContext
-                ).value();
+                    resumeTypeContext
+                );
+                var explicitResumeType = typeCheckResult.value().value();
+                resumeTypeContext = typeCheckResult.context();
+
                 checkIsSubType(
                     explicitResumeType,
                     expectedResumeType,
                     untypedResumeTypeNode.source(),
-                    initialContext
+                    resumeTypeContext
                 );
-                return explicitResumeType;
-            })
-            .orElse(expectedResumeType);
+
+                return resumeTypeContext.toResult(explicitResumeType);
+            },
+            resumeTypeContext -> resumeTypeContext.toResult(expectedResumeType)
+        );
+        var resumeType = typeCheckResumeTypeResult.value();
+        context = typeCheckResumeTypeResult.context();
 
         var partialSignature = new PartialFunctionSignature(
             Optional.empty(),
@@ -2606,10 +2740,12 @@ public class TypeChecker {
         OrderedPSet<StructType> variantTypes,
         TypeCheckerContext context
     ) {
-        var typedCaseType = typeCheckTypeLevelExpression(
+        var typeCheckCaseTypeResult = typeCheckTypeLevelExpression(
             untyped.caseType(),
             context
         );
+        var typedCaseType = typeCheckCaseTypeResult.value();
+        context = typeCheckCaseTypeResult.context();
 
         var variant = switch (typedCaseType.value()) {
             case StructType variantStructType ->
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerResult.java b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerResult.java
index 8c4fa2b1..effe9ba4 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerResult.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerResult.java
@@ -1,5 +1,9 @@
 package org.shedlang.compiler.typechecker;
 
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Optional;
+import java.util.function.BiFunction;
 import java.util.function.Function;
 
 public class TypeCheckerResult<T> {
@@ -11,6 +15,47 @@ public class TypeCheckerResult<T> {
         this.context = context;
     }
 
+    public static <T, R> TypeCheckerResult<Optional<R>> map(
+        Optional<T> value,
+        TypeCheckerContext context,
+        BiFunction<T, TypeCheckerContext, TypeCheckerResult<R>> func
+    ) {
+        if (value.isPresent()) {
+            return func.apply(value.get(), context).map(Optional::of);
+        } else {
+            return context.toResult(Optional.empty());
+        }
+    }
+
+    public static <T, R> TypeCheckerResult<R> mapOr(
+        Optional<T> value,
+        TypeCheckerContext context,
+        BiFunction<T, TypeCheckerContext, TypeCheckerResult<R>> func,
+        Function<TypeCheckerContext, TypeCheckerResult<R>> orElse
+    ) {
+        if (value.isPresent()) {
+            return func.apply(value.get(), context);
+        } else {
+            return orElse.apply(context);
+        }
+    }
+
+    public static <T, R> TypeCheckerResult<List<R>> map(
+        List<T> values,
+        TypeCheckerContext context,
+        BiFunction<T, TypeCheckerContext, TypeCheckerResult<R>> func
+    ) {
+        var result = new ArrayList<R>();
+
+        for (var value : values) {
+            var valueResult = func.apply(value, context);
+            result.add(valueResult.value);
+            context = valueResult.context;
+        }
+
+        return context.toResult(result);
+    }
+
     public <U> TypeCheckerResult<U> map(Function<T, U> func) {
         return new TypeCheckerResult<>(func.apply(value), context);
     }