Type-check effect in function types

2 files changed, 13 insertions, 2 deletions

diff --git a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionTypeNode.java b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionTypeNode.java
index ca8404e3..c2529ad2 100644
--- a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionTypeNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionTypeNode.java
@@ -37,7 +37,16 @@ public record UntypedFunctionTypeNode(
             );
         }
 
-        public Builder withReturnType(UntypedTypeLevelReferenceNode returnType) {
+        public Builder withEffect(UntypedTypeLevelExpressionNode effect) {
+            return new Builder(
+                this.positionalParams,
+                this.namedParams,
+                Optional.of(effect),
+                this.returnType
+            );
+        }
+
+        public Builder withReturnType(UntypedTypeLevelExpressionNode returnType) {
             return new Builder(
                 this.positionalParams,
                 this.namedParams,
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index aeed2edd..a3be369a 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -650,7 +650,9 @@ public class TypeChecker {
             new FunctionType(
                 List.of(),
                 List.of(),
-                Effects.NONE,
+                untyped.effect()
+                    .map(untypedEffect -> typeCheckTypeLevelExpression(TypeSets.EFFECT, untypedEffect, context).value())
+                    .orElse(Effects.NONE),
                 typeCheckTypeLevelExpression(TypeSets.META_TYPE, untyped.returnType(), context)
                     .value()
             ),