@@ -163,6 +163,10 @@ newtype TArgumentPosition =
163163 // position, we need to ensure we make these available (these are specified as
164164 // parameters in the flow-summary spec)
165165 FlowSummaryImpl:: ParsePositions:: isParsedPositionalParameterPosition ( _, index )
166+ or
167+ // the generated function inside a comprehension has a positional argument at index 0
168+ exists ( Comp c ) and
169+ index = 0
166170 } or
167171 TKeywordArgumentPosition ( string name ) {
168172 exists ( any ( CallNode c ) .getArgByName ( name ) )
@@ -314,12 +318,11 @@ newtype TDataFlowCallable =
314318 * class instantiations, and (in the future) special methods.
315319 */
316320 TFunction ( Function func ) {
317- // For generators/list-comprehensions we create a synthetic function. In the
318- // points-to call-graph these were not considered callable, and instead we added
319- // data-flow steps (read/write) for these. As an easy solution for now, we do the
320- // same to keep things easy to reason about (and therefore exclude things that do
321- // not have a definition)
321+ // Functions with an explicit definition
322322 exists ( func .getDefinition ( ) )
323+ or
324+ // For generators/list-comprehensions we create a synthetic function.
325+ exists ( Comp c | c .getFunction ( ) = func )
323326 } or
324327 /** see QLDoc for `DataFlowModuleScope` for why we need this. */
325328 TModule ( Module m ) or
@@ -1381,6 +1384,8 @@ private predicate sameEnclosingCallable(Node node1, Node node2) {
13811384// =============================================================================
13821385newtype TDataFlowCall =
13831386 TNormalCall ( CallNode call , Function target , CallType type ) { resolveCall ( call , target , type ) } or
1387+ /** A call to the generated function inside a comprehension */
1388+ TComprehensionCall ( Comp c ) or
13841389 TPotentialLibraryCall ( CallNode call ) or
13851390 /** A synthesized call inside a summarized callable */
13861391 TSummaryCall (
@@ -1467,6 +1472,34 @@ class NormalCall extends ExtractedDataFlowCall, TNormalCall {
14671472 CallType getCallType ( ) { result = type }
14681473}
14691474
1475+ /** A call to the generated function inside a comprhension */
1476+ class ComprehensionCall extends ExtractedDataFlowCall , TComprehensionCall {
1477+ Comp c ;
1478+ Function target ;
1479+
1480+ ComprehensionCall ( ) {
1481+ this = TComprehensionCall ( c ) and
1482+ target = c .getFunction ( )
1483+ }
1484+
1485+ Comp getComprehension ( ) { result = c }
1486+
1487+ override string toString ( ) { result = "comprehension call" }
1488+
1489+ override ControlFlowNode getNode ( ) { result .getNode ( ) = c }
1490+
1491+ override Scope getScope ( ) { result = c .getScope ( ) }
1492+
1493+ override DataFlowCallable getCallable ( ) { result .( DataFlowFunction ) .getScope ( ) = target }
1494+
1495+ override ArgumentNode getArgument ( ArgumentPosition apos ) {
1496+ result .asExpr ( ) = c .getIterable ( ) and
1497+ apos .isPositional ( 0 )
1498+ }
1499+
1500+ override Location getLocation ( ) { result = c .getLocation ( ) }
1501+ }
1502+
14701503/**
14711504 * A potential call to a summarized callable, a `LibraryCallable`.
14721505 *
@@ -1697,6 +1730,47 @@ class CapturedVariablesArgumentNode extends CfgNode, ArgumentNode {
16971730 }
16981731}
16991732
1733+ /** A synthetic node representing the values of variables captured by a comprehension. */
1734+ class SynthCompCapturedVariablesArgumentNode extends Node , TSynthCompCapturedVariablesArgumentNode ,
1735+ ArgumentNode
1736+ {
1737+ Comp comp ;
1738+
1739+ SynthCompCapturedVariablesArgumentNode ( ) { this = TSynthCompCapturedVariablesArgumentNode ( comp ) }
1740+
1741+ override string toString ( ) { result = "Capturing closure argument (comp)" }
1742+
1743+ override Scope getScope ( ) { result = comp .getScope ( ) }
1744+
1745+ override Location getLocation ( ) { result = comp .getLocation ( ) }
1746+
1747+ Comp getComprehension ( ) { result = comp }
1748+
1749+ override predicate argumentOf ( DataFlowCall call , ArgumentPosition pos ) {
1750+ call .( ComprehensionCall ) .getComprehension ( ) = comp and
1751+ pos .isLambdaSelf ( )
1752+ }
1753+ }
1754+
1755+ /** A synthetic node representing the values of variables captured by a comprehension after the output has been computed. */
1756+ class SynthCompCapturedVariablesArgumentPostUpdateNode extends PostUpdateNodeImpl ,
1757+ TSynthCompCapturedVariablesArgumentPostUpdateNode
1758+ {
1759+ Comp comp ;
1760+
1761+ SynthCompCapturedVariablesArgumentPostUpdateNode ( ) {
1762+ this = TSynthCompCapturedVariablesArgumentPostUpdateNode ( comp )
1763+ }
1764+
1765+ override string toString ( ) { result = "[post] Capturing closure argument (comp)" }
1766+
1767+ override Scope getScope ( ) { result = comp .getScope ( ) }
1768+
1769+ override Location getLocation ( ) { result = comp .getLocation ( ) }
1770+
1771+ override Node getPreUpdateNode ( ) { result = TSynthCompCapturedVariablesArgumentNode ( comp ) }
1772+ }
1773+
17001774/** Gets a viable run-time target for the call `call`. */
17011775DataFlowCallable viableCallable ( DataFlowCall call ) {
17021776 call instanceof ExtractedDataFlowCall and
@@ -1736,7 +1810,10 @@ abstract class ReturnNode extends Node {
17361810/** A data flow node that represents a value returned by a callable. */
17371811class ExtractedReturnNode extends ReturnNode , CfgNode {
17381812 // See `TaintTrackingImplementation::returnFlowStep`
1739- ExtractedReturnNode ( ) { node = any ( Return ret ) .getValue ( ) .getAFlowNode ( ) }
1813+ ExtractedReturnNode ( ) {
1814+ node = any ( Return ret ) .getValue ( ) .getAFlowNode ( ) or
1815+ node = any ( Yield yield ) .getAFlowNode ( )
1816+ }
17401817
17411818 override ReturnKind getKind ( ) { any ( ) }
17421819}
0 commit comments