Add WASM code gen for push-string instruction

6 files changed, 171 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 8efb7617..d65d25b5 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGen.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGen.java
@@ -37,9 +37,14 @@ public class WasmCodeGen {
             case PushIntInstruction pushInt ->
                 context.addInstruction(new WasmI32Const(pushInt.value().intValueExact()));
 
-            case PushStringInstruction pushStringInstruction -> {
-                throw new UnsupportedOperationException("TODO");
+            case PushStringInstruction pushString -> {
+                var defineStringResult = WasmCodeGenStrings.defineString(pushString.value(), context);
+                var offset = defineStringResult.value();
+                context = defineStringResult.context();
+
+                yield context.addInstruction(new WasmI32Const(offset));
             }
+
             case PushUnitInstruction pushUnit ->
                 context.addInstruction(new WasmI32Const(0));
 
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 8f0fb1d8..10f55c68 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenContext.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenContext.java
@@ -7,30 +7,46 @@ 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 org.shedlang.compiler.backends.wasm.structure.ast.modules.*;
 
 import java.util.List;
 import java.util.Optional;
 
 public record WasmCodeGenContext(
     PVector<WasmFunction> functions,
+    WasmMemoryBuilder memory,
     PVector<WasmInstruction> instructions,
     OrderedPMap<String, WasmLocal> locals,
     PMap<String, WasmLocalIndex> shedLocalNameToLocalIndex
 ) {
     public static final WasmCodeGenContext EMPTY = new WasmCodeGenContext(
         Empty.vector(),
+        WasmMemoryBuilder.EMPTY,
         Empty.vector(),
         OrderedPMap.empty(),
         Empty.map()
     );
 
+    public WasmCodeGenResult<Integer> addData(
+        WasmCodeGenGlobalSpace space,
+        byte[] data,
+        WasmAlignment alignment
+    ) {
+        var result = memory.add(space, data, alignment);
+
+        return new WasmCodeGenContext(
+            functions,
+            result.memory(),
+            instructions,
+            locals,
+            shedLocalNameToLocalIndex
+        ).toResult(result.offset());
+    }
+
     public WasmCodeGenContext addInstruction(WasmInstruction instruction) {
         return new WasmCodeGenContext(
             functions,
+            memory,
             instructions.plus(instruction),
             locals,
             shedLocalNameToLocalIndex
@@ -47,6 +63,7 @@ public record WasmCodeGenContext(
         );
         return new WasmCodeGenContext(
             functions.plus(function),
+            memory,
             Empty.vector(),
             OrderedPMap.empty(),
             Empty.map()
@@ -60,7 +77,7 @@ public record WasmCodeGenContext(
             module = module.addFunction(function);
         }
 
-        return module;
+        return memory.toModule(module);
     }
 
     public WasmCodeGenResult<WasmLocalIndex> addShedLocal(String name) {
@@ -70,9 +87,10 @@ public record WasmCodeGenContext(
 
         return new WasmCodeGenContext(
             context.functions,
+            context.memory,
             context.instructions,
             context.locals,
-            shedLocalNameToLocalIndex.plus(name, localIndex)
+            context.shedLocalNameToLocalIndex.plus(name, localIndex)
         ).toResult(localIndex);
     }
 
@@ -95,6 +113,7 @@ public record WasmCodeGenContext(
         );
         return new WasmCodeGenContext(
             functions,
+            memory,
             instructions,
             locals.plus(localName, new WasmLocal(localName, WasmCodeGenTypes.VALUE_TYPE)),
             shedLocalNameToLocalIndex
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenGlobalSpace.java b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenGlobalSpace.java
new file mode 100644
index 00000000..4b1b5293
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenGlobalSpace.java
@@ -0,0 +1,16 @@
+package org.shedlang.compiler.backends.wasm;
+
+public enum WasmCodeGenGlobalSpace {
+    DATA_STRING("string");
+
+    private final String name;
+
+    WasmCodeGenGlobalSpace(String name) {
+        this.name = name;
+    }
+
+    @Override
+    public String toString() {
+        return name;
+    }
+}
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenStrings.java b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenStrings.java
new file mode 100644
index 00000000..a2431e60
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenStrings.java
@@ -0,0 +1,27 @@
+package org.shedlang.compiler.backends.wasm;
+
+import org.shedlang.compiler.backends.wasm.structure.ast.modules.WasmAlignment;
+
+import java.nio.ByteBuffer;
+import java.nio.charset.StandardCharsets;
+
+public class WasmCodeGenStrings {
+    private WasmCodeGenStrings() {
+    }
+
+    public static WasmCodeGenResult<Integer> defineString(
+        String value,
+        WasmCodeGenContext context
+    ) {
+        var codeUnits = value.getBytes(StandardCharsets.UTF_8);
+        var data = ByteBuffer.allocate(4 + codeUnits.length);
+        data.putInt(codeUnits.length);
+        data.put(codeUnits);
+
+        return context.addData(
+            WasmCodeGenGlobalSpace.DATA_STRING,
+            data.array(),
+            WasmAlignment.ALIGN_32_BIT
+        );
+    }
+}
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/WasmMemoryBuilder.java b/src/main/java/org/shedlang/compiler/backends/wasm/WasmMemoryBuilder.java
new file mode 100644
index 00000000..579fc586
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/WasmMemoryBuilder.java
@@ -0,0 +1,92 @@
+package org.shedlang.compiler.backends.wasm;
+
+import org.pcollections.OrderedPMap;
+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.WasmI32Const;
+import org.shedlang.compiler.backends.wasm.structure.ast.linking.WasmDataSegmentIndex;
+import org.shedlang.compiler.backends.wasm.structure.ast.linking.WasmDataSymbol;
+import org.shedlang.compiler.backends.wasm.structure.ast.modules.WasmAlignment;
+import org.shedlang.compiler.backends.wasm.structure.ast.modules.WasmDataSegment;
+import org.shedlang.compiler.backends.wasm.structure.ast.modules.WasmMemory;
+import org.shedlang.compiler.backends.wasm.structure.ast.modules.WasmModule;
+import org.shedlang.compiler.backends.wasm.structure.ast.types.WasmLimits;
+import org.shedlang.compiler.backends.wasm.structure.ast.types.WasmMemoryType;
+
+import java.nio.ByteBuffer;
+import java.util.List;
+import java.util.Optional;
+
+public class WasmMemoryBuilder {
+    public static final WasmMemoryBuilder EMPTY = new WasmMemoryBuilder(OrderedPMap.empty(), 0);
+
+    private final OrderedPMap<String, Entity> entities;
+    private final int size;
+
+    private WasmMemoryBuilder(OrderedPMap<String, Entity> entities, int size) {
+        this.entities = entities;
+        this.size = size;
+    }
+
+    public WasmModule toModule(WasmModule module) {
+        if (entities.size() > 0) {
+            var data = ByteBuffer.allocate(size);
+
+            for (var entity : entities.values()) {
+                data.put(entity.offset, entity.data);
+            }
+
+            // TODO: we assume that we're in complete control of memory and data segments.
+            // Either remove this assumption or document it properly.
+            module = module
+                .addMemory(new WasmMemory(new WasmMemoryType(WasmLimits.min(size))))
+                .addDataSegment(new WasmDataSegment(
+                    ".data",
+                    data.array(),
+                    new WasmExpression(List.of(new WasmI32Const(0))),
+                    WasmAlignment.ALIGN_32_BIT,
+                    entities.values().stream()
+                        .map(entity -> new WasmDataSymbol(
+                            entity.name,
+                            new WasmDataSegmentIndex(0),
+                            entity.offset,
+                            entity.data.length
+                        ))
+                        .toList()
+                ));
+        }
+
+        return module;
+    }
+
+    public record Result(int offset, WasmMemoryBuilder memory) {
+    }
+
+    public Result add(
+        WasmCodeGenGlobalSpace space,
+        byte[] data,
+        WasmAlignment alignment
+    ) {
+        var fullName = NameAllocator.generateName(
+            x -> entities.containsKey(x),
+            space,
+            Optional.empty()
+        );
+
+        var offset = size;
+        if (offset % alignment.bytes() != 0) {
+            offset += alignment.bytes() - (offset % alignment.bytes());
+        }
+
+        var builder = new WasmMemoryBuilder(
+            entities.plus(fullName, new Entity(fullName, offset, data, alignment)),
+            offset + data.length
+        );
+        // TODO: return relocated offset
+        return new Result(offset, builder);
+    }
+
+    private record Entity(String name, int offset, byte[] data, WasmAlignment alignment) {
+
+    }
+}
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/modules/WasmAlignment.java b/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/modules/WasmAlignment.java
index 81befcaf..5ca736b9 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/modules/WasmAlignment.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/modules/WasmAlignment.java
@@ -8,4 +8,8 @@ package org.shedlang.compiler.backends.wasm.structure.ast.modules;
 public record WasmAlignment(int alignment) {
     public static final WasmAlignment NONE = new WasmAlignment(0);
     public static final WasmAlignment ALIGN_32_BIT = new WasmAlignment(2);
+
+    public int bytes() {
+        return 1 << alignment;
+    }
 }