Generate Patio code for defining an effect

5 files changed, 195 insertions, 2 deletions

diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGen.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGen.java
index e5dcc654..80f5b1c8 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGen.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGen.java
@@ -195,7 +195,24 @@ public class PatioCodeGen {
             }
 
             case DefineEffectInstruction defineEffect -> {
-                throw new UnsupportedOperationException("TODO");
+                var targetResult = context.allocateVal(
+                    PatioCodeGenLocalSpace.USER_DEFINED_EFFECT,
+                    defineEffect.effect().name(),
+                    PatioCodeGenTypes.VALUE_TYPE
+                );
+                var target = targetResult.value();
+                context = targetResult.context();
+
+                context = PatioCodeGenEffects.createUserDefinedEffect(
+                    target,
+                    defineEffect.effect(),
+                    defineEffect.source(),
+                    context
+                );
+
+                context = context.temporaryPush(target);
+
+                yield context;
             }
 
             case DefineFunctionInstruction defineFunction -> {
diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenClosures.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenClosures.java
index f4f5f6be..ad5ae9f2 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenClosures.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenClosures.java
@@ -21,7 +21,8 @@ import java.util.stream.Stream;
 import static org.shedlang.compiler.util.Iterables.sortedBy;
 
 class PatioCodeGenClosures {
-    public static final PatioType CLOSURE_POINTER_TYPE = PatioTypes.DATA_POINTER;
+    private static final PatioType CLOSURE_POINTER_TYPE = PatioTypes.DATA_POINTER;
+    public static final PatioType CLOSURE_TYPE = PatioCodeGenTypes.VALUE_TYPE;
 
     private PatioCodeGenClosures() {}
 
diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenEffects.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenEffects.java
new file mode 100644
index 00000000..dd1f414b
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenEffects.java
@@ -0,0 +1,167 @@
+package org.shedlang.compiler.patio;
+
+import org.shedlang.compiler.ast.typed.FreeVariables;
+import org.shedlang.compiler.patio.ast.PatioFunctionDeclaration;
+import org.shedlang.compiler.patio.ast.PatioLocal;
+import org.shedlang.compiler.patio.ast.instructions.*;
+import org.shedlang.compiler.patio.ast.operands.PatioInt;
+import org.shedlang.compiler.patio.ast.operands.PatioLocalRef;
+import org.shedlang.compiler.patio.ast.types.PatioDataLayout;
+import org.shedlang.compiler.patio.ast.types.PatioDataLayouts;
+import org.shedlang.compiler.patio.ast.types.PatioType;
+import org.shedlang.compiler.patio.ast.types.PatioTypes;
+import org.shedlang.compiler.sources.Source;
+import org.shedlang.compiler.stackir.instructions.DefineFunctionInstruction;
+import org.shedlang.compiler.types.UserDefinedEffect;
+import org.shedlang.compiler.util.Lists;
+
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.stream.IntStream;
+
+public class PatioCodeGenEffects {
+    private static final PatioType EFFECT_TAG_TYPE = PatioTypes.INT;
+    private static final PatioType EFFECT_METHOD_INDEX_TYPE = PatioTypes.INT;
+
+    private PatioCodeGenEffects() {
+    }
+
+    public static PatioCodeGenContext createUserDefinedEffect(
+        PatioLocalRef target,
+        UserDefinedEffect effect,
+        Source source,
+        PatioCodeGenContext context
+    ) {
+        var effectPointerResult = context.allocateVal(
+            PatioCodeGenLocalSpace.USER_DEFINED_EFFECT_POINTER,
+            effect.name(),
+            PatioTypes.DATA_POINTER
+        );
+        var effectPointer = effectPointerResult.value();
+        context = effectPointerResult.context();
+
+        var effectLayout = userDefinedEffectLayout(effect, context);
+
+        context = context.addInstruction(new PatioAllocHeap(
+            target,
+            effectLayout,
+            source
+        ));
+
+        var effectInfo = context.userDefinedEffectInfo(effect);
+        for (var method : effectInfo.methods()) {
+            var effectMethodResult = context.allocateVal(
+                PatioCodeGenLocalSpace.USER_DEFINED_EFFECT_METHOD,
+                method.name(),
+                PatioCodeGenClosures.CLOSURE_TYPE
+            );
+            var effectMethod = effectMethodResult.value();
+            context = effectMethodResult.context();
+
+            var positionalParams = IntStream.range(0, method.type().positionalParams().size())
+                // TODO: ensure names of positional params doesn't clash with named params
+                .mapToObj(paramIndex -> new DefineFunctionInstruction.Param("arg" + paramIndex, source))
+                .toList();
+
+            context = PatioCodeGenClosures.create(
+                effectMethod,
+                method.name(),
+                positionalParams,
+                method.type().namedParams().stream()
+                    .map(namedParam -> new DefineFunctionInstruction.Param(namedParam.name(), source))
+                    .toList(),
+                FreeVariables.empty(),
+                bodyContext -> {
+                    var runtimeFindHandlerMethodName = "shed__runtime__effects_find_handler_method";
+
+                    // TODO: check if element already exists?
+                    bodyContext = bodyContext.addCompilationUnitElement(new PatioFunctionDeclaration(
+                        runtimeFindHandlerMethodName,
+                        List.of(
+                            new PatioLocal.Param("effectTag", EFFECT_TAG_TYPE),
+                            new PatioLocal.Param("methodIndex", EFFECT_METHOD_INDEX_TYPE)
+                        ),
+                        PatioCodeGenTypes.VALUE_TYPE,
+                        source
+                    ));
+
+                    var handlerMethodResult = bodyContext.allocateVal(
+                        PatioCodeGenLocalSpace.EFFECT_HANDLER_METHOD,
+                        PatioCodeGenClosures.CLOSURE_TYPE
+                    );
+                    var handlerMethod = handlerMethodResult.value();
+                    bodyContext = handlerMethodResult.context();
+
+                    bodyContext = bodyContext.addInstruction(new PatioCall(
+                        Optional.of(handlerMethod),
+                        runtimeFindHandlerMethodName,
+                        List.of(
+                            new PatioInt(bodyContext.effectTag(effect)),
+                            new PatioInt(method.fieldIndex())
+                        ),
+                        source
+                    ));
+
+                    var returnValueResult = bodyContext.allocateVal(
+                        PatioCodeGenLocalSpace.TEMPORARY,
+                        PatioCodeGenTypes.VALUE_TYPE
+                    );
+                    var returnValue = returnValueResult.value();
+                    bodyContext = returnValueResult.context();
+
+                    var bodyContextArgs = bodyContext;
+                    bodyContext = PatioCodeGenClosures.call(
+                        returnValue,
+                        handlerMethod,
+                        positionalParams.stream()
+                            .map(param -> bodyContextArgs.stackIrLocalNameToLocal(param.name()))
+                            .toList(),
+                        method.type().namedParams().stream()
+                            .map(param -> Map.entry(
+                                param.name(),
+                                bodyContextArgs.stackIrLocalNameToLocal(param.name())
+                            ))
+                            .toList(),
+                        source,
+                        bodyContext
+                    );
+
+                    bodyContext = bodyContext.addInstruction(new PatioReturn(returnValue, source));
+
+                    return bodyContext;
+                },
+                source,
+                context
+            );
+
+            context = context
+                .addInstruction(new PatioMemoryStore(
+                    effectPointer,
+                    effectLayout,
+                    List.of(method.fieldIndex()),
+                    effectMethod,
+                    source
+                ));
+        }
+
+        context = context.addInstruction(new PatioDataPtrToInt(
+            target,
+            effectPointer,
+            source
+        ));
+
+        return context;
+    }
+
+    private static PatioDataLayout userDefinedEffectLayout(
+        UserDefinedEffect effect,
+        PatioCodeGenContext context
+    ) {
+        var effectInfo = context.userDefinedEffectInfo(effect);
+        return PatioDataLayouts.struct(Lists.repeat(
+            PatioDataLayouts.fromType(PatioCodeGenClosures.CLOSURE_TYPE),
+            effectInfo.methodCount()
+        ));
+    }
+}
diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenLocalSpace.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenLocalSpace.java
index cbe6c9b5..255a69c5 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenLocalSpace.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenLocalSpace.java
@@ -7,6 +7,7 @@ public enum PatioCodeGenLocalSpace {
     CLOSURE("closure"),
     CLOSURE_POINTER("closurePointer"),
     CLOSURE_FUNCTION_POINTER("closureFunctionPointer"),
+    EFFECT_HANDLER_METHOD("effectHandlerMethod"),
     MAIN("main"),
     NAMESPACE_IS_INITED_POINTER("namespaceIsInitedPointer"),
     NAMESPACE_POINTER("namespacePointer"),
@@ -23,6 +24,9 @@ public enum PatioCodeGenLocalSpace {
     TUPLE("tuple"),
     TUPLE_ELEMENT("tupleElement"),
     TUPLE_POINTER("tuplePointer"),
+    USER_DEFINED_EFFECT("userDefinedEffect"),
+    USER_DEFINED_EFFECT_METHOD("userDefinedMethod"),
+    USER_DEFINED_EFFECT_POINTER("userDefinedEffectPointer"),
     VAR_VALUE("varValue");
 
     private final String name;
diff --git a/src/main/java/org/shedlang/compiler/types/UserDefinedEffectInfo.java b/src/main/java/org/shedlang/compiler/types/UserDefinedEffectInfo.java
index d5276e9f..fe0ef3dd 100644
--- a/src/main/java/org/shedlang/compiler/types/UserDefinedEffectInfo.java
+++ b/src/main/java/org/shedlang/compiler/types/UserDefinedEffectInfo.java
@@ -3,6 +3,10 @@ package org.shedlang.compiler.types;
 import java.util.Optional;
 
 public record UserDefinedEffectInfo(FieldsInfo<FunctionType> methods) {
+    public int methodCount() {
+        return methods.fieldCount();
+    }
+
     public Optional<FunctionType> methodType(String methodName) {
         return methods.fieldType(methodName);
     }