Fix generation of if instruction for while loop

1 files changed, 9 insertions, 2 deletions

diff --git a/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java b/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
index f4580e69..570ce62a 100644
--- a/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
+++ b/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
@@ -418,11 +418,18 @@ 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(
-            List.of(),
-            List.of(new LoopBreakInstruction(node.source())),
+            List.of(
+                new PushUnitInstruction(node.source())
+            ),
+            List.of(
+                new LoopBreakInstruction(node.source()),
+                new PushUnitInstruction(node.source())
+            ),
             node.source()
         ));
+        context = context.addInstruction(new DiscardInstruction(node.source()));
 
         context = context.addInstructions(ifTrueInstructions);