Support mutual recursion

5 files changed, 157 insertions, 37 deletions

diff --git a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGenClosures.java b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGenClosures.java
index b32d69b7..9ec1b97a 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGenClosures.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGenClosures.java
@@ -350,11 +350,15 @@ public class LlvmCodeGenClosures {
         var closureEnvironmentPointer = closureEnvironmentPointerResult.value();
         context = closureEnvironmentPointerResult.context();
 
-        return context.addInstruction(new LlvmStore(
-            LlvmCodeGenTypes.VALUE_TYPE,
-            context.shedLocal(freeVariable.name()),
-            closureEnvironmentPointer
-        ));
+        // TODO: test behaviour when local is defined later
+        return context.onceShedLocalDefined(
+            freeVariable.name(),
+            c -> c.addInstruction(new LlvmStore(
+                LlvmCodeGenTypes.VALUE_TYPE,
+                c.shedLocal(freeVariable.name()),
+                closureEnvironmentPointer
+            ))
+        );
     }
 
     private static LlvmCodeGenContext loadClosureFreeVariable(
diff --git a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGenContext.java b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGenContext.java
index eef011f4..b73d7138 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGenContext.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGenContext.java
@@ -11,6 +11,7 @@ import org.shedlang.compiler.backends.llvm.ir.ast.operands.LlvmOperand;
 import org.shedlang.compiler.backends.llvm.ir.ast.operands.LlvmVariableLocal;
 import org.shedlang.compiler.backends.llvm.ir.ast.types.LlvmType;
 import org.shedlang.compiler.typechecker.TypeInfo;
+import org.shedlang.compiler.util.P;
 import org.shedlang.compiler.util.POrderedMap;
 
 import java.util.List;
@@ -66,6 +67,7 @@ public record LlvmCodeGenContext(
     record FunctionContext(
         PStack<LlvmOperand> stack,
         PMap<String, LlvmOperand> shedLocalNameToOperand,
+        PMap<String, PVector<Function<LlvmCodeGenContext, LlvmCodeGenContext>>> waitingForShedLocal,
         PVector<LlvmBasicBlock> basicBlocks,
         Optional<String> basicBlockLabel,
         PVector<LlvmInstruction> instructions,
@@ -74,6 +76,7 @@ public record LlvmCodeGenContext(
         public static final FunctionContext EMPTY = new FunctionContext(
             Empty.stack(),
             Empty.map(),
+            P.map(),
             Empty.vector(),
             Optional.empty(),
             Empty.vector(),
@@ -84,6 +87,7 @@ public record LlvmCodeGenContext(
             return new FunctionContext(
                 stack.plus(operand),
                 shedLocalNameToOperand,
+                waitingForShedLocal,
                 basicBlocks,
                 basicBlockLabel,
                 instructions,
@@ -95,6 +99,7 @@ public record LlvmCodeGenContext(
             return new FunctionContext(
                 stack.minus(0),
                 shedLocalNameToOperand,
+                waitingForShedLocal,
                 basicBlocks,
                 basicBlockLabel,
                 instructions,
@@ -102,15 +107,22 @@ public record LlvmCodeGenContext(
             ).toResult(stack.get(0));
         }
 
-        public FunctionContext addShedLocal(String name, LlvmOperand value) {
+        public Result<PVector<Function<LlvmCodeGenContext, LlvmCodeGenContext>>> addShedLocal(String name, LlvmOperand value) {
+            var funcs = this.waitingForShedLocal.getOrDefault(name, P.vector());
+
             return new FunctionContext(
                 stack,
                 shedLocalNameToOperand.plus(name, value),
+                waitingForShedLocal.minus(name),
                 basicBlocks,
                 basicBlockLabel,
                 instructions,
                 localVariableNames
-            );
+            ).toResult(funcs);
+        }
+
+        public Optional<LlvmOperand> shedLocalOrNone(String name) {
+            return Optional.ofNullable(shedLocalNameToOperand.get(name));
         }
 
         public LlvmOperand shedLocal(String name) {
@@ -123,6 +135,21 @@ public record LlvmCodeGenContext(
             }
         }
 
+        public FunctionContext addWaitingForShedLocal(String name, Function<LlvmCodeGenContext, LlvmCodeGenContext> func) {
+            return new FunctionContext(
+                stack,
+                shedLocalNameToOperand,
+                waitingForShedLocal.plus(
+                    name,
+                    waitingForShedLocal.getOrDefault(name, P.vector()).plus(func)
+                ),
+                basicBlocks,
+                basicBlockLabel,
+                instructions,
+                localVariableNames
+            );
+        }
+
         public FunctionContext setBasicBlockLabel(String basicBlockLabel) {
             if (this.basicBlockLabel.isPresent()) {
                 // TODO: better error
@@ -132,6 +159,7 @@ public record LlvmCodeGenContext(
             return new FunctionContext(
                 stack,
                 shedLocalNameToOperand,
+                waitingForShedLocal,
                 basicBlocks,
                 Optional.of(basicBlockLabel),
                 instructions,
@@ -143,6 +171,7 @@ public record LlvmCodeGenContext(
             return new FunctionContext(
                 stack,
                 shedLocalNameToOperand,
+                waitingForShedLocal,
                 basicBlocks,
                 basicBlockLabel,
                 instructions.plus(instruction),
@@ -156,6 +185,7 @@ public record LlvmCodeGenContext(
             return new FunctionContext(
                 stack,
                 shedLocalNameToOperand,
+                waitingForShedLocal,
                 basicBlocks.plus(new LlvmBasicBlock(
                     basicBlockLabel,
                     instructions,
@@ -175,6 +205,7 @@ public record LlvmCodeGenContext(
             return new FunctionContext(
                 stack,
                 shedLocalNameToOperand,
+                waitingForShedLocal,
                 basicBlocks,
                 basicBlockLabel,
                 instructions,
@@ -287,10 +318,34 @@ public record LlvmCodeGenContext(
     }
 
     public LlvmCodeGenContext addShedLocal(String name, LlvmOperand value) {
-        return new LlvmCodeGenContext(
+        var result = this.functionContext.addShedLocal(name, value);
+
+        var context = new LlvmCodeGenContext(
             this.moduleContext,
-            this.functionContext.addShedLocal(name, value)
+            result.context
         );
+
+        for (var func : result.value) {
+            context = func.apply(context);
+        }
+
+        return context;
+    }
+
+    public LlvmCodeGenContext onceShedLocalDefined(
+        String name,
+        Function<LlvmCodeGenContext, LlvmCodeGenContext> func
+    ) {
+        var local = this.functionContext.shedLocalOrNone(name);
+
+        if (local.isEmpty()) {
+            return new LlvmCodeGenContext(
+                this.moduleContext,
+                this.functionContext.addWaitingForShedLocal(name, func)
+            );
+        } else {
+            return func.apply(this);
+        }
     }
 
     public LlvmOperand shedLocal(String name) {
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGen.java b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGen.java
index ac8c8d62..254d0b86 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGen.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGen.java
@@ -215,10 +215,10 @@ public class WasmCodeGen {
                 context.addInstruction(new WasmLocalGet(context.shedLocal(localLoad.name())));
 
             case LocalStoreInstruction localStore -> {
-                var localResult = context.addShedLocal(localStore.name());
-                var local = localResult.value();
-                context = localResult.context();
-                yield context.addInstruction(new WasmLocalSet(local));
+                yield context.addShedLocal(
+                    localStore.name(),
+                    (local, c) -> c.addInstruction(new WasmLocalSet(local))
+                );
             }
 
             case NamespaceFieldStoreInstruction namespaceFieldStore -> {
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenClosures.java b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenClosures.java
index 28fe631c..be7406bb 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenClosures.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenClosures.java
@@ -56,19 +56,17 @@ public class WasmCodeGenClosures {
             bodyContext = closurePointerResult.context();
 
             for (var shedParam : shedParams) {
-                var paramResult = bodyContext.addShedLocal(shedParam.name());
-                bodyContext = paramResult.context();
+                bodyContext = bodyContext.addShedArg(shedParam.name());
             }
 
             for (var freeVariable : freeVariables) {
-                var freeVariableResult = bodyContext.addShedLocal(freeVariable.name());
-                var freeVariableIndex = freeVariableResult.value();
-                bodyContext = freeVariableResult.context();
-
-                bodyContext = bodyContext
-                    .addInstruction(new WasmLocalGet(closurePointer))
-                    .addInstruction(new WasmI32Load(freeVariableMemoryArg(freeVariable)))
-                    .addInstruction(new WasmLocalSet(freeVariableIndex));
+                bodyContext = bodyContext.addShedLocal(
+                    freeVariable.name(),
+                    (freeVariableIndex, c) -> c
+                        .addInstruction(new WasmLocalGet(closurePointer))
+                        .addInstruction(new WasmI32Load(freeVariableMemoryArg(freeVariable)))
+                        .addInstruction(new WasmLocalSet(freeVariableIndex))
+                );
             }
 
             bodyContext = compileBody.apply(bodyContext);
@@ -91,10 +89,14 @@ public class WasmCodeGenClosures {
             .addInstruction(new WasmI32Store(functionPointerMemoryArg()));
 
         for (var freeVariable : freeVariables) {
-            context = context
-                .addInstruction(new WasmLocalGet(closurePointer))
-                .addInstruction(new WasmLocalGet(context.shedLocal(freeVariable.name())))
-                .addInstruction(new WasmI32Store(freeVariableMemoryArg(freeVariable)));
+            // TODO: test behaviour when local is defined later
+            context = context.onceShedLocalDefined(
+                freeVariable.name(),
+                c -> c
+                    .addInstruction(new WasmLocalGet(closurePointer))
+                    .addInstruction(new WasmLocalGet(c.shedLocal(freeVariable.name())))
+                    .addInstruction(new WasmI32Store(freeVariableMemoryArg(freeVariable)))
+            );
         }
 
         context = context.addInstruction(new WasmLocalGet(closurePointer));
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenContext.java b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenContext.java
index 1cb0e2e2..afce597f 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenContext.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenContext.java
@@ -15,6 +15,7 @@ import org.shedlang.compiler.util.P;
 
 import java.util.List;
 import java.util.Optional;
+import java.util.function.BiFunction;
 import java.util.function.Function;
 
 public record WasmCodeGenContext(
@@ -171,7 +172,8 @@ public record WasmCodeGenContext(
         PVector<WasmInstruction> instructions,
         PStack<PVector<WasmInstruction>> blockStack,
         OrderedPMap<String, WasmLocal> locals,
-        PMap<String, WasmLocalIndex> shedLocalNameToLocalIndex
+        PMap<String, WasmLocalIndex> shedLocalNameToLocalIndex,
+        PMap<String, PVector<Function<WasmCodeGenContext, WasmCodeGenContext>>> waitingForShedLocal
     ) {
         public static final FunctionContext EMPTY = forFuncType(
             new WasmFuncType(List.of(), List.of())
@@ -183,6 +185,7 @@ public record WasmCodeGenContext(
                 Empty.vector(),
                 Empty.stack(),
                 OrderedPMap.empty(),
+                Empty.map(),
                 Empty.map()
             );
         }
@@ -193,7 +196,8 @@ public record WasmCodeGenContext(
                 this.instructions.plus(instruction),
                 this.blockStack,
                 this.locals,
-                this.shedLocalNameToLocalIndex
+                this.shedLocalNameToLocalIndex,
+                this.waitingForShedLocal
             );
         }
 
@@ -203,7 +207,8 @@ public record WasmCodeGenContext(
                 P.vector(),
                 this.blockStack.plus(this.instructions),
                 this.locals,
-                this.shedLocalNameToLocalIndex
+                this.shedLocalNameToLocalIndex,
+                this.waitingForShedLocal
             );
         }
 
@@ -213,7 +218,8 @@ public record WasmCodeGenContext(
                 this.blockStack.get(0),
                 this.blockStack.minus(0),
                 this.locals,
-                this.shedLocalNameToLocalIndex
+                this.shedLocalNameToLocalIndex,
+                this.waitingForShedLocal
             ).toResult(this.instructions);
         }
 
@@ -232,22 +238,30 @@ public record WasmCodeGenContext(
                     localName,
                     new WasmLocal(localName, WasmCodeGenTypes.VALUE_TYPE)
                 ),
-                this.shedLocalNameToLocalIndex
+                this.shedLocalNameToLocalIndex,
+                this.waitingForShedLocal
             ).toResult(localIndex);
         }
 
-        public Result<WasmLocalIndex> addShedLocal(String name) {
+        public Result<PVector<Function<WasmCodeGenContext, WasmCodeGenContext>>> addShedLocal(String name) {
             var localIndexResult = generateLocal(WasmCodeGenLocalSpace.LOCAL, Optional.of(name));
             var localIndex = localIndexResult.value();
             var context = localIndexResult.context();
 
+            var funcs = this.waitingForShedLocal.getOrDefault(name, P.vector());
+
             return new FunctionContext(
                 context.funcType,
                 context.instructions,
                 context.blockStack,
                 context.locals,
-                context.shedLocalNameToLocalIndex.plus(name, localIndex)
-            ).toResult(localIndex);
+                context.shedLocalNameToLocalIndex.plus(name, localIndex),
+                this.waitingForShedLocal.minus(name)
+            ).toResult(funcs);
+        }
+
+        public Optional<WasmLocalIndex> shedLocalOrNone(String name) {
+            return Optional.ofNullable(shedLocalNameToLocalIndex.get(name));
         }
 
         public WasmLocalIndex shedLocal(String name) {
@@ -260,6 +274,23 @@ public record WasmCodeGenContext(
             }
         }
 
+        public FunctionContext addWaitingForShedLocal(
+            String name,
+            Function<WasmCodeGenContext, WasmCodeGenContext> func
+        ) {
+            return new FunctionContext(
+                this.funcType,
+                this.instructions,
+                this.blockStack,
+                this.locals,
+                this.shedLocalNameToLocalIndex,
+                this.waitingForShedLocal.plus(
+                    name,
+                    this.waitingForShedLocal.getOrDefault(name, P.vector()).plus(func)
+                )
+            );
+        }
+
         public WasmFunction toFunction(String name) {
             return new WasmFunction(
                 name,
@@ -374,10 +405,38 @@ public record WasmCodeGenContext(
         return this.functionContext.toFunction(name);
     }
 
-    public WasmCodeGenResult<WasmLocalIndex> addShedLocal(String name) {
+    public WasmCodeGenContext addShedArg(String name) {
+        return addShedLocal(name, (l, c) -> c);
+    }
+
+    public WasmCodeGenContext addShedLocal(
+        String name,
+        BiFunction<WasmLocalIndex, WasmCodeGenContext, WasmCodeGenContext> setLocal
+    ) {
         var result = this.functionContext.addShedLocal(name);
 
-        return new WasmCodeGenContext(moduleContext, result.context).toResult(result.value);
+        var context = new WasmCodeGenContext(moduleContext, result.context);
+
+        context = setLocal.apply(context.shedLocal(name), context);
+
+        for (var func : result.value) {
+            context = func.apply(context);
+        }
+
+        return context;
+    }
+
+    public WasmCodeGenContext onceShedLocalDefined(String name, Function<WasmCodeGenContext, WasmCodeGenContext> func) {
+        var local = this.functionContext.shedLocalOrNone(name);
+
+        if (local.isEmpty()) {
+            return new WasmCodeGenContext(
+                this.moduleContext,
+                this.functionContext.addWaitingForShedLocal(name, func)
+            );
+        } else {
+            return func.apply(this);
+        }
     }
 
     public WasmLocalIndex shedLocal(String name) {