Compile while loop into Stack IR

1 files changed, 60 insertions, 5 deletions

diff --git a/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java b/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
index cd38d602..f4580e69 100644
--- a/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
+++ b/src/main/java/org/shedlang/compiler/stackir/StackCompiler.java
@@ -44,7 +44,7 @@ public class StackCompiler {
         return context;
     }
 
-    static Context compileBlock(TypedBlockNode node, Context context) {
+    static Context compileBlockWithYield(TypedBlockNode node, Context context) {
         for (var statement : node.statements()) {
             context = compileFunctionStatement(statement, context);
         }
@@ -62,6 +62,24 @@ public class StackCompiler {
         return context;
     }
 
+    static Context compileBlockWithoutYield(TypedBlockNode node, Context context) {
+        for (var statement : node.statements()) {
+            context = compileFunctionStatement(statement, context);
+        }
+
+        switch (node.behaviour()) {
+            case TypedFunctionStatementBehaviour.BlockReturn blockReturn -> {
+                context = context.addInstruction(new DiscardInstruction(node.source()));
+            }
+            case TypedFunctionStatementBehaviour.Continue continue_ -> {
+            }
+            case TypedFunctionStatementBehaviour.Terminate terminate -> {
+            }
+        }
+
+        return context;
+    }
+
     private static Context compileBlockReturn(TypedBlockReturnNode blockReturn, Context context) {
         return compileExpression(blockReturn.expression(), context);
     }
@@ -150,7 +168,7 @@ public class StackCompiler {
             ))
             .toList();
 
-        context = compileBlock(node.body(), context);
+        context = compileBlockWithYield(node.body(), context);
 
         switch (node.body().behaviour()) {
             case TypedFunctionStatementBehaviour.BlockReturn blockReturn -> {
@@ -228,7 +246,7 @@ public class StackCompiler {
                 compileVar(var, context);
 
             case TypedWhileNode while_ ->
-                throw new UnsupportedOperationException("TODO");
+                compileWhile(while_, context);
         };
     }
 
@@ -247,7 +265,7 @@ public class StackCompiler {
     ) {
         if (conditionalBranches.isEmpty()) {
             context = context.enterBlock(elseBody.boundVariableNames());
-            context = compileBlock(elseBody, context);
+            context = compileBlockWithYield(elseBody, context);
             var elseInstructions = context.instructions;
             context = context.exitBlock();
             for (var instruction : elseInstructions) {
@@ -258,7 +276,7 @@ public class StackCompiler {
             context = compileExpression(conditionalBranch.condition(), context);
 
             context = context.enterBlock(conditionalBranch.body().boundVariableNames());
-            context = compileBlock(conditionalBranch.body(), context);
+            context = compileBlockWithYield(conditionalBranch.body(), context);
             var ifTrueInstructions = context.instructions;
             context = context.exitBlock();
 
@@ -391,6 +409,33 @@ public class StackCompiler {
         return context;
     }
 
+    private static Context compileWhile(TypedWhileNode node, Context context) {
+        context = context.enterBlock(node.body().boundVariableNames());
+        context = compileBlockWithoutYield(node.body(), context);
+        var ifTrueInstructions = context.instructions;
+        context = context.exitBlock();
+
+        context = context.enterBlock(List.of());
+
+        context = compileExpression(node.condition(), context);
+        context = context.addInstruction(new IfInstruction(
+            List.of(),
+            List.of(new LoopBreakInstruction(node.source())),
+            node.source()
+        ));
+
+        context = context.addInstructions(ifTrueInstructions);
+
+        context = context.addInstruction(new LoopContinueInstruction(node.source()));
+
+        var bodyInstructions = context.instructions;
+        context = context.exitBlock();
+
+        context = context.addInstruction(new LoopInstruction(bodyInstructions, node.source()));
+
+        return context;
+    }
+
     public static class Context {
         public static final Context EMPTY = new Context(
             P.stack(),
@@ -456,6 +501,16 @@ public class StackCompiler {
             );
         }
 
+        public Context addInstructions(PVector<Instruction> instructions) {
+            return new Context(
+                this.instructionsStack,
+                this.instructions.plusAll(instructions),
+                this.localNames,
+                this.shedVariablesByNameStack,
+                this.shedVariablesByName
+            );
+        }
+
         public String shedVariableToLocal(String variableName) {
             return shedVariable(variableName).stackIrLocalName;
         }