Generate effect tags

2 files changed, 31 insertions, 13 deletions

diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java
index f7876225..155e391f 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java
@@ -13,10 +13,7 @@ import org.shedlang.compiler.patio.ast.types.PatioDataLayout;
 import org.shedlang.compiler.patio.ast.types.PatioType;
 import org.shedlang.compiler.sources.Source;
 import org.shedlang.compiler.typechecker.TypeInfos;
-import org.shedlang.compiler.types.NamespaceName;
-import org.shedlang.compiler.types.NamespaceTypeInfo;
-import org.shedlang.compiler.types.StructType;
-import org.shedlang.compiler.types.StructTypeInfo;
+import org.shedlang.compiler.types.*;
 import org.shedlang.compiler.util.P;
 
 import java.util.ArrayList;
@@ -33,25 +30,42 @@ public record PatioCodeGenContext(
 ) {
     public record GlobalContext(
         TypeInfos typeInfos,
-        PMap<StructType, Integer> structTags
+        PMap<StructType, Integer> structTags,
+        PMap<UserDefinedEffect, Integer> effectTags
     ) {
         public static final GlobalContext EMPTY = forTypeInfos(TypeInfos.EMPTY);
 
         public static GlobalContext forTypeInfos(TypeInfos typeInfos) {
-            var structTags = P.<StructType, Integer>map();
-            var structTypes = typeInfos.structTypes();
-            var nextTag = 0;
-            for (var structType : structTypes) {
-                structTags = structTags.plus(structType, nextTag);
-                nextTag++;
-            }
+            var structTags = generateStructTags(typeInfos);
+            var effectTags = generateEffectTags(typeInfos);
 
             return new GlobalContext(
                 typeInfos,
-                structTags
+                structTags,
+                effectTags
             );
         }
 
+        private static PMap<StructType, Integer> generateStructTags(TypeInfos typeInfos) {
+            var tags = P.<StructType, Integer>map();
+            var nextTag = 0;
+            for (var structType : typeInfos.structTypes()) {
+                tags = tags.plus(structType, nextTag);
+                nextTag++;
+            }
+            return tags;
+        }
+
+        private static PMap<UserDefinedEffect, Integer> generateEffectTags(TypeInfos typeInfos) {
+            var tags = P.<UserDefinedEffect, Integer>map();
+            var nextTag = 0;
+            for (var effect : typeInfos.userDefinedEffects()) {
+                tags = tags.plus(effect, nextTag);
+                nextTag++;
+            }
+            return tags;
+        }
+
         public NamespaceTypeInfo namespaceTypeInfo(NamespaceName namespaceName) {
             return typeInfos.namespaceTypeInfo(namespaceName)
                 // TODO: throw better error
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeInfos.java b/src/main/java/org/shedlang/compiler/typechecker/TypeInfos.java
index 77bb5493..8de1f0de 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeInfos.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeInfos.java
@@ -90,4 +90,8 @@ public class TypeInfos {
     public Optional<UserDefinedEffectInfo> userDefinedEffectInfo(UserDefinedEffect effect) {
         return Optional.ofNullable(userDefinedEffectInfos.get(effect));
     }
+
+    public Set<UserDefinedEffect> userDefinedEffects() {
+        return this.userDefinedEffectInfos.keySet();
+    }
 }