Generate stack IR for when expressions

2 files changed, 31 insertions, 1 deletions

diff --git a/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java b/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
index e4f3d23f..5746fa34 100644
--- a/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
+++ b/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
@@ -8,6 +8,7 @@ import org.shedlang.compiler.types.StructType;
 import org.shedlang.compiler.util.Iterables;
 import org.shedlang.compiler.util.P;
 
+import java.util.ArrayList;
 import java.util.List;
 
 public class StackCompiler {
@@ -142,7 +143,7 @@ public class StackCompiler {
                 compileTuple(typedTupleNode, context);
 
             case TypedWhenNode typedWhenNode ->
-                throw new UnsupportedOperationException("TODO");
+                compileWhen(typedWhenNode, context);
         };
     }
 
@@ -527,6 +528,33 @@ public class StackCompiler {
         return context.addInstruction(instruction);
     }
 
+    private static Context compileWhen(TypedWhenNode node, Context context) {
+        context = compileExpression(node.expression(), context);
+
+        var switchCases = new ArrayList<SwitchInstruction.Case>();
+
+        for (var whenCase : node.cases()) {
+            context = context.enterBlock(whenCase.body().boundVariableNames());
+            context = compileBlockWithYield(whenCase.body(), context);
+            var caseBodyInstructions = context.instructions;
+            context = context.exitBlock();
+
+            var switchCase = new SwitchInstruction.Case(
+                whenCase.caseType().value(),
+                caseBodyInstructions,
+                whenCase.source()
+            );
+            switchCases.add(switchCase);
+        }
+
+        context = context.addInstruction(new SwitchInstruction(
+            switchCases,
+            node.source()
+        ));
+
+        return context;
+    }
+
     private static Context compileWhile(TypedWhileNode node, Context context) {
         context = context.enterBlock(node.body().boundVariableNames());
         context = compileBlockWithoutYield(node.body(), context);
diff --git a/src/main/java/org/shedlang/compiler/stackir/instructions/SwitchInstruction.java b/src/main/java/org/shedlang/compiler/stackir/instructions/SwitchInstruction.java
index 85f292e8..d978df98 100644
--- a/src/main/java/org/shedlang/compiler/stackir/instructions/SwitchInstruction.java
+++ b/src/main/java/org/shedlang/compiler/stackir/instructions/SwitchInstruction.java
@@ -9,6 +9,8 @@ import java.util.List;
  * Pop a value from the stack, find the case corresponding to that value's type,
  * and execute the instructions in the case's body. If there is no such case,
  * the behaviour of this instruction is not defined.
+ * <p>
+ * Each case body should push exactly one value onto the stack.
  */
 public record SwitchInstruction(
     List<Case> cases,