Add WASM code gen for local-store instruction

10 files changed, 95 insertions, 8 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 482d308e..67b80ec7 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGen.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGen.java
@@ -2,6 +2,7 @@ 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.stackir.instructions.*;
 
@@ -23,9 +24,13 @@ public class WasmCodeGen {
             case LocalLoadInstruction localLoadInstruction -> {
                 throw new UnsupportedOperationException("TODO");
             }
-            case LocalStoreInstruction localStoreInstruction -> {
-                throw new UnsupportedOperationException("TODO");
+            case LocalStoreInstruction localStore -> {
+                var localResult = context.generateLocal(WasmCodeGenLocalSpace.LOCAL, localStore.name());
+                var local = localResult.value();
+                context = localResult.context();
+                yield context.addInstruction(new WasmLocalSet(local));
             }
+
             case NamespaceFieldStoreInstruction namespaceFieldStoreInstruction -> {
                 throw new UnsupportedOperationException("TODO");
             }
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 125d125c..9513b4b9 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenContext.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenContext.java
@@ -1,27 +1,35 @@
 package org.shedlang.compiler.backends.wasm;
 
 import org.pcollections.Empty;
+import org.pcollections.OrderedPMap;
 import org.pcollections.PVector;
+import org.shedlang.compiler.backends.llvm.NameAllocator;
 import org.shedlang.compiler.backends.wasm.structure.ast.instructions.WasmExpression;
 import org.shedlang.compiler.backends.wasm.structure.ast.instructions.WasmInstruction;
 import org.shedlang.compiler.backends.wasm.structure.ast.modules.WasmFunction;
+import org.shedlang.compiler.backends.wasm.structure.ast.modules.WasmLocal;
+import org.shedlang.compiler.backends.wasm.structure.ast.modules.WasmLocalIndex;
 import org.shedlang.compiler.backends.wasm.structure.ast.modules.WasmModule;
 
 import java.util.List;
+import java.util.Optional;
 
 public record WasmCodeGenContext(
     PVector<WasmFunction> functions,
-    PVector<WasmInstruction> instructions
+    PVector<WasmInstruction> instructions,
+    OrderedPMap<String, WasmLocal> locals
 ) {
     public static final WasmCodeGenContext EMPTY = new WasmCodeGenContext(
         Empty.vector(),
-        Empty.vector()
+        Empty.vector(),
+        OrderedPMap.empty()
     );
 
     public WasmCodeGenContext addInstruction(WasmInstruction instruction) {
         return new WasmCodeGenContext(
             functions,
-            instructions.plus(instruction)
+            instructions.plus(instruction),
+            locals
         );
     }
 
@@ -29,13 +37,14 @@ public record WasmCodeGenContext(
         var function = new WasmFunction(
             "f",
             List.of(),
-            List.of(),
+            locals.values().stream().toList(),
             List.of(),
             new WasmExpression(instructions)
         );
         return new WasmCodeGenContext(
             functions.plus(function),
-            instructions
+            Empty.vector(),
+            OrderedPMap.empty()
         );
     }
 
@@ -48,4 +57,22 @@ public record WasmCodeGenContext(
 
         return module;
     }
+
+    public WasmCodeGenResult<WasmLocalIndex> generateLocal(WasmCodeGenLocalSpace space, String name) {
+        var localIndex = new WasmLocalIndex(locals.size());
+        var localName = NameAllocator.generateName(
+            x -> locals.containsKey(x),
+            space,
+            Optional.of(name)
+        );
+        return new WasmCodeGenContext(
+            functions,
+            instructions,
+            locals.plus(localName, new WasmLocal(localName, WasmCodeGenTypes.VALUE_TYPE))
+        ).toResult(localIndex);
+    }
+
+    public <T> WasmCodeGenResult<T> toResult(T value) {
+        return new WasmCodeGenResult<>(value, this);
+    }
 }
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenLocalSpace.java b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenLocalSpace.java
new file mode 100644
index 00000000..196a5304
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenLocalSpace.java
@@ -0,0 +1,16 @@
+package org.shedlang.compiler.backends.wasm;
+
+public enum WasmCodeGenLocalSpace {
+    LOCAL("local");
+
+    private final String name;
+
+    WasmCodeGenLocalSpace(String name) {
+        this.name = name;
+    }
+
+    @Override
+    public String toString() {
+        return name;
+    }
+}
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenResult.java b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenResult.java
new file mode 100644
index 00000000..c20e8409
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenResult.java
@@ -0,0 +1,4 @@
+package org.shedlang.compiler.backends.wasm;
+
+public record WasmCodeGenResult<T>(T value, WasmCodeGenContext context) {
+}
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenTypes.java b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenTypes.java
new file mode 100644
index 00000000..5822f1b6
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenTypes.java
@@ -0,0 +1,10 @@
+package org.shedlang.compiler.backends.wasm;
+
+import org.shedlang.compiler.backends.wasm.structure.ast.types.WasmValueType;
+
+public class WasmCodeGenTypes {
+    private WasmCodeGenTypes() {
+    }
+
+    public static final WasmValueType VALUE_TYPE = WasmValueType.I32;
+}
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 5cdd7a23..c50e29fe 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, WasmReturn {
+public sealed interface WasmInstruction permits WasmDrop, WasmI32Const, WasmI64Const, WasmLocalSet, WasmReturn {
 }
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/instructions/WasmLocalSet.java b/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/instructions/WasmLocalSet.java
new file mode 100644
index 00000000..90e12a9a
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/instructions/WasmLocalSet.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 WasmLocalSet(
+    WasmLocalIndex localIndex
+) implements WasmInstruction {
+}
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/modules/WasmLocalIndex.java b/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/modules/WasmLocalIndex.java
new file mode 100644
index 00000000..7ee0ef50
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/modules/WasmLocalIndex.java
@@ -0,0 +1,4 @@
+package org.shedlang.compiler.backends.wasm.structure.ast.modules;
+
+public record WasmLocalIndex(int value) {
+}
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 988d8c9e..a2b11ac7 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
@@ -5,6 +5,8 @@ import org.shedlang.compiler.backends.wasm.structure.ast.instructions.*;
 
 import java.io.IOException;
 
+import static org.shedlang.compiler.backends.wasm.structure.binary.WasmBinaryFormatModules.writeLocalIndex;
+
 public class WasmBinaryFormatInstructions {
     // Corresponds to https://webassembly.github.io/spec/core/binary/instructions.html
 
@@ -26,6 +28,13 @@ public class WasmBinaryFormatInstructions {
                 writer.writeByte(0x1a);
             }
 
+            // == Variable Instructions ==
+
+            case WasmLocalSet localSet -> {
+                writer.writeByte(0x21);
+                writeLocalIndex(localSet.localIndex(), writer);
+            }
+
             // == Numeric Instructions ==
 
             case WasmI32Const i32Const -> {
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/WasmBinaryFormatModules.java b/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/WasmBinaryFormatModules.java
index 1b993461..f8a91ae3 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/WasmBinaryFormatModules.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/WasmBinaryFormatModules.java
@@ -40,6 +40,10 @@ public class WasmBinaryFormatModules {
         Leb128Encoding.encodeUnsignedInt32(globalIndex.value(), writer);
     }
 
+    static void writeLocalIndex(WasmLocalIndex localIndex, BinaryWriter writer) {
+        Leb128Encoding.encodeUnsignedInt32(localIndex.value(), writer);
+    }
+
     // == Sections ==
 
     interface Write {