WASM binary format: support element segments

3 files changed, 63 insertions, 0 deletions

diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/modules/WasmElementSegment.java b/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/modules/WasmElementSegment.java
new file mode 100644
index 00000000..e389e7f0
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/modules/WasmElementSegment.java
@@ -0,0 +1,11 @@
+package org.shedlang.compiler.backends.wasm.structure.ast.modules;
+
+import java.util.List;
+
+public record WasmElementSegment(
+    List<WasmFuncIndex> elements
+) {
+    public static WasmElementSegment funcIndicesActive(List<WasmFuncIndex> elements) {
+        return new WasmElementSegment(elements);
+    }
+}
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 0a15cc70..6772e6ae 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
@@ -13,6 +13,7 @@ public record WasmModule(
     PVector<WasmFunction> functions,
     PVector<WasmTable> tables,
     PVector<WasmMemory> memories,
+    PVector<WasmElementSegment> elementSegments,
     PVector<WasmImport> imports
 ) {
     public static final WasmModule EMPTY = new WasmModule(
@@ -20,6 +21,7 @@ public record WasmModule(
         P.vector(),
         P.vector(),
         P.vector(),
+        P.vector(),
         P.vector()
     );
 
@@ -46,6 +48,7 @@ public record WasmModule(
             functions,
             tables,
             memories,
+            elementSegments,
             imports
         );
     }
@@ -56,6 +59,7 @@ public record WasmModule(
             functions.plus(function),
             tables,
             memories,
+            elementSegments,
             imports
         );
     }
@@ -66,6 +70,7 @@ public record WasmModule(
             functions,
             tables.plus(table),
             memories,
+            elementSegments,
             imports
         );
     }
@@ -76,6 +81,18 @@ public record WasmModule(
             functions,
             tables,
             memories.plus(memory),
+            elementSegments,
+            imports
+        );
+    }
+
+    public WasmModule addElementSegment(WasmElementSegment elementSegment) {
+        return new WasmModule(
+            types,
+            functions,
+            tables,
+            memories,
+            elementSegments.plus(elementSegment),
             imports
         );
     }
@@ -93,6 +110,7 @@ public record WasmModule(
             functions,
             tables,
             memories,
+            elementSegments,
             newImports
         );
     }
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 faed48d5..603319d8 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
@@ -2,12 +2,15 @@ package org.shedlang.compiler.backends.wasm.structure.binary;
 
 import org.pcollections.OrderedPSet;
 import org.shedlang.compiler.backends.wasm.structure.BinaryWriter;
+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.modules.*;
 import org.shedlang.compiler.backends.wasm.structure.ast.types.WasmFuncType;
 
 import java.io.IOException;
 import java.io.OutputStream;
 import java.util.HexFormat;
+import java.util.List;
 
 import static org.shedlang.compiler.backends.wasm.structure.binary.WasmBinaryFormatConventions.writeVec;
 import static org.shedlang.compiler.backends.wasm.structure.binary.WasmBinaryFormatInstructions.writeExpression;
@@ -190,6 +193,36 @@ public class WasmBinaryFormatModules {
         writeTableType(table.type(), writer);
     }
 
+    // == Element Section ==
+
+    private static void writeElementSection(
+        WasmModule module,
+        BinaryWriter writer
+    ) throws IOException {
+        if (module.elementSegments().size() > 0) {
+            writeSection(SectionId.ELEMENT, writer, () -> {
+                writeVec(
+                    module.elementSegments(),
+                    elementSegment -> writeElementSegment(elementSegment, writer),
+                    writer
+                );
+            });
+        }
+    }
+
+    private static void writeElementSegment(
+        WasmElementSegment elementSegment,
+        BinaryWriter writer
+    ) throws IOException {
+        writer.writeByte(0x00);
+        writeExpression(new WasmExpression(List.of(new WasmI32Const(0))), writer);
+        writeVec(
+            elementSegment.elements(),
+            element -> writeFuncIndex(element, writer),
+            writer
+        );
+    }
+
     // == Code Section ==
 
     private static void writeCodeSection(WasmModule module, BinaryWriter writer) throws IOException {
@@ -256,6 +289,7 @@ public class WasmBinaryFormatModules {
         writeFunctionSection(module, indexTable, writer);
         writeTableSection(module, writer);
         writeMemorySection(module, writer);
+        writeElementSection(module, writer);
         writeCodeSection(module, writer);
         writeNameSection(module, indexTable, writer);