Support variable result count from if instruction

6 files changed, 44 insertions, 22 deletions

diff --git a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGen.java b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGen.java
index ae361d25..36898220 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGen.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGen.java
@@ -141,6 +141,10 @@ public class LlvmCodeGen {
             }
 
             case IfInstruction if_ -> {
+                if (if_.resultCount() != 1) {
+                    throw new UnsupportedOperationException("TODO");
+                }
+
                 var conditionResult = context.pop();
                 var condition = conditionResult.value();
                 context = conditionResult.context();
diff --git a/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGen.java b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGen.java
index 84c2052a..431ebc1d 100644
--- a/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGen.java
+++ b/src/main/java/org/shedlang/compiler/backends/wasm/WasmCodeGen.java
@@ -9,6 +9,7 @@ import org.shedlang.compiler.stackir.StackIrNamespace;
 import org.shedlang.compiler.stackir.instructions.*;
 import org.shedlang.compiler.typechecker.TypeInfo;
 import org.shedlang.compiler.types.NamespaceName;
+import org.shedlang.compiler.util.Lists;
 
 import java.util.ArrayList;
 import java.util.Collections;
@@ -193,7 +194,7 @@ public class WasmCodeGen {
                 context = falseBodyResult.context();
 
                 yield context.addInstruction(new WasmIf(
-                    new WasmBlockType(List.of(WasmCodeGenTypes.VALUE_TYPE)),
+                    new WasmBlockType(Lists.repeat(WasmCodeGenTypes.VALUE_TYPE, if_.resultCount())),
                     trueBody,
                     falseBody
                 ));
diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGen.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGen.java
index 948bfd6a..42a606de 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGen.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGen.java
@@ -5,6 +5,7 @@ import org.shedlang.compiler.patio.ast.PatioVisibility;
 import org.shedlang.compiler.patio.ast.instructions.*;
 import org.shedlang.compiler.patio.ast.operands.PatioInt;
 import org.shedlang.compiler.patio.ast.operands.PatioLocalRef;
+import org.shedlang.compiler.patio.ast.operands.PatioOperand;
 import org.shedlang.compiler.patio.ast.types.PatioTypes;
 import org.shedlang.compiler.sources.Source;
 import org.shedlang.compiler.sources.Sources;
@@ -256,30 +257,38 @@ public class PatioCodeGen {
                 var conditionBool = conditionBoolResult.value();
                 context = conditionBoolResult.context();
 
-                var valueResult = context.allocateVal(
-                    PatioCodeGenLocalSpace.TEMPORARY,
-                    PatioCodeGenTypes.VALUE_TYPE
-                );
-                var value = valueResult.value();
-                context = valueResult.context();
+                var targets = new ArrayList<PatioLocalRef>();
+
+                for (var i = 0; i < if_.resultCount(); i++) {
+                    var targetResult = context.allocateVal(
+                        PatioCodeGenLocalSpace.TEMPORARY,
+                        PatioCodeGenTypes.VALUE_TYPE
+                    );
+                    var target = targetResult.value();
+                    context = targetResult.context();
 
-                var trueBodyResult = generateBlock(if_.trueBody(), if_.source(), context);
+                    targets.add(target);
+                }
+
+                var trueBodyResult = generateBlock(if_.resultCount(), if_.trueBody(), if_.source(), context);
                 var trueBody = trueBodyResult.value();
                 context = trueBodyResult.context();
 
-                var falseBodyResult = generateBlock(if_.falseBody(), if_.source(), context);
+                var falseBodyResult = generateBlock(if_.resultCount(), if_.falseBody(), if_.source(), context);
                 var falseBody = falseBodyResult.value();
                 context = falseBodyResult.context();
 
                 context = context.addInstruction(new PatioIf(
-                    List.of(value),
+                    targets,
                     conditionBool,
                     trueBody,
                     falseBody,
                     if_.source()
                 ));
 
-                context = context.temporaryPush(value);
+                for (var target : targets) {
+                    context = context.temporaryPush(target);
+                }
 
                 yield context;
             }
@@ -566,6 +575,7 @@ public class PatioCodeGen {
     }
 
     private static PatioCodeGenResult<PatioBlock> generateBlock(
+        int resultCount,
         List<Instruction> stackIrInstructions,
         Source source,
         PatioCodeGenContext context
@@ -573,9 +583,15 @@ public class PatioCodeGen {
         context = context.enterBlock();
         context = generateInstructions(stackIrInstructions, context);
 
-        var targetResult = context.temporaryPop();
-        var target = targetResult.value();
-        context = targetResult.context();
+        var targets = new ArrayList<PatioOperand>();
+        for (var i = 0; i < resultCount; i++) {
+            var targetResult = context.temporaryPop();
+            var target = targetResult.value();
+            context = targetResult.context();
+
+            targets.add(target);
+        }
+        Collections.reverse(targets);
 
         var instructionsResult = context.exitBlock();
         var instructions = instructionsResult.value();
@@ -583,7 +599,7 @@ public class PatioCodeGen {
 
         return context.toResult(new PatioBlock(
             instructions,
-            List.of(target),
+            targets,
             source
         ));
     }
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