From 760c89832ee4ac09ddaae70a780ebe11f39d86dc Mon Sep 17 00:00:00 2001
From: Advik Kabra <64316822+advikkabra@users.noreply.github.com>
Date: Thu, 4 Apr 2024 09:25:33 +0530
Subject: [PATCH] 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
---
 integration_tests/loop_10.py   | 95 +++++++++++++++++++++++++++++-----
 src/libasr/pass/while_else.cpp |  2 +
 2 files changed, 84 insertions(+), 13 deletions(-)

diff --git a/integration_tests/loop_10.py b/integration_tests/loop_10.py
index 00c91d58ed..6f9f0defb9 100644
--- a/integration_tests/loop_10.py
+++ b/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,44 +10,109 @@ 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()
@@ -55,3 +120,7 @@ def nested_loop_for_for():
 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()
diff --git a/src/libasr/pass/while_else.cpp b/src/libasr/pass/while_else.cpp
index df44951087..07b2744a85 100644
--- a/src/libasr/pass/while_else.cpp
+++ b/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);