Compile while loop into Stack IR

3 files changed, 100 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 c73e77a0..2fb3c66c 100644
--- a/src/test/java/org/shedlang/compiler/ast/typed/Typed.java
+++ b/src/test/java/org/shedlang/compiler/ast/typed/Typed.java
@@ -209,4 +209,11 @@ public class Typed {
     ) {
         return new TypedVarNode(target, expression, Sources.NULL);
     }
+
+    public static TypedWhileNode while_(
+        TypedReferenceNode condition,
+        TypedBlockNode body
+    ) {
+        return new TypedWhileNode(condition, body, Sources.NULL);
+    }
 }
diff --git a/src/test/java/org/shedlang/compiler/stackir/StackCompilerWhileTests.java b/src/test/java/org/shedlang/compiler/stackir/StackCompilerWhileTests.java
new file mode 100644
index 00000000..c8fd1cdc
--- /dev/null
+++ b/src/test/java/org/shedlang/compiler/stackir/StackCompilerWhileTests.java
@@ -0,0 +1,76 @@
+package org.shedlang.compiler.stackir;
+
+import org.junit.jupiter.api.Test;
+import org.shedlang.compiler.ast.typed.Typed;
+
+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 StackCompilerWhileTests {
+    @Test
+    public void whileIsCompiledToLoop() {
+        var node = Typed.while_(
+            Typed.reference("c"),
+            Typed.block(List.of(
+                Typed.expressionStatement(Typed.intLiteral(42))
+            ))
+        );
+        var context = StackCompiler.Context.EMPTY
+            .addShedVariable("c");
+
+        var result = StackCompiler.compileFunctionStatement(node, context);
+
+        assertThat(result.instructions(), isSequence(
+            isLoopInstruction(
+                isSequence(
+                    isLocalLoadInstruction("c__0"),
+                    isIfInstruction(
+                        isSequence(),
+                        isSequence(
+                            isLoopBreakInstruction()
+                        )
+                    ),
+                    isPushIntInstruction(42),
+                    isDiscardInstruction(),
+                    isLoopContinueInstruction()
+                )
+            )
+        ));
+    }
+
+    @Test
+    public void whenBodyYieldsValueThenValueIsDiscarded() {
+        // In principle, the type checker should prevent the body from yielding
+        // a value, but we check we handle this case anyway.
+        var node = Typed.while_(
+            Typed.reference("c"),
+            Typed.block(List.of(
+                Typed.blockReturn(Typed.intLiteral(42))
+            ))
+        );
+        var context = StackCompiler.Context.EMPTY
+            .addShedVariable("c");
+
+        var result = StackCompiler.compileFunctionStatement(node, context);
+
+        assertThat(result.instructions(), isSequence(
+            isLoopInstruction(
+                isSequence(
+                    isLocalLoadInstruction("c__0"),
+                    isIfInstruction(
+                        isSequence(),
+                        isSequence(
+                            isLoopBreakInstruction()
+                        )
+                    ),
+                    isPushIntInstruction(42),
+                    isDiscardInstruction(),
+                    isLoopContinueInstruction()
+                )
+            )
+        ));
+    }
+}
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 a95680d0..b70ff8dd 100644
--- a/src/test/java/org/shedlang/compiler/stackir/instructions/InstructionMatchers.java
+++ b/src/test/java/org/shedlang/compiler/stackir/instructions/InstructionMatchers.java
@@ -74,6 +74,23 @@ public class InstructionMatchers {
         );
     }
 
+    public static Matcher<Instruction> isLoopInstruction(
+        Matcher<Iterable<? extends Instruction>> body
+    ) {
+        return instanceOf(
+            LoopInstruction.class,
+            has("body", x -> x.body(), body)
+        );
+    }
+
+    public static Matcher<Instruction> isLoopBreakInstruction() {
+        return instanceOf(LoopBreakInstruction.class);
+    }
+
+    public static Matcher<Instruction> isLoopContinueInstruction() {
+        return instanceOf(LoopContinueInstruction.class);
+    }
+
     public static Matcher<Instruction> isNamespaceFieldStore(NamespaceName namespaceName, String fieldName) {
         return instanceOf(
             NamespaceFieldStoreInstruction.class,