From 405fc623d679c0d40a8e4fc031dd4e5e32d5f35f Mon Sep 17 00:00:00 2001
From: Michael Williamson <mike@zwobble.org>
Date: Tue, 24 Oct 2023 21:21:29 +0100
Subject: Type-check effect in function types

---
 .../ast/untyped/UntypedFunctionTypeNode.java       | 11 ++++-
 .../shedlang/compiler/typechecker/TypeChecker.java |  4 +-
 .../typechecker/TypeCheckerFunctionTypeTests.java  | 49 ++++++++++++++++++++++
 3 files changed, 62 insertions(+), 2 deletions(-)

(limited to 'src')

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()
             ),
diff --git a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionTypeTests.java b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionTypeTests.java
index 608e3caf..d1c7c2f2 100644
--- a/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionTypeTests.java
+++ b/src/test/java/org/shedlang/compiler/typechecker/TypeCheckerFunctionTypeTests.java
@@ -4,6 +4,7 @@ import org.junit.jupiter.api.Test;
 import org.shedlang.compiler.ast.untyped.Untyped;
 import org.shedlang.compiler.ast.untyped.UntypedFunctionTypeNode;
 import org.shedlang.compiler.typechecker.typesets.TypeSets;
+import org.shedlang.compiler.types.Effects;
 import org.shedlang.compiler.types.FunctionType;
 import org.shedlang.compiler.types.Types;
 import org.shedlang.compiler.util.P;
@@ -36,4 +37,52 @@ public class TypeCheckerFunctionTypeTests {
             ))
         );
     }
+
+    @Test
+    public void whenNodeHasNoEffectThenEffectIsNone() {
+        var untyped = UntypedFunctionTypeNode.builder()
+            .withReturnType(Untyped.typeLevelReference("String"))
+            .build();
+        var context = TypeCheckerContext.EMPTY
+            .enterGlobal(P.map(
+                Map.entry("String", Types.metaType(Types.STRING))
+            ));
+
+        var typed = TypeChecker.typeCheckTypeLevelExpression(TypeSets.META_TYPE, untyped, context);
+
+        assertThat(
+            typed,
+            isTypedFunctionTypeNode(equalTo(
+                FunctionType.builder()
+                    .withEffect(Effects.NONE)
+                    .withReturnType(Types.STRING)
+                    .build()
+            ))
+        );
+    }
+
+    @Test
+    public void effectIsTypeChecked() {
+        var untyped = UntypedFunctionTypeNode.builder()
+            .withEffect(Untyped.typeLevelReference("Io"))
+            .withReturnType(Untyped.typeLevelReference("String"))
+            .build();
+        var context = TypeCheckerContext.EMPTY
+            .enterGlobal(P.map(
+                Map.entry("Io", Types.effectType(Effects.IO)),
+                Map.entry("String", Types.metaType(Types.STRING))
+            ));
+
+        var typed = TypeChecker.typeCheckTypeLevelExpression(TypeSets.META_TYPE, untyped, context);
+
+        assertThat(
+            typed,
+            isTypedFunctionTypeNode(equalTo(
+                FunctionType.builder()
+                    .withEffect(Effects.IO)
+                    .withReturnType(Types.STRING)
+                    .build()
+            ))
+        );
+    }
 }
-- 
cgit v1.2.3