Support while loops in LLVM Patio backend

2 files changed, 86 insertions, 13 deletions

diff --git a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGenLocalVariableSpace.java b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGenLocalVariableSpace.java
index ba5019c2..8e2967df 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGenLocalVariableSpace.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmCodeGenLocalVariableSpace.java
@@ -17,6 +17,8 @@ public enum LlvmCodeGenLocalVariableSpace {
     IF_TRUE_START("ifTrueStart"),
     IF_TRUE_END("ifTrueEnd"),
     LOCAL("local"),
+    LOOP_START("loopStart"),
+    LOOP_END("loopEnd"),
     NAMESPACE("namespace"),
     NAMESPACE_STRUCT_POINTER("namespaceStructPointer"),
     PARAM("param"),
diff --git a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmPatioCodeGen.java b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmPatioCodeGen.java
index a9919af8..ed704b85 100644
--- a/src/main/java/org/shedlang/compiler/backends/llvm/LlvmPatioCodeGen.java
+++ b/src/main/java/org/shedlang/compiler/backends/llvm/LlvmPatioCodeGen.java
@@ -1,6 +1,7 @@
 package org.shedlang.compiler.backends.llvm;
 
 import org.pcollections.PMap;
+import org.pcollections.PStack;
 import org.pcollections.PVector;
 import org.pcollections.TreePVector;
 import org.shedlang.compiler.backends.llvm.ir.ast.*;
@@ -424,14 +425,37 @@ public class LlvmPatioCodeGen {
                     patioOperandToLlvmOperand(patioLocalStore.value(), context)
                 );
 
-            case PatioLoop patioLoop ->
-                throw new UnsupportedOperationException("TODO");
+            case PatioLoop patioLoop -> {
+                var loopStartLabelResult = context.allocateLocalName(
+                    LlvmCodeGenLocalVariableSpace.LOOP_START
+                );
+                var loopStartLabel = loopStartLabelResult.value();
+                context = loopStartLabelResult.context();
+
+                var loopEndLabelResult = context.allocateLocalName(
+                    LlvmCodeGenLocalVariableSpace.LOOP_END
+                );
+                var loopEndLabel = loopEndLabelResult.value();
+                context = loopEndLabelResult.context();
+                
+                context = context.addTerminatorInstruction(new LlvmBrUnconditional(loopStartLabel.name()));
+
+                context = context.setBasicBlockLabel(loopStartLabel.name());
+
+                context = context.enterLoop(loopStartLabel, loopEndLabel);
+                context = generateInstructions(patioLoop.body(), context);
+                context = context.exitLoop();
+
+                context = context.setBasicBlockLabel(loopEndLabel.name());
+
+                yield context;
+            }
 
             case PatioLoopBreak patioLoopBreak ->
-                throw new UnsupportedOperationException("TODO");
+                context.addTerminatorInstruction(new LlvmBrUnconditional(context.loopEndLabel().name()));
 
             case PatioLoopContinue patioLoopContinue ->
-                throw new UnsupportedOperationException("TODO");
+                context.addTerminatorInstruction(new LlvmBrUnconditional(context.loopStartLabel().name()));
 
             case PatioMemoryLoad patioMemoryLoad -> {
                 var targetResult = context.addPatioLocalToLlvmLocal(patioMemoryLoad.target());
@@ -862,7 +886,8 @@ public class LlvmPatioCodeGen {
                 P.vector(),
                 P.map(),
                 P.map(),
-                NameAllocator.empty()
+                NameAllocator.empty(),
+                P.stack()
             );
         }
 
@@ -873,6 +898,7 @@ public class LlvmPatioCodeGen {
         private final PMap<PatioLocalRef, LlvmType> patioLocalToType;
         private final PMap<PatioLocalRef, LlvmOperand> patioLocalToOperand;
         private final NameAllocator<LlvmCodeGenLocalVariableSpace> localNameAllocator;
+        private final PStack<LoopLabels> loopLabels;
 
         private FunctionContext(
             PMap<String, LlvmType> functionNameToReturnType,
@@ -881,7 +907,8 @@ public class LlvmPatioCodeGen {
             PVector<LlvmInstruction> instructions,
             PMap<PatioLocalRef, LlvmType> patioLocalToType,
             PMap<PatioLocalRef, LlvmOperand> patioLocalToOperand,
-            NameAllocator<LlvmCodeGenLocalVariableSpace> localNameAllocator
+            NameAllocator<LlvmCodeGenLocalVariableSpace> localNameAllocator,
+            PStack<LoopLabels> loopLabels
         ) {
             this.functionNameToReturnType = functionNameToReturnType;
             this.basicBlocks = basicBlocks;
@@ -890,6 +917,7 @@ public class LlvmPatioCodeGen {
             this.patioLocalToType = patioLocalToType;
             this.patioLocalToOperand = patioLocalToOperand;
             this.localNameAllocator = localNameAllocator;
+            this.loopLabels = loopLabels;
         }
 
         private LlvmType functionNameToReturnType(String functionName) {
@@ -908,7 +936,8 @@ public class LlvmPatioCodeGen {
                 P.vector(),
                 this.patioLocalToType,
                 this.patioLocalToOperand,
-                this.localNameAllocator
+                this.localNameAllocator,
+                this.loopLabels
             );
         }
 
@@ -921,7 +950,8 @@ public class LlvmPatioCodeGen {
                     this.instructions,
                     this.patioLocalToType,
                     this.patioLocalToOperand,
-                    this.localNameAllocator
+                    this.localNameAllocator,
+                    this.loopLabels
                 );
             } else {
                 throw new RuntimeException("basic block label is already set");
@@ -953,7 +983,8 @@ public class LlvmPatioCodeGen {
                 this.instructions,
                 this.patioLocalToType.plus(patioLocal, llvmType),
                 this.patioLocalToOperand,
-                this.localNameAllocator
+                this.localNameAllocator,
+                this.loopLabels
             );
         }
 
@@ -965,7 +996,8 @@ public class LlvmPatioCodeGen {
                 this.instructions.plus(instruction),
                 this.patioLocalToType,
                 this.patioLocalToOperand,
-                this.localNameAllocator
+                this.localNameAllocator,
+                this.loopLabels
             );
         }
 
@@ -979,7 +1011,8 @@ public class LlvmPatioCodeGen {
                 this.instructions,
                 this.patioLocalToType,
                 this.patioLocalToOperand,
-                result.allocator()
+                result.allocator(),
+                this.loopLabels
             ).toResult(new LlvmVariableLocal(result.name()));
         }
 
@@ -993,7 +1026,8 @@ public class LlvmPatioCodeGen {
                 this.instructions,
                 this.patioLocalToType,
                 this.patioLocalToOperand,
-                result.allocator()
+                result.allocator(),
+                this.loopLabels
             ).toResult(new LlvmVariableLocal(result.name()));
         }
 
@@ -1003,6 +1037,40 @@ public class LlvmPatioCodeGen {
             return result.context.addPatioLocal(patioLocal, result.value).toResult(result.value);
         }
 
+        public FunctionContext enterLoop(LlvmVariableLocal startLabel, LlvmVariableLocal endLabel) {
+            return new FunctionContext(
+                this.functionNameToReturnType,
+                this.basicBlocks,
+                this.basicBlockLabel,
+                this.instructions,
+                this.patioLocalToType,
+                this.patioLocalToOperand,
+                this.localNameAllocator,
+                this.loopLabels.plus(new LoopLabels(startLabel, endLabel))
+            );
+        }
+
+        public FunctionContext exitLoop() {
+            return new FunctionContext(
+                this.functionNameToReturnType,
+                this.basicBlocks,
+                this.basicBlockLabel,
+                this.instructions,
+                this.patioLocalToType,
+                this.patioLocalToOperand,
+                this.localNameAllocator,
+                this.loopLabels.minus(0)
+            );
+        }
+
+        public LlvmVariableLocal loopStartLabel() {
+            return this.loopLabels.get(0).start;
+        }
+
+        public LlvmVariableLocal loopEndLabel() {
+            return this.loopLabels.get(0).end;
+        }
+
         public FunctionContext addPatioLocal(PatioLocalRef patioLocal, LlvmOperand llvmOperand) {
             return new FunctionContext(
                 this.functionNameToReturnType,
@@ -1011,7 +1079,8 @@ public class LlvmPatioCodeGen {
                 this.instructions,
                 this.patioLocalToType,
                 this.patioLocalToOperand.plus(patioLocal, llvmOperand),
-                this.localNameAllocator
+                this.localNameAllocator,
+                this.loopLabels
             );
         }
 
@@ -1022,4 +1091,6 @@ public class LlvmPatioCodeGen {
         public record Result<T>(T value, FunctionContext context) {
         }
     }
+
+    private record LoopLabels(LlvmVariableLocal start, LlvmVariableLocal end) {}
 }