diff --git a/artiq/compiler/transforms/artiq_ir_generator.py b/artiq/compiler/transforms/artiq_ir_generator.py index 330ff0ecf..4575c381c 100644 --- a/artiq/compiler/transforms/artiq_ir_generator.py +++ b/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)) diff --git a/artiq/test/lit/integration/for.py b/artiq/test/lit/integration/for.py index 9c305f5da..7156bee5d 100644 --- a/artiq/test/lit/integration/for.py +++ b/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 diff --git a/artiq/test/lit/integration/while.py b/artiq/test/lit/integration/while.py index 7cd2ac626..e7d03c0d7 100644 --- a/artiq/test/lit/integration/while.py +++ b/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