Apply suggestions from code review

yoff · RasmusWL · web-flow · commit c395d2d9573c · 2023-12-15T09:58:27.000+01:00
Co-authored-by: Rasmus Wriedt Larsen &lt;rasmuswriedtlarsen@gmail.com&gt;
diff --git a/python/ql/lib/semmle/python/dataflow/new/internal/DataFlowDispatch.qll b/python/ql/lib/semmle/python/dataflow/new/internal/DataFlowDispatch.qll
@@ -1531,7 +1531,11 @@ abstract class ParameterNodeImpl extends Node {
 class LambdaSelfReferenceNode extends ParameterNodeImpl, TLambdaSelfReferenceNode {
   private Function callable;
 
-  LambdaSelfReferenceNode() { this = TLambdaSelfReferenceNode(callable) }
+  LambdaSelfReferenceNode() { 
+    this = TLambdaSelfReferenceNode(callable) and
+    // TODO: Remove this restriction when adding proper support for captured variables in the body of the function we generate for comprehensions
+    exists(TFunction(callable))
+  }
 
   final Function getCallable() { result = callable }
 
diff --git a/python/ql/lib/semmle/python/dataflow/new/internal/DataFlowPrivate.qll b/python/ql/lib/semmle/python/dataflow/new/internal/DataFlowPrivate.qll
@@ -392,7 +392,7 @@ module VariableCapture {
     ExprCfgNode() { isExpressionNode(this) }
   }
 
-  private predicate closureFlowStep(ExprCfgNode e1, ExprCfgNode e2) {
+  private predicate closureFlowStep(ExprCfgNode fromNode, ExprCfgNode toNode) {
     // simpleAstFlowStep(e1, e2)
     // or
     exists(SsaVariable def |
@@ -689,7 +689,7 @@ predicate summaryFlowSteps(Node nodeFrom, Node nodeTo) {
   IncludePostUpdateFlow<PhaseDependentFlow<summaryLocalStep/2>::step/2>::step(nodeFrom, nodeTo)
 }
 
-predicate variableCaptureFlowStep(Node nodeFrom, Node nodeTo) {
+predicate variableCaptureLocalFlowStep(Node nodeFrom, Node nodeTo) {
   // Blindly applying use-use flow can result in a node that steps to itself, for
   // example in while-loops. To uphold dataflow consistency checks, we don't want
   // that. However, we do want to allow `[post] n` to `n` (to handle while loops), so
diff --git a/python/ql/lib/semmle/python/dataflow/new/internal/DataFlowPublic.qll b/python/ql/lib/semmle/python/dataflow/new/internal/DataFlowPublic.qll
@@ -116,7 +116,7 @@ newtype TNode =
     exists(ParameterPosition ppos | ppos.isKeyword(_) | exists(callable.getParameter(ppos)))
   } or
   /** A synthetic node representing a captured variable. */
-  TCaptureNode(VariableCapture::Flow::SynthesizedCaptureNode cn) or
+  TSynthCaptureNode(VariableCapture::Flow::SynthesizedCaptureNode cn) or
   /** A synthetic node representing the heap of a function. Used for variable capture. */
   TLambdaSelfReferenceNode(Function f) {
     f = any(VariableCapture::CapturedVariable v).getACapturingScope()
@@ -489,10 +489,8 @@ class CaptureNode extends Node, TCaptureNode {
   /** Gets the `SynthesizedCaptureNode` that this node represents. */
   VariableCapture::Flow::SynthesizedCaptureNode getSynthesizedCaptureNode() { result = cn }
 
-  override DataFlowCallable getEnclosingCallable() {
-    result = TFunction(cn.getEnclosingCallable())
-    or
-    result = TModule(cn.getEnclosingCallable())
+  override Scope getScope() {
+    result = cn.getEnclosingCallable()
   }
 
   override Location getLocation() { result = cn.getLocation() }
