From 4b4805265d148e4c36aeca3b5c46e0c3b8a3584f Mon Sep 17 00:00:00 2001
From: whitequark <whitequark@whitequark.org>
Date: Sat, 13 Jun 2015 12:07:46 +0300
Subject: [PATCH] Add support for Break and Continue.

---
 artiq/py2llvm/typing.py                       | 46 +++++++++++++++----
 lit-test/py2llvm/typing/error_control_flow.py | 19 ++++++++
 2 files changed, 56 insertions(+), 9 deletions(-)
 create mode 100644 lit-test/py2llvm/typing/error_control_flow.py

diff --git a/artiq/py2llvm/typing.py b/artiq/py2llvm/typing.py
index f7f754195..87d68e797 100644
--- a/artiq/py2llvm/typing.py
+++ b/artiq/py2llvm/typing.py
@@ -284,9 +284,7 @@ class ASTTypedRewriter(algorithm.Transformer):
 
     # stmt
     visit_Assert = visit_unsupported
-    visit_Break = visit_unsupported
     visit_ClassDef = visit_unsupported
-    visit_Continue = visit_unsupported
     visit_Delete = visit_unsupported
     visit_Import = visit_unsupported
     visit_ImportFrom = visit_unsupported
@@ -298,8 +296,8 @@ class ASTTypedRewriter(algorithm.Transformer):
 class Inferencer(algorithm.Visitor):
     def __init__(self, engine):
         self.engine = engine
-        # currently visited function, for Return inference
-        self.function = None
+        self.function = None # currently visited function, for Return inference
+        self.in_loop = False
 
     def _unify(self, typea, typeb, loca, locb, makenotes=None):
         try:
@@ -385,11 +383,6 @@ class Inferencer(algorithm.Visitor):
                     node.operand.loc)
                 self.engine.process(diag)
 
-    def visit_FunctionDefT(self, node):
-        old_function, self.function = self.function, node
-        node = self.generic_visit(node)
-        self.function = old_function
-
     def visit_Assign(self, node):
         self.generic_visit(node)
         if len(node.targets) > 1:
@@ -405,12 +398,47 @@ class Inferencer(algorithm.Visitor):
                     node.target.loc, node.value.loc)
 
     def visit_For(self, node):
+        old_in_loop, self.in_loop = self.in_loop, True
         self.generic_visit(node)
+        self.in_loop = old_in_loop
         # TODO: support more than just lists
         self._unify(builtins.TList(node.target.type), node.iter.type,
                     node.target.loc, node.iter.loc)
 
+    def visit_While(self, node):
+        old_in_loop, self.in_loop = self.in_loop, True
+        self.generic_visit(node)
+        self.in_loop = old_in_loop
+
+    def visit_Break(self, node):
+        if not self.in_loop:
+            diag = diagnostic.Diagnostic("error",
+                "break statement outside of a loop", {},
+                node.keyword_loc)
+            self.engine.process(diag)
+
+    def visit_Continue(self, node):
+        if not self.in_loop:
+            diag = diagnostic.Diagnostic("error",
+                "continue statement outside of a loop", {},
+                node.keyword_loc)
+            self.engine.process(diag)
+
+    def visit_FunctionDefT(self, node):
+        old_function, self.function = self.function, node
+        old_in_loop, self.in_loop = self.in_loop, False
+        self.generic_visit(node)
+        self.function = old_function
+        self.in_loop = old_in_loop
+
     def visit_Return(self, node):
+        if not self.function:
+            diag = diagnostic.Diagnostic("error",
+                "return statement outside of a function", {},
+                node.keyword_loc)
+            self.engine.process(diag)
+            return
+
         self.generic_visit(node)
         def makenotes(printer, typea, typeb, loca, locb):
             return [
diff --git a/lit-test/py2llvm/typing/error_control_flow.py b/lit-test/py2llvm/typing/error_control_flow.py
new file mode 100644
index 000000000..ea72347dd
--- /dev/null
+++ b/lit-test/py2llvm/typing/error_control_flow.py
@@ -0,0 +1,19 @@
+# RUN: %python -m artiq.py2llvm.typing +diag %s >%t
+# RUN: OutputCheck %s --file-to-check=%t
+
+# CHECK-L: ${LINE:+1}: error: return statement outside of a function
+return
+
+# CHECK-L: ${LINE:+1}: error: break statement outside of a loop
+break
+
+# CHECK-L: ${LINE:+1}: error: continue statement outside of a loop
+continue
+
+while True:
+    def f():
+        # CHECK-L: ${LINE:+1}: error: break statement outside of a loop
+        break
+
+        # CHECK-L: ${LINE:+1}: error: continue statement outside of a loop
+        continue