Else blocks for while statements and breaking loops

Author: Omnixerion

compiler/src/main/java/dev/ultreon/pythonc/JvmWriter.java

@@ -22,6 +22,19 @@ public Context getContext() {
         return context;
     }
 
+    public LoopContext getLoopContext() {
+        for (int i = pc.contextStack.size() - 1; i >= 0; i--) {
+            Context context = pc.contextStack.get(i);
+            if (context instanceof FunctionContext) {
+                throw new CompilerException("Not inside a loop or loop is outside of function.");
+            }
+            if (context instanceof LoopContext) {
+                return (LoopContext) context;
+            }
+        }
+        throw new RuntimeException("No loop context found");
+    }
+
     public JvmClass getClassSymbol(String className) {
         return pc.getClassSymbol(className);
     }

compiler/src/main/java/dev/ultreon/pythonc/LoopContext.java

@@ -0,0 +1,10 @@
+package dev.ultreon.pythonc;
+
+import org.objectweb.asm.Label;
+
+public interface LoopContext extends Context {
+
+    Label getContinuationLabel();
+
+    Label getBreakLabel();
+}

compiler/src/main/java/dev/ultreon/pythonc/PythonCompiler.java

@@ -62,7 +62,7 @@ public class PythonCompiler extends PythonParserBaseVisitor<Object> {
     private PyClass curPyClass = null;
     private final List<CompilerException> compileErrors = new ArrayList<>();
     private Label elifLabel;
-    private final Stack<Context> contextStack = new Stack<>();
+    final Stack<Context> contextStack = new Stack<>();
 
     public final JvmWriter writer = new JvmWriter(this);
 
@@ -318,19 +318,41 @@ public Object visitCompound_stmt(PythonParser.Compound_stmtContext ctx) {
             return visit(whileStmtContext);
         }
 
+        PythonParser.For_stmtContext forStmtContext = ctx.for_stmt();
+        if (forStmtContext != null) {
+            return visit(forStmtContext);
+        }
+
         throw new RuntimeException("No supported matching compound_stmt found for:\n" + ctx.getText());
     }
 
+    @Override
+    public Object visitFor_stmt(PythonParser.For_stmtContext ctx) {
+//        PythonParser.Named_expressionContext namedExpressionContext = ctx.named_expression();
+//        if (namedExpressionContext != null) {
+//            Object visit = visit(namedExpressionContext);
+//            loadExpr(ctx, visit);
+//        } else {
+//            throw new RuntimeException("No supported matching named_expression found for:\n" + ctx.getText());
+//        }
+
+        throw new RuntimeException("No supported matching for_stmt found for:\n" + ctx.getText());
+    }
+
     @Override
     public Object visitWhile_stmt(PythonParser.While_stmtContext ctx) {
         Label loopStart = new Label();
         Label loopEnd = new Label();
+        Label elseBlock = null;
+        if (ctx.else_block() != null) {
+            elseBlock = new Label();
+        }
 
         // Loop start:
         mv.visitLabel(loopStart);
 
         // region Comparison(named_expression)
-        pushContext(new WhileConditionContext(loopEnd));
+        pushContext(new WhileConditionContext(loopEnd, elseBlock));
         PythonParser.Named_expressionContext namedExpressionContext = ctx.named_expression();
         if (namedExpressionContext != null) {
             Object visit = visit(namedExpressionContext);
@@ -343,18 +365,35 @@ public Object visitWhile_stmt(PythonParser.While_stmtContext ctx) {
         // endregion Comparison
 
         // region Loop(block)
+        Context loopContext = new WhileLoopContext(loopStart, loopEnd);
+        pushContext(loopContext);
         PythonParser.BlockContext block = ctx.block();
         if (block != null) {
             visit(block);
         }
         if (context.needsPop()) {
             throw new RuntimeException("Still values on stack for:\n" + ctx.getText());
         }
+        popContext();
         // endregion Loop
 
         // Jump to loop start
         mv.visitJumpInsn(GOTO, loopStart);
 
+        if (elseBlock != null) {
+            mv.visitLabel(elseBlock);
+
+            // region Else(block)
+            PythonParser.BlockContext elseBlockContext = ctx.else_block().block();
+            if (elseBlockContext != null) {
+                visit(elseBlockContext);
+            }
+            if (context.needsPop()) {
+                throw new RuntimeException("Still values on stack for:\n" + ctx.getText());
+            }
+            // endregion Else
+        }
+
         // Loop end:
         mv.visitLabel(loopEnd);
 
@@ -980,6 +1019,21 @@ public Object visitSimple_stmt(PythonParser.Simple_stmtContext ctx) {
             writer.pop();
             return visit;
         }
+
+        TerminalNode aBreak = ctx.BREAK();
+        if (aBreak != null) {
+            Context context = writer.getLoopContext();
+            if (context instanceof WhileLoopContext whileConditionContext) {
+                mv.visitJumpInsn(GOTO, whileConditionContext.endLabel);
+                return Unit.Instance;
+//            } else if (writer.getContext() instanceof ForConditionContext) {
+//                ForConditionContext forConditionContext = (ForConditionContext) writer.getContext();
+//                mv.visitJumpInsn(GOTO, forConditionContext.loopEnd);
+//                return Unit.Instance;
+            } else {
+                throw new RuntimeException("Not in a loop:\n" + ctx.getText());
+            }
+        }
         throw new RuntimeException("No supported matching simple_stmt found of type " + ctx.getClass().getSimpleName() + " for:\n" + ctx.getText());
     }
 

compiler/src/main/java/dev/ultreon/pythonc/WhileConditionContext.java

@@ -5,9 +5,11 @@
 
 public class WhileConditionContext extends AbstractContext implements ConditionContext {
     public final Label loopEnd;
+    public final Label elseBlock;
 
-    public WhileConditionContext(Label loopEnd) {
+    public WhileConditionContext(Label loopEnd, Label elseBlock) {
         this.loopEnd = loopEnd;
+        this.elseBlock = elseBlock;
     }
 
     @Override
@@ -17,6 +19,6 @@ public WhileConditionContext(Label loopEnd) {
 
     @Override
     public @Nullable Label ifFalse() {
-        return loopEnd;
+        return elseBlock != null ? elseBlock : loopEnd;
     }
 }

compiler/src/main/java/dev/ultreon/pythonc/WhileLoopContext.java

@@ -0,0 +1,23 @@
+package dev.ultreon.pythonc;
+
+import org.objectweb.asm.Label;
+
+public class WhileLoopContext extends AbstractContext implements LoopContext {
+    public final Label startLabel;
+    public final Label endLabel;
+
+    public WhileLoopContext(Label startLabel, Label endLabel) {
+        this.startLabel = startLabel;
+        this.endLabel = endLabel;
+    }
+
+    @Override
+    public Label getContinuationLabel() {
+        return startLabel;
+    }
+
+    @Override
+    public Label getBreakLabel() {
+        return endLabel;
+    }
+}

src/main/python/example/hello.py

@@ -127,5 +127,18 @@ def init():
         print(i)
         print("Hello, world!")
         i = i + 1
+    else:
+        print("Done!")
+
+    i = 0
+    while i < 10:
+        print(i)
+        print("Hello, world!")
+        i = i + 1
+        if i == 5:
+            print("Break!")
+            break
+    else:
+        print("Done!")