Add WASM code gen for local-load instruction

5 files changed, 53 insertions, 14 deletions

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 67b80ec7..8efb7617 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGen.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGen.java
@@ -1,9 +1,6 @@
 package org.shedlang.compiler.backends.wasm;
 
-import org.shedlang.compiler.backends.wasm.structure.ast.instructions.WasmDrop;
-import org.shedlang.compiler.backends.wasm.structure.ast.instructions.WasmI32Const;
-import org.shedlang.compiler.backends.wasm.structure.ast.instructions.WasmLocalSet;
-import org.shedlang.compiler.backends.wasm.structure.ast.instructions.WasmReturn;
+import org.shedlang.compiler.backends.wasm.structure.ast.instructions.*;
 import org.shedlang.compiler.stackir.instructions.*;
 
 public class WasmCodeGen {
@@ -21,11 +18,11 @@ public class WasmCodeGen {
             case DiscardInstruction discard ->
                 context.addInstruction(new WasmDrop());
 
-            case LocalLoadInstruction localLoadInstruction -> {
-                throw new UnsupportedOperationException("TODO");
-            }
+            case LocalLoadInstruction localLoad ->
+                context.addInstruction(new WasmLocalGet(context.shedLocal(localLoad.name())));
+
             case LocalStoreInstruction localStore -> {
-                var localResult = context.generateLocal(WasmCodeGenLocalSpace.LOCAL, localStore.name());
+                var localResult = context.addShedLocal(localStore.name());
                 var local = localResult.value();
                 context = localResult.context();
                 yield context.addInstruction(new WasmLocalSet(local));
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 9513b4b9..8f0fb1d8 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenContext.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenContext.java
@@ -2,6 +2,7 @@ package org.shedlang.compiler.backends.wasm;
 
 import org.pcollections.Empty;
 import org.pcollections.OrderedPMap;
+import org.pcollections.PMap;
 import org.pcollections.PVector;
 import org.shedlang.compiler.backends.llvm.NameAllocator;
 import org.shedlang.compiler.backends.wasm.structure.ast.instructions.WasmExpression;
@@ -17,19 +18,22 @@ import java.util.Optional;
 public record WasmCodeGenContext(
     PVector<WasmFunction> functions,
     PVector<WasmInstruction> instructions,
-    OrderedPMap<String, WasmLocal> locals
+    OrderedPMap<String, WasmLocal> locals,
+    PMap<String, WasmLocalIndex> shedLocalNameToLocalIndex
 ) {
     public static final WasmCodeGenContext EMPTY = new WasmCodeGenContext(
         Empty.vector(),
         Empty.vector(),
-        OrderedPMap.empty()
+        OrderedPMap.empty(),
+        Empty.map()
     );
 
     public WasmCodeGenContext addInstruction(WasmInstruction instruction) {
         return new WasmCodeGenContext(
             functions,
             instructions.plus(instruction),
-            locals
+            locals,
+            shedLocalNameToLocalIndex
         );
     }
 
@@ -44,7 +48,8 @@ public record WasmCodeGenContext(
         return new WasmCodeGenContext(
             functions.plus(function),
             Empty.vector(),
-            OrderedPMap.empty()
+            OrderedPMap.empty(),
+            Empty.map()
         );
     }
 
@@ -58,6 +63,29 @@ public record WasmCodeGenContext(
         return module;
     }
 
+    public WasmCodeGenResult<WasmLocalIndex> addShedLocal(String name) {
+        var localIndexResult = generateLocal(WasmCodeGenLocalSpace.LOCAL, name);
+        var localIndex = localIndexResult.value();
+        var context = localIndexResult.context();
+
+        return new WasmCodeGenContext(
+            context.functions,
+            context.instructions,
+            context.locals,
+            shedLocalNameToLocalIndex.plus(name, localIndex)
+        ).toResult(localIndex);
+    }
+
+    public WasmLocalIndex shedLocal(String name) {
+        var localIndex = shedLocalNameToLocalIndex.get(name);
+        if (localIndex == null) {
+            // TODO: throw better error
+            throw new RuntimeException("unknown local: " + name);
+        } else {
+            return localIndex;
+        }
+    }
+
     public WasmCodeGenResult<WasmLocalIndex> generateLocal(WasmCodeGenLocalSpace space, String name) {
         var localIndex = new WasmLocalIndex(locals.size());
         var localName = NameAllocator.generateName(
@@ -68,7 +96,8 @@ public record WasmCodeGenContext(
         return new WasmCodeGenContext(
             functions,
             instructions,
-            locals.plus(localName, new WasmLocal(localName, WasmCodeGenTypes.VALUE_TYPE))
+            locals.plus(localName, new WasmLocal(localName, WasmCodeGenTypes.VALUE_TYPE)),
+            shedLocalNameToLocalIndex
         ).toResult(localIndex);
     }
 
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/instructions/WasmInstruction.java b/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/instructions/WasmInstruction.java
index c50e29fe..9467a9e7 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/instructions/WasmInstruction.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/instructions/WasmInstruction.java
@@ -1,4 +1,4 @@
 package org.shedlang.compiler.backends.wasm.structure.ast.instructions;
 
-public sealed interface WasmInstruction permits WasmDrop, WasmI32Const, WasmI64Const, WasmLocalSet, WasmReturn {
+public sealed interface WasmInstruction permits WasmDrop, WasmI32Const, WasmI64Const, WasmLocalGet, WasmLocalSet, WasmReturn {
 }
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/instructions/WasmLocalGet.java b/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/instructions/WasmLocalGet.java
new file mode 100644
index 00000000..862bfbed
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/instructions/WasmLocalGet.java
@@ -0,0 +1,8 @@
+package org.shedlang.compiler.backends.wasm.structure.ast.instructions;
+
+import org.shedlang.compiler.backends.wasm.structure.ast.modules.WasmLocalIndex;
+
+public record WasmLocalGet(
+    WasmLocalIndex localIndex
+) implements WasmInstruction {
+}
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/WasmBinaryFormatInstructions.java b/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/WasmBinaryFormatInstructions.java
index a2b11ac7..09161143 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/WasmBinaryFormatInstructions.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/WasmBinaryFormatInstructions.java
@@ -30,6 +30,11 @@ public class WasmBinaryFormatInstructions {
 
             // == Variable Instructions ==
 
+            case WasmLocalGet localGet -> {
+                writer.writeByte(0x20);
+                writeLocalIndex(localGet.localIndex(), writer);
+            }
+
             case WasmLocalSet localSet -> {
                 writer.writeByte(0x21);
                 writeLocalIndex(localSet.localIndex(), writer);