Add nested loop tests for else clause in loops (#2630)

* Add nested loop tests for else clause in loops

* Add nested else test

* Fix bug in for else implementation

* Update tests to use assert

* Add print statements in tests

Author: advikkabra

integration_tests/loop_10.py

@@ -1,7 +1,7 @@
 def with_break_for():
-    i: i32
+    i: i32 = 0
     for i in range(4):
-        print(i)
+        i += 1
         break
     else:
         assert False
@@ -10,48 +10,117 @@ def with_break_for():
 def with_break_while():
     i: i32 = 0
     while i < 4:
-        print(i)
+        i += 1
         break
     else:
-        print(10)
         assert False
 
 
 def no_break_for():
     i: i32
+    j: i32 = 0
     for i in range(2):
-        print(i)
+        j += 1
     else:
-        print(10)
+        print(j)
+        assert j == 2
         return
     assert False
 
 def break_in_if_for():
     i: i32
+    j: i32 = 0
     for i in range(2):
-        print(i)
+        j += 1
         if i == 1:
             break
     else:
-        print(10)
         assert False
+    print(j)
+    assert j == 2
 
 def nested_loop_for_for():
     i: i32
     j: i32
+    m: i32 = 0
     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
+        m = 10
+    print(m)
+    assert m == 10
+
+def nested_loop_for_while():
+    i: i32
+    j: i32 = 10
+    m: i32 = 0
+    for i in range(2):
+        while j < 20:
+            break
+    else:
+        m = 10
+    print(m)
+    assert m == 10
+
+def nested_loop_while_for():
+    i: i32 = 0
+    j: i32
+    m: i32 = 0
+    while i < 2:
+        i += 1
+        for j in range(10, 20):
+            break
+    else:
+        print(i)
+        assert i == 2
+        m = 10
+    print(m)
+    assert m == 10
+
+def nested_loop_while_while():
+    i: i32 = 0
+    j: i32 = 10
+    m: i32 = 0
+    while i < 2:
+        i += 1
+        while j < 20:
+            break
+    else:
+        print(i)
+        assert i == 2
+        m = 10
+    print(m)
+    assert m == 10
+
+def nested_loop_else():
+    i: i32
+    j: i32
+    l: i32 = 0
+    m: i32 = 0
+    for i in range(2):
+        l += 1
+        m: i32 = 0
+        for j in range(10, 12):
+            m += 1
+        else:
+            print(m)
+            assert m == 2
+            l += 10
+    else:
+        print(l)
+        assert l == 22
+        m = 10
+    print(m)
+    assert m == 10
 
 
 with_break_for()
 with_break_while()
 no_break_for()
 break_in_if_for()
 nested_loop_for_for()
+nested_loop_for_while()
+nested_loop_while_for()
+nested_loop_while_while()
+nested_loop_else()

src/libasr/pass/while_else.cpp

@@ -73,6 +73,8 @@ class WhileLoopVisitor : public ASR::StatementWalkVisitor<WhileLoopVisitor>
         Creating a flag variable in case of a while-else loop
         Creates an if statement after the loop to check if the flag was changed
         */
+        ASR::WhileLoop_t &xx = const_cast<ASR::WhileLoop_t&>(x);
+        transform_stmts(xx.m_body, xx.n_body);
         if (x.n_orelse > 0) {
             Vec<ASR::stmt_t*> result;
             result.reserve(al, 3);