Support init funcs in WASM binary format

4 files changed, 95 insertions, 14 deletions

diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/linking/WasmInitFunc.java b/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/linking/WasmInitFunc.java
new file mode 100644
index 00000000..1613a96f
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/linking/WasmInitFunc.java
@@ -0,0 +1,9 @@
+package org.shedlang.compiler.backends.wasm.structure.ast.linking;
+
+import org.shedlang.compiler.backends.wasm.structure.ast.modules.WasmFuncIndex;
+
+public record WasmInitFunc(
+    int priority,
+    WasmFuncIndex funcIndex
+) {
+}
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/linking/WasmSymbolIndex.java b/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/linking/WasmSymbolIndex.java
new file mode 100644
index 00000000..b2baac2c
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/linking/WasmSymbolIndex.java
@@ -0,0 +1,4 @@
+package org.shedlang.compiler.backends.wasm.structure.ast.linking;
+
+public record WasmSymbolIndex(int value) {
+}
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 5875a591..a4284202 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
@@ -2,6 +2,7 @@ package org.shedlang.compiler.backends.wasm.structure.ast.modules;
 
 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.util.P;
 
@@ -15,7 +16,8 @@ public record WasmModule(
     PVector<WasmMemory> memories,
     PVector<WasmGlobal> globals,
     PVector<WasmElementSegment> elementSegments,
-    PVector<WasmImport> imports
+    PVector<WasmImport> imports,
+    PVector<WasmInitFunc> initFuncs
 ) {
     public static final WasmModule EMPTY = new WasmModule(
         P.orderedSet(),
@@ -24,6 +26,7 @@ public record WasmModule(
         P.vector(),
         P.vector(),
         P.vector(),
+        P.vector(),
         P.vector()
     );
 
@@ -124,7 +127,8 @@ public record WasmModule(
             memories,
             globals,
             elementSegments,
-            imports
+            imports,
+            initFuncs
         );
     }
 
@@ -136,7 +140,8 @@ public record WasmModule(
             memories,
             globals,
             elementSegments,
-            imports
+            imports,
+            initFuncs
         );
     }
 
@@ -148,7 +153,8 @@ public record WasmModule(
             memories,
             globals,
             elementSegments,
-            imports
+            imports,
+            initFuncs
         );
     }
 
@@ -160,7 +166,8 @@ public record WasmModule(
             memories.plus(memory),
             globals,
             elementSegments,
-            imports
+            imports,
+            initFuncs
         );
     }
 
@@ -172,7 +179,8 @@ public record WasmModule(
             memories,
             globals.plus(global),
             elementSegments,
-            imports
+            imports,
+            initFuncs
         );
     }
 
@@ -184,7 +192,8 @@ public record WasmModule(
             memories,
             globals,
             elementSegments.plus(elementSegment),
-            imports
+            imports,
+            initFuncs
         );
     }
 
@@ -203,7 +212,21 @@ public record WasmModule(
             memories,
             globals,
             elementSegments,
-            newImports
+            newImports,
+            initFuncs
+        );
+    }
+
+    public WasmModule addInitFunc(WasmInitFunc initFunc) {
+        return new WasmModule(
+            types,
+            functions,
+            tables,
+            memories,
+            globals,
+            elementSegments,
+            imports,
+            initFuncs.plus(initFunc)
         );
     }
 }
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 c92baba2..94130edb 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
@@ -1,6 +1,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.WasmFuncIndex;
 import org.shedlang.compiler.backends.wasm.structure.ast.modules.WasmGlobalIndex;
 import org.shedlang.compiler.backends.wasm.structure.ast.modules.WasmModule;
@@ -32,7 +33,14 @@ public class WasmBinaryFormatLinking {
     ) throws IOException {
         writeCustomSection("linking", writer, () -> {
             Leb128Encoding.encodeUnsignedInt32(VERSION, writer);
-            writeSymbolTableSubsection(module, indexTable, writer);
+
+            var symbolTable = generateSymbolTable(module);
+            
+            writeSymbolTableSubsection(symbolTable, writer);
+            
+            if (module.initFuncs().size() > 0) {
+                writeInitFuncsSubsection(module, symbolTable, writer);
+            }
         });
     }
 
@@ -58,14 +66,31 @@ public class WasmBinaryFormatLinking {
         }
     }
 
-    private static void writeSymbolTableSubsection(
+    private static void writeInitFuncsSubsection(
         WasmModule module,
-        WasmIndexTable indexTable,
+        SymbolTable symbolTable,
+        BinaryWriter writer
+    ) throws IOException {
+        writeSubsection(SubsectionType.WASM_INIT_FUNCS, writer, () -> {
+            writeVec(
+                module.initFuncs(),
+                initFunc -> {
+                    Leb128Encoding.encodeUnsignedInt32(initFunc.priority(), writer);
+                    var symbolIndex = symbolTable.funcIndexToSymbolIndex(initFunc.funcIndex());
+                    writeSymbolIndex(symbolIndex, writer);
+                },
+                writer
+            );
+        });
+    }
+
+    private static void writeSymbolTableSubsection(
+        SymbolTable symbolTable,
         BinaryWriter writer
     ) throws IOException {
         writeSubsection(SubsectionType.WASM_SYMBOL_TABLE, writer, () -> {
             writeVec(
-                symbolInfos(module),
+                symbolTable.symbolInfos,
                 symbolInfo -> writeSymbolInfo(symbolInfo, writer),
                 writer
             );
@@ -103,8 +128,23 @@ public class WasmBinaryFormatLinking {
             this.id = (byte) id;
         }
     }
+
+    private static class SymbolTable {
+        private final List<SymbolInfo> symbolInfos;
+        private final int functionSymbolIndexOffset;
+
+        public SymbolTable(List<SymbolInfo> symbolInfos, int functionSymbolIndexOffset) {
+            this.symbolInfos = symbolInfos;
+            this.functionSymbolIndexOffset = functionSymbolIndexOffset;
+        }
+
+        public WasmSymbolIndex funcIndexToSymbolIndex(WasmFuncIndex funcIndex) {
+            return new WasmSymbolIndex(funcIndex.value() + functionSymbolIndexOffset);
+        }
+    }
     
-    private static List<SymbolInfo> symbolInfos(WasmModule module) {
+    private static SymbolTable generateSymbolTable(WasmModule module) {
+        var functionSymbolIndexOffset = 0;
         var functionSymbols = module.indexedFunctions().stream()
             .map(function -> {
                 var flags = switch (function) {
@@ -137,7 +177,8 @@ public class WasmBinaryFormatLinking {
                 );
             });
 
-        return Stream.<SymbolInfo>concat(functionSymbols, globalSymbols).toList();
+        var symbolInfos = Stream.<SymbolInfo>concat(functionSymbols, globalSymbols).toList();
+        return new SymbolTable(symbolInfos, functionSymbolIndexOffset);
     }
 
     private sealed interface SymbolInfo {
@@ -173,4 +214,8 @@ public class WasmBinaryFormatLinking {
 
         return asInt;
     }
+
+    private static void writeSymbolIndex(WasmSymbolIndex symbolIndex, BinaryWriter writer) {
+        Leb128Encoding.encodeUnsignedInt32(symbolIndex.value(), writer);
+    }
 }