Generate separate WASM data segments

5 files changed, 28 insertions, 46 deletions

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 5f80dbb9..8e147ad6 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenContext.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenContext.java
@@ -125,10 +125,6 @@ public record WasmCodeGenContext(
             ).toResult(result.offset());
         }
 
-        public WasmMemoryAddress memoryAddress(String name) {
-            return memory.memoryAddress(name);
-        }
-
         public WasmModule toModule() {
             var module = WasmModule.EMPTY;
 
@@ -315,10 +311,6 @@ 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/WasmCodeGenNamespaces.java b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenNamespaces.java
index 281c2f73..9704c787 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenNamespaces.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGenNamespaces.java
@@ -97,7 +97,7 @@ public class WasmCodeGenNamespaces {
                             new WasmGlobalSet(new WasmGlobalName(globalNameNamespaceIsInited(namespaceName)))
                         )
                     ))
-                    .addInstruction(new WasmMemoryAddressConst(namespaceMemoryAddress(namespaceName, bodyContext)));
+                    .addInstruction(new WasmMemoryAddressConst(namespaceMemoryAddress(namespaceName)));
             }
         );
 
@@ -110,7 +110,7 @@ public class WasmCodeGenNamespaces {
         WasmLocalIndex fieldValue,
         WasmCodeGenContext context
     ) {
-        var namespaceAddress = namespaceMemoryAddress(namespaceName, context);
+        var namespaceAddress = namespaceMemoryAddress(namespaceName);
 
         return context
             .addInstruction(new WasmMemoryAddressConst(namespaceAddress))
@@ -137,9 +137,9 @@ public class WasmCodeGenNamespaces {
             )));
     }
 
-    private static WasmMemoryAddress namespaceMemoryAddress(NamespaceName namespaceName, WasmCodeGenContext context) {
+    private static WasmMemoryAddress namespaceMemoryAddress(NamespaceName namespaceName) {
         var dataName = globalNameNamespaceValue(namespaceName);
-        return context.memoryAddress(dataName);
+        return new WasmMemoryAddress(dataName);
     }
 
     private static WasmMemoryArg fieldMemoryArg(
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 6b69874e..df72afc4 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/WasmMemoryBuilder.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/WasmMemoryBuilder.java
@@ -10,14 +10,12 @@ import org.shedlang.compiler.backends.wasm.structure.ast.modules.*;
 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;
 
 public class WasmMemoryBuilder {
     public static final WasmMemoryBuilder EMPTY = new WasmMemoryBuilder(OrderedPMap.empty(), 0);
 
     private static final int PAGE_SIZE = 65536;
-    private static final WasmDataSegmentIndex DATA_SEGMENT_INDEX = new WasmDataSegmentIndex(0);
 
     private final OrderedPMap<String, Entity> entities;
     private final int size;
@@ -30,6 +28,8 @@ public class WasmMemoryBuilder {
     public WasmModule toModule(WasmModule module) {
         var pageCount = (size + PAGE_SIZE - 1) / PAGE_SIZE;
 
+        // TODO: we assume that we're in complete control of memory and data segments.
+        // Either remove this assumption or document it properly.
         module = module
             .addImport(new WasmImport(
                 "env",
@@ -38,29 +38,27 @@ public class WasmMemoryBuilder {
             ));
 
         if (entities.size() > 0) {
-            var data = ByteBuffer.allocate(size);
-
+            var dataSegmentIndex = module.dataSegments().size();
             for (var entity : entities.values()) {
-                data.put(entity.offset, entity.data);
+                module = module
+                    .addDataSegment(new WasmDataSegment(
+                        ".data." + entity.name,
+                        entity.data,
+                        new WasmExpression(List.of(new WasmI32Const(0))),
+                        WasmAlignment.ALIGN_32_BIT,
+                        List.of(
+                            new WasmDataSymbol(
+                                entity.name,
+                                // TODO: remove the data segment index from data symbol (its redundant)
+                                new WasmDataSegmentIndex(dataSegmentIndex),
+                                0,
+                                entity.data.length
+                            )
+                        )
+                    ));
+
+                dataSegmentIndex++;
             }
-
-            // TODO: we assume that we're in complete control of memory and data segments.
-            // Either remove this assumption or document it properly.
-            module = module
-                .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,
-                            DATA_SEGMENT_INDEX,
-                            entity.offset,
-                            entity.data.length
-                        ))
-                        .toList()
-                ));
         }
 
         return module;
@@ -83,12 +81,7 @@ public class WasmMemoryBuilder {
             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);
+        return new Result(new WasmMemoryAddress(name), 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/linking/WasmMemoryAddress.java b/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/linking/WasmMemoryAddress.java
index 263d67b5..a3d2c4dd 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/linking/WasmMemoryAddress.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/linking/WasmMemoryAddress.java
@@ -1,7 +1,4 @@
 package org.shedlang.compiler.backends.wasm.structure.ast.linking;
 
-public record WasmMemoryAddress(
-    WasmDataSegmentIndex dataSegmentIndex,
-    int dataSegmentOffset
-) {
+public record WasmMemoryAddress(String dataSegmentName) {
 }
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/WasmBinaryFormatLinking.java b/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/WasmBinaryFormatLinking.java
index f8813bb5..de46dd3f 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/WasmBinaryFormatLinking.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/WasmBinaryFormatLinking.java
@@ -321,7 +321,7 @@ public class WasmBinaryFormatLinking {
                 var symbolIndex = new WasmSymbolIndex(symbolInfos.size());
                 symbolInfos.add(symbolInfo);
                 memoryAddressToSymbolIndex.put(
-                    new WasmMemoryAddress(dataSymbol.dataSegmentIndex(), dataSymbol.offset()),
+                    new WasmMemoryAddress(dataSymbol.name()),
                     symbolIndex
                 );
             }