compiler: Fix break/continue targets in loop else blocks

Previously, they would still target the respective labels in the
just-exited loop.

Signed-off-by: David Nadlinger <code@klickverbot.at>

artiq/compiler/transforms/artiq_ir_generator.py

@@ -452,30 +452,30 @@ class ARTIQIRGenerator(algorithm.Visitor):
             self.current_block = body
             self.visit(node.body)
             post_body = self.current_block
-
-            if any(node.orelse):
-                else_tail = self.add_block("while.else")
-                self.current_block = else_tail
-                self.visit(node.orelse)
-                post_else_tail = self.current_block
-
-            tail = self.add_block("while.tail")
-            self.current_block = tail
-
-            if any(node.orelse):
-                if not post_else_tail.is_terminated():
-                    post_else_tail.append(ir.Branch(tail))
-            else:
-                else_tail = tail
-
-            post_head.append(ir.BranchIf(cond, body, else_tail))
-            if not post_body.is_terminated():
-                post_body.append(ir.Branch(head))
-            break_block.append(ir.Branch(tail))
         finally:
             self.break_target = old_break
             self.continue_target = old_continue
 
+        if any(node.orelse):
+            else_tail = self.add_block("while.else")
+            self.current_block = else_tail
+            self.visit(node.orelse)
+            post_else_tail = self.current_block
+
+        tail = self.add_block("while.tail")
+        self.current_block = tail
+
+        if any(node.orelse):
+            if not post_else_tail.is_terminated():
+                post_else_tail.append(ir.Branch(tail))
+        else:
+            else_tail = tail
+
+        post_head.append(ir.BranchIf(cond, body, else_tail))
+        if not post_body.is_terminated():
+            post_body.append(ir.Branch(head))
+        break_block.append(ir.Branch(tail))
+
     def iterable_len(self, value, typ=_size_type):
         if builtins.is_listish(value.type):
             if isinstance(value, ir.Constant):
@@ -541,33 +541,33 @@ class ARTIQIRGenerator(algorithm.Visitor):
                 self.current_assign = None
             self.visit(node.body)
             post_body = self.current_block
-
-            if any(node.orelse):
-                else_tail = self.add_block("for.else")
-                self.current_block = else_tail
-                self.visit(node.orelse)
-                post_else_tail = self.current_block
-
-            tail = self.add_block("for.tail")
-            self.current_block = tail
-
-            if any(node.orelse):
-                if not post_else_tail.is_terminated():
-                    post_else_tail.append(ir.Branch(tail))
-            else:
-                else_tail = tail
-
-            if node.trip_count is not None:
-                head.append(ir.Loop(node.trip_count, phi, cond, body, else_tail))
-            else:
-                head.append(ir.BranchIf(cond, body, else_tail))
-            if not post_body.is_terminated():
-                post_body.append(ir.Branch(continue_block))
-            break_block.append(ir.Branch(tail))
         finally:
             self.break_target = old_break
             self.continue_target = old_continue
 
+        if any(node.orelse):
+            else_tail = self.add_block("for.else")
+            self.current_block = else_tail
+            self.visit(node.orelse)
+            post_else_tail = self.current_block
+
+        tail = self.add_block("for.tail")
+        self.current_block = tail
+
+        if any(node.orelse):
+            if not post_else_tail.is_terminated():
+                post_else_tail.append(ir.Branch(tail))
+        else:
+            else_tail = tail
+
+        if node.trip_count is not None:
+            head.append(ir.Loop(node.trip_count, phi, cond, body, else_tail))
+        else:
+            head.append(ir.BranchIf(cond, body, else_tail))
+        if not post_body.is_terminated():
+            post_body.append(ir.Branch(continue_block))
+        break_block.append(ir.Branch(tail))
+
     def visit_Break(self, node):
         self.append(ir.Branch(self.break_target))
 

artiq/test/lit/integration/for.py

@@ -27,3 +27,27 @@ for x in range(10):
     assert False
 else:
     assert False
+
+# Verify continue target is reset in else block.
+cond = False
+while True:
+    if cond:
+        break
+    cond = True
+    for _ in range(1):
+        pass
+    else:
+        continue
+    assert False
+else:
+    assert False
+
+# Verify for target is reset in else block.
+while True:
+    for _ in range(1):
+        pass
+    else:
+        break
+    assert False
+else:
+    assert False

artiq/test/lit/integration/while.py

@@ -29,3 +29,27 @@ while cond:
     cond = False
     continue
     assert False
+
+# Verify continue target is reset in else block.
+cond = False
+while True:
+    if cond:
+        break
+    cond = True
+    while False:
+        assert False
+    else:
+        continue
+    assert False
+else:
+    assert False
+
+# Verify break target is reset in else block.
+while True:
+    while False:
+        assert False
+    else:
+        break
+    assert False
+else:
+    assert False