Declare variables up-front

4 files changed, 183 insertions, 66 deletions

diff --git a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedWhenNode.java b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedWhenNode.java
index e8c2d348..75ff89fe 100644
--- a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedWhenNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedWhenNode.java
@@ -9,4 +9,5 @@ public record UntypedWhenNode(
     List<UntypedWhenCaseNode> cases,
     Source source
 ) implements UntypedExpressionNode {
+    // TODO: should implement UntypedBranchingExpressionNode?
 }
diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGen.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGen.java
index dc4733f4..26928a24 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGen.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGen.java
@@ -1,5 +1,7 @@
 package org.shedlang.compiler.patio;
 
+import org.pcollections.OrderedPSet;
+import org.shedlang.compiler.ast.typed.BoundVariable;
 import org.shedlang.compiler.patio.ast.PatioCompilationUnit;
 import org.shedlang.compiler.patio.ast.PatioVisibility;
 import org.shedlang.compiler.patio.ast.instructions.*;
@@ -12,6 +14,7 @@ import org.shedlang.compiler.sources.Sources;
 import org.shedlang.compiler.stackir.StackIrNamespace;
 import org.shedlang.compiler.stackir.instructions.*;
 import org.shedlang.compiler.types.NamespaceName;
+import org.shedlang.compiler.util.P;
 
 import java.util.*;
 
@@ -99,8 +102,11 @@ public class PatioCodeGen {
     ) {
         context = PatioCodeGenNamespaces.define(
             stackIrNamespace.namespaceName(),
-            initContext ->
-                generateInstructions(stackIrNamespace.instructions(), initContext),
+            initContext -> {
+                initContext = declareStackIrLocals(stackIrNamespace.instructions(), initContext);
+
+                return generateInstructions(stackIrNamespace.instructions(), initContext);
+            },
             context
         );
 
@@ -227,8 +233,13 @@ public class PatioCodeGen {
                     defineFunction.positionalParams(),
                     defineFunction.namedParams(),
                     defineFunction.freeVariables(),
-                    functionContext ->
-                        generateInstructions(defineFunction.body(), functionContext),
+                    functionContext -> {
+                        // TODO: explicitly store locals on a function?
+
+                        functionContext = declareStackIrLocals(defineFunction.body(), functionContext);
+
+                        return generateInstructions(defineFunction.body(), functionContext);
+                    },
                     defineFunction.source(),
                     context
                 );
@@ -421,12 +432,17 @@ public class PatioCodeGen {
                 yield context.temporaryPush(local);
             }
 
-            case LocalStoreInstruction localStore ->
-                context.addStackIrVal(
-                    localStore.name(),
-                    (local, c) -> c.temporaryPopIntoLocal(local, localStore.source())
+            case LocalStoreInstruction localStore -> {
+                context = context.temporaryPopIntoLocal(
+                    context.stackIrLocalNameToLocal(localStore.name()),
+                    localStore.source()
                 );
 
+                context = context.initStackIrLocal(localStore.name());
+
+                yield context;
+            }
+
             case LoopBreakInstruction loopBreak ->
                 context.addInstruction(new PatioLoopBreak(loopBreak.source()));
 
@@ -691,8 +707,6 @@ public class PatioCodeGen {
             }
 
             case VarStoreInstruction varStore -> {
-                context = context.addStackIrVarIfMissing(varStore.name());
-
                 var valueResult = context.temporaryPop();
                 var value = valueResult.value();
                 context = valueResult.context();
@@ -838,6 +852,97 @@ 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()) {
+                context = context.defineStackIrVar(boundVariable.name());
+            } else {
+                context = context.declareStackIrVal(boundVariable.name()).context();
+            }
+        }
+        return context;
+    }
+
+    private static OrderedPSet<BoundVariable> findBoundVariables(List<Instruction> body, OrderedPSet<BoundVariable> boundVariables) {
+        for (var bodyInstruction : body) {
+            switch (bodyInstruction) {
+                case CallInstruction callInstruction -> {
+                }
+                case DefineEffectInstruction defineEffectInstruction -> {
+                }
+                case DefineFunctionInstruction defineFunctionInstruction -> {
+                }
+                case DefineStructInstruction defineStructInstruction -> {
+                }
+                case DiscardInstruction discardInstruction -> {
+                }
+                case DuplicateInstruction duplicateInstruction -> {
+                }
+                case FieldLoadInstruction fieldLoadInstruction -> {
+                }
+                case IfInstruction ifInstruction -> {
+                    boundVariables = findBoundVariables(ifInstruction.trueBody(), boundVariables);
+                    boundVariables = findBoundVariables(ifInstruction.falseBody(), boundVariables);
+                }
+                case IntAddInstruction intAddInstruction -> {
+                }
+                case IntEqualsInstruction intEqualsInstruction -> {
+                }
+                case IntNotEqualInstruction intNotEqualInstruction -> {
+                }
+                case IntSubInstruction intSubInstruction -> {
+                }
+                case LocalLoadInstruction localLoadInstruction -> {
+                }
+                case LocalStoreInstruction localStoreInstruction -> {
+                    boundVariables = boundVariables.plus(BoundVariable.immutable(localStoreInstruction.name()));
+                }
+                case LoopBreakInstruction loopBreakInstruction -> {
+                }
+                case LoopContinueInstruction loopContinueInstruction -> {
+                }
+                case LoopInstruction loopInstruction -> {
+                    boundVariables = findBoundVariables(loopInstruction.body(), boundVariables);
+                }
+                case NamespaceFieldStoreInstruction namespaceFieldStoreInstruction -> {
+                }
+                case NamespaceImportInstruction namespaceImportInstruction -> {
+                }
+                case PushBoolInstruction pushBoolInstruction -> {
+                }
+                case PushIntInstruction pushIntInstruction -> {
+                }
+                case PushStringInstruction pushStringInstruction -> {
+                }
+                case PushUnitInstruction pushUnitInstruction -> {
+                }
+                case ReturnInstruction returnInstruction -> {
+                }
+                case StringAddInstruction stringAddInstruction -> {
+                }
+                case SwitchInstruction switchInstruction -> {
+                    for (var switchCase : switchInstruction.cases()) {
+                        boundVariables = findBoundVariables(switchCase.body(), boundVariables);
+                    }
+                }
+                case TryInstruction tryInstruction -> {
+                    boundVariables = findBoundVariables(tryInstruction.body(), boundVariables);
+                }
+                case TupleCreateInstruction tupleCreateInstruction -> {
+                }
+                case TupleGetInstruction tupleGetInstruction -> {
+                }
+                case VarLoadInstruction varLoadInstruction -> {
+                }
+                case VarStoreInstruction varStoreInstruction -> {
+                    boundVariables = boundVariables.plus(BoundVariable.mutable(varStoreInstruction.name()));
+                }
+            }
+        }
+
+        return boundVariables;
+    }
+
     private static PatioCodeGenContext pushValue(int value, Source source, PatioCodeGenContext context) {
         var localResult = context.allocateVal(PatioCodeGenLocalSpace.TEMPORARY, PatioCodeGenTypes.VALUE_TYPE);
         var local = localResult.value();
diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenClosures.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenClosures.java
index ad5ae9f2..cd2f1555 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenClosures.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenClosures.java
@@ -76,20 +76,23 @@ class PatioCodeGenClosures {
             var stackIrParams = Stream.concat(positionalParams.stream(), orderedNamedParams.stream()).toList();
 
             for (var stackIrParam : stackIrParams) {
-                bodyContext = bodyContext.addStackIrParam(stackIrParam.name());
+                bodyContext = bodyContext.defineStackIrParam(stackIrParam.name());
             }
 
             for (var freeVariable : freeVariables) {
-                bodyContext = bodyContext.addStackIrVal(
-                    freeVariable.name(),
-                    (freeVariableLocal, c) -> c.addInstruction(new PatioMemoryLoad(
-                        freeVariableLocal,
-                        closurePointer,
-                        layout,
-                        freeVariableOffset(freeVariable),
-                        source
-                    ))
-                );
+                var freeVariableLocalResult = bodyContext.declareStackIrVal(freeVariable.name());
+                var freeVariableLocal = freeVariableLocalResult.value();
+                bodyContext = freeVariableLocalResult.context();
+
+                bodyContext = bodyContext.addInstruction(new PatioMemoryLoad(
+                    freeVariableLocal,
+                    closurePointer,
+                    layout,
+                    freeVariableOffset(freeVariable),
+                    source
+                ));
+
+                bodyContext = bodyContext.initStackIrLocal(freeVariable.name());
             }
 
             bodyContext = generateBody.apply(bodyContext);
diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java
index a1089405..3f0b43bb 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java
@@ -1,9 +1,6 @@
 package org.shedlang.compiler.patio;
 
-import org.pcollections.OrderedPMap;
-import org.pcollections.PMap;
-import org.pcollections.PStack;
-import org.pcollections.PVector;
+import org.pcollections.*;
 import org.shedlang.compiler.backends.NameAllocator;
 import org.shedlang.compiler.patio.ast.*;
 import org.shedlang.compiler.patio.ast.instructions.PatioInstruction;
@@ -21,7 +18,6 @@ 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(
@@ -251,9 +247,10 @@ public record PatioCodeGenContext(
         }
 
         @CheckReturnValue
-        public Result<PVector<Function<PatioCodeGenContext, PatioCodeGenContext>>> addStackIrLocal(
+        public FunctionContext.Result<PatioLocalRef> declareStackIrLocal(
             String name,
-            PatioLocalFactory createLocal
+            PatioType type,
+            PatioLocalType localType
         ) {
             var context = this;
 
@@ -275,10 +272,23 @@ public record PatioCodeGenContext(
                 this.instructions,
                 this.localsByName.plus(
                     fullName,
-                    createLocal.create(fullName, PatioCodeGenTypes.VALUE_TYPE)
+                    new PatioLocal(fullName, type, localType)
                 ),
                 // TODO: the mapping is redundant given we fix the mapping
                 this.stackIrNameToLocal.plus(name, localRef),
+                this.waitingForStackIrLocal.plus(name, P.vector())
+            );
+
+            return context.toResult(localRef);
+        }
+
+        @CheckReturnValue
+        public FunctionContext.Result<PVector<Function<PatioCodeGenContext, PatioCodeGenContext>>> initStackIrLocal(String name) {
+            var context = new FunctionContext(
+                this.temporaryStack,
+                this.instructions,
+                this.localsByName,
+                this.stackIrNameToLocal,
                 this.waitingForStackIrLocal.minus(name)
             );
 
@@ -305,6 +315,11 @@ public record PatioCodeGenContext(
         }
 
         @CheckReturnValue
+        public boolean isStackIrLocalDefined(String name) {
+            return !this.waitingForStackIrLocal.containsKey(name);
+        }
+
+        @CheckReturnValue
         public FunctionContext addWaitingForShedLocal(
             String name,
             Function<PatioCodeGenContext, PatioCodeGenContext> func
@@ -490,56 +505,46 @@ public record PatioCodeGenContext(
     }
 
     @CheckReturnValue
-    public PatioCodeGenContext addStackIrParam(
-        String name
-    ) {
-        return addStackIrLocal(
-            name,
-            (l, c) -> c,
-            (patioName, type) -> new PatioLocal(patioName, type, PatioLocalType.PARAM)
-        );
+    public PatioCodeGenContext defineStackIrParam(String name) {
+        return declareStackIrParam(name).initStackIrLocal(name);
     }
 
     @CheckReturnValue
-    public PatioCodeGenContext addStackIrVal(
-        String name,
-        BiFunction<PatioLocalRef, PatioCodeGenContext, PatioCodeGenContext> setLocal
-    ) {
-        return addStackIrLocal(
-            name,
-            setLocal,
-            (patioName, type) -> new PatioLocal(patioName, type, PatioLocalType.VAL)
-        );
+    private PatioCodeGenContext declareStackIrParam(String name) {
+        return declareStackIrLocal(name, PatioLocalType.PARAM).context();
     }
 
     @CheckReturnValue
-    public PatioCodeGenContext addStackIrVarIfMissing(
-        String name
-    ) {
-        var local = this.functionContext.stackIrLocalNameToLocalOrNone(name);
-        if (local.isEmpty()) {
-            return addStackIrLocal(
-                name,
-                (l, c) -> c,
-                (patioName, type) -> new PatioLocal(patioName, type, PatioLocalType.VAR)
-            );
-        } else {
-            return this;
-        }
+    public PatioCodeGenResult<PatioLocalRef> declareStackIrVal(String name) {
+        return declareStackIrLocal(name, PatioLocalType.VAL);
     }
 
     @CheckReturnValue
-    public PatioCodeGenContext addStackIrLocal(
+    public PatioCodeGenContext defineStackIrVar(String name) {
+        return declareStackIrVar(name).initStackIrLocal(name);
+    }
+
+    @CheckReturnValue
+    private PatioCodeGenContext declareStackIrVar(String name) {
+        return declareStackIrLocal(name, PatioLocalType.VAR).context();
+    }
+
+    @CheckReturnValue
+    private PatioCodeGenResult<PatioLocalRef> declareStackIrLocal(
         String name,
-        BiFunction<PatioLocalRef, PatioCodeGenContext, PatioCodeGenContext> setLocal,
-        PatioLocalFactory createLocal
+        PatioLocalType patioLocalType
     ) {
-        var result = this.functionContext.addStackIrLocal(name, createLocal);
+        return withFunctionContextResult(this.functionContext.declareStackIrLocal(
+            name,
+            PatioCodeGenTypes.VALUE_TYPE,
+            patioLocalType
+        ));
+    }
 
+    @CheckReturnValue
+    public PatioCodeGenContext initStackIrLocal(String name) {
+        var result = this.functionContext.initStackIrLocal(name);
         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);
@@ -556,6 +561,9 @@ public record PatioCodeGenContext(
         var local = this.functionContext.stackIrLocalNameToLocalOrNone(name);
 
         if (local.isEmpty()) {
+            // TODO: better error
+            throw new RuntimeException("Unknown stack IR local: " + name);
+        } else if (!this.functionContext.isStackIrLocalDefined(name)) {
             return withFunctionContext(
                 this.functionContext.addWaitingForShedLocal(name, func)
             );