Type check effect in function definition

5 files changed, 49 insertions, 6 deletions

diff --git a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionDefinitionNode.java b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionDefinitionNode.java
index 57c02568..a3c5447a 100644
--- a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionDefinitionNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionDefinitionNode.java
@@ -79,6 +79,17 @@ public record UntypedFunctionDefinitionNode(
             );
         }
 
+        public Builder withEffect(UntypedTypeLevelExpressionNode effect) {
+            return new Builder(
+                name,
+                positionalParams,
+                namedParams,
+                Optional.of(effect),
+                returnType,
+                body
+            );
+        }
+
         public Builder withReturnType(UntypedTypeLevelExpressionNode returnType) {
             return new Builder(
                 name,
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index c6875d1b..cbaf366a 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -6,10 +6,7 @@ import org.shedlang.compiler.sources.Source;
 import org.shedlang.compiler.typechecker.errors.*;
 import org.shedlang.compiler.typechecker.typesets.TypeSetTypeLevelValue;
 import org.shedlang.compiler.typechecker.typesets.TypeSets;
-import org.shedlang.compiler.types.FunctionType;
-import org.shedlang.compiler.types.Type;
-import org.shedlang.compiler.types.TypeLevelValue;
-import org.shedlang.compiler.types.TypeLevelValueType;
+import org.shedlang.compiler.types.*;
 
 import java.util.HashSet;
 import java.util.Optional;
@@ -131,6 +128,13 @@ public class TypeChecker {
         );
     }
 
+    private static TypedTypeLevelExpressionNode<Effect> typeCheckEffect(
+        UntypedTypeLevelExpressionNode untyped,
+        TypeCheckerContext context
+    ) {
+        return typeCheckTypeLevelExpression(TypeSets.EFFECT, untyped, context);
+    }
+
     public static TypedExpressionNode typeCheckExpression(
         UntypedExpressionNode untyped,
         TypeCheckerContext context
@@ -169,6 +173,13 @@ public class TypeChecker {
     ) {
         var typedParams = typeCheckParams(untyped.params(), context);
 
+        Optional<TypedTypeLevelExpressionNode<Effect>> typedEffect;
+        if (untyped.effect().isPresent()) {
+            typedEffect = Optional.of(typeCheckEffect(untyped.effect().get(), context));
+        } else {
+            typedEffect = Optional.empty();
+        }
+
         var typedReturnType = typeCheckTypeLevelExpression(TypeSets.META_TYPE, untyped.returnType(), context);
 
         var typedBody = untyped.body().stream()
@@ -178,7 +189,7 @@ public class TypeChecker {
         return new TypedFunctionDefinitionNode(
             untyped.name(),
             typedParams,
-            Optional.empty(),
+            typedEffect,
             typedReturnType,
             typedBody,
             untyped.source()
diff --git a/src/main/java/org/shedlang/compiler/typechecker/typesets/TypeSetEffect.java b/src/main/java/org/shedlang/compiler/typechecker/typesets/TypeSetEffect.java
new file mode 100644
index 00000000..6c250a7a
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/typechecker/typesets/TypeSetEffect.java
@@ -0,0 +1,20 @@
+package org.shedlang.compiler.typechecker.typesets;
+
+import org.shedlang.compiler.types.Effect;
+
+public final class TypeSetEffect implements TypeSetTypeLevelValue<Effect> {
+    private TypeSetEffect() {
+    }
+
+    static final TypeSetEffect INSTANCE = new TypeSetEffect();
+
+    @Override
+    public String describe() {
+        return "effect";
+    }
+
+    @Override
+    public Class<Effect> valueClass() {
+        return Effect.class;
+    }
+}
diff --git a/src/main/java/org/shedlang/compiler/typechecker/typesets/TypeSetTypeLevelValue.java b/src/main/java/org/shedlang/compiler/typechecker/typesets/TypeSetTypeLevelValue.java
index c4389aee..4007eb2c 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/typesets/TypeSetTypeLevelValue.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/typesets/TypeSetTypeLevelValue.java
@@ -2,6 +2,6 @@ package org.shedlang.compiler.typechecker.typesets;
 
 import org.shedlang.compiler.types.TypeLevelValue;
 
-public sealed interface TypeSetTypeLevelValue<T extends TypeLevelValue> extends TypeSet permits TypeSetMetaType {
+public sealed interface TypeSetTypeLevelValue<T extends TypeLevelValue> extends TypeSet permits TypeSetEffect, TypeSetMetaType {
     Class<T> valueClass();
 }
diff --git a/src/main/java/org/shedlang/compiler/typechecker/typesets/TypeSets.java b/src/main/java/org/shedlang/compiler/typechecker/typesets/TypeSets.java
index 633bff86..3dd47cae 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/typesets/TypeSets.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/typesets/TypeSets.java
@@ -6,6 +6,7 @@ public class TypeSets {
     private TypeSets() {
     }
 
+    public static final TypeSetEffect EFFECT = TypeSetEffect.INSTANCE;
     public static final TypeSet FUNCTION = TypeSetFunction.INSTANCE;
     public static final TypeSetMetaType META_TYPE = TypeSetMetaType.INSTANCE;