Support variable result count from if instruction

6 files changed, 58 insertions, 14 deletions

diff --git a/src/test/java/org/shedlang/compiler/backends/llvm/LlvmCodeGenIfTests.java b/src/test/java/org/shedlang/compiler/backends/llvm/LlvmCodeGenIfTests.java
index 3a04ce3f..0243ae90 100644
--- a/src/test/java/org/shedlang/compiler/backends/llvm/LlvmCodeGenIfTests.java
+++ b/src/test/java/org/shedlang/compiler/backends/llvm/LlvmCodeGenIfTests.java
@@ -15,6 +15,7 @@ public class LlvmCodeGenIfTests {
     @Test
     public void ifIsCompiled(Snapshotter snapshotter) throws IOException {
         var instruction = Instructions.if_(
+            1,
             List.of(
                 Instructions.pushInt(1),
                 Instructions.pushInt(2),
@@ -38,6 +39,7 @@ public class LlvmCodeGenIfTests {
     @Test
     public void canCompileNestedIf(Snapshotter snapshotter) throws IOException {
         var instruction = Instructions.if_(
+            1,
             List.of(
                 Instructions.pushInt(1),
                 Instructions.pushInt(2),
@@ -46,6 +48,7 @@ public class LlvmCodeGenIfTests {
             List.of(
                 Instructions.pushBool(true),
                 Instructions.if_(
+                    1,
                     List.of(
                         Instructions.pushInt(101),
                         Instructions.pushInt(102),
diff --git a/src/test/java/org/shedlang/compiler/backends/wasm/WasmCodeGenIfTests.java b/src/test/java/org/shedlang/compiler/backends/wasm/WasmCodeGenIfTests.java
index 37cf28af..76eb4a1a 100644
--- a/src/test/java/org/shedlang/compiler/backends/wasm/WasmCodeGenIfTests.java
+++ b/src/test/java/org/shedlang/compiler/backends/wasm/WasmCodeGenIfTests.java
@@ -16,6 +16,7 @@ public class WasmCodeGenIfTests {
     @Test
     public void ifIsCompiledToWasmIf(Snapshotter snapshotter) throws IOException, InterruptedException {
         var instruction = Instructions.if_(
+            1,
             List.of(
                 Instructions.pushInt(1),
                 Instructions.pushInt(2),
diff --git a/src/test/java/org/shedlang/compiler/patio/PatioCodeGenIfTests.java b/src/test/java/org/shedlang/compiler/patio/PatioCodeGenIfTests.java
index f8fea8da..c06f0c5e 100644
--- a/src/test/java/org/shedlang/compiler/patio/PatioCodeGenIfTests.java
+++ b/src/test/java/org/shedlang/compiler/patio/PatioCodeGenIfTests.java
@@ -15,8 +15,31 @@ import static org.shedlang.compiler.patio.PatioTesting.serialize;
 @ExtendWith({SnapshotResolver.class})
 public class PatioCodeGenIfTests {
     @Test
-    public void canGenerateCodeForIf(Snapshotter snapshotter) throws IOException {
+    public void canGenerateCodeForIfWithNoResults(Snapshotter snapshotter) throws IOException {
         var instruction = Instructions.if_(
+            0,
+            List.of(
+                Instructions.pushInt(1),
+                Instructions.discard()
+            ),
+            List.of(
+                Instructions.pushInt(101),
+                Instructions.discard()
+            )
+        );
+        var context = PatioCodeGenContext.EMPTY
+            .temporaryPush(new PatioLocalRef("other"))
+            .temporaryPush(new PatioLocalRef("condition"));
+
+        var result = PatioCodeGen.generateInstruction(instruction, context);
+
+        snapshotter.assertSnapshot(serialize(result));
+    }
+
+    @Test
+    public void canGenerateCodeForIfWithSingleResult(Snapshotter snapshotter) throws IOException {
+        var instruction = Instructions.if_(
+            1,
             List.of(
                 Instructions.pushInt(1),
                 Instructions.pushInt(2),
@@ -36,4 +59,26 @@ public class PatioCodeGenIfTests {
 
         snapshotter.assertSnapshot(serialize(result));
     }
+
+    @Test
+    public void canGenerateCodeForIfWithMultipleResults(Snapshotter snapshotter) throws IOException {
+        var instruction = Instructions.if_(
+            2,
+            List.of(
+                Instructions.pushInt(1),
+                Instructions.pushInt(2)
+            ),
+            List.of(
+                Instructions.pushInt(101),
+                Instructions.pushInt(102)
+            )
+        );
+        var context = PatioCodeGenContext.EMPTY
+            .temporaryPush(new PatioLocalRef("other"))
+            .temporaryPush(new PatioLocalRef("condition"));
+
+        var result = PatioCodeGen.generateInstruction(instruction, context);
+
+        snapshotter.assertSnapshot(serialize(result));
+    }
 }
diff --git a/src/test/java/org/shedlang/compiler/stackir/InstructionSerializerTests.java b/src/test/java/org/shedlang/compiler/stackir/InstructionSerializerTests.java
index 407621fb..b8a45a0b 100644
--- a/src/test/java/org/shedlang/compiler/stackir/InstructionSerializerTests.java
+++ b/src/test/java/org/shedlang/compiler/stackir/InstructionSerializerTests.java
@@ -78,6 +78,7 @@ public class InstructionSerializerTests {
     @Test
     public void canSerializeIf() {
         var instruction = Instructions.if_(
+            1,
             List.of(
                 Instructions.pushUnit()
             ),
@@ -93,6 +94,7 @@ public class InstructionSerializerTests {
         assertThat(string, equalTo("""
             (
                 if
+                result-count=1
                 (
                     (push-unit)
                 )
diff --git a/src/test/java/org/shedlang/compiler/stackir/StackCompilerWhileTests.java b/src/test/java/org/shedlang/compiler/stackir/StackCompilerWhileTests.java
index 398a9827..c8fd1cdc 100644
--- a/src/test/java/org/shedlang/compiler/stackir/StackCompilerWhileTests.java
+++ b/src/test/java/org/shedlang/compiler/stackir/StackCompilerWhileTests.java
@@ -28,15 +28,11 @@ public class StackCompilerWhileTests {
                 isSequence(
                     isLocalLoadInstruction("c__0"),
                     isIfInstruction(
+                        isSequence(),
                         isSequence(
-                            isPushUnitInstruction()
-                        ),
-                        isSequence(
-                            isLoopBreakInstruction(),
-                            isPushUnitInstruction()
+                            isLoopBreakInstruction()
                         )
                     ),
-                    isDiscardInstruction(),
                     isPushIntInstruction(42),
                     isDiscardInstruction(),
                     isLoopContinueInstruction()
@@ -65,15 +61,11 @@ public class StackCompilerWhileTests {
                 isSequence(
                     isLocalLoadInstruction("c__0"),
                     isIfInstruction(
+                        isSequence(),
                         isSequence(
-                            isPushUnitInstruction()
-                        ),
-                        isSequence(
-                            isLoopBreakInstruction(),
-                            isPushUnitInstruction()
+                            isLoopBreakInstruction()
                         )
                     ),
-                    isDiscardInstruction(),
                     isPushIntInstruction(42),
                     isDiscardInstruction(),
                     isLoopContinueInstruction()
diff --git a/src/test/java/org/shedlang/compiler/stackir/instructions/Instructions.java b/src/test/java/org/shedlang/compiler/stackir/instructions/Instructions.java
index 4bc3dc50..ee48a1bd 100644
--- a/src/test/java/org/shedlang/compiler/stackir/instructions/Instructions.java
+++ b/src/test/java/org/shedlang/compiler/stackir/instructions/Instructions.java
@@ -23,10 +23,11 @@ public class Instructions {
     }
 
     public static Instruction if_(
+        int resultCount,
         List<Instruction> trueBody,
         List<Instruction> falseBody
     ) {
-        return new IfInstruction(trueBody, falseBody, Sources.NULL);
+        return new IfInstruction(resultCount, trueBody, falseBody, Sources.NULL);
     }
 
     public static Instruction intAdd() {