Support import of cuntions in WASM binary format

5 files changed, 86 insertions, 3 deletions

diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/structure/BinaryWriter.java b/src/main/java/org/shedlang/compiler/backends/wasm/structure/BinaryWriter.java
index ade47714..2fef5255 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/structure/BinaryWriter.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/structure/BinaryWriter.java
@@ -10,6 +10,10 @@ public class BinaryWriter {
         buffer.put(value);
     }
 
+    public void writeByte(int value) {
+        writeByte((byte) value);
+    }
+
     public void writeBytes(byte[] bytes) {
         ensureSpareCapacity(bytes.length);
         buffer.put(bytes, 0, bytes.length);
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/structure/WasmBinaryFormat.java b/src/main/java/org/shedlang/compiler/backends/wasm/structure/WasmBinaryFormat.java
index c87a86e5..f9fceab5 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/structure/WasmBinaryFormat.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/structure/WasmBinaryFormat.java
@@ -2,14 +2,19 @@ package org.shedlang.compiler.backends.wasm.structure;
 
 import org.pcollections.PVector;
 import org.shedlang.compiler.backends.wasm.structure.ast.WasmModule;
+import org.shedlang.compiler.backends.wasm.structure.ast.imports.WasmImport;
+import org.shedlang.compiler.backends.wasm.structure.ast.imports.WasmImportDescriptor;
 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.Action;
 
 import java.io.IOException;
 import java.io.OutputStream;
+import java.nio.charset.StandardCharsets;
+import java.util.HashMap;
 import java.util.HexFormat;
 import java.util.List;
+import java.util.Map;
 
 public class WasmBinaryFormat {
     private static final byte[] WASM_MAGIC = HexFormat.of().parseHex("0061736D");
@@ -43,6 +48,7 @@ public class WasmBinaryFormat {
         writer.writeBytes(WASM_MAGIC);
         writer.writeBytes(WASM_VERSION);
         writeTypeSection(module.types(), writer);
+        writeImportSection(module, writer);
 
         outputStream.write(writer.toBytes());
     }
@@ -61,6 +67,42 @@ public class WasmBinaryFormat {
         }
     }
 
+    private static void writeImportSection(
+        WasmModule module,
+        BinaryWriter writer
+    ) throws IOException {
+        if (module.imports().size() > 0) {
+            writeSection(SectionType.IMPORT, writer, () -> {
+                Leb128Encoding.encodeUnsignedInt32(module.imports().size(), writer);
+                
+                var funcTypeToFuncIndex = new HashMap<WasmFuncType, Integer>();
+                for (var funcIndex = 0; funcIndex < module.types().size(); funcIndex++) {
+                    funcTypeToFuncIndex.put(module.types().get(funcIndex), funcIndex);
+                }
+                
+                for (var import_ : module.imports()) {
+                    writeImport(import_, funcTypeToFuncIndex, writer);
+                }
+            });
+        }
+    }
+
+    private static void writeImport(
+        WasmImport import_,
+        Map<WasmFuncType, Integer> funcTypeToFuncIndex,
+        BinaryWriter writer
+    ) {
+        writeString(import_.moduleName(), writer);
+        writeString(import_.entityName(), writer);
+        switch (import_.descriptor()) {
+            case WasmImportDescriptor.Func func -> {
+                writer.writeByte(0x00);
+                var funcIndex = funcTypeToFuncIndex.get(func.type());
+                Leb128Encoding.encodeUnsignedInt32(funcIndex, writer);
+            }
+        }
+    }
+
     private static void writeFuncType(WasmFuncType type, BinaryWriter writer) {
         writer.writeByte((byte) 0x60);
         writeValueTypeVector(type.params(), writer);
@@ -109,4 +151,10 @@ public class WasmBinaryFormat {
         Leb128Encoding.encodePaddedUnsignedInt32(size, writer);
         writer.seek(contentEndPosition);
     }
+
+    private static void writeString(String value, BinaryWriter writer) {
+        var bytes = value.getBytes(StandardCharsets.UTF_8);
+        Leb128Encoding.encodeUnsignedInt32(bytes.length, writer);
+        writer.writeBytes(bytes);
+    }
 }
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/WasmModule.java b/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/WasmModule.java
index ed0d88f2..7bffd4c6 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/WasmModule.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/WasmModule.java
@@ -1,15 +1,30 @@
 package org.shedlang.compiler.backends.wasm.structure.ast;
 
 import org.pcollections.PVector;
+import org.shedlang.compiler.backends.wasm.structure.ast.imports.WasmImport;
 import org.shedlang.compiler.backends.wasm.structure.ast.types.WasmFuncType;
 import org.shedlang.compiler.util.P;
 
 public record WasmModule(
-    PVector<WasmFuncType> types
+    PVector<WasmFuncType> types,
+    PVector<WasmImport> imports
 ) {
-    public static final WasmModule EMPTY = new WasmModule(P.vector());
+    public static final WasmModule EMPTY = new WasmModule(
+        P.vector(),
+        P.vector()
+    );
 
     public WasmModule addType(WasmFuncType type) {
-        return new WasmModule(types.plus(type));
+        return new WasmModule(
+            types.plus(type),
+            imports
+        );
+    }
+
+    public WasmModule addImport(WasmImport import_) {
+        return new WasmModule(
+            types,
+            imports.plus(import_)
+        );
     }
 }
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/imports/WasmImport.java b/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/imports/WasmImport.java
new file mode 100644
index 00000000..5f7c7768
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/imports/WasmImport.java
@@ -0,0 +1,9 @@
+package org.shedlang.compiler.backends.wasm.structure.ast.imports;
+
+public record WasmImport(
+    String moduleName,
+    String entityName,
+    String identifier,
+    WasmImportDescriptor descriptor
+) {
+}
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/imports/WasmImportDescriptor.java b/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/imports/WasmImportDescriptor.java
new file mode 100644
index 00000000..10acc19f
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/imports/WasmImportDescriptor.java
@@ -0,0 +1,7 @@
+package org.shedlang.compiler.backends.wasm.structure.ast.imports;
+
+import org.shedlang.compiler.backends.wasm.structure.ast.types.WasmFuncType;
+
+public sealed interface WasmImportDescriptor {
+    record Func(WasmFuncType type) implements WasmImportDescriptor {}
+}