Generate stack IR for when expressions

3 files changed, 96 insertions, 0 deletions

diff --git a/src/test/java/org/shedlang/compiler/ast/typed/Typed.java b/src/test/java/org/shedlang/compiler/ast/typed/Typed.java
index 06c481ba..7c5aa0fc 100644
--- a/src/test/java/org/shedlang/compiler/ast/typed/Typed.java
+++ b/src/test/java/org/shedlang/compiler/ast/typed/Typed.java
@@ -252,6 +252,21 @@ public class Typed {
         return new TypedAssignNode(targetVar(targetName), expression, Sources.NULL);
     }
 
+    public static TypedWhenNode when(
+        TypedExpressionNode expression,
+        List<TypedWhenCaseNode> cases,
+        Type type
+    ) {
+        return new TypedWhenNode(expression, cases, type, Sources.NULL);
+    }
+
+    public static TypedWhenCaseNode whenCase(
+        TypedTypeLevelExpressionNode<StructType> caseType,
+        TypedBlockNode body
+    ) {
+        return new TypedWhenCaseNode(caseType, body, Sources.NULL);
+    }
+
     public static TypedWhileNode while_(
         TypedReferenceNode condition,
         TypedBlockNode body
diff --git a/src/test/java/org/shedlang/compiler/stackir/StackCompilerWhenTests.java b/src/test/java/org/shedlang/compiler/stackir/StackCompilerWhenTests.java
new file mode 100644
index 00000000..0f85d48c
--- /dev/null
+++ b/src/test/java/org/shedlang/compiler/stackir/StackCompilerWhenTests.java
@@ -0,0 +1,62 @@
+package org.shedlang.compiler.stackir;
+
+import org.junit.jupiter.api.Test;
+import org.shedlang.compiler.ast.typed.Typed;
+import org.shedlang.compiler.types.NamespaceName;
+import org.shedlang.compiler.types.Types;
+
+import java.util.List;
+
+import static org.shedlang.compiler.stackir.instructions.InstructionMatchers.*;
+import static org.zwobble.precisely.AssertThat.assertThat;
+import static org.zwobble.precisely.Matchers.isSequence;
+
+public class StackCompilerWhenTests {
+    private final static NamespaceName NAMESPACE_NAME = NamespaceName.of("Example", "Project");
+
+    @Test
+    public void whenIsCompiledToSwitchInstruction() {
+        var node = Typed.when(
+            Typed.reference("x"),
+            List.of(
+                Typed.whenCase(
+                    Typed.typeLevelReference("A", Types.struct(NAMESPACE_NAME, "A")),
+                    Typed.block(List.of(
+                        Typed.blockReturn(Typed.stringLiteral("a"))
+                    ))
+                ),
+                Typed.whenCase(
+                    Typed.typeLevelReference("B", Types.struct(NAMESPACE_NAME, "B")),
+                    Typed.block(List.of(
+                        Typed.blockReturn(Typed.stringLiteral("b"))
+                    ))
+                )
+            ),
+            Types.STRING
+        );
+        var context = StackCompiler.Context.EMPTY
+            .addShedVariable("x");
+
+        var result = StackCompiler.compileExpression(node, context);
+
+        assertThat(result.instructions(), isSequence(
+            isLocalLoadInstruction("x__0"),
+            isSwitchInstruction(
+                isSequence(
+                    isSwitchInstructionCase(
+                        Types.struct(NAMESPACE_NAME, "A"),
+                        isSequence(
+                            isPushStringInstruction("a")
+                        )
+                    ),
+                    isSwitchInstructionCase(
+                        Types.struct(NAMESPACE_NAME, "B"),
+                        isSequence(
+                            isPushStringInstruction("b")
+                        )
+                    )
+                )
+            )
+        ));
+    }
+}
diff --git a/src/test/java/org/shedlang/compiler/stackir/instructions/InstructionMatchers.java b/src/test/java/org/shedlang/compiler/stackir/instructions/InstructionMatchers.java
index 3d1689e1..10a9cb95 100644
--- a/src/test/java/org/shedlang/compiler/stackir/instructions/InstructionMatchers.java
+++ b/src/test/java/org/shedlang/compiler/stackir/instructions/InstructionMatchers.java
@@ -162,6 +162,25 @@ public class InstructionMatchers {
         return instanceOf(StringAddInstruction.class);
     }
 
+    public static Matcher<Instruction> isSwitchInstruction(
+        Matcher<Iterable<? extends SwitchInstruction.Case>> cases
+    ) {
+        return instanceOf(
+            SwitchInstruction.class,
+            has("cases", x -> x.cases(), cases)
+        );
+    }
+
+    public static Matcher<SwitchInstruction.Case> isSwitchInstructionCase(
+        StructType caseType,
+        Matcher<Iterable<? extends Instruction>> body
+    ) {
+        return allOf(
+            has("caseType", x -> x.caseType(), equalTo(caseType)),
+            has("body", x -> x.body(), body)
+        );
+    }
+
     public static Matcher<Instruction> isTupleCreateInstruction(int elementCount) {
         return instanceOf(
             TupleCreateInstruction.class,