Add datasegments to WASM binary format

6 files changed, 106 insertions, 3 deletions

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
new file mode 100644
index 00000000..81befcaf
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/modules/WasmAlignment.java
@@ -0,0 +1,11 @@
+package org.shedlang.compiler.backends.wasm.structure.ast.modules;
+
+/**
+ * The alignment of data in memory.
+ *
+ * @param alignment The alignment as a power of two.
+ */
+public record WasmAlignment(int alignment) {
+    public static final WasmAlignment NONE = new WasmAlignment(0);
+    public static final WasmAlignment ALIGN_32_BIT = new WasmAlignment(2);
+}
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/modules/WasmDataSegment.java b/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/modules/WasmDataSegment.java
new file mode 100644
index 00000000..7227f54b
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/modules/WasmDataSegment.java
@@ -0,0 +1,11 @@
+package org.shedlang.compiler.backends.wasm.structure.ast.modules;
+
+import org.shedlang.compiler.backends.wasm.structure.ast.instructions.WasmExpression;
+
+public record WasmDataSegment(
+    String name,
+    byte[] init,
+    WasmExpression offset,
+    WasmAlignment alignment
+) {
+}
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/modules/WasmModule.java b/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/modules/WasmModule.java
index a4284202..5365b5ee 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/modules/WasmModule.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/modules/WasmModule.java
@@ -16,6 +16,7 @@ public record WasmModule(
     PVector<WasmMemory> memories,
     PVector<WasmGlobal> globals,
     PVector<WasmElementSegment> elementSegments,
+    PVector<WasmDataSegment> dataSegments,
     PVector<WasmImport> imports,
     PVector<WasmInitFunc> initFuncs
 ) {
@@ -27,6 +28,7 @@ public record WasmModule(
         P.vector(),
         P.vector(),
         P.vector(),
+        P.vector(),
         P.vector()
     );
 
@@ -127,6 +129,7 @@ public record WasmModule(
             memories,
             globals,
             elementSegments,
+            dataSegments,
             imports,
             initFuncs
         );
@@ -140,6 +143,7 @@ public record WasmModule(
             memories,
             globals,
             elementSegments,
+            dataSegments,
             imports,
             initFuncs
         );
@@ -153,6 +157,7 @@ public record WasmModule(
             memories,
             globals,
             elementSegments,
+            dataSegments,
             imports,
             initFuncs
         );
@@ -166,6 +171,7 @@ public record WasmModule(
             memories.plus(memory),
             globals,
             elementSegments,
+            dataSegments,
             imports,
             initFuncs
         );
@@ -179,6 +185,7 @@ public record WasmModule(
             memories,
             globals.plus(global),
             elementSegments,
+            dataSegments,
             imports,
             initFuncs
         );
@@ -192,6 +199,21 @@ public record WasmModule(
             memories,
             globals,
             elementSegments.plus(elementSegment),
+            dataSegments,
+            imports,
+            initFuncs
+        );
+    }
+
+    public WasmModule addDataSegment(WasmDataSegment dataSegment) {
+        return new WasmModule(
+            types,
+            functions,
+            tables,
+            memories,
+            globals,
+            elementSegments,
+            dataSegments.plus(dataSegment),
             imports,
             initFuncs
         );
@@ -212,6 +234,7 @@ public record WasmModule(
             memories,
             globals,
             elementSegments,
+            dataSegments,
             newImports,
             initFuncs
         );
@@ -225,6 +248,7 @@ public record WasmModule(
             memories,
             globals,
             elementSegments,
+            dataSegments,
             imports,
             initFuncs.plus(initFunc)
         );
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/WasmBinaryFormatConventions.java b/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/WasmBinaryFormatConventions.java
index 3607e447..cdce10c6 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/WasmBinaryFormatConventions.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/WasmBinaryFormatConventions.java
@@ -43,4 +43,12 @@ public class WasmBinaryFormatConventions {
             writeElement.write(elementIndex, vec.get(elementIndex));
         }
     }
+
+    static void writeVec(
+        byte[] vec,
+        BinaryWriter writer
+    ) throws IOException {
+        Leb128Encoding.encodeUnsignedInt32(vec.length, writer);
+        writer.writeBytes(vec);
+    }
 }
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 94130edb..1023206e 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
@@ -2,6 +2,7 @@ package org.shedlang.compiler.backends.wasm.structure.binary;
 
 import org.shedlang.compiler.backends.wasm.structure.BinaryWriter;
 import org.shedlang.compiler.backends.wasm.structure.ast.linking.WasmSymbolIndex;
+import org.shedlang.compiler.backends.wasm.structure.ast.modules.WasmAlignment;
 import org.shedlang.compiler.backends.wasm.structure.ast.modules.WasmFuncIndex;
 import org.shedlang.compiler.backends.wasm.structure.ast.modules.WasmGlobalIndex;
 import org.shedlang.compiler.backends.wasm.structure.ast.modules.WasmModule;
@@ -35,6 +36,8 @@ public class WasmBinaryFormatLinking {
             Leb128Encoding.encodeUnsignedInt32(VERSION, writer);
 
             var symbolTable = generateSymbolTable(module);
+
+            writeSegmentInfoSubsection(module, writer);
             
             writeSymbolTableSubsection(symbolTable, writer);
             
@@ -66,6 +69,25 @@ public class WasmBinaryFormatLinking {
         }
     }
 
+    private static void writeSegmentInfoSubsection(
+        WasmModule module,
+        BinaryWriter writer
+    ) throws IOException {
+        if (module.dataSegments().size() > 0) {
+            writeSubsection(SubsectionType.WASM_SEGMENT_INFO, writer, () -> {
+                writeVec(
+                    module.dataSegments(),
+                    dataSegment -> {
+                        writeString(dataSegment.name(), writer);
+                        writeAlignment(dataSegment.alignment(), writer);
+                        writeFlags(List.of(), writer);
+                    },
+                    writer
+                );
+            });
+        }
+    }
+
     private static void writeInitFuncsSubsection(
         WasmModule module,
         SymbolTable symbolTable,
@@ -101,13 +123,13 @@ public class WasmBinaryFormatLinking {
         switch (symbolInfo) {
             case SymbolInfo.Function function -> {
                 writer.writeByte(SymbolType.FUNCTION.id);
-                Leb128Encoding.encodeUnsignedInt32(flagsToInt(function.flags), writer);
+                writeFlags(function.flags, writer);
                 writeFuncIndex(function.funcIndex, writer);
                 writeString(function.name(), writer);
             }
             case SymbolInfo.Global global -> {
                 writer.writeByte(SymbolType.GLOBAL.id);
-                Leb128Encoding.encodeUnsignedInt32(flagsToInt(global.flags), writer);
+                writeFlags(global.flags, writer);
                 writeGlobalIndex(global.globalIndex, writer);
                 writeString(global.name(), writer);
             }
@@ -215,7 +237,15 @@ public class WasmBinaryFormatLinking {
         return asInt;
     }
 
+    private static void writeFlags(List<SymbolFlag> flags, BinaryWriter writer) {
+        Leb128Encoding.encodeUnsignedInt32(flagsToInt(flags), writer);
+    }
+
     private static void writeSymbolIndex(WasmSymbolIndex symbolIndex, BinaryWriter writer) {
         Leb128Encoding.encodeUnsignedInt32(symbolIndex.value(), writer);
     }
+
+    private static void writeAlignment(WasmAlignment alignment, BinaryWriter writer) {
+        Leb128Encoding.encodeUnsignedInt32(alignment.alignment(), writer);
+    }
 }
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 2f424d98..1b993461 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
@@ -275,7 +275,7 @@ public class WasmBinaryFormatModules {
         WasmElementSegment elementSegment,
         BinaryWriter writer
     ) throws IOException {
-        writer.writeByte(0x00);
+        Leb128Encoding.encodeUnsignedInt32(0x00, writer);
         writeExpression(new WasmExpression(List.of(new WasmI32Const(0))), writer);
         writeVec(
             elementSegment.elements(),
@@ -316,6 +316,24 @@ public class WasmBinaryFormatModules {
         writeExpression(function.body(), writer);
     }
 
+    // == Data Section ==
+
+    private static void writeDataSection(WasmModule module, BinaryWriter writer) throws IOException {
+        if (module.dataSegments().size() > 0) {
+            writeSection(SectionId.DATA, writer, () -> {
+                writeVec(
+                    module.dataSegments(),
+                    dataSegment -> {
+                        Leb128Encoding.encodeUnsignedInt32(0x00, writer);
+                        writeExpression(dataSegment.offset(), writer);
+                        writeVec(dataSegment.init(), writer);
+                    },
+                    writer
+                );
+            });
+        }
+    }
+
     // == Modules ==
 
     private static final byte[] WASM_MAGIC = HexFormat.of().parseHex("0061736D");
@@ -335,6 +353,7 @@ public class WasmBinaryFormatModules {
         writeGlobalSection(module, writer);
         writeElementSection(module, writer);
         writeCodeSection(module, writer);
+        writeDataSection(module, writer);
         writeNameSection(module, indexTable, writer);
         writeLinkingSection(module, indexTable, writer);