Support Wasm block types with multiple results

3 files changed, 51 insertions, 22 deletions

diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/PatioToWasm.java b/src/main/java/org/shedlang/compiler/backends/wasm/PatioToWasm.java
index 11e6c22f..c4e73446 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/PatioToWasm.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/PatioToWasm.java
@@ -213,7 +213,9 @@ public class PatioToWasm {
             case PatioIf patioIf -> {
                 context = generateOperand(patioIf.condition(), context);
 
-                var wasmBlockType = generateBlockType(patioIf.destinations(), context);
+                var wasmBlockTypeResult = generateBlockType(patioIf.destinations(), context);
+                var wasmBlockType = wasmBlockTypeResult.value();
+                context = wasmBlockTypeResult.context();
 
                 context = context.enterIfInstruction(wasmBlockType);
                 context = generateBlockBody(patioIf.trueBody(), context);
@@ -359,7 +361,9 @@ public class PatioToWasm {
 
             case PatioSwitch patioSwitch -> {
                 // TODO: use br_table?
-                var wasmBlockType = generateBlockType(patioSwitch.destinations(), context);
+                var wasmBlockTypeResult = generateBlockType(patioSwitch.destinations(), context);
+                var wasmBlockType = wasmBlockTypeResult.value();
+                context = wasmBlockTypeResult.context();
 
                 for (var patioSwitchCase : patioSwitch.cases()) {
                     context = generateOperand(patioSwitch.value(), context);
@@ -397,7 +401,7 @@ public class PatioToWasm {
             case PatioTry try_ -> {
                 // TODO: don't assume we always have a Patio int
                 context = context.enterTryInstruction(
-                    new WasmBlockType(List.of(target.patioIntType()))
+                    WasmBlockType.of(target.patioIntType())
                 );
 
                 context = generateBlockBody(try_.body(), context);
@@ -437,15 +441,25 @@ public class PatioToWasm {
         };
     }
 
-    private static WasmBlockType generateBlockType(
+    private static Context.Result<WasmBlockType> generateBlockType(
         List<PatioLocalRef> destinations,
         Context context
     ) {
-        return new WasmBlockType(
-            destinations.stream()
+        if (destinations.isEmpty()) {
+            return context.toResult(WasmBlockType.EMPTY);
+        } else if (destinations.size() == 1) {
+            var valueType = patioOperandToWasmValueType(destinations.get(0), context);
+            return context.toResult(WasmBlockType.of(valueType));
+        } else {
+            var resultTypes = destinations.stream()
                 .map(patioTarget -> patioOperandToWasmValueType(patioTarget, context))
-                .toList()
-        );
+                .toList();
+
+            var funcType = new WasmFuncType(List.of(), resultTypes);
+
+            return context.addType(funcType).toResult(WasmBlockType.of(funcType));
+        }
+
     }
 
     private static Context generateBlockBody(
@@ -1104,6 +1118,12 @@ public class PatioToWasm {
 
             throw new IllegalStateException("cannot find label index");
         }
+
+        record Result<T>(T value, Context context) {}
+
+        public <T> Result<T> toResult(T value) {
+            return new Result<>(value, this);
+        }
     }
 
     private interface WasmStructuredInstructionBuilder {
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/instructions/WasmBlockType.java b/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/instructions/WasmBlockType.java
index 731deb06..0721ccbf 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/instructions/WasmBlockType.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/structure/ast/instructions/WasmBlockType.java
@@ -1,14 +1,20 @@
 package org.shedlang.compiler.backends.wasm.structure.ast.instructions;
 
+import org.shedlang.compiler.backends.wasm.structure.ast.types.WasmFuncType;
 import org.shedlang.compiler.backends.wasm.structure.ast.types.WasmValueType;
 
-import java.util.Arrays;
-import java.util.List;
+public sealed interface WasmBlockType {
+    WasmBlockType EMPTY = new WasmBlockType.Empty();
 
-public record WasmBlockType(List<WasmValueType> resultTypes) {
-    public static final WasmBlockType EMPTY = new WasmBlockType(List.of());
+    static WasmBlockType of(WasmValueType valueType) {
+        return new WasmBlockType.ValueType(valueType);
+    }
 
-    public static WasmBlockType of(WasmValueType... resultTypes) {
-        return new WasmBlockType(Arrays.asList(resultTypes));
+    static WasmBlockType of(WasmFuncType funcType) {
+        return new WasmBlockType.FuncType(funcType);
     }
+
+    record Empty() implements WasmBlockType {}
+    record ValueType(WasmValueType valueType) implements WasmBlockType {}
+    record FuncType(WasmFuncType funcType) implements WasmBlockType {}
 }
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
index 3a712e33..5d6b18d9 100644
--- 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
@@ -224,14 +224,17 @@ public class WasmBinaryFormatInstructions {
 
     // == Control Instructions ==
 
-    private static void writeBlockType(WasmBlockType blockType, WasmBinaryWriter writer) {
-        if (blockType.resultTypes().isEmpty()) {
-            writer.writeByte(0x40);
-        } else if (blockType.resultTypes().size() == 1) {
-            writeValueType(blockType.resultTypes().get(0), writer);
-        } else {
-            // TODO: implement this
-            throw new RuntimeException("not implemented");
+    private static void writeBlockType(WasmBlockType blockType, WasmBinaryWriter writer) throws IOException {
+        switch (blockType) {
+            case WasmBlockType.Empty empty -> {
+                writer.writeByte(0x40);
+            }
+            case WasmBlockType.ValueType valueType -> {
+                writeValueType(valueType.valueType(), writer);
+            }
+            case WasmBlockType.FuncType funcType -> {
+                writer.writeTypeIndex(funcType.funcType());
+            }
         }
     }