From 4a5e39885752ec0190e67cd4c7d0955147c89cbc Mon Sep 17 00:00:00 2001
From: Sebastien Bourdeauducq <sb@m-labs.hk>
Date: Sat, 27 Sep 2014 00:27:11 +0800
Subject: [PATCH] py2llvm: support continue statement

---
 artiq/py2llvm/ast_body.py | 50 ++++++++++++++++++++++++++-------------
 1 file changed, 34 insertions(+), 16 deletions(-)

diff --git a/artiq/py2llvm/ast_body.py b/artiq/py2llvm/ast_body.py
index 9570bbf09..a3fd7b234 100644
--- a/artiq/py2llvm/ast_body.py
+++ b/artiq/py2llvm/ast_body.py
@@ -43,6 +43,7 @@ class Visitor:
         self.ns = ns
         self.builder = builder
         self._break_stack = []
+        self._continue_stack = []
         self._active_exception_stack = []
         self._exception_level_stack = [0]
 
@@ -210,31 +211,38 @@ class Visitor:
 
         self.builder.position_at_end(merge_block)
 
-    def _enter_loop_body(self, break_block):
+    def _enter_loop_body(self, break_block, continue_block):
         self._break_stack.append(break_block)
+        self._continue_stack.append(continue_block)
         self._exception_level_stack.append(0)
 
     def _leave_loop_body(self):
         self._exception_level_stack.pop()
+        self._continue_stack.pop()
+        self._break_stack.pop()
 
     def _visit_stmt_While(self, node):
         function = self.builder.basic_block.function
+
         body_block = function.append_basic_block("w_body")
         else_block = function.append_basic_block("w_else")
-        merge_block = function.append_basic_block("w_merge")
-
         condition = self.visit_expression(node.test).o_bool(self.builder)
         self.builder.cbranch(
             condition.auto_load(self.builder), body_block, else_block)
 
+        continue_block = function.append_basic_block("w_continue")
+        merge_block = function.append_basic_block("w_merge")
         self.builder.position_at_end(body_block)
-        self._enter_loop_body(merge_block)
+        self._enter_loop_body(merge_block, continue_block)
         self.visit_statements(node.body)
         self._leave_loop_body()
         if not self._bb_terminated():
-            condition = self.visit_expression(node.test).o_bool(self.builder)
-            self.builder.cbranch(
-                condition.auto_load(self.builder), body_block, merge_block)
+            self.builder.branch(continue_block)
+
+        self.builder.position_at_end(continue_block)
+        condition = self.visit_expression(node.test).o_bool(self.builder)
+        self.builder.cbranch(
+            condition.auto_load(self.builder), body_block, merge_block)
 
         self.builder.position_at_end(else_block)
         self.visit_statements(node.orelse)
@@ -245,27 +253,31 @@ class Visitor:
 
     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")
 
         it = self.visit_expression(node.iter)
         target = self.visit_expression(node.target)
         itval = it.get_value_ptr()
 
+        body_block = function.append_basic_block("f_body")
+        else_block = function.append_basic_block("f_else")
         cont = it.o_next(self.builder)
         self.builder.cbranch(
             cont.auto_load(self.builder), body_block, else_block)
 
+        continue_block = function.append_basic_block("f_continue")
+        merge_block = function.append_basic_block("f_merge")
         self.builder.position_at_end(body_block)
         target.set_value(self.builder, itval)
-        self._enter_loop_body(merge_block)
+        self._enter_loop_body(merge_block, continue_block)
         self.visit_statements(node.body)
         self._leave_loop_body()
         if not self._bb_terminated():
-            cont = it.o_next(self.builder)
-            self.builder.cbranch(
-                cont.auto_load(self.builder), body_block, merge_block)
+            self.builder.branch(continue_block)
+
+        self.builder.position_at_end(continue_block)
+        cont = it.o_next(self.builder)
+        self.builder.cbranch(
+            cont.auto_load(self.builder), body_block, merge_block)
 
         self.builder.position_at_end(else_block)
         self.visit_statements(node.orelse)
@@ -274,11 +286,17 @@ class Visitor:
 
         self.builder.position_at_end(merge_block)
 
-    def _visit_stmt_Break(self, node):
+    def _break_loop_body(self, target_block):
         exception_levels = self._exception_level_stack[-1]
         if exception_levels:
             self.env.build_pop(self.builder, exception_levels)
-        self.builder.branch(self._break_stack[-1])
+        self.builder.branch(target_block)
+
+    def _visit_stmt_Break(self, node):
+        self._break_loop_body(self._break_stack[-1])
+
+    def _visit_stmt_Continue(self, node):
+        self._break_loop_body(self._continue_stack[-1])
 
     def _visit_stmt_Return(self, node):
         if node.value is None: