diff --git a/artiq/py2llvm/ast_body.py b/artiq/py2llvm/ast_body.py index bd69ed62d..5a44b1f6e 100644 --- a/artiq/py2llvm/ast_body.py +++ b/artiq/py2llvm/ast_body.py @@ -41,6 +41,7 @@ class Visitor: self.env = env self.ns = ns self.builder = builder + self._break_stack = [] # builder can be None for visit_expression def visit_expression(self, node): @@ -206,6 +207,7 @@ class Visitor: body_block = function.append_basic_block("w_body") else_block = function.append_basic_block("w_else") merge_block = function.append_basic_block("w_merge") + self._break_stack.append(merge_block) condition = self.visit_expression(node.test).o_bool(self.builder) self.builder.cbranch( @@ -224,12 +226,14 @@ class Visitor: self.builder.branch(merge_block) self.builder.position_at_end(merge_block) + self._break_stack.pop() def _visit_stmt_For(self, node): function = self.builder.basic_block.function body_block = function.append_basic_block("f_body") else_block = function.append_basic_block("f_else") merge_block = function.append_basic_block("f_merge") + self._break_stack.append(merge_block) it = self.visit_expression(node.iter) target = self.visit_expression(node.target) @@ -253,6 +257,10 @@ class Visitor: self.builder.branch(merge_block) self.builder.position_at_end(merge_block) + self._break_stack.pop() + + def _visit_stmt_Break(self, node): + self.builder.branch(self._break_stack[-1]) def _visit_stmt_Return(self, node): if node.value is None: diff --git a/examples/coredev_test.py b/examples/coredev_test.py index 127648a25..9512e689d 100644 --- a/examples/coredev_test.py +++ b/examples/coredev_test.py @@ -21,6 +21,7 @@ class CompilerTest(AutoContext): while d*d <= x: if x % d == 0: prime = False + break d += 1 if prime: self.output(x)