Add WASM code gen for namespace field stores

6 files changed, 122 insertions, 11 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 2a380d35..88d46346 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGen.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGen.java
@@ -85,8 +85,22 @@ public class WasmCodeGen {
                 yield context.addInstruction(new WasmLocalSet(local));
             }
 
-            case NamespaceFieldStoreInstruction namespaceFieldStoreInstruction -> {
-                throw new UnsupportedOperationException("TODO");
+            case NamespaceFieldStoreInstruction namespaceFieldStore -> {
+                var fieldValueResult = context.generateLocal(
+                    WasmCodeGenLocalSpace.FIELD_VALUE,
+                    namespaceFieldStore.fieldName()
+                );
+                var fieldValue = fieldValueResult.value();
+                context = fieldValueResult.context();
+
+                context = context.addInstruction(new WasmLocalSet(fieldValue));
+
+                yield WasmCodeGenNamespaces.fieldStore(
+                    namespaceFieldStore.namespaceName(),
+                    namespaceFieldStore.fieldName(),
+                    fieldValue,
+                    context
+                );
             }
             case PushBoolInstruction pushBool ->
                 context.addInstruction(new WasmI32Const(pushBool.value() ? 1 : 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 bee0bc98..b31ae005 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenContext.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenContext.java
@@ -10,6 +10,7 @@ import org.shedlang.compiler.backends.wasm.structure.ast.types.WasmFuncType;
 import org.shedlang.compiler.backends.wasm.structure.ast.types.WasmLimits;
 import org.shedlang.compiler.backends.wasm.structure.ast.types.WasmReferenceType;
 import org.shedlang.compiler.backends.wasm.structure.ast.types.WasmTableType;
+import org.shedlang.compiler.typechecker.TypeInfo;
 
 import java.util.List;
 import java.util.Optional;
@@ -20,6 +21,7 @@ public record WasmCodeGenContext(
     FunctionContext functionContext
 ) {
     record ModuleContext(
+        TypeInfo typeInfo,
         OrderedPSet<WasmFuncType> types,
         OrderedPSet<WasmImport> imports,
         PVector<WasmFunction> functions,
@@ -27,6 +29,7 @@ public record WasmCodeGenContext(
         WasmMemoryBuilder memory
     ) {
         public static final ModuleContext EMPTY = new ModuleContext(
+            TypeInfo.EMPTY,
             OrderedPSet.empty(),
             OrderedPSet.empty(),
             Empty.vector(),
@@ -34,8 +37,20 @@ public record WasmCodeGenContext(
             WasmMemoryBuilder.EMPTY
         );
 
+        public static ModuleContext fromTypeInfo(TypeInfo typeInfo) {
+            return new ModuleContext(
+                typeInfo,
+                OrderedPSet.empty(),
+                OrderedPSet.empty(),
+                Empty.vector(),
+                Empty.set(),
+                WasmMemoryBuilder.EMPTY
+            );
+        }
+
         public ModuleContext addType(WasmFuncType funcType) {
             return new ModuleContext(
+                this.typeInfo,
                 this.types.plus(funcType),
                 this.imports,
                 this.functions,
@@ -46,6 +61,7 @@ public record WasmCodeGenContext(
 
         public ModuleContext addImport(WasmImport import_) {
             return new ModuleContext(
+                this.typeInfo,
                 this.types,
                 this.imports.plus(import_),
                 this.functions,
@@ -56,6 +72,7 @@ public record WasmCodeGenContext(
 
         public ModuleContext addFunction(WasmFunction function) {
             return new ModuleContext(
+                this.typeInfo,
                 this.types.plus(function.type()),
                 this.imports,
                 this.functions.plus(function),
@@ -80,6 +97,7 @@ public record WasmCodeGenContext(
             var result = memory.add(name, data, alignment);
 
             return new ModuleContext(
+                this.typeInfo,
                 this.types,
                 this.imports,
                 this.functions,
@@ -88,6 +106,10 @@ public record WasmCodeGenContext(
             ).toResult(result.offset());
         }
 
+        public WasmMemoryAddress memoryAddress(String name) {
+            return memory.memoryAddress(name);
+        }
+
         public WasmModule toModule() {
             var module = WasmModule.EMPTY;
 
@@ -210,6 +232,17 @@ public record WasmCodeGenContext(
         FunctionContext.EMPTY
     );
 
+    public static WasmCodeGenContext fromTypeInfo(TypeInfo typeInfo) {
+        return new WasmCodeGenContext(
+            ModuleContext.fromTypeInfo(typeInfo),
+            FunctionContext.EMPTY
+        );
+    }
+
+    public TypeInfo typeInfo() {
+        return moduleContext.typeInfo;
+    }
+
     public WasmCodeGenContext addType(WasmFuncType funcType) {
         return new WasmCodeGenContext(
             this.moduleContext.addType(funcType),
@@ -252,6 +285,10 @@ public record WasmCodeGenContext(
         ).toResult(result.value);
     }
 
+    public WasmMemoryAddress memoryAddress(String name) {
+        return moduleContext.memoryAddress(name);
+    }
+
     public WasmModule toModule() {
         return moduleContext.toModule();
     }
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenGlobalSpace.java b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenGlobalSpace.java
index 97e51b68..18c265e2 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenGlobalSpace.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenGlobalSpace.java
@@ -2,6 +2,7 @@ package org.shedlang.compiler.backends.wasm;
 
 public enum WasmCodeGenGlobalSpace {
     FUNCTION("function"),
+    DATA_NAMESPACE("namespace"),
     DATA_STRING("string");
 
     private final String name;
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenLocalSpace.java b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenLocalSpace.java
index 83bf9459..aa130318 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenLocalSpace.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenLocalSpace.java
@@ -5,6 +5,7 @@ public enum WasmCodeGenLocalSpace {
     CALL_ARG_NAMED("callArgNamed"),
     CLOSURE_POINTER("closurePointer"),
     CLOSURE_FREE_VARIABLE("closureFreeVariable"),
+    FIELD_VALUE("fieldValue"),
     LOCAL("local"),
     RECEIVER("receiver");
 
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenNamespaces.java b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenNamespaces.java
new file mode 100644
index 00000000..38f57f52
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenNamespaces.java
@@ -0,0 +1,55 @@
+package org.shedlang.compiler.backends.wasm;
+
+import org.shedlang.compiler.backends.wasm.structure.ast.instructions.WasmI32Store;
+import org.shedlang.compiler.backends.wasm.structure.ast.instructions.WasmLocalGet;
+import org.shedlang.compiler.backends.wasm.structure.ast.instructions.WasmMemoryAddressConst;
+import org.shedlang.compiler.backends.wasm.structure.ast.instructions.WasmMemoryArg;
+import org.shedlang.compiler.backends.wasm.structure.ast.linking.WasmMemoryAddress;
+import org.shedlang.compiler.backends.wasm.structure.ast.modules.WasmAlignment;
+import org.shedlang.compiler.backends.wasm.structure.ast.modules.WasmLocalIndex;
+import org.shedlang.compiler.types.NamespaceName;
+
+public class WasmCodeGenNamespaces {
+    private WasmCodeGenNamespaces() {
+    }
+
+    public static WasmCodeGenContext fieldStore(
+        NamespaceName namespaceName,
+        String fieldName,
+        WasmLocalIndex fieldValue,
+        WasmCodeGenContext context
+    ) {
+        var namespaceTypeInfo = context.typeInfo().namespaceTypeInfo(namespaceName);
+
+        var namespaceAddress = namespaceMemoryAddress(namespaceName, context);
+
+        return context
+            .addInstruction(new WasmMemoryAddressConst(namespaceAddress))
+            .addInstruction(new WasmLocalGet(fieldValue))
+            .addInstruction(new WasmI32Store(new WasmMemoryArg(
+                // TODO: remove magic number
+                namespaceTypeInfo.fieldIndex(fieldName) * 4,
+                WasmAlignment.ALIGN_32_BIT
+            )));
+    }
+
+    private static WasmMemoryAddress namespaceMemoryAddress(NamespaceName namespaceName, WasmCodeGenContext context) {
+        var dataName = dataNameNamespaceValue(namespaceName);
+        return context.memoryAddress(dataName);
+
+//        context.addData(
+//            dataName,
+//            // TODO: remove magic number
+//            new byte[4 * ],
+//            WasmAlignment.ALIGN_32_BIT
+//        );
+    }
+
+    private static String dataNameNamespaceValue(NamespaceName namespaceName) {
+        return "shed__namespace__" + namespaceNameToWasmIdentifier(namespaceName);
+    }
+
+    private static String namespaceNameToWasmIdentifier(NamespaceName namespaceName) {
+        return String.join("_", namespaceName.parts());
+    }
+}
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/WasmMemoryBuilder.java b/src/main/java/org/shedlang/compiler/backends/wasm/WasmMemoryBuilder.java
index c963e732..eac98b5c 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/WasmMemoryBuilder.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/WasmMemoryBuilder.java
@@ -1,7 +1,6 @@
 package org.shedlang.compiler.backends.wasm;
 
-import org.pcollections.Empty;
-import org.pcollections.PVector;
+import org.pcollections.OrderedPMap;
 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;
@@ -15,15 +14,14 @@ import java.nio.ByteBuffer;
 import java.util.List;
 
 public class WasmMemoryBuilder {
-    public static final WasmMemoryBuilder EMPTY = new WasmMemoryBuilder(Empty.vector(), 0);
+    public static final WasmMemoryBuilder EMPTY = new WasmMemoryBuilder(OrderedPMap.empty(), 0);
 
     private static final WasmDataSegmentIndex DATA_SEGMENT_INDEX = new WasmDataSegmentIndex(0);
 
-    // TODO: rename to elements?
-    private final PVector<Entity> entities;
+    private final OrderedPMap<String, Entity> entities;
     private final int size;
 
-    private WasmMemoryBuilder(PVector<Entity> entities, int size) {
+    private WasmMemoryBuilder(OrderedPMap<String, Entity> entities, int size) {
         this.entities = entities;
         this.size = size;
     }
@@ -39,7 +37,7 @@ public class WasmMemoryBuilder {
         if (entities.size() > 0) {
             var data = ByteBuffer.allocate(size);
 
-            for (var entity : entities) {
+            for (var entity : entities.values()) {
                 data.put(entity.offset, entity.data);
             }
 
@@ -51,7 +49,7 @@ public class WasmMemoryBuilder {
                     data.array(),
                     new WasmExpression(List.of(new WasmI32Const(0))),
                     WasmAlignment.ALIGN_32_BIT,
-                    entities.stream()
+                    entities.values().stream()
                         .map(entity -> new WasmDataSymbol(
                             entity.name,
                             DATA_SEGMENT_INDEX,
@@ -79,12 +77,17 @@ public class WasmMemoryBuilder {
         }
 
         var builder = new WasmMemoryBuilder(
-            entities.plus(new Entity(name, offset, data, alignment)),
+            entities.plus(name, new Entity(name, offset, data, alignment)),
             offset + data.length
         );
         return new Result(new WasmMemoryAddress(DATA_SEGMENT_INDEX, offset), builder);
     }
 
+    public WasmMemoryAddress memoryAddress(String name) {
+        var offset = entities.get(name).offset;
+        return new WasmMemoryAddress(DATA_SEGMENT_INDEX, offset);
+    }
+
     private record Entity(String name, int offset, byte[] data, WasmAlignment alignment) {
 
     }