From 53fb03d1bf7324c38931bb9ffa7dc4dff5ec17d2 Mon Sep 17 00:00:00 2001
From: whitequark <whitequark@whitequark.org>
Date: Thu, 16 Jul 2015 14:52:41 +0300
Subject: [PATCH] Restrict comprehensions to single for and no if clauses.

---
 .../compiler/transforms/asttyped_rewriter.py  | 23 +++++++++++++++----
 artiq/compiler/transforms/inferencer.py       | 12 ++++++++++
 .../inferencer/error_comprehension.py         |  8 +++++++
 3 files changed, 39 insertions(+), 4 deletions(-)
 create mode 100644 lit-test/compiler/inferencer/error_comprehension.py

diff --git a/artiq/compiler/transforms/asttyped_rewriter.py b/artiq/compiler/transforms/asttyped_rewriter.py
index de72fce3e..e89e7ac9c 100644
--- a/artiq/compiler/transforms/asttyped_rewriter.py
+++ b/artiq/compiler/transforms/asttyped_rewriter.py
@@ -57,6 +57,25 @@ class LocalExtractor(algorithm.Visitor):
         for if_ in node.ifs:
             self.visit(node.ifs)
 
+    def visit_generator(self, node):
+        if self.in_root:
+            return
+        self.in_root = True
+        self.visit(list(reversed(node.generators)))
+        self.visit(node.elt)
+
+    visit_ListComp     = visit_generator
+    visit_SetComp      = visit_generator
+    visit_GeneratorExp = visit_generator
+
+    def visit_DictComp(self, node):
+        if self.in_root:
+            return
+        self.in_root = True
+        self.visit(list(reversed(node.generators)))
+        self.visit(node.key)
+        self.visit(node.value)
+
     def visit_root(self, node):
         if self.in_root:
             return
@@ -66,10 +85,6 @@ class LocalExtractor(algorithm.Visitor):
     visit_Module       = visit_root # don't look at inner scopes
     visit_ClassDef     = visit_root
     visit_Lambda       = visit_root
-    visit_DictComp     = visit_root
-    visit_ListComp     = visit_root
-    visit_SetComp      = visit_root
-    visit_GeneratorExp = visit_root
 
     def visit_FunctionDef(self, node):
         if self.in_root:
diff --git a/artiq/compiler/transforms/inferencer.py b/artiq/compiler/transforms/inferencer.py
index 7fbf09813..dbb9aeb04 100644
--- a/artiq/compiler/transforms/inferencer.py
+++ b/artiq/compiler/transforms/inferencer.py
@@ -390,11 +390,23 @@ class Inferencer(algorithm.Visitor):
                     node.loc, None)
 
     def visit_ListCompT(self, node):
+        if len(node.generators) > 1:
+            diag = diagnostic.Diagnostic("error",
+                "multiple for clauses in comprehensions are not supported", {},
+                node.generators[1].for_loc)
+            self.engine.process(diag)
+
         self.generic_visit(node)
         self._unify(node.type, builtins.TList(node.elt.type),
                     node.loc, None)
 
     def visit_comprehension(self, node):
+        if any(node.ifs):
+            diag = diagnostic.Diagnostic("error",
+                "if clauses in comprehensions are not supported", {},
+                node.if_locs[0])
+            self.engine.process(diag)
+
         self.generic_visit(node)
         self._unify_iterable(element=node.target, collection=node.iter)
 
diff --git a/lit-test/compiler/inferencer/error_comprehension.py b/lit-test/compiler/inferencer/error_comprehension.py
new file mode 100644
index 000000000..d586dd657
--- /dev/null
+++ b/lit-test/compiler/inferencer/error_comprehension.py
@@ -0,0 +1,8 @@
+# RUN: %python -m artiq.compiler.testbench.inferencer +diag %s >%t
+# RUN: OutputCheck %s --file-to-check=%t
+
+# CHECK-L: ${LINE:+1}: error: if clauses in comprehensions are not supported
+[x for x in [] if x]
+
+# CHECK-L: ${LINE:+1}: error: multiple for clauses in comprehensions are not supported
+[(x, y) for x in [] for y in []]