Extract WasmIndexTable

2 files changed, 39 insertions, 13 deletions

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 5bab03c2..f03cdd32 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
@@ -7,9 +7,7 @@ import org.shedlang.compiler.backends.wasm.structure.ast.types.WasmFuncType;
 
 import java.io.IOException;
 import java.io.OutputStream;
-import java.util.HashMap;
 import java.util.HexFormat;
-import java.util.Map;
 
 import static org.shedlang.compiler.backends.wasm.structure.binary.WasmBinaryFormatConventions.writeVec;
 import static org.shedlang.compiler.backends.wasm.structure.binary.WasmBinaryFormatNameSection.writeNameSection;
@@ -108,18 +106,14 @@ public class WasmBinaryFormatModules {
 
     private static void writeImportSection(
         WasmModule module,
+        WasmIndexTable indexTable,
         BinaryWriter writer
     ) throws IOException {
         if (module.imports().size() > 0) {
             writeSection(SectionId.IMPORT, writer, () -> {
-                var funcTypeToTypeIndex = new HashMap<WasmFuncType, Integer>();
-                for (var funcType : module.types()) {
-                    funcTypeToTypeIndex.put(funcType, funcTypeToTypeIndex.size());
-                }
-
                 writeVec(
                     module.imports(),
-                    import_ -> writeImport(import_, funcTypeToTypeIndex, writer),
+                    import_ -> writeImport(import_, indexTable, writer),
                     writer
                 );
             });
@@ -128,23 +122,23 @@ public class WasmBinaryFormatModules {
 
     private static void writeImport(
         WasmImport import_,
-        Map<WasmFuncType, Integer> funcTypeToTypeIndex,
+        WasmIndexTable indexTable,
         BinaryWriter writer
     ) {
         writeString(import_.moduleName(), writer);
         writeString(import_.entityName(), writer);
-        writeImportDesciptor(import_.descriptor(), funcTypeToTypeIndex, writer);
+        writeImportDesciptor(import_.descriptor(), indexTable, writer);
     }
 
     private static void writeImportDesciptor(
         WasmImportDescriptor descriptor,
-        Map<WasmFuncType, Integer> funcTypeToTypeIndex,
+        WasmIndexTable indexTable,
         BinaryWriter writer
     ) {
         switch (descriptor) {
             case WasmImportDescriptor.Func func -> {
                 writer.writeByte(0x00);
-                var typeIndex = funcTypeToTypeIndex.get(func.type());
+                var typeIndex = indexTable.funcIndex(func.type());
                 writeTypeIndex(typeIndex, writer);
             }
             case WasmImportDescriptor.Table table -> {
@@ -207,11 +201,13 @@ public class WasmBinaryFormatModules {
     private static final byte[] WASM_VERSION = HexFormat.of().parseHex("01000000");
 
     public static void serializeModule(WasmModule module, OutputStream outputStream) throws IOException {
+        var indexTable = WasmIndexTable.fromModule(module);
+
         var writer = new BinaryWriter();
         writer.writeBytes(WASM_MAGIC);
         writer.writeBytes(WASM_VERSION);
         writeTypeSection(module.types(), writer);
-        writeImportSection(module, writer);
+        writeImportSection(module, indexTable, writer);
         writeTableSection(module, writer);
         writeMemorySection(module, writer);
         writeNameSection(module, writer);
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/WasmIndexTable.java b/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/WasmIndexTable.java
new file mode 100644
index 00000000..0b2c5c56
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/WasmIndexTable.java
@@ -0,0 +1,30 @@
+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.types.WasmFuncType;
+
+import java.util.HashMap;
+import java.util.Map;
+
+public class WasmIndexTable {
+    public static WasmIndexTable fromModule(WasmModule module) {
+        var funcTypeToTypeIndex = new HashMap<WasmFuncType, Integer>();
+        for (var funcType : module.types()) {
+            funcTypeToTypeIndex.put(funcType, funcTypeToTypeIndex.size());
+        }
+
+        return new WasmIndexTable(
+            funcTypeToTypeIndex
+        );
+    }
+
+    private final Map<WasmFuncType, Integer> funcTypeToTypeIndex;
+
+    private WasmIndexTable(Map<WasmFuncType, Integer> funcTypeToTypeIndex) {
+        this.funcTypeToTypeIndex = funcTypeToTypeIndex;
+    }
+
+    public int funcIndex(WasmFuncType type) {
+        return funcTypeToTypeIndex.get(type);
+    }
+}