Merge pull request #4132 from yoff/SharedDataflow_NestedComprehensions

Python: Shared dataflow, nested comprehensions

Author: RasmusWL

python/ql/src/experimental/dataflow/internal/DataFlowPrivate.qll

@@ -461,6 +461,10 @@ predicate comprehensionStoreStep(CfgNode nodeFrom, Content c, CfgNode nodeTo) {
   // Dictionary
   nodeTo.getNode().getNode().(DictComp).getElt() = nodeFrom.getNode().getNode() and
   c instanceof DictionaryElementAnyContent
+  or
+  // Generator
+  nodeTo.getNode().getNode().(GeneratorExp).getElt() = nodeFrom.getNode().getNode() and
+  c instanceof ListElementContent
 }
 
 /**
@@ -538,33 +542,23 @@ predicate comprehensionReadStep(CfgNode nodeFrom, Content c, EssaNode nodeTo) {
   //   nodeFrom is `l`, cfg node
   //   nodeTo is `x`, essa var
   //   c denotes element of list or set
-  exists(For f, Comp comp |
-    f = getCompFor(comp) and
-    nodeFrom.getNode().getNode() = getCompIter(comp) and
+  exists(Comp comp |
+    // outermost for
+    nodeFrom.getNode().getNode() = comp.getIterable() and
     nodeTo.getVar().getDefinition().(AssignmentDefinition).getDefiningNode().getNode() =
-      f.getTarget() and
-    (
-      c instanceof ListElementContent
-      or
-      c instanceof SetElementContent
+      comp.getIterationVariable(0).getAStore()
+    or
+    // an inner for
+    exists(int n | n > 0 |
+      nodeFrom.getNode().getNode() = comp.getNthInnerLoop(n).getIter() and
+      nodeTo.getVar().getDefinition().(AssignmentDefinition).getDefiningNode().getNode() =
+        comp.getNthInnerLoop(n).getTarget()
     )
-  )
-}
-
-/** This seems to compensate for extractor shortcomings */
-For getCompFor(Comp c) {
-  c.contains(result) and
-  c.getFunction() = result.getScope()
-}
-
-/** This seems to compensate for extractor shortcomings */
-AstNode getCompIter(Comp c) {
-  c.contains(result) and
-  c.getScope() = result.getScope() and
-  not result = c.getFunction() and
-  not exists(AstNode between |
-    c.contains(between) and
-    between.contains(result)
+  ) and
+  (
+    c instanceof ListElementContent
+    or
+    c instanceof SetElementContent
   )
 }
 

python/ql/src/semmle/python/Comprehensions.qll

@@ -4,18 +4,22 @@ import python
 abstract class Comp extends Expr {
   abstract Function getFunction();
 
-  /** Gets the iteration variable for the nth innermost generator of this list comprehension */
+  /** Gets the iterable of this set comprehension. */
+  abstract Expr getIterable();
+
+  /** Gets the iteration variable for the nth innermost generator of this comprehension. */
   Variable getIterationVariable(int n) {
     result.getAnAccess() = this.getNthInnerLoop(n).getTarget()
   }
 
-  private For getNthInnerLoop(int n) {
+  /** Gets the nth innermost For expression of this comprehension. */
+  For getNthInnerLoop(int n) {
     n = 0 and result = this.getFunction().getStmt(0)
     or
     result = this.getNthInnerLoop(n - 1).getStmt(0)
   }
 
-  /** Gets the iteration variable for a generator of this list comprehension */
+  /** Gets the iteration variable for a generator of this list comprehension. */
   Variable getAnIterationVariable() { result = this.getIterationVariable(_) }
 
   /** Gets the scope in which the body of this list comprehension evaluates. */
@@ -62,6 +66,8 @@ class ListComp extends ListComp_, Comp {
 
   override Function getFunction() { result = ListComp_.super.getFunction() }
 
+  override Expr getIterable() { result = ListComp_.super.getIterable() }
+
   override string toString() { result = ListComp_.super.toString() }
 
   override Expr getElt() { result = Comp.super.getElt() }
@@ -79,6 +85,8 @@ class SetComp extends SetComp_, Comp {
   override predicate hasSideEffects() { any() }
 
   override Function getFunction() { result = SetComp_.super.getFunction() }
+
+  override Expr getIterable() { result = SetComp_.super.getIterable() }
 }
 
 /** A dictionary comprehension, such as `{ k:v for k, v in enumerate("0123456789") }` */
@@ -93,6 +101,8 @@ class DictComp extends DictComp_, Comp {
   override predicate hasSideEffects() { any() }
 
   override Function getFunction() { result = DictComp_.super.getFunction() }
+
+  override Expr getIterable() { result = DictComp_.super.getIterable() }
 }
 
 /** A generator expression, such as `(var for var in iterable)` */
@@ -107,4 +117,6 @@ class GeneratorExp extends GeneratorExp_, Comp {
   override predicate hasSideEffects() { any() }
 
   override Function getFunction() { result = GeneratorExp_.super.getFunction() }
+
+  override Expr getIterable() { result = GeneratorExp_.super.getIterable() }
 }