WASM binary format: support functions

8 files changed, 168 insertions, 2 deletions

diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/instructions/WasmI32Const.java b/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/instructions/WasmI32Const.java
new file mode 100644
index 00000000..e9815e06
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/instructions/WasmI32Const.java
@@ -0,0 +1,4 @@
+package org.shedlang.compiler.backends.wasm.structure.ast.instructions;
+
+public record WasmI32Const(int value) implements WasmInstruction {
+}
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/instructions/WasmI64Const.java b/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/instructions/WasmI64Const.java
new file mode 100644
index 00000000..b6f87aaa
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/instructions/WasmI64Const.java
@@ -0,0 +1,4 @@
+package org.shedlang.compiler.backends.wasm.structure.ast.instructions;
+
+public record WasmI64Const(long value) implements WasmInstruction {
+}
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/instructions/WasmInstruction.java b/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/instructions/WasmInstruction.java
new file mode 100644
index 00000000..4f33a003
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/instructions/WasmInstruction.java
@@ -0,0 +1,4 @@
+package org.shedlang.compiler.backends.wasm.structure.ast.instructions;
+
+public sealed interface WasmInstruction permits WasmI32Const, WasmI64Const {
+}
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/modules/WasmFunction.java b/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/modules/WasmFunction.java
new file mode 100644
index 00000000..50ad805b
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/modules/WasmFunction.java
@@ -0,0 +1,14 @@
+package org.shedlang.compiler.backends.wasm.structure.ast.modules;
+
+import org.shedlang.compiler.backends.wasm.structure.ast.instructions.WasmInstruction;
+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 WasmFunction(
+    WasmFuncType type,
+    List<WasmValueType> locals,
+    List<WasmInstruction> body
+) {
+}
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 12fbb80e..7b19cee9 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
@@ -7,6 +7,7 @@ import org.shedlang.compiler.util.P;
 
 public record WasmModule(
     OrderedPSet<WasmFuncType> types,
+    PVector<WasmFunction> functions,
     PVector<WasmTable> tables,
     PVector<WasmMemory> memories,
     PVector<WasmImport> imports
@@ -15,12 +16,24 @@ public record WasmModule(
         P.orderedSet(),
         P.vector(),
         P.vector(),
+        P.vector(),
         P.vector()
     );
 
     public WasmModule addType(WasmFuncType type) {
         return new WasmModule(
             types.plus(type),
+            functions,
+            tables,
+            memories,
+            imports
+        );
+    }
+
+    public WasmModule addFunction(WasmFunction function) {
+        return new WasmModule(
+            types.plus(function.type()),
+            functions.plus(function),
             tables,
             memories,
             imports
@@ -30,6 +43,7 @@ public record WasmModule(
     public WasmModule addTable(WasmTable table) {
         return new WasmModule(
             types,
+            functions,
             tables.plus(table),
             memories,
             imports
@@ -39,6 +53,7 @@ public record WasmModule(
     public WasmModule addMemory(WasmMemory memory) {
         return new WasmModule(
             types,
+            functions,
             tables,
             memories.plus(memory),
             imports
@@ -55,6 +70,7 @@ public record WasmModule(
 
         return new WasmModule(
             newTypes,
+            functions,
             tables,
             memories,
             newImports
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/Leb128Encoding.java b/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/Leb128Encoding.java
index e0c4734d..87d91520 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/Leb128Encoding.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/Leb128Encoding.java
@@ -64,4 +64,28 @@ public class Leb128Encoding {
             writer.writeByte((byte) (value & 0x7f));
         }
     }
+
+    static void encodeSignedInt64(long value, BinaryWriter writer) {
+        if (value >= 0) {
+            while (true) {
+                if ((value >> 6) == 0) {
+                    writer.writeByte((byte) value);
+                    return;
+                } else {
+                    writer.writeByte((byte) ((value & 0x7f) | 0x80));
+                    value >>= 7;
+                }
+            }
+        } else {
+            while (true) {
+                if ((value >> 6) == -1) {
+                    writer.writeByte((byte) (value & 0x7f));
+                    return;
+                } else {
+                    writer.writeByte((byte) ((value & 0x7f) | 0x80));
+                    value >>= 7;
+                }
+            }
+        }
+    }
 }
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/WasmBinaryFormatInstructions.java b/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/WasmBinaryFormatInstructions.java
new file mode 100644
index 00000000..58d3d28c
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/structure/binary/WasmBinaryFormatInstructions.java
@@ -0,0 +1,44 @@
+package org.shedlang.compiler.backends.wasm.structure.binary;
+
+import org.shedlang.compiler.backends.wasm.structure.BinaryWriter;
+import org.shedlang.compiler.backends.wasm.structure.ast.instructions.WasmI32Const;
+import org.shedlang.compiler.backends.wasm.structure.ast.instructions.WasmI64Const;
+import org.shedlang.compiler.backends.wasm.structure.ast.instructions.WasmInstruction;
+
+import java.io.IOException;
+import java.util.List;
+
+public class WasmBinaryFormatInstructions {
+    // Corresponds to https://webassembly.github.io/spec/core/binary/instructions.html
+
+    private WasmBinaryFormatInstructions() {
+    }
+
+    private static void writeInstruction(WasmInstruction instruction, BinaryWriter writer) {
+        switch (instruction) {
+
+            // == Numeric Instructions ==
+
+            case WasmI32Const i32Const -> {
+                writer.writeByte(0x41);
+                Leb128Encoding.encodeSignedInt32(i32Const.value(), writer);
+            }
+            case WasmI64Const i64Const -> {
+                writer.writeByte(0x42);
+                Leb128Encoding.encodeSignedInt64(i64Const.value(), writer);
+            }
+        }
+    }
+
+    // == Expressions ==
+
+    static void writeExpression(
+        List<WasmInstruction> instructions,
+        BinaryWriter writer
+    ) throws IOException  {
+        for (var instruction : instructions) {
+            writeInstruction(instruction, writer);
+        }
+        writer.writeByte(0x0b);
+    }
+}
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 cddfb026..58ae08a8 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
@@ -10,6 +10,7 @@ import java.io.OutputStream;
 import java.util.HexFormat;
 
 import static org.shedlang.compiler.backends.wasm.structure.binary.WasmBinaryFormatConventions.writeVec;
+import static org.shedlang.compiler.backends.wasm.structure.binary.WasmBinaryFormatInstructions.writeExpression;
 import static org.shedlang.compiler.backends.wasm.structure.binary.WasmBinaryFormatNameSection.writeNameSection;
 import static org.shedlang.compiler.backends.wasm.structure.binary.WasmBinaryFormatSize.writeSize;
 import static org.shedlang.compiler.backends.wasm.structure.binary.WasmBinaryFormatTypes.*;
@@ -118,10 +119,10 @@ public class WasmBinaryFormatModules {
     ) {
         writeString(import_.moduleName(), writer);
         writeString(import_.entityName(), writer);
-        writeImportDesciptor(import_.descriptor(), indexTable, writer);
+        writeImportDescriptor(import_.descriptor(), indexTable, writer);
     }
 
-    private static void writeImportDesciptor(
+    private static void writeImportDescriptor(
         WasmImportDescriptor descriptor,
         WasmIndexTable indexTable,
         BinaryWriter writer
@@ -147,6 +148,27 @@ public class WasmBinaryFormatModules {
         }
     }
 
+    // == Function Section ==
+
+    private static void writeFunctionSection(
+        WasmModule module,
+        WasmIndexTable indexTable,
+        BinaryWriter writer
+    ) throws IOException {
+        if (module.functions().size() > 0) {
+            writeSection(SectionId.FUNCTION, writer, () -> {
+                writeVec(
+                    module.functions(),
+                    function -> writeTypeIndex(
+                        indexTable.funcIndex(function.type()),
+                        writer
+                    ),
+                    writer
+                );
+            });
+        }
+    }
+
     // == Table Section ==
 
     private static void writeTableSection(
@@ -168,6 +190,38 @@ public class WasmBinaryFormatModules {
         writeTableType(table.type(), writer);
     }
 
+    // == Code Section ==
+
+    private static void writeCodeSection(WasmModule module, BinaryWriter writer) throws IOException {
+        if (module.functions().size() > 0) {
+            writeSection(SectionId.CODE, writer, () -> {
+                writeVec(
+                    module.functions(),
+                    function -> writeCode(function, writer),
+                    writer
+                );
+            });
+        }
+    }
+
+    private static void writeCode(WasmFunction function, BinaryWriter writer) throws IOException {
+        writeSize(writer, () -> {
+            writeFunction(function, writer);
+        });
+    }
+
+    private static void writeFunction(WasmFunction function, BinaryWriter writer) throws IOException {
+        writeVec(
+            function.locals(),
+            local -> {
+                Leb128Encoding.encodeUnsignedInt32(1, writer);
+                writeValueType(local, writer);
+            },
+            writer
+        );
+        writeExpression(function.body(), writer);
+    }
+
     // == Memory Section ==
 
     private static void writeMemorySection(WasmModule module, BinaryWriter writer) throws IOException {
@@ -199,8 +253,10 @@ public class WasmBinaryFormatModules {
         writer.writeBytes(WASM_VERSION);
         writeTypeSection(module.types(), writer);
         writeImportSection(module, indexTable, writer);
+        writeFunctionSection(module, indexTable, writer);
         writeTableSection(module, writer);
         writeMemorySection(module, writer);
+        writeCodeSection(module, writer);
         writeNameSection(module, writer);
 
         outputStream.write(writer.toBytes());