Use runtime type tags as effect tags in effect handler

2 files changed, 73 insertions, 19 deletions

diff --git a/src/main/java/org/shedlang/compiler/shidtopatio/ShidToPatioContext.java b/src/main/java/org/shedlang/compiler/shidtopatio/ShidToPatioContext.java
index cab33a07..bdfacb9a 100644
--- a/src/main/java/org/shedlang/compiler/shidtopatio/ShidToPatioContext.java
+++ b/src/main/java/org/shedlang/compiler/shidtopatio/ShidToPatioContext.java
@@ -38,21 +38,18 @@ public record ShidToPatioContext(
 
         @CheckReturnValue
         public static GlobalContext forTypeInfos(TypeInfos typeInfos) {
-            var effectTags = generateEffectTags(typeInfos);
-            var runtimeTypeTags = P.<RuntimeType, Integer>orderedPMap();
-
             var context = new GlobalContext(
                 typeInfos,
                 P.map(),
-                effectTags,
-                runtimeTypeTags
+                P.map(),
+                P.orderedPMap()
             );
 
             // Ensure that scalars have tag of 0: this is relied on in the
             // runtime (for instance, when allocating ints and strings).
             context = context.addRuntimeTypeTag(ScalarRuntimeType.INSTANCE);
-
             context = generateStructTags(context);
+            context = generateEffectTags(context);
 
             return context;
         }
@@ -82,14 +79,27 @@ public record ShidToPatioContext(
         }
 
         @CheckReturnValue
-        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++;
+        private static GlobalContext generateEffectTags(GlobalContext context) {
+            var effectTags = P.<UserDefinedEffect, Integer>map();
+
+            for (var effect : context.typeInfos.userDefinedEffects()) {
+                // TODO: get effects with infos to avoid potential throw.
+                var effectInfo = context.typeInfos.userDefinedEffectInfo(effect).orElseThrow();
+                var effectHandlerRuntimeType = ShidToPatioEffects.effectHandlerRuntimeType(effect, effectInfo);
+
+                var runtimeTypeTagResult = context.runtimeTypeTag(effectHandlerRuntimeType);
+                var runtimeTypeTag = runtimeTypeTagResult.value();
+                context = runtimeTypeTagResult.context();
+
+                effectTags = effectTags.plus(effect, runtimeTypeTag);
             }
-            return tags;
+
+            return new GlobalContext(
+                context.typeInfos,
+                context.structTags,
+                effectTags,
+                context.runtimeTypeTags
+            );
         }
 
         @CheckReturnValue
diff --git a/src/main/java/org/shedlang/compiler/shidtopatio/ShidToPatioEffects.java b/src/main/java/org/shedlang/compiler/shidtopatio/ShidToPatioEffects.java
index 9f165a2d..8cf27d7d 100644
--- a/src/main/java/org/shedlang/compiler/shidtopatio/ShidToPatioEffects.java
+++ b/src/main/java/org/shedlang/compiler/shidtopatio/ShidToPatioEffects.java
@@ -16,6 +16,7 @@ import org.shedlang.compiler.shid.instructions.ShidDefineFunction;
 import org.shedlang.compiler.sources.Source;
 import org.shedlang.compiler.sources.Sources;
 import org.shedlang.compiler.types.UserDefinedEffect;
+import org.shedlang.compiler.types.UserDefinedEffectInfo;
 import org.shedlang.compiler.util.Lists;
 
 import java.util.List;
@@ -83,6 +84,42 @@ public class ShidToPatioEffects {
         }
     }
 
+    public static EffectHandlerRuntimeType effectHandlerRuntimeType(
+        UserDefinedEffect effect,
+        UserDefinedEffectInfo effectInfo
+    ) {
+        return new EffectHandlerRuntimeType(effect, effectInfo.methodCount());
+    }
+
+    public record EffectHandlerRuntimeType(
+        UserDefinedEffect effect,
+        int methodCount
+    ) implements RuntimeType {
+        @Override
+        public PatioDataLayout layout() {
+            return ShidToPatioGc.gcObjectDataLayout(PatioDataLayouts.struct(List.of(
+                PatioDataLayouts.fromType(EFFECT_HANDLER_TYPE),
+                PatioDataLayouts.struct(Lists.repeat(
+                    PatioDataLayouts.fromType(EFFECT_HANDLER_METHOD_POINTER_TYPE),
+                    methodCount
+                ))
+            )));
+        }
+
+        private static final List<Integer> GC_CHILDREN_OFFSET =
+            ShidToPatioGc.gcObjectOffset(List.of(1));
+
+        @Override
+        public List<Integer> gcChildrenOffset() {
+            return GC_CHILDREN_OFFSET;
+        }
+
+        @Override
+        public int gcChildrenCount() {
+            return methodCount;
+        }
+    }
+
     private static final PatioFunctionDeclaration FIND_HANDLER_METHOD_DECLARATION = new PatioFunctionDeclaration(
         "shed__runtime__effects_find_handler_method",
         List.of(
@@ -110,10 +147,9 @@ public class ShidToPatioEffects {
     private static final PatioFunctionDeclaration PUSH_HANDLER_DECLARATION = new PatioFunctionDeclaration(
         "shed__runtime__effects_push_handler",
         List.of(
-            PatioLocal.param("effectTag", EFFECT_TAG_TYPE),
-            PatioLocal.param("methodCount", PatioTypes.INT)
+            PatioLocal.param("effectHandler", EFFECT_HANDLER_TYPE)
         ),
-        EFFECT_HANDLER_TYPE,
+        PatioTypes.VOID,
         Sources.RUNTIME
     );
 
@@ -321,12 +357,20 @@ public class ShidToPatioEffects {
         var effectHandler = effectHandlerResult.value();
         context = effectHandlerResult.context();
 
+        var effectHandlerRuntimeType = effectHandlerRuntimeType(effect, effectInfo);
+
+        context = ShidToPatioGc.alloc(
+            effectHandler,
+            effectHandlerRuntimeType,
+            source,
+            context
+        );
+
         context = context.addInstruction(new PatioCall(
-            Optional.of(effectHandler),
+            Optional.empty(),
             PUSH_HANDLER_DECLARATION.name(),
             List.of(
-                PatioIntConst.of(context.effectTag(effect)),
-                PatioIntConst.of(effectInfo.methodCount())
+                effectHandler
             ),
             source
         ));