Support vars in stack compiler

1 files changed, 78 insertions, 24 deletions

diff --git a/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java b/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
index 8863253c..5b58d975 100644
--- a/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
+++ b/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
@@ -3,15 +3,23 @@ package org.shedlang.compiler.stackir;
 import org.pcollections.*;
 import org.shedlang.compiler.ShedBuiltins;
 import org.shedlang.compiler.ast.typed.*;
+import org.shedlang.compiler.sources.Source;
 import org.shedlang.compiler.stackir.instructions.*;
 import org.shedlang.compiler.util.P;
 
 import java.util.List;
+import java.util.function.BiFunction;
 
 public class StackCompiler {
     private StackCompiler() {
     }
 
+    private static Context compileAssign(TypedAssignNode node, Context context) {
+        context = compileExpression(node.expression(), context);
+        context = compileTarget(node.target(), VarStoreInstruction::new, context);
+        return context;
+    }
+
     private static Context compileBinaryOperation(
         TypedBinaryOperationNode node,
         Context context
@@ -202,7 +210,7 @@ public class StackCompiler {
     ) {
         return switch (node) {
             case TypedAssignNode assign ->
-                throw new UnsupportedOperationException();
+                compileAssign(assign, context);
 
             case TypedBlockReturnNode blockReturn ->
                 compileBlockReturn(blockReturn, context);
@@ -217,7 +225,7 @@ public class StackCompiler {
                 compileVal(val, context);
 
             case TypedVarNode var ->
-                throw new UnsupportedOperationException();
+                compileVar(var, context);
         };
     }
 
@@ -319,9 +327,19 @@ public class StackCompiler {
     }
 
     private static Context compileReference(TypedReferenceNode node, Context context) {
-        var localName = context.shedVariableToLocal(node.name());
+        var shedVariable = context.shedVariable(node.name());
 
-        return context.addInstruction(new LocalLoadInstruction(localName, node.source()));
+        if (shedVariable.isVar) {
+            return context.addInstruction(new VarLoadInstruction(
+                shedVariable.stackIrLocalName,
+                node.source()
+            ));
+        } else {
+            return context.addInstruction(new LocalLoadInstruction(
+                shedVariable.stackIrLocalName,
+                node.source()
+            ));
+        }
     }
 
     private static Context compileStringLiteral(
@@ -332,11 +350,15 @@ public class StackCompiler {
         return context.addInstruction(instruction);
     }
 
-    private static Context compileTarget(TypedTargetNode target, Context context) {
+    private static Context compileTarget(
+        TypedTargetNode target,
+        BiFunction<String, Source, Instruction> store,
+        Context context
+    ) {
         return switch (target) {
             case TypedTargetNode.Variable variable -> {
                 context = context.addInstruction(
-                    new LocalStoreInstruction(
+                    store.apply(
                         context.shedVariableToLocal(variable.name()),
                         variable.source()
                     )
@@ -349,7 +371,20 @@ public class StackCompiler {
 
     private static Context compileVal(TypedValNode val, Context context) {
         context = compileExpression(val.expression(), context);
-        context = compileTarget(val.target(), context);
+        context = compileTarget(val.target(), LocalStoreInstruction::new, context);
+        return context;
+    }
+
+    private static Context compileVar(TypedVarNode var, Context context) {
+        context = compileExpression(var.expression(), context);
+        context = compileTarget(var.target(), VarStoreInstruction::new, context);
+
+        // TODO: better to store this information in the block node with the bound variables?
+        switch (var.target()) {
+            case TypedTargetNode.Variable variable ->
+                context = context.markShedVariableAsVar(variable.name());
+        }
+
         return context;
     }
 
@@ -365,20 +400,21 @@ public class StackCompiler {
         private final PStack<PVector<Instruction>> instructionsStack;
         private final PVector<Instruction> instructions;
         private final PSet<String> localNames;
-        private final PStack<PMap<String, String>> shedVariableNameToLocalNameStack;
-        private final PMap<String, String> shedVariableNameToLocalName;
+        private final PStack<PMap<String, ShedVariable>> shedVariablesByNameStack;
+        private final PMap<String, ShedVariable> shedVariablesByName;
 
         private Context(
             PStack<PVector<Instruction>> instructionsStack,
             PVector<Instruction> instructions,
             PSet<String> localNames,
-            PStack<PMap<String, String>> shedVariableNameToLocalNameStack, PMap<String, String> shedVariableNameToLocalName
+            PStack<PMap<String, ShedVariable>> shedVariablesByNameStack,
+            PMap<String, ShedVariable> shedVariablesByName
         ) {
             this.instructionsStack = instructionsStack;
             this.instructions = instructions;
             this.localNames = localNames;
-            this.shedVariableNameToLocalNameStack = shedVariableNameToLocalNameStack;
-            this.shedVariableNameToLocalName = shedVariableNameToLocalName;
+            this.shedVariablesByNameStack = shedVariablesByNameStack;
+            this.shedVariablesByName = shedVariablesByName;
         }
 
         public Context enterBlock(List<String> shedVariableNames) {
@@ -386,8 +422,8 @@ public class StackCompiler {
                 this.instructionsStack.plus(this.instructions),
                 P.vector(),
                 this.localNames,
-                this.shedVariableNameToLocalNameStack.plus(this.shedVariableNameToLocalName),
-                this.shedVariableNameToLocalName
+                this.shedVariablesByNameStack.plus(this.shedVariablesByName),
+                this.shedVariablesByName
             );
 
             for (var shedVariableName : shedVariableNames) {
@@ -402,8 +438,8 @@ public class StackCompiler {
                 this.instructionsStack.minus(0),
                 this.instructionsStack.get(0),
                 this.localNames,
-                this.shedVariableNameToLocalNameStack.minus(0),
-                this.shedVariableNameToLocalNameStack.get(0)
+                this.shedVariablesByNameStack.minus(0),
+                this.shedVariablesByNameStack.get(0)
             );
         }
 
@@ -412,19 +448,23 @@ public class StackCompiler {
                 this.instructionsStack,
                 this.instructions.plus(instruction),
                 this.localNames,
-                this.shedVariableNameToLocalNameStack,
-                this.shedVariableNameToLocalName
+                this.shedVariablesByNameStack,
+                this.shedVariablesByName
             );
         }
 
         public String shedVariableToLocal(String variableName) {
-            var localName = this.shedVariableNameToLocalName.get(variableName);
+            return shedVariable(variableName).stackIrLocalName;
+        }
 
-            if (localName == null) {
-                throw new RuntimeException("No local name for Shed variable: " + variableName);
+        public ShedVariable shedVariable(String variableName) {
+            var shedVariable = this.shedVariablesByName.get(variableName);
+
+            if (shedVariable == null) {
+                throw new RuntimeException("Unknown Shed variable: " + variableName);
             }
 
-            return localName;
+            return shedVariable;
         }
 
         public Context addShedVariable(String variableName) {
@@ -434,8 +474,20 @@ public class StackCompiler {
                 this.instructionsStack,
                 this.instructions,
                 this.localNames.plus(localName),
-                this.shedVariableNameToLocalNameStack,
-                this.shedVariableNameToLocalName.plus(variableName, localName)
+                this.shedVariablesByNameStack,
+                this.shedVariablesByName.plus(variableName, new ShedVariable(localName, false))
+            );
+        }
+
+        public Context markShedVariableAsVar(String variableName) {
+            var localName = shedVariableToLocal(variableName);
+
+            return new Context(
+                this.instructionsStack,
+                this.instructions,
+                this.localNames,
+                this.shedVariablesByNameStack,
+                this.shedVariablesByName.plus(variableName, new ShedVariable(localName, true))
             );
         }
 
@@ -455,4 +507,6 @@ public class StackCompiler {
             return this.instructions;
         }
     }
+
+    private record ShedVariable(String stackIrLocalName, boolean isVar) {}
 }