Support variable result count from if instruction

3 files changed, 8 insertions, 7 deletions

diff --git a/src/main/java/org/shedlang/compiler/stackir/InstructionSerializer.java b/src/main/java/org/shedlang/compiler/stackir/InstructionSerializer.java
index 1ae8f04d..cc4d2390 100644
--- a/src/main/java/org/shedlang/compiler/stackir/InstructionSerializer.java
+++ b/src/main/java/org/shedlang/compiler/stackir/InstructionSerializer.java
@@ -86,6 +86,7 @@ public class InstructionSerializer {
             case IfInstruction if_ ->
                 treeListIndented(
                     treeIdentifier("if"),
+                    treeNamed("result-count", treeInt(if_.resultCount())),
                     instructionsToTree(if_.trueBody()),
                     instructionsToTree(if_.falseBody())
                 );
diff --git a/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java b/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
index 570ce62a..1621a50c 100644
--- a/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
+++ b/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
@@ -290,6 +290,7 @@ public class StackCompiler {
             context = context.exitBlock();
 
             context = context.addInstruction(new IfInstruction(
+                1,
                 ifTrueInstructions,
                 ifFalseInstructions,
                 conditionalBranch.source()
@@ -418,18 +419,14 @@ public class StackCompiler {
         context = context.enterBlock(List.of());
 
         context = compileExpression(node.condition(), context);
-        // TODO: support no yielding from if instruction?
         context = context.addInstruction(new IfInstruction(
+            0,
+            List.of(),
             List.of(
-                new PushUnitInstruction(node.source())
-            ),
-            List.of(
-                new LoopBreakInstruction(node.source()),
-                new PushUnitInstruction(node.source())
+                new LoopBreakInstruction(node.source())
             ),
             node.source()
         ));
-        context = context.addInstruction(new DiscardInstruction(node.source()));
 
         context = context.addInstructions(ifTrueInstructions);
 
diff --git a/src/main/java/org/shedlang/compiler/stackir/instructions/IfInstruction.java b/src/main/java/org/shedlang/compiler/stackir/instructions/IfInstruction.java
index df09859f..9595d7f7 100644
--- a/src/main/java/org/shedlang/compiler/stackir/instructions/IfInstruction.java
+++ b/src/main/java/org/shedlang/compiler/stackir/instructions/IfInstruction.java
@@ -9,8 +9,11 @@ import java.util.List;
  * Pop a value from a stack, and assume it is a Bool. If it is true, execute the
  * instructions in the "true" body, otherwise execute the instructions in the
  * "false" body.
+ *
+ * Each branch body should push <code>resultCount</code> values onto the stack.
  */
 public record IfInstruction(
+    int resultCount,
     List<Instruction> trueBody,
     List<Instruction> falseBody,
     Source source