Support capture of vars in effect handler methods

3 files changed, 71 insertions, 19 deletions

diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGen.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGen.java
index 7420b69b..1fb32e5c 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGen.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGen.java
@@ -2,6 +2,7 @@ package org.shedlang.compiler.patio;
 
 import org.pcollections.OrderedPSet;
 import org.shedlang.compiler.ast.typed.BoundVariable;
+import org.shedlang.compiler.ast.typed.FreeVariables;
 import org.shedlang.compiler.patio.ast.PatioCompilationUnit;
 import org.shedlang.compiler.patio.ast.PatioVisibility;
 import org.shedlang.compiler.patio.ast.instructions.*;
@@ -18,6 +19,7 @@ import org.shedlang.compiler.types.NamespaceName;
 import org.shedlang.compiler.util.P;
 
 import java.util.*;
+import java.util.stream.Collectors;
 
 public class PatioCodeGen {
     private PatioCodeGen() {}
@@ -104,7 +106,11 @@ public class PatioCodeGen {
         context = PatioCodeGenNamespaces.define(
             stackIrNamespace.namespaceName(),
             initContext -> {
-                initContext = declareStackIrLocals(stackIrNamespace.instructions(), initContext);
+                initContext = declareStackIrLocals(
+                    stackIrNamespace.instructions(),
+                    FreeVariables.empty(),
+                    initContext
+                );
 
                 return generateInstructions(stackIrNamespace.instructions(), initContext);
             },
@@ -237,7 +243,11 @@ public class PatioCodeGen {
                     functionContext -> {
                         // TODO: explicitly store locals on a function?
 
-                        functionContext = declareStackIrLocals(defineFunction.body(), functionContext);
+                        functionContext = declareStackIrLocals(
+                            defineFunction.body(),
+                            defineFunction.freeVariables(),
+                            functionContext
+                        );
 
                         return generateInstructions(defineFunction.body(), functionContext);
                     },
@@ -857,9 +867,21 @@ public class PatioCodeGen {
         return context.toResult(conditionBool);
     }
 
-    private static PatioCodeGenContext declareStackIrLocals(List<Instruction> body, PatioCodeGenContext context) {
-        for (var boundVariable : findBoundVariables(body, P.orderedSet())) {
-            if (boundVariable.isMutable()) {
+    private static PatioCodeGenContext declareStackIrLocals(
+        List<Instruction> body,
+        FreeVariables freeVariables,
+        PatioCodeGenContext context
+    ) {
+        var freeVariableNames = freeVariables.stream()
+            .map(freeVariable -> freeVariable.name())
+            .collect(Collectors.toSet());
+
+        var boundVariables = findAssignedVariables(body, P.orderedSet());
+
+        for (var boundVariable : boundVariables) {
+            if (freeVariableNames.contains(boundVariable.name())) {
+                // This is a terrible hack. Another reason to collect variables up-front.
+            } else if (boundVariable.isMutable()) {
                 var patioLocalResult = context.declareStackIrVal(
                     boundVariable.name(),
                     PatioTypes.DATA_POINTER
@@ -885,7 +907,10 @@ public class PatioCodeGen {
         return context;
     }
 
-    private static OrderedPSet<BoundVariable> findBoundVariables(List<Instruction> body, OrderedPSet<BoundVariable> boundVariables) {
+    private static OrderedPSet<BoundVariable> findAssignedVariables(
+        List<Instruction> body,
+        OrderedPSet<BoundVariable> boundVariables
+    ) {
         for (var bodyInstruction : body) {
             switch (bodyInstruction) {
                 case CallInstruction callInstruction -> {
@@ -903,8 +928,8 @@ public class PatioCodeGen {
                 case FieldLoadInstruction fieldLoadInstruction -> {
                 }
                 case IfInstruction ifInstruction -> {
-                    boundVariables = findBoundVariables(ifInstruction.trueBody(), boundVariables);
-                    boundVariables = findBoundVariables(ifInstruction.falseBody(), boundVariables);
+                    boundVariables = findAssignedVariables(ifInstruction.trueBody(), boundVariables);
+                    boundVariables = findAssignedVariables(ifInstruction.falseBody(), boundVariables);
                 }
                 case IntAddInstruction intAddInstruction -> {
                 }
@@ -924,7 +949,7 @@ public class PatioCodeGen {
                 case LoopContinueInstruction loopContinueInstruction -> {
                 }
                 case LoopInstruction loopInstruction -> {
-                    boundVariables = findBoundVariables(loopInstruction.body(), boundVariables);
+                    boundVariables = findAssignedVariables(loopInstruction.body(), boundVariables);
                 }
                 case NamespaceFieldStoreInstruction namespaceFieldStoreInstruction -> {
                 }
@@ -944,11 +969,11 @@ public class PatioCodeGen {
                 }
                 case SwitchInstruction switchInstruction -> {
                     for (var switchCase : switchInstruction.cases()) {
-                        boundVariables = findBoundVariables(switchCase.body(), boundVariables);
+                        boundVariables = findAssignedVariables(switchCase.body(), boundVariables);
                     }
                 }
                 case TryInstruction tryInstruction -> {
-                    boundVariables = findBoundVariables(tryInstruction.body(), boundVariables);
+                    boundVariables = findAssignedVariables(tryInstruction.body(), boundVariables);
                 }
                 case TupleCreateInstruction tupleCreateInstruction -> {
                 }
diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenClosures.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenClosures.java
index 5ff66e0f..961eeef0 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenClosures.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenClosures.java
@@ -10,7 +10,6 @@ import org.shedlang.compiler.patio.ast.operands.PatioOperand;
 import org.shedlang.compiler.patio.ast.types.*;
 import org.shedlang.compiler.sources.Source;
 import org.shedlang.compiler.stackir.instructions.DefineFunctionInstruction;
-import org.shedlang.compiler.util.Lists;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -34,13 +33,30 @@ class PatioCodeGenClosures {
         ));
     }
 
-    private static PatioDataLayout closureDataLayout(int freeVariableCount) {
+    private static PatioDataLayout closureDataLayout(
+        FreeVariables freeVariables,
+        PatioCodeGenContext context
+    ) {
         return PatioDataLayouts.struct(List.of(
             PatioDataLayouts.FUNCTION_POINTER,
-            PatioDataLayouts.struct(Lists.repeat(PatioDataLayouts.INT, freeVariableCount))
+            PatioDataLayouts.struct(
+                freeVariables.stream()
+                    .map(freeVariable -> PatioDataLayouts.fromType(
+                        freeVariableType(freeVariable, context)
+                    ))
+                    .toList()
+            )
         ));
     }
 
+    private static PatioType freeVariableType(
+        FreeVariable freeVariable,
+        PatioCodeGenContext context
+    ) {
+        var local = context.stackIrLocalNameToLocal(freeVariable.name());
+        return context.localType(local);
+    }
+
     private static List<Integer> freeVariableOffset(FreeVariable freeVariable) {
         return List.of(1, freeVariable.index());
     }
@@ -55,7 +71,7 @@ class PatioCodeGenClosures {
         Source source,
         PatioCodeGenContext context
     ) {
-        var layout = closureDataLayout(freeVariables.size());
+        var layout = closureDataLayout(freeVariables, context);
 
         var fullNameResult = context.allocateCompilationUnitElementName(
             PatioCodeGenCompilationUnitElementSpace.FUNCTION,
@@ -64,6 +80,7 @@ class PatioCodeGenClosures {
         var fullName = fullNameResult.value();
         context = fullNameResult.context();
 
+        var bodyContextParent = context;
         context = context.createFunction(fullName, PatioVisibility.PRIVATE, PatioCodeGenTypes.VALUE_TYPE, source, bodyContext -> {
             var closurePointerResult = bodyContext.allocateParam(
                 PatioCodeGenLocalSpace.CLOSURE_POINTER,
@@ -82,8 +99,7 @@ class PatioCodeGenClosures {
             for (var freeVariable : freeVariables) {
                 var freeVariableLocalResult = bodyContext.declareStackIrVal(
                     freeVariable.name(),
-                    // TODO: vary based on mutability
-                    PatioCodeGenTypes.VALUE_TYPE
+                    freeVariableType(freeVariable, bodyContextParent)
                 );
                 var freeVariableLocal = freeVariableLocalResult.value();
                 bodyContext = freeVariableLocalResult.context();
diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java
index 9509036c..cfb075c3 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java
@@ -323,6 +323,12 @@ public record PatioCodeGenContext(
         }
 
         @CheckReturnValue
+        public PatioType localType(PatioLocalRef local){
+            // TODO: handle missing key.
+            return this.localsByName.get(local.name()).type();
+        }
+
+        @CheckReturnValue
         public FunctionContext addWaitingForShedLocal(
             String name,
             Function<PatioCodeGenContext, PatioCodeGenContext> func
@@ -531,7 +537,7 @@ public record PatioCodeGenContext(
     }
 
     @CheckReturnValue
-    private PatioCodeGenResult<PatioLocalRef> declareStackIrLocal(
+    public PatioCodeGenResult<PatioLocalRef> declareStackIrLocal(
         String name,
         PatioType type,
         PatioLocalType patioLocalType
@@ -580,6 +586,11 @@ public record PatioCodeGenContext(
     }
 
     @CheckReturnValue
+    public PatioType localType(PatioLocalRef local) {
+        return this.functionContext.localType(local);
+    }
+
+    @CheckReturnValue
     public PatioCodeGenContext createFunction(
         String name,
         PatioVisibility visibility,
@@ -682,7 +693,7 @@ public record PatioCodeGenContext(
         );
     }
 
-    private interface PatioLocalFactory {
+    public interface PatioLocalFactory {
         PatioLocal create(String name, PatioType type);
     }
 }