support `for..else` and `while..else`

include/pocketpy/codeobject.h

@@ -42,9 +42,15 @@ struct CodeBlock {
     int for_loop_depth; // this is used for exception handling
     int start;          // start index of this block in codes, inclusive
     int end;            // end index of this block in codes, exclusive
+    int end2;           // ...
 
     CodeBlock(CodeBlockType type, int parent, int for_loop_depth, int start):
-        type(type), parent(parent), for_loop_depth(for_loop_depth), start(start), end(-1) {}
+        type(type), parent(parent), for_loop_depth(for_loop_depth), start(start), end(-1), end2(-1) {}
+
+    int get_break_end() const{
+        if(end2 != -1) return end2;
+        return end;
+    }
 };
 
 struct CodeObject;

include/pocketpy/expr.h

@@ -53,8 +53,8 @@ struct CodeEmitContext{
     bool is_compiling_class = false;
     int for_loop_depth = 0;
 
-    bool is_curr_block_loop() const;
-    void enter_block(CodeBlockType type);
+    int get_loop() const;
+    CodeBlock* enter_block(CodeBlockType type);
     void exit_block();
     void emit_expr();   // clear the expression stack and generate bytecode
     std::string _log_s_expr();

src/ceval.cpp

@@ -482,10 +482,10 @@ __NEXT_STEP:;
         } else POP();                       // [b]
         DISPATCH();
     TARGET(LOOP_CONTINUE)
-        frame->jump_abs(co_blocks[byte.block].start);
+        frame->jump_abs(co_blocks[byte.arg].start);
         DISPATCH();
     TARGET(LOOP_BREAK)
-        frame->jump_abs_break(co_blocks[byte.block].end);
+        frame->jump_abs_break(co_blocks[byte.arg].get_break_end());
         DISPATCH();
     TARGET(GOTO) {
         _name = StrName(byte.arg);

src/compiler.cpp

@@ -559,13 +559,18 @@ __SUBSCR_END:
     }
 
     void Compiler::compile_while_loop() {
-        ctx()->enter_block(WHILE_LOOP);
+        CodeBlock* block = ctx()->enter_block(WHILE_LOOP);
         EXPR(false);   // condition
         int patch = ctx()->emit(OP_POP_JUMP_IF_FALSE, BC_NOARG, prev().line);
         compile_block_body();
-        ctx()->emit(OP_LOOP_CONTINUE, BC_NOARG, BC_KEEPLINE);
+        ctx()->emit(OP_LOOP_CONTINUE, ctx()->get_loop(), BC_KEEPLINE);
         ctx()->patch_jump(patch);
         ctx()->exit_block();
+        // optional else clause
+        if (match(TK("else"))) {
+            compile_block_body();
+            block->end2 = ctx()->co->codes.size();
+        }
     }
 
     void Compiler::compile_for_loop() {
@@ -573,13 +578,18 @@ __SUBSCR_END:
         consume(TK("in"));
         EXPR_TUPLE(false);
         ctx()->emit(OP_GET_ITER, BC_NOARG, BC_KEEPLINE);
-        ctx()->enter_block(FOR_LOOP);
+        CodeBlock* block = ctx()->enter_block(FOR_LOOP);
         ctx()->emit(OP_FOR_ITER, BC_NOARG, BC_KEEPLINE);
         bool ok = vars->emit_store(ctx());
         if(!ok) SyntaxError();  // this error occurs in `vars` instead of this line, but...nevermind
         compile_block_body();
-        ctx()->emit(OP_LOOP_CONTINUE, BC_NOARG, BC_KEEPLINE);
+        ctx()->emit(OP_LOOP_CONTINUE, ctx()->get_loop(), BC_KEEPLINE);
         ctx()->exit_block();
+        // optional else clause
+        if (match(TK("else"))) {
+            compile_block_body();
+            block->end2 = ctx()->co->codes.size();
+        }
     }
 
     void Compiler::compile_try_except() {
@@ -668,15 +678,16 @@ __SUBSCR_END:
     void Compiler::compile_stmt() {
         advance();
         int kw_line = prev().line;  // backup line number
+        int curr_loop_block = ctx()->get_loop();
         switch(prev().type){
             case TK("break"):
-                if (!ctx()->is_curr_block_loop()) SyntaxError("'break' outside loop");
-                ctx()->emit(OP_LOOP_BREAK, BC_NOARG, kw_line);
+                if (curr_loop_block < 0) SyntaxError("'break' outside loop");
+                ctx()->emit(OP_LOOP_BREAK, curr_loop_block, kw_line);
                 consume_end_stmt();
                 break;
             case TK("continue"):
-                if (!ctx()->is_curr_block_loop()) SyntaxError("'continue' not properly in loop");
-                ctx()->emit(OP_LOOP_CONTINUE, BC_NOARG, kw_line);
+                if (curr_loop_block < 0) SyntaxError("'continue' not properly in loop");
+                ctx()->emit(OP_LOOP_CONTINUE, curr_loop_block, kw_line);
                 consume_end_stmt();
                 break;
             case TK("yield"): 
@@ -696,7 +707,7 @@ __SUBSCR_END:
                 ctx()->enter_block(FOR_LOOP);
                 ctx()->emit(OP_FOR_ITER, BC_NOARG, BC_KEEPLINE);
                 ctx()->emit(OP_YIELD_VALUE, BC_NOARG, BC_KEEPLINE);
-                ctx()->emit(OP_LOOP_CONTINUE, BC_NOARG, BC_KEEPLINE);
+                ctx()->emit(OP_LOOP_CONTINUE, ctx()->get_loop(), BC_KEEPLINE);
                 ctx()->exit_block();
                 consume_end_stmt();
                 break;

src/expr.cpp

@@ -2,16 +2,23 @@
 
 namespace pkpy{
 
-    bool CodeEmitContext::is_curr_block_loop() const {
-        return co->blocks[curr_block_i].type == FOR_LOOP || co->blocks[curr_block_i].type == WHILE_LOOP;
+    int CodeEmitContext::get_loop() const {
+        int index = curr_block_i;
+        while(index >= 0){
+            if(co->blocks[index].type == FOR_LOOP) break;
+            if(co->blocks[index].type == WHILE_LOOP) break;
+            index = co->blocks[index].parent;
+        }
+        return index;
     }
 
-    void CodeEmitContext::enter_block(CodeBlockType type){
+    CodeBlock* CodeEmitContext::enter_block(CodeBlockType type){
         if(type == FOR_LOOP) for_loop_depth++;
         co->blocks.push_back(CodeBlock(
             type, curr_block_i, for_loop_depth, (int)co->codes.size()
         ));
         curr_block_i = co->blocks.size()-1;
+        return &co->blocks[curr_block_i];
     }
 
     void CodeEmitContext::exit_block(){
@@ -338,7 +345,7 @@ namespace pkpy{
             expr->emit(ctx);
             ctx->emit(op1(), BC_NOARG, BC_KEEPLINE);
         }
-        ctx->emit(OP_LOOP_CONTINUE, BC_NOARG, BC_KEEPLINE);
+        ctx->emit(OP_LOOP_CONTINUE, ctx->get_loop(), BC_KEEPLINE);
         ctx->exit_block();
     }
 

tests/20_controlflow.py

@@ -71,3 +71,64 @@ assert d == 1
 d = 1 if 2 < 1 else 2
 assert d == 2
 
+t = 0
+for i in range(5):
+    try:
+        break
+    except:
+        pass
+    t = 1
+assert t == 0
+
+t = 0
+for i in range(5):
+    if True and 1:
+        break
+    t = 1
+assert t == 0
+
+for i in range(5):
+    break
+else:
+    assert False
+
+for i in range(5):
+    if i==3:
+        break
+else:
+    assert False
+
+flag = False
+for i in range(5):
+    if i==6:
+        break
+else:
+    flag = True
+assert flag is True
+
+while True:
+    break
+else:
+    assert False
+
+flag = False
+while False:
+    assert False
+else:
+    flag = True
+assert flag is True
+
+x = 1
+while 0:
+    while True:
+        break
+else:
+    x = 2
+assert x == 2
+
+if x == 2:
+    while 0:
+        pass
+else:
+    x = 3
+assert x == 2