Implementing for-else and while-else loops (#2555)

* Added for-else functionality
* Added support for nested loops

Author: advikkabra

integration_tests/CMakeLists.txt

@@ -518,6 +518,8 @@ RUN(NAME loop_05             LABELS cpython llvm c)
 RUN(NAME loop_06             LABELS cpython llvm c NOFAST)
 RUN(NAME loop_07             LABELS cpython llvm c)
 RUN(NAME loop_08             LABELS cpython llvm c)
+RUN(NAME loop_09             LABELS cpython llvm)
+RUN(NAME loop_10             LABELS cpython llvm)
 RUN(NAME if_01               LABELS cpython llvm c wasm wasm_x86 wasm_x64)
 RUN(NAME if_02               LABELS cpython llvm c wasm wasm_x86 wasm_x64)
 RUN(NAME if_03    FAIL       LABELS cpython llvm c NOFAST)

integration_tests/loop_09.py

@@ -0,0 +1,45 @@
+def with_break_for():
+    i: i32
+    s: i32 = 0
+    for i in range(4):
+        s += i
+        break
+    else:
+        s += 10
+    assert s == 0
+
+def with_break_while():
+    i: i32 = 0
+    s: i32 = 0
+    while i < 4:
+        s += i
+        break
+    else:
+        s += 10
+    assert s == 0
+
+def no_break_for():
+    i: i32
+    s: i32 = 0
+    for i in range(2):
+        s += i
+    else:
+        s += 10
+    assert s == 11
+
+def no_break_while():
+    i: i32 = 0
+    s: i32 = 0
+    while i < 2:
+        s += i
+        i += 1
+    else:
+        s += 10
+    assert s == 11
+
+no_break_for()
+no_break_while()
+
+with_break_for()
+with_break_while()
+

integration_tests/loop_10.py

@@ -0,0 +1,57 @@
+def with_break_for():
+    i: i32
+    for i in range(4):
+        print(i)
+        break
+    else:
+        assert False
+
+
+def with_break_while():
+    i: i32 = 0
+    while i < 4:
+        print(i)
+        break
+    else:
+        print(10)
+        assert False
+
+
+def no_break_for():
+    i: i32
+    for i in range(2):
+        print(i)
+    else:
+        print(10)
+        return
+    assert False
+
+def break_in_if_for():
+    i: i32
+    for i in range(2):
+        print(i)
+        if i == 1:
+            break
+    else:
+        print(10)
+        assert False
+
+def nested_loop_for_for():
+    i: i32
+    j: i32
+    for i in range(2):
+        print("outer: " + str(i))
+        for j in range(10, 20):
+            print("  inner: " + str(j))
+            break
+    else:
+        print("no break in outer loop")
+        return
+    assert False
+
+
+with_break_for()
+with_break_while()
+no_break_for()
+break_in_if_for()
+nested_loop_for_for()

run_tests.py

@@ -39,7 +39,7 @@ def is_included(backend):
                            "dead_code_removal", "loop_vectorise", "print_list_tuple",
                            "class_constructor"]
 
-    if pass_ and (pass_ not in ["do_loops", "global_stmts"] and
+    if pass_ and (pass_ not in ["do_loops", "global_stmts", "while_else"] and
                   pass_ not in optimization_passes):
         raise Exception(f"Unknown pass: {pass_}")
     if no_color:

src/libasr/ASR.asdl

@@ -185,7 +185,7 @@ stmt
     -- deallocates if allocated otherwise does nothing
     | ImplicitDeallocate(expr* vars)
     | DoConcurrentLoop(do_loop_head head, stmt* body)
-    | DoLoop(identifier? name, do_loop_head head, stmt* body)
+    | DoLoop(identifier? name, do_loop_head head, stmt* body, stmt* orelse)
     | ErrorStop(expr? code)
     | Exit(identifier? stmt_name)
     | ForAllSingle(do_loop_head head, stmt assign_stmt)
@@ -219,7 +219,7 @@ stmt
     | Assert(expr test, expr? msg)
     | SubroutineCall(symbol name, symbol? original_name, call_arg* args, expr? dt)
     | Where(expr test, stmt* body, stmt* orelse)
-    | WhileLoop(identifier? name, expr test, stmt* body)
+    | WhileLoop(identifier? name, expr test, stmt* body, stmt* orelse)
     | Nullify(symbol* vars)
     | Flush(int label, expr unit, expr? err, expr? iomsg, expr? iostat)
     | ListAppend(expr a, expr ele)

src/libasr/CMakeLists.txt

@@ -34,6 +34,7 @@ set(SRC
     pass/param_to_const.cpp
     pass/do_loops.cpp
     pass/for_all.cpp
+    pass/while_else.cpp
     pass/global_stmts.cpp
     pass/select_case.cpp
     pass/init_expr.cpp

src/libasr/codegen/asr_to_julia.cpp

@@ -1237,7 +1237,7 @@ class ASRToJuliaVisitor : public ASR::BaseVisitor<ASRToJuliaVisitor>
 
     void visit_DoConcurrentLoop(const ASR::DoConcurrentLoop_t& x)
     {
-        const ASR::DoLoop_t do_loop = ASR::DoLoop_t{ x.base, nullptr, x.m_head, x.m_body, x.n_body };
+        const ASR::DoLoop_t do_loop = ASR::DoLoop_t{ x.base, nullptr, x.m_head, x.m_body, x.n_body , nullptr, 0};
         visit_DoLoop(do_loop, true);
     }
 

src/libasr/gen_pass.py

@@ -8,6 +8,7 @@
         "replace_flip_sign",
         "replace_fma",
         "replace_for_all",
+        "while_else",
         "wrap_global_stmts",
         "replace_implied_do_loops",
         "replace_init_expr",

src/libasr/pass/arr_slice.cpp

@@ -155,7 +155,7 @@ class ReplaceArraySection: public ASR::BaseExprReplacer<ReplaceArraySection> {
                                         const_1, int_type, nullptr));
             ASR::stmt_t* assign_stmt = ASRUtils::STMT(ASR::make_Assignment_t(al, x->base.base.loc, idx_vars_target[i], inc_expr, nullptr));
             doloop_body.push_back(al, assign_stmt);
-            doloop = ASRUtils::STMT(ASR::make_DoLoop_t(al, x->base.base.loc, nullptr, head, doloop_body.p, doloop_body.size()));
+            doloop = ASRUtils::STMT(ASR::make_DoLoop_t(al, x->base.base.loc, nullptr, head, doloop_body.p, doloop_body.size(), nullptr, 0));
         }
         int a_kind = ASRUtils::extract_kind_from_ttype_t(ASRUtils::expr_type(idx_vars_target[0]));
         ASR::ttype_t* int_type = ASRUtils::TYPE(ASR::make_Integer_t(al, x->base.base.loc, a_kind));

src/libasr/pass/array_op.cpp

@@ -171,7 +171,7 @@ class ReplaceArrayOp: public ASR::BaseExprReplacer<ReplaceArrayOp> {
                         doloop_body.push_back(al, assign_stmt2);
                     }
                 }
-                doloop = ASRUtils::STMT(ASR::make_DoLoop_t(al, loc, nullptr, head, doloop_body.p, doloop_body.size()));
+                doloop = ASRUtils::STMT(ASR::make_DoLoop_t(al, loc, nullptr, head, doloop_body.p, doloop_body.size(), nullptr, 0));
             }
             if( var_rank > 0 ) {
                 ASR::expr_t* idx_lb = PassUtils::get_bound(op_expr1, 1, "lbound", al);
@@ -207,7 +207,7 @@ class ReplaceArrayOp: public ASR::BaseExprReplacer<ReplaceArrayOp> {
                 } else {
                     doloop_body.push_back(al, doloop);
                 }
-                doloop = ASRUtils::STMT(ASR::make_DoLoop_t(al, loc, nullptr, head, doloop_body.p, doloop_body.size()));
+                doloop = ASRUtils::STMT(ASR::make_DoLoop_t(al, loc, nullptr, head, doloop_body.p, doloop_body.size(), nullptr, 0));
             }
             pass_result.push_back(al, doloop);
         }
@@ -460,7 +460,7 @@ class ReplaceArrayOp: public ASR::BaseExprReplacer<ReplaceArrayOp> {
                     al, loc, idx_vars_value[i], inc_expr, nullptr));
                 doloop_body.push_back(al, assign_stmt);
             }
-            doloop = ASRUtils::STMT(ASR::make_DoLoop_t(al, loc, nullptr, head, doloop_body.p, doloop_body.size()));
+            doloop = ASRUtils::STMT(ASR::make_DoLoop_t(al, loc, nullptr, head, doloop_body.p, doloop_body.size(), nullptr, 0));
         }
         if( ASRUtils::is_array(ASRUtils::expr_type(op_expr)) ) {
             ASR::expr_t* idx_lb = PassUtils::get_bound(op_expr, 1, "lbound", al);
@@ -505,7 +505,7 @@ class ReplaceArrayOp: public ASR::BaseExprReplacer<ReplaceArrayOp> {
             } else {
                 doloop_body.push_back(al, doloop);
             }
-            doloop = ASRUtils::STMT(ASR::make_DoLoop_t(al, loc, nullptr, head, doloop_body.p, doloop_body.size()));
+            doloop = ASRUtils::STMT(ASR::make_DoLoop_t(al, loc, nullptr, head, doloop_body.p, doloop_body.size(), nullptr, 0));
         }
         pass_result.push_back(al, doloop);
     }
@@ -1780,7 +1780,7 @@ class ArrayOpVisitor : public ASR::CallReplacerOnExpressionsVisitor<ArrayOpVisit
                             al, loc, idx_vars_value[i], inc_expr, nullptr));
                         doloop_body.push_back(al, assign_stmt);
                     }
-                    doloop = ASRUtils::STMT(ASR::make_DoLoop_t(al, loc, nullptr, head, doloop_body.p, doloop_body.size()));
+                    doloop = ASRUtils::STMT(ASR::make_DoLoop_t(al, loc, nullptr, head, doloop_body.p, doloop_body.size(), nullptr, 0));
                 }
                 if( var_rank > 0 ) {
                     ASR::expr_t* idx_lb = PassUtils::get_bound(op_expr, 1, "lbound", al);
@@ -1810,7 +1810,7 @@ class ArrayOpVisitor : public ASR::CallReplacerOnExpressionsVisitor<ArrayOpVisit
                     } else {
                         doloop_body.push_back(al, doloop);
                     }
-                    doloop = ASRUtils::STMT(ASR::make_DoLoop_t(al, loc, nullptr, head, doloop_body.p, doloop_body.size()));
+                    doloop = ASRUtils::STMT(ASR::make_DoLoop_t(al, loc, nullptr, head, doloop_body.p, doloop_body.size(), nullptr, 0));
                 }
                 pass_result.push_back(al, doloop);
             }