Automatically add function types to module

3 files changed, 18 insertions, 9 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 d8757402..12fbb80e 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
@@ -1,17 +1,18 @@
 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.types.WasmFuncType;
 import org.shedlang.compiler.util.P;
 
 public record WasmModule(
-    PVector<WasmFuncType> types,
+    OrderedPSet<WasmFuncType> types,
     PVector<WasmTable> tables,
     PVector<WasmMemory> memories,
     PVector<WasmImport> imports
 ) {
     public static final WasmModule EMPTY = new WasmModule(
-        P.vector(),
+        P.orderedSet(),
         P.vector(),
         P.vector(),
         P.vector()
@@ -45,11 +46,18 @@ public record WasmModule(
     }
 
     public WasmModule addImport(WasmImport import_) {
+        var newImports = imports.plus(import_);
+
+        var newTypes = types;
+        if (import_.descriptor() instanceof WasmImportDescriptor.Func descriptor) {
+            newTypes = newTypes.plus(descriptor.type());
+        }
+
         return new WasmModule(
-            types,
+            newTypes,
             tables,
             memories,
-            imports.plus(import_)
+            newImports
         );
     }
 }
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 9dc25ffa..3607e447 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
@@ -3,6 +3,7 @@ package org.shedlang.compiler.backends.wasm.structure.binary;
 import org.shedlang.compiler.backends.wasm.structure.BinaryWriter;
 
 import java.io.IOException;
+import java.util.Collection;
 import java.util.List;
 
 public class WasmBinaryFormatConventions {
@@ -18,7 +19,7 @@ public class WasmBinaryFormatConventions {
     }
 
     static <T> void writeVec(
-        List<T> vec,
+        Collection<T> vec,
         WriteElement<T> writeElement,
         BinaryWriter writer
     ) throws IOException {
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 8945d7ed..5bab03c2 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
@@ -1,6 +1,6 @@
 package org.shedlang.compiler.backends.wasm.structure.binary;
 
-import org.pcollections.PVector;
+import org.pcollections.OrderedPSet;
 import org.shedlang.compiler.backends.wasm.structure.BinaryWriter;
 import org.shedlang.compiler.backends.wasm.structure.ast.modules.*;
 import org.shedlang.compiler.backends.wasm.structure.ast.types.WasmFuncType;
@@ -90,7 +90,7 @@ public class WasmBinaryFormatModules {
     // == Type Section ==
 
     private static void writeTypeSection(
-        PVector<WasmFuncType> types,
+        OrderedPSet<WasmFuncType> types,
         BinaryWriter writer
     ) throws IOException {
         if (types.size() > 0) {
@@ -113,8 +113,8 @@ public class WasmBinaryFormatModules {
         if (module.imports().size() > 0) {
             writeSection(SectionId.IMPORT, writer, () -> {
                 var funcTypeToTypeIndex = new HashMap<WasmFuncType, Integer>();
-                for (var funcIndex = 0; funcIndex < module.types().size(); funcIndex++) {
-                    funcTypeToTypeIndex.put(module.types().get(funcIndex), funcIndex);
+                for (var funcType : module.types()) {
+                    funcTypeToTypeIndex.put(funcType, funcTypeToTypeIndex.size());
                 }
 
                 writeVec(