Support mutual-recursion example in Patio

4 files changed, 100 insertions, 36 deletions

diff --git a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmPatioCodeGen.java b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmPatioCodeGen.java
index e5a3ac15..aa389dc2 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmPatioCodeGen.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmPatioCodeGen.java
@@ -697,8 +697,7 @@ public class LlvmPatioCodeGen {
     private static LlvmOperand patioDataValueToLlvmOperand(PatioDataValue value) {
         return switch (value) {
             case PatioDataValueArray patioDataValueArray -> {
-                var patioElement = patioDataValueArray.elements().get(0);
-                var llvmType = patioDataLayoutToLlvmType(patioElement.layout());
+                var llvmType = patioDataLayoutToLlvmType(patioDataValueArray.elementLayout());
 
                 var llvmElements = patioDataValueArray.elements().stream()
                     .map(element -> patioDataValueToLlvmOperand(element))
diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGen.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGen.java
index 45554f4c..11189e7b 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGen.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGen.java
@@ -360,13 +360,11 @@ public class PatioCodeGen {
                 yield context.temporaryPush(local);
             }
 
-            case LocalStoreInstruction localStore -> {
-                var localResult = context.addStackIrLocal(localStore.name());
-                var local = localResult.value();
-                context = localResult.context();
-
-                yield context.temporaryPopIntoLocal(local, localStore.source());
-            }
+            case LocalStoreInstruction localStore ->
+                context.addStackIrLocal(
+                    localStore.name(),
+                    (local, c) -> c.temporaryPopIntoLocal(local, localStore.source())
+                );
 
             case NamespaceFieldStoreInstruction namespaceFieldStore -> {
                 var fieldValueResult = context.temporaryPop();
diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenClosures.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenClosures.java
index 9920eeaa..db864db0 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenClosures.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenClosures.java
@@ -75,25 +75,20 @@ class PatioCodeGenClosures {
             var stackIrParams = Stream.concat(positionalParams.stream(), namedParams.stream()).toList();
 
             for (var stackIrParam : stackIrParams) {
-                var paramResult = bodyContext.addStackIrLocal(stackIrParam.name());
-                var param = paramResult.value();
-                bodyContext = paramResult.context();
-
-                bodyContext = bodyContext.addArg(param);
+                bodyContext = bodyContext.addStackIrArg(stackIrParam.name());
             }
 
             for (var freeVariable : freeVariables) {
-                var freeVariableLocalResult = bodyContext.addStackIrLocal(freeVariable.name());
-                var freeVariableLocal = freeVariableLocalResult.value();
-                bodyContext = freeVariableLocalResult.context();
-
-                bodyContext = bodyContext.addInstruction(new PatioMemoryLoad(
-                    freeVariableLocal,
-                    closurePointer,
-                    layout,
-                    freeVariableOffset(freeVariable),
-                    source
-                ));
+                bodyContext = bodyContext.addStackIrLocal(
+                    freeVariable.name(),
+                    (freeVariableLocal, c) -> c.addInstruction(new PatioMemoryLoad(
+                        freeVariableLocal,
+                        closurePointer,
+                        layout,
+                        freeVariableOffset(freeVariable),
+                        source
+                    ))
+                );
             }
 
             bodyContext = generateBody.apply(bodyContext);
@@ -126,13 +121,17 @@ class PatioCodeGenClosures {
         ));
 
         for (var freeVariable : freeVariables) {
-            context = context.addInstruction(new PatioMemoryStore(
-                closurePointer,
-                layout,
-                freeVariableOffset(freeVariable),
-                context.stackIrLocalNameToLocal(freeVariable.name()),
-                source
-            ));
+            context = context.onceStackIrLocalDefined(
+                freeVariable.name(),
+                c -> c.addInstruction(new PatioMemoryStore(
+                    closurePointer,
+                    layout,
+                    freeVariableOffset(freeVariable),
+                    c.stackIrLocalNameToLocal(freeVariable.name()),
+                    source
+                ))
+            );
+
         }
 
         return context;
diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java
index 59f307ce..17052735 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java
@@ -20,6 +20,7 @@ import org.shedlang.compiler.util.P;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Optional;
+import java.util.function.BiFunction;
 import java.util.function.Function;
 
 public record PatioCodeGenContext(
@@ -86,6 +87,7 @@ public record PatioCodeGenContext(
         PVector<PatioInstruction> instructions,
         OrderedPMap<String, PatioLocal> localsByName,
         PMap<String, PatioLocalRef> stackIrNameToLocal,
+        PMap<String, PVector<Function<PatioCodeGenContext, PatioCodeGenContext>>> waitingForStackIrLocal,
         PVector<PatioLocalRef> args
     ) {
         public static final FunctionContext EMPTY = new FunctionContext(
@@ -93,6 +95,7 @@ public record PatioCodeGenContext(
             P.vector(),
             OrderedPMap.empty(),
             P.map(),
+            P.map(),
             P.vector()
         );
 
@@ -102,6 +105,7 @@ public record PatioCodeGenContext(
                 this.instructions,
                 this.localsByName,
                 this.stackIrNameToLocal,
+                this.waitingForStackIrLocal,
                 this.args
             ).toResult(this.temporaryStack.get(0));
         }
@@ -112,6 +116,7 @@ public record PatioCodeGenContext(
                 this.instructions,
                 this.localsByName,
                 this.stackIrNameToLocal,
+                this.waitingForStackIrLocal,
                 this.args
             );
         }
@@ -122,6 +127,7 @@ public record PatioCodeGenContext(
                 this.instructions.plus(instruction),
                 this.localsByName,
                 this.stackIrNameToLocal,
+                this.waitingForStackIrLocal,
                 this.args
             );
         }
@@ -142,11 +148,12 @@ public record PatioCodeGenContext(
                 this.instructions,
                 this.localsByName.plus(fullName, new PatioLocal(fullName, type)),
                 this.stackIrNameToLocal,
+                this.waitingForStackIrLocal,
                 this.args
             ).toResult(new PatioLocalRef(fullName));
         }
 
-        public Result<PatioLocalRef> addStackIrLocal(String name) {
+        public Result<PVector<Function<PatioCodeGenContext, PatioCodeGenContext>>> addStackIrLocal(String name) {
             var context = this;
 
             if (this.stackIrNameToLocal.containsKey(name)) {
@@ -171,10 +178,11 @@ public record PatioCodeGenContext(
                 ),
                 // TODO: the mapping is redundant given we fix the mapping
                 this.stackIrNameToLocal.plus(name, localRef),
+                this.waitingForStackIrLocal.minus(name),
                 this.args
             );
 
-            return context.toResult(localRef);
+            return context.toResult(this.waitingForStackIrLocal.getOrDefault(name, P.vector()));
         }
 
         public PatioLocalRef stackIrLocalNameToLocal(String name) {
@@ -188,12 +196,36 @@ public record PatioCodeGenContext(
             return local;
         }
 
+        public Optional<PatioLocalRef> stackIrLocalNameToLocalOrNone(String name) {
+            var local = this.stackIrNameToLocal.get(name);
+
+            return Optional.ofNullable(local);
+        }
+
+        public FunctionContext addWaitingForShedLocal(
+            String name,
+            Function<PatioCodeGenContext, PatioCodeGenContext> func
+        ) {
+            return new FunctionContext(
+                this.temporaryStack,
+                this.instructions,
+                this.localsByName,
+                this.stackIrNameToLocal,
+                this.waitingForStackIrLocal.plus(
+                    name,
+                    this.waitingForStackIrLocal.getOrDefault(name, P.vector()).plus(func)
+                ),
+                this.args
+            );
+        }
+
         public FunctionContext addArg(PatioLocalRef local) {
             return new FunctionContext(
                 this.temporaryStack,
                 this.instructions,
                 this.localsByName,
                 this.stackIrNameToLocal,
+                this.waitingForStackIrLocal,
                 this.args.plus(local)
             );
         }
@@ -301,10 +333,44 @@ public record PatioCodeGenContext(
         return withFunctionContextResult(result);
     }
 
-    public PatioCodeGenResult<PatioLocalRef> addStackIrLocal(String name) {
+    public PatioCodeGenContext addStackIrArg(String name) {
+        return addStackIrLocal(
+            name,
+            (local, c) -> c.addArg(local)
+        );
+    }
+
+    public PatioCodeGenContext addStackIrLocal(
+        String name,
+        BiFunction<PatioLocalRef, PatioCodeGenContext, PatioCodeGenContext> setLocal
+    ) {
         var result = this.functionContext.addStackIrLocal(name);
 
-        return withFunctionContextResult(result);
+        var context = withFunctionContext(result.context);
+        var patioLocal = context.functionContext.stackIrLocalNameToLocal(name);
+
+        context = setLocal.apply(patioLocal, context);
+
+        for (var func : result.value) {
+            context = func.apply(context);
+        }
+
+        return context;
+    }
+
+    public PatioCodeGenContext onceStackIrLocalDefined(
+        String name,
+        Function<PatioCodeGenContext, PatioCodeGenContext> func
+    ) {
+        var local = this.functionContext.stackIrLocalNameToLocalOrNone(name);
+
+        if (local.isEmpty()) {
+            return withFunctionContext(
+                this.functionContext.addWaitingForShedLocal(name, func)
+            );
+        } else {
+            return func.apply(this);
+        }
     }
 
     public PatioLocalRef stackIrLocalNameToLocal(String name) {
@@ -365,6 +431,7 @@ public record PatioCodeGenContext(
                 P.vector(),
                 this.functionContext.localsByName,
                 this.functionContext.stackIrNameToLocal,
+                this.functionContext.waitingForStackIrLocal,
                 P.vector()
             )
         );
@@ -387,6 +454,7 @@ public record PatioCodeGenContext(
                 this.blockStack.get(0).instructions,
                 this.functionContext.localsByName,
                 this.functionContext.stackIrNameToLocal,
+                this.functionContext.waitingForStackIrLocal,
                 this.blockStack.get(0).args
             )
         ).toResult(this.functionContext.instructions);