Add effect info for user defined effects

9 files changed, 96 insertions, 14 deletions

diff --git a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedEffectDefinitionNode.java b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedEffectDefinitionNode.java
index 5731aefe..f6b5c2ce 100644
--- a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedEffectDefinitionNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedEffectDefinitionNode.java
@@ -43,5 +43,13 @@ public record UntypedEffectDefinitionNode(
                 this.functions
             );
         }
+
+        public Builder addFunction(String name, UntypedFunctionTypeNode type) {
+            return new Builder(
+                this.name,
+                this.visibility,
+                this.functions.plus(new UntypedEffectDefinitionFunctionNode(name, type, Sources.NULL))
+            );
+        }
     }
 }
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 5c104fe1..c6c57866 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -409,6 +409,23 @@ public class TypeChecker {
             }
 
             @Override
+            public TypeCheckerContext defineTypes(TypeCheckerContext context) {
+                var effect = context.getLexicalState(UserDefinedEffect.class, untyped.name());
+
+                var functions = untyped.functions().stream()
+                    .map(functionNode -> {
+                        var typedFunctionTypeNode = typeCheckFunctionType(functionNode.type(), context);
+
+                        return Map.entry(functionNode.name(), typedFunctionTypeNode.value());
+                    })
+                    .collect(PCollectors.toOrderedMap());
+
+                var effectInfo = new UserDefinedEffectInfo(new FieldsInfo<>(functions));
+
+                return context.addUserDefinedEffectInfo(effect, effectInfo);
+            }
+
+            @Override
             public TypeCheckerResult<TypedEffectDefinitionNode> finish(TypeCheckerContext context) {
                 var effect = context.getLexicalState(UserDefinedEffect.class, untyped.name());
 
@@ -1272,7 +1289,7 @@ public class TypeChecker {
 
                 var type = context.getLexicalState(StructType.class, untyped.name());
 
-                context = context.addStructTypeInfo(type, new StructTypeInfo(new FieldsInfo(fieldTypes)));
+                context = context.addStructTypeInfo(type, new StructTypeInfo(new FieldsInfo<>(fieldTypes)));
 
                 var typedNode = new TypedStructDefinitionNode(
                     type,
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
index 832c5137..de989863 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
@@ -400,7 +400,7 @@ public class TypeCheckerContext {
         NamespaceName namespaceName,
         POrderedMap<String, Type> fields
     ) {
-        var namespaceTypeInfo = new NamespaceTypeInfo(new FieldsInfo(fields));
+        var namespaceTypeInfo = new NamespaceTypeInfo(new FieldsInfo<>(fields));
 
         return new TypeCheckerContext(
             this.typeInfos.addNamespaceTypeInfo(namespaceName, namespaceTypeInfo),
@@ -443,6 +443,21 @@ public class TypeCheckerContext {
             .orElseThrow(() -> new MissingSumTypeInfoError(type, source));
     }
 
+    public TypeCheckerContext addUserDefinedEffectInfo(
+        UserDefinedEffect effect,
+        UserDefinedEffectInfo effectInfo
+    ) {
+        return new TypeCheckerContext(
+            this.typeInfos.addUserDefinedEffectInfo(effect, effectInfo),
+            this.stack
+        );
+    }
+
+    public UserDefinedEffectInfo userDefinedEffectInfo(UserDefinedEffect effect, Source source) {
+        return this.typeInfos.userDefinedEffectInfo(effect)
+            .orElseThrow(() -> new MissingUserDefinedEffectInfoError(effect, source));
+    }
+
     public <T> TypeCheckerResult<T> toResult(T value) {
         return new TypeCheckerResult<>(value, this);
     }
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeInfos.java b/src/main/java/org/shedlang/compiler/typechecker/TypeInfos.java
index f18b8836..77bb5493 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeInfos.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeInfos.java
@@ -10,20 +10,23 @@ import java.util.Optional;
 import java.util.Set;
 
 public class TypeInfos {
-    public static final TypeInfos EMPTY = new TypeInfos(Empty.map(), P.orderedPMap(), Empty.map());
+    public static final TypeInfos EMPTY = new TypeInfos(Empty.map(), P.orderedPMap(), Empty.map(), Empty.map());
 
     private final PMap<NamespaceName, NamespaceTypeInfo> namespaceTypeInfos;
     private final OrderedPMap<StructType, StructTypeInfo> structTypeInfos;
     private final PMap<SumType, SumTypeInfo> sumTypeInfos;
+    private final PMap<UserDefinedEffect, UserDefinedEffectInfo> userDefinedEffectInfos;
 
     public TypeInfos(
         PMap<NamespaceName, NamespaceTypeInfo> namespaceTypeInfos,
         OrderedPMap<StructType, StructTypeInfo> structTypeInfos,
-        PMap<SumType, SumTypeInfo> sumTypeInfos
+        PMap<SumType, SumTypeInfo> sumTypeInfos,
+        PMap<UserDefinedEffect, UserDefinedEffectInfo> userDefinedEffectInfos
     ) {
         this.namespaceTypeInfos = namespaceTypeInfos;
         this.structTypeInfos = structTypeInfos;
         this.sumTypeInfos = sumTypeInfos;
+        this.userDefinedEffectInfos = userDefinedEffectInfos;
     }
 
     public TypeInfos addNamespaceTypeInfo(
@@ -33,7 +36,8 @@ public class TypeInfos {
         return new TypeInfos(
             this.namespaceTypeInfos.plus(namespaceName, namespaceTypeInfo),
             this.structTypeInfos,
-            this.sumTypeInfos
+            this.sumTypeInfos,
+            this.userDefinedEffectInfos
         );
     }
 
@@ -48,7 +52,8 @@ public class TypeInfos {
         return new TypeInfos(
             this.namespaceTypeInfos,
             this.structTypeInfos.plus(type, structTypeInfo),
-            this.sumTypeInfos
+            this.sumTypeInfos,
+            this.userDefinedEffectInfos
         );
     }
 
@@ -64,11 +69,25 @@ public class TypeInfos {
         return new TypeInfos(
             this.namespaceTypeInfos,
             this.structTypeInfos,
-            this.sumTypeInfos.plus(sumType, sumTypeInfo)
+            this.sumTypeInfos.plus(sumType, sumTypeInfo),
+            this.userDefinedEffectInfos
         );
     }
 
     public Optional<SumTypeInfo> sumTypeInfo(SumType type) {
         return Optional.ofNullable(sumTypeInfos.get(type));
     }
+
+    public TypeInfos addUserDefinedEffectInfo(UserDefinedEffect effect, UserDefinedEffectInfo effectInfo) {
+        return new TypeInfos(
+            this.namespaceTypeInfos,
+            this.structTypeInfos,
+            this.sumTypeInfos,
+            this.userDefinedEffectInfos.plus(effect, effectInfo)
+        );
+    }
+
+    public Optional<UserDefinedEffectInfo> userDefinedEffectInfo(UserDefinedEffect effect) {
+        return Optional.ofNullable(userDefinedEffectInfos.get(effect));
+    }
 }
diff --git a/src/main/java/org/shedlang/compiler/typechecker/errors/MissingUserDefinedEffectInfoError.java b/src/main/java/org/shedlang/compiler/typechecker/errors/MissingUserDefinedEffectInfoError.java
new file mode 100644
index 00000000..ffc0b5e3
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/typechecker/errors/MissingUserDefinedEffectInfoError.java
@@ -0,0 +1,18 @@
+package org.shedlang.compiler.typechecker.errors;
+
+import org.shedlang.compiler.errors.InternalCompilerError;
+import org.shedlang.compiler.sources.Source;
+import org.shedlang.compiler.types.UserDefinedEffect;
+
+public class MissingUserDefinedEffectInfoError extends InternalCompilerError {
+    private final UserDefinedEffect effect;
+
+    public MissingUserDefinedEffectInfoError(UserDefinedEffect effect, Source source) {
+        super("Missing effect info for " + effect.describe(), source);
+        this.effect = effect;
+    }
+
+    public UserDefinedEffect effect() {
+        return effect;
+    }
+}
diff --git a/src/main/java/org/shedlang/compiler/types/FieldsInfo.java b/src/main/java/org/shedlang/compiler/types/FieldsInfo.java
index 7723b7dc..867e09ac 100644
--- a/src/main/java/org/shedlang/compiler/types/FieldsInfo.java
+++ b/src/main/java/org/shedlang/compiler/types/FieldsInfo.java
@@ -6,7 +6,7 @@ import java.util.Map;
 import java.util.Optional;
 import java.util.stream.Stream;
 
-public record FieldsInfo(POrderedMap<String, Type> fields) {
+public record FieldsInfo<T extends Type>(POrderedMap<String, T> fields) {
     public int fieldCount() {
         return fields.size();
     }
@@ -21,12 +21,12 @@ public record FieldsInfo(POrderedMap<String, Type> fields) {
         }
     }
 
-    public Optional<Type> fieldType(String fieldName) {
+    public Optional<T> fieldType(String fieldName) {
         return fields.get(fieldName);
     }
 
     // TODO: better element type
-    public Stream<Map.Entry<String, Type>> stream() {
+    public Stream<Map.Entry<String, T>> stream() {
         return fields.entryStream();
     }
 }
diff --git a/src/main/java/org/shedlang/compiler/types/NamespaceTypeInfo.java b/src/main/java/org/shedlang/compiler/types/NamespaceTypeInfo.java
index cd2dc1ff..5cfb2aa8 100644
--- a/src/main/java/org/shedlang/compiler/types/NamespaceTypeInfo.java
+++ b/src/main/java/org/shedlang/compiler/types/NamespaceTypeInfo.java
@@ -4,9 +4,9 @@ import org.shedlang.compiler.util.POrderedMap;
 
 import java.util.Optional;
 
-public record NamespaceTypeInfo(FieldsInfo fields) {
+public record NamespaceTypeInfo(FieldsInfo<Type> fields) {
     public static NamespaceTypeInfo fromFields(POrderedMap<String, Type> fields) {
-        return new NamespaceTypeInfo(new FieldsInfo(fields));
+        return new NamespaceTypeInfo(new FieldsInfo<>(fields));
     }
 
     public int fieldCount() {
diff --git a/src/main/java/org/shedlang/compiler/types/StructTypeInfo.java b/src/main/java/org/shedlang/compiler/types/StructTypeInfo.java
index 17436827..15338cd6 100644
--- a/src/main/java/org/shedlang/compiler/types/StructTypeInfo.java
+++ b/src/main/java/org/shedlang/compiler/types/StructTypeInfo.java
@@ -4,9 +4,9 @@ import org.shedlang.compiler.util.POrderedMap;
 
 import java.util.Optional;
 
-public record StructTypeInfo(FieldsInfo fields) {
+public record StructTypeInfo(FieldsInfo<Type> fields) {
     public static StructTypeInfo fromFields(POrderedMap<String, Type> fields) {
-        return new StructTypeInfo(new FieldsInfo(fields));
+        return new StructTypeInfo(new FieldsInfo<>(fields));
     }
 
     public int fieldCount() {
diff --git a/src/main/java/org/shedlang/compiler/types/UserDefinedEffectInfo.java b/src/main/java/org/shedlang/compiler/types/UserDefinedEffectInfo.java
new file mode 100644
index 00000000..eb05cf1d
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/types/UserDefinedEffectInfo.java
@@ -0,0 +1,5 @@
+package org.shedlang.compiler.types;
+
+public record UserDefinedEffectInfo(FieldsInfo<FunctionType> functions) {
+
+}