Add var-load and var-store to Stack IR

13 files changed, 184 insertions, 2 deletions

diff --git a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGen.java b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGen.java
index 8763d6f8..60aa2152 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGen.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGen.java
@@ -306,6 +306,12 @@ public class LlvmCodeGen {
                         LlvmCodeGenStrings.add(target, left, right, operationContext),
                     context
                 );
+
+            case VarLoadInstruction varLoad ->
+                throw new UnsupportedOperationException();
+
+            case VarStoreInstruction varStore ->
+                throw new UnsupportedOperationException();
         };
     }
 
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 87cc5ad5..a85a49c3 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGen.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGen.java
@@ -264,6 +264,12 @@ public class WasmCodeGen {
 
             case StringAddInstruction stringAdd ->
                 WasmCodeGenStrings.add(context);
+
+            case VarLoadInstruction varLoad ->
+                throw new UnsupportedOperationException();
+
+            case VarStoreInstruction varStore ->
+                throw new UnsupportedOperationException();
         };
     }
 
diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGen.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGen.java
index e11f5dd1..57a08d69 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGen.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGen.java
@@ -458,6 +458,42 @@ public class PatioCodeGen {
                         ),
                     context
                 );
+
+            case VarLoadInstruction varLoad -> {
+                var varValueResult = context.allocateVal(
+                    PatioCodeGenLocalSpace.VAR_VALUE,
+                    varLoad.name(),
+                    PatioCodeGenTypes.VALUE_TYPE
+                );
+                var varValue = varValueResult.value();
+                context = varValueResult.context();
+
+                context = context.addInstruction(new PatioVarLoad(
+                    varValue,
+                    context.stackIrLocalNameToLocal(varLoad.name()),
+                    varLoad.source()
+                ));
+
+                context = context.temporaryPush(varValue);
+
+                yield context;
+            }
+
+            case VarStoreInstruction varStore -> {
+                context = context.addStackIrVarIfMissing(varStore.name());
+
+                var valueResult = context.temporaryPop();
+                var value = valueResult.value();
+                context = valueResult.context();
+
+                context = context.addInstruction(new PatioVarStore(
+                    context.stackIrLocalNameToLocal(varStore.name()),
+                    value,
+                    varStore.source()
+                ));
+
+                yield context;
+            }
         };
     }
 
diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java
index ecec3aee..034d945d 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java
@@ -370,6 +370,17 @@ public record PatioCodeGenContext(
         return addStackIrLocal(name, setLocal, PatioLocal.Val::new);
     }
 
+    public PatioCodeGenContext addStackIrVarIfMissing(
+        String name
+    ) {
+        var local = this.functionContext.stackIrLocalNameToLocalOrNone(name);
+        if (local.isEmpty()) {
+            return addStackIrLocal(name, (l, c) -> c, PatioLocal.Var::new);
+        } else {
+            return this;
+        }
+    }
+
     public PatioCodeGenContext addStackIrLocal(
         String name,
         BiFunction<PatioLocalRef, PatioCodeGenContext, PatioCodeGenContext> setLocal,
diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenLocalSpace.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenLocalSpace.java
index fb46b165..74e394bb 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenLocalSpace.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenLocalSpace.java
@@ -14,7 +14,8 @@ public enum PatioCodeGenLocalSpace {
     LOCAL("local"),
     RECEIVER("receiver"),
     STRING("string"),
-    TEMPORARY("tmp");
+    TEMPORARY("tmp"),
+    VAR_VALUE("varValue");
 
     private final String name;
 
diff --git a/src/main/java/org/shedlang/compiler/stackir/InstructionSerializer.java b/src/main/java/org/shedlang/compiler/stackir/InstructionSerializer.java
index 0f6881ae..807caa57 100644
--- a/src/main/java/org/shedlang/compiler/stackir/InstructionSerializer.java
+++ b/src/main/java/org/shedlang/compiler/stackir/InstructionSerializer.java
@@ -157,6 +157,18 @@ public class InstructionSerializer {
 
             case StringAddInstruction stringAddInstruction ->
                 treeList(treeIdentifier("string-add"));
+
+            case VarLoadInstruction varLoadInstruction ->
+                treeList(
+                    treeIdentifier("var-load"),
+                    treeString(varLoadInstruction.name())
+                );
+
+            case VarStoreInstruction varStoreInstruction ->
+                treeList(
+                    treeIdentifier("var-store"),
+                    treeString(varStoreInstruction.name())
+                );
         };
     }
 }
diff --git a/src/main/java/org/shedlang/compiler/stackir/instructions/Instruction.java b/src/main/java/org/shedlang/compiler/stackir/instructions/Instruction.java
index 80593264..b4a59136 100644
--- a/src/main/java/org/shedlang/compiler/stackir/instructions/Instruction.java
+++ b/src/main/java/org/shedlang/compiler/stackir/instructions/Instruction.java
@@ -3,7 +3,7 @@ package org.shedlang.compiler.stackir.instructions;
 import org.shedlang.compiler.sources.Source;
 
 public sealed interface Instruction
-    permits CallInstruction, DefineFunctionInstruction, DiscardInstruction, FieldLoadInstruction, IfInstruction, IntAddInstruction, IntEqualsInstruction, IntNotEqualInstruction, IntSubInstruction, LocalLoadInstruction, LocalStoreInstruction, NamespaceFieldStoreInstruction, NamespaceImportInstruction, PushBoolInstruction, PushIntInstruction, PushStringInstruction, PushUnitInstruction, ReturnInstruction, StringAddInstruction {
+    permits CallInstruction, DefineFunctionInstruction, DiscardInstruction, FieldLoadInstruction, IfInstruction, IntAddInstruction, IntEqualsInstruction, IntNotEqualInstruction, IntSubInstruction, LocalLoadInstruction, LocalStoreInstruction, NamespaceFieldStoreInstruction, NamespaceImportInstruction, PushBoolInstruction, PushIntInstruction, PushStringInstruction, PushUnitInstruction, ReturnInstruction, StringAddInstruction, VarLoadInstruction, VarStoreInstruction {
 
     Source source();
 }
diff --git a/src/main/java/org/shedlang/compiler/stackir/instructions/VarLoadInstruction.java b/src/main/java/org/shedlang/compiler/stackir/instructions/VarLoadInstruction.java
new file mode 100644
index 00000000..2ca49d90
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/stackir/instructions/VarLoadInstruction.java
@@ -0,0 +1,10 @@
+package org.shedlang.compiler.stackir.instructions;
+
+import org.shedlang.compiler.sources.Source;
+
+/**
+ * Load the value stored in the var with the given name and push it onto the
+ * stack.
+ */
+public record VarLoadInstruction(String name, Source source) implements Instruction {
+}
diff --git a/src/main/java/org/shedlang/compiler/stackir/instructions/VarStoreInstruction.java b/src/main/java/org/shedlang/compiler/stackir/instructions/VarStoreInstruction.java
new file mode 100644
index 00000000..b355d39c
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/stackir/instructions/VarStoreInstruction.java
@@ -0,0 +1,9 @@
+package org.shedlang.compiler.stackir.instructions;
+
+import org.shedlang.compiler.sources.Source;
+
+/**
+ * Pop a value from the stack and store it in the var with the given name.
+ */
+public record VarStoreInstruction(String name, Source source) implements Instruction {
+}
diff --git a/src/test/java/org/shedlang/compiler/patio/PatioCodeGenVarLoadTests.java b/src/test/java/org/shedlang/compiler/patio/PatioCodeGenVarLoadTests.java
new file mode 100644
index 00000000..ac51b8f1
--- /dev/null
+++ b/src/test/java/org/shedlang/compiler/patio/PatioCodeGenVarLoadTests.java
@@ -0,0 +1,25 @@
+package org.shedlang.compiler.patio;
+
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.shedlang.compiler.stackir.instructions.Instructions;
+import org.shedlang.compiler.testing.snapshots.SnapshotResolver;
+import org.shedlang.compiler.testing.snapshots.Snapshotter;
+
+import java.io.IOException;
+
+import static org.shedlang.compiler.patio.PatioTesting.serialize;
+
+@ExtendWith({SnapshotResolver.class})
+public class PatioCodeGenVarLoadTests {
+    @Test
+    public void canGenerateCodeForVarLoad(Snapshotter snapshotter) throws IOException {
+        var stackIr = Instructions.varLoad("x");
+        var context = PatioCodeGenContext.EMPTY
+            .addStackIrVarIfMissing("x");
+
+        var result = PatioCodeGen.generateInstruction(stackIr, context);
+
+        snapshotter.assertSnapshot(serialize(result));
+    }
+}
diff --git a/src/test/java/org/shedlang/compiler/patio/PatioCodeGenVarStoreTests.java b/src/test/java/org/shedlang/compiler/patio/PatioCodeGenVarStoreTests.java
new file mode 100644
index 00000000..837be982
--- /dev/null
+++ b/src/test/java/org/shedlang/compiler/patio/PatioCodeGenVarStoreTests.java
@@ -0,0 +1,40 @@
+package org.shedlang.compiler.patio;
+
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.shedlang.compiler.patio.ast.operands.PatioLocalRef;
+import org.shedlang.compiler.stackir.instructions.Instructions;
+import org.shedlang.compiler.testing.snapshots.SnapshotResolver;
+import org.shedlang.compiler.testing.snapshots.Snapshotter;
+
+import java.io.IOException;
+
+import static org.shedlang.compiler.patio.PatioTesting.serialize;
+
+@ExtendWith({SnapshotResolver.class})
+public class PatioCodeGenVarStoreTests {
+    @Test
+    public void canGenerateCodeForVarStore(Snapshotter snapshotter) throws IOException {
+        var stackIr = Instructions.varStore("x");
+        var context = PatioCodeGenContext.EMPTY
+            .temporaryPush(new PatioLocalRef("a"))
+            .temporaryPush(new PatioLocalRef("b"));
+
+        var result = PatioCodeGen.generateInstruction(stackIr, context);
+
+        snapshotter.assertSnapshot(serialize(result));
+    }
+
+    @Test
+    public void canStoreIntoVarMultipleTimes(Snapshotter snapshotter) throws IOException {
+        var stackIr = Instructions.varStore("x");
+        var context = PatioCodeGenContext.EMPTY
+            .temporaryPush(new PatioLocalRef("a"))
+            .temporaryPush(new PatioLocalRef("b"));
+
+        var result = PatioCodeGen.generateInstruction(stackIr, context);
+        result = PatioCodeGen.generateInstruction(stackIr, result);
+
+        snapshotter.assertSnapshot(serialize(result));
+    }
+}
diff --git a/src/test/java/org/shedlang/compiler/stackir/InstructionSerializerTests.java b/src/test/java/org/shedlang/compiler/stackir/InstructionSerializerTests.java
index e4893c31..aff45481 100644
--- a/src/test/java/org/shedlang/compiler/stackir/InstructionSerializerTests.java
+++ b/src/test/java/org/shedlang/compiler/stackir/InstructionSerializerTests.java
@@ -244,6 +244,24 @@ public class InstructionSerializerTests {
         assertThat(string, equalTo("(string-add)"));
     }
 
+    @Test
+    public void canSerializeVarLoad() {
+        var instruction = Instructions.varLoad("x");
+
+        var string = serialize(instruction);
+
+        assertThat(string, equalTo("(var-load \"x\")"));
+    }
+
+    @Test
+    public void canSerializeVarStore() {
+        var instruction = Instructions.varStore("x");
+
+        var string = serialize(instruction);
+
+        assertThat(string, equalTo("(var-store \"x\")"));
+    }
+
     private String serialize(Instruction instruction) {
         var builder = new StringBuilder();
         InstructionSerializer.serializeInstruction(instruction, builder);
diff --git a/src/test/java/org/shedlang/compiler/stackir/instructions/Instructions.java b/src/test/java/org/shedlang/compiler/stackir/instructions/Instructions.java
index 131207b8..18526f4a 100644
--- a/src/test/java/org/shedlang/compiler/stackir/instructions/Instructions.java
+++ b/src/test/java/org/shedlang/compiler/stackir/instructions/Instructions.java
@@ -88,4 +88,12 @@ public class Instructions {
     public static Instruction stringAdd() {
         return new StringAddInstruction(Sources.NULL);
     }
+
+    public static Instruction varLoad(String name) {
+        return new VarLoadInstruction(name, Sources.NULL);
+    }
+
+    public static Instruction varStore(String name) {
+        return new VarStoreInstruction(name, Sources.NULL);
+    }
 }