Add tags to WASM binary format

6 files changed, 140 insertions, 29 deletions

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 2175e46c..5fdbbbe7 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
@@ -4,6 +4,7 @@ import org.pcollections.OrderedPSet;
 import org.pcollections.PVector;
 import org.shedlang.compiler.backends.wasm.structure.ast.linking.WasmInitFunc;
 import org.shedlang.compiler.backends.wasm.structure.ast.types.WasmFuncType;
+import org.shedlang.compiler.backends.wasm.structure.ast.types.WasmValueType;
 import org.shedlang.compiler.util.P;
 
 import java.util.ArrayList;
@@ -18,7 +19,8 @@ public record WasmModule(
     PVector<WasmElementSegment> elementSegments,
     PVector<WasmDataSegment> dataSegments,
     PVector<WasmImport> imports,
-    PVector<WasmInitFunc> initFuncs
+    PVector<WasmInitFunc> initFuncs,
+    PVector<WasmTag> tags
 ) {
     public static final WasmModule EMPTY = new WasmModule(
         P.orderedSet(),
@@ -29,6 +31,7 @@ public record WasmModule(
         P.vector(),
         P.vector(),
         P.vector(),
+        P.vector(),
         P.vector()
     );
 
@@ -116,7 +119,8 @@ public record WasmModule(
             elementSegments,
             dataSegments,
             imports,
-            initFuncs
+            initFuncs,
+            tags
         );
     }
 
@@ -130,7 +134,8 @@ public record WasmModule(
             elementSegments,
             dataSegments,
             imports,
-            initFuncs
+            initFuncs,
+            tags
         );
     }
 
@@ -144,7 +149,8 @@ public record WasmModule(
             elementSegments,
             dataSegments,
             imports,
-            initFuncs
+            initFuncs,
+            tags
         );
     }
 
@@ -158,7 +164,8 @@ public record WasmModule(
             elementSegments,
             dataSegments,
             imports,
-            initFuncs
+            initFuncs,
+            tags
         );
     }
 
@@ -172,7 +179,8 @@ public record WasmModule(
             elementSegments,
             dataSegments,
             imports,
-            initFuncs
+            initFuncs,
+            tags
         );
     }
 
@@ -186,7 +194,8 @@ public record WasmModule(
             elementSegments.plus(elementSegment),
             dataSegments,
             imports,
-            initFuncs
+            initFuncs,
+            tags
         );
     }
 
@@ -200,7 +209,8 @@ public record WasmModule(
             elementSegments,
             dataSegments.plus(dataSegment),
             imports,
-            initFuncs
+            initFuncs,
+            tags
         );
     }
 
@@ -221,7 +231,8 @@ public record WasmModule(
             elementSegments,
             dataSegments,
             newImports,
-            initFuncs
+            initFuncs,
+            tags
         );
     }
 
@@ -235,7 +246,28 @@ public record WasmModule(
             elementSegments,
             dataSegments,
             imports,
-            initFuncs.plus(initFunc)
+            initFuncs.plus(initFunc),
+            tags
+        );
+    }
+
+    public WasmModule addTag(
+        WasmTagType tagType,
+        String name,
+        List<WasmValueType> fieldTypes
+    ) {
+        var tag = new WasmTag(tagType, name, fieldTypes);
+        return new WasmModule(
+            types.plus(tag.funcType()),
+            functions,
+            tables,
+            memories,
+            globals,
+            elementSegments,
+            dataSegments,
+            imports,
+            initFuncs,
+            tags.plus(tag)
         );
     }
 }
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/modules/WasmTag.java b/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/modules/WasmTag.java
new file mode 100644
index 00000000..020fc754
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/modules/WasmTag.java
@@ -0,0 +1,12 @@
+package org.shedlang.compiler.backends.wasm.structure.ast.modules;
+
+import org.shedlang.compiler.backends.wasm.structure.ast.types.WasmFuncType;
+import org.shedlang.compiler.backends.wasm.structure.ast.types.WasmValueType;
+
+import java.util.List;
+
+public record WasmTag(WasmTagType tagType, String name, List<WasmValueType> fieldTypes) {
+    public WasmFuncType funcType() {
+        return new WasmFuncType(fieldTypes, List.of());
+    }
+}
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/modules/WasmTagType.java b/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/modules/WasmTagType.java
new file mode 100644
index 00000000..d82e4039
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/modules/WasmTagType.java
@@ -0,0 +1,15 @@
+package org.shedlang.compiler.backends.wasm.structure.ast.modules;
+
+public enum WasmTagType {
+    EXCEPTION(0);
+
+    private final byte value;
+
+    WasmTagType(int value) {
+        this.value = (byte) value;
+    }
+
+    public byte id() {
+        return value;
+    }
+}
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/WasmBinaryFormatExceptions.java b/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/WasmBinaryFormatExceptions.java
new file mode 100644
index 00000000..29b0d9ed
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/WasmBinaryFormatExceptions.java
@@ -0,0 +1,29 @@
+package org.shedlang.compiler.backends.wasm.structure.binary;
+
+import org.shedlang.compiler.backends.wasm.structure.ast.modules.WasmModule;
+import org.shedlang.compiler.backends.wasm.structure.ast.modules.WasmTag;
+
+import java.io.IOException;
+
+public class WasmBinaryFormatExceptions {
+    // Corresponds to https://github.com/WebAssembly/exception-handling/blob/main/proposals/exception-handling/Exceptions.md
+
+    private WasmBinaryFormatExceptions() {
+    }
+
+    public static void writeTagSection(WasmModule module, WasmBinaryWriter writer) throws IOException {
+        if (module.tags().size() > 0) {
+            writer.writeSection(WasmSectionId.TAG, () -> {
+                writer.writeVec(
+                    module.tags(),
+                    tag -> writeTag(tag, writer)
+                );
+            });
+        }
+    }
+
+    private static void writeTag(WasmTag tag, WasmBinaryWriter writer) throws IOException {
+        writer.writeByte(tag.tagType().id());
+        writer.writeTypeIndex(tag.funcType());
+    }
+}
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 1709f78d..863a5d0f 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
@@ -12,6 +12,7 @@ import java.io.OutputStream;
 import java.util.HexFormat;
 import java.util.List;
 
+import static org.shedlang.compiler.backends.wasm.structure.binary.WasmBinaryFormatExceptions.writeTagSection;
 import static org.shedlang.compiler.backends.wasm.structure.binary.WasmBinaryFormatInstructions.writeExpression;
 import static org.shedlang.compiler.backends.wasm.structure.binary.WasmBinaryFormatLinking.*;
 import static org.shedlang.compiler.backends.wasm.structure.binary.WasmBinaryFormatNameSection.writeNameSection;
@@ -335,6 +336,7 @@ public class WasmBinaryFormatModules {
         writeFunctionSection(module, indexTable, writer);
         writeTableSection(module, writer);
         writeMemorySection(module, writer);
+        writeTagSection(module, writer);
         writeGlobalSection(module, writer);
         writeElementSection(module, writer);
         writeCodeSection(module, writer);
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/WasmBinaryWriter.java b/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/WasmBinaryWriter.java
index 3ac5d2bc..bb6798b9 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/WasmBinaryWriter.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/WasmBinaryWriter.java
@@ -15,6 +15,7 @@ public class WasmBinaryWriter {
     private final BinaryWriter writer = new BinaryWriter();
     private final Map<Integer, List<WasmRelocationEntry>> relocationEntries = new HashMap<>();
     private final WasmIndexTable indexTable;
+    private Optional<WasmSectionId> currentSection = Optional.empty();
     private int sectionIndex = -1;
     private int sectionContentsStart = -1;
 
@@ -58,7 +59,7 @@ public class WasmBinaryWriter {
         var funcIndex = indexTable.funcIndex(funcName);
         writeRelocatableIndex(
             WasmRelocationType.FUNCTION_INDEX_LEB,
-            () -> writePaddedUnsignedInt32(funcIndex.value()),
+            funcIndex.value(),
             new WasmRelocatableIndex.FuncIndex(funcName)
         );
     }
@@ -67,7 +68,7 @@ public class WasmBinaryWriter {
         var globalIndex = indexTable.globalIndex(globalName);
         writeRelocatableIndex(
             WasmRelocationType.GLOBAL_INDEX_LEB,
-            () -> writePaddedUnsignedInt32(globalIndex.value()),
+            globalIndex.value(),
             new WasmRelocatableIndex.GlobalIndex(globalName)
         );
     }
@@ -75,7 +76,8 @@ public class WasmBinaryWriter {
     public void writeMemoryAddressConst(WasmMemoryAddress index) throws IOException {
         writeRelocatableIndex(
             WasmRelocationType.MEMORY_ADDR_SLEB,
-            () -> writePaddedSignedInt32(0),
+            // TODO: use actual value? Or throw if not in relocatable section?
+            0,
             new WasmRelocatableIndex.MemoryAddress(index)
         );
     }
@@ -87,47 +89,66 @@ public class WasmBinaryWriter {
     public void writeTableElementIndexConst(WasmFuncName funcName) throws IOException {
         writeRelocatableIndex(
             WasmRelocationType.TABLE_INDEX_SLEB,
-            () -> writePaddedSignedInt32(0),
+            // TODO: use actual value? Or throw if not in relocatable section?
+            0,
             new WasmRelocatableIndex.TableElementIndex(funcName)
         );
     }
 
     public void writeTypeIndex(WasmFuncType funcType) throws IOException {
+        var typeIndex = indexTable.typeIndex(funcType);
         writeRelocatableIndex(
             WasmRelocationType.TYPE_INDEX_LEB,
-            () -> writePaddedUnsignedInt32(0),
+            typeIndex.value(),
             new WasmRelocatableIndex.TypeIndex(funcType)
         );
     }
 
     private void writeRelocatableIndex(
         WasmRelocationType relocationType,
-        WasmBinaryFormatModules.Write writePlaceholder,
+        int unrelocatedIndex,
         WasmRelocatableIndex index
     ) throws IOException {
-        // TODO: handle relocatable values in non-CODE sections
-
-        var sectionContentsOffset = writer.position() - sectionContentsStart;
-        writePlaceholder.write();
-
-        if (!relocationEntries.containsKey(sectionIndex)) {
-            relocationEntries.put(sectionIndex, new ArrayList<>());
+        // TODO: handle custom sections?
+        var isRelocatable = currentSection
+            .map(c -> c == WasmSectionId.CODE || c == WasmSectionId.DATA)
+            .orElse(false);
+
+        if (isRelocatable) {
+            var sectionContentsOffset = writer.position() - sectionContentsStart;
+            if (relocationType.isSigned()) {
+                writePaddedSignedInt32(unrelocatedIndex);
+            } else {
+                writePaddedUnsignedInt32(unrelocatedIndex);
+            }
+
+            if (!relocationEntries.containsKey(sectionIndex)) {
+                relocationEntries.put(sectionIndex, new ArrayList<>());
+            }
+            relocationEntries.get(sectionIndex).add(new WasmRelocationEntry(
+                sectionIndex,
+                relocationType,
+                sectionContentsOffset,
+                index
+            ));
+        } else {
+            if (relocationType.isSigned()) {
+                writeSignedInt32(unrelocatedIndex);
+            } else {
+                writeUnsignedInt32(unrelocatedIndex);
+            }
         }
-        relocationEntries.get(sectionIndex).add(new WasmRelocationEntry(
-            sectionIndex,
-            relocationType,
-            sectionContentsOffset,
-            index
-        ));
     }
 
     // TODO: extract write interface
     public void writeSection(WasmSectionId sectionId, WasmBinaryFormatModules.Write writeContents) throws IOException {
+        currentSection = Optional.of(sectionId);
         sectionIndex++;
         WasmBinaryFormatModules.writeSection(sectionId, writer, () -> {
             sectionContentsStart = writer.position();
             writeContents.write();
         });
+        currentSection = Optional.empty();
     }
 
     public <T> void writeVec(