cached stages iteration 3

python/ql/lib/semmle/python/AstExtended.qll

@@ -37,6 +37,7 @@ abstract class AstNode extends AstNode_ {
    * Expr.getASubExpression(), Stmt.getASubStatement(), Stmt.getASubExpression() or
    * Scope.getAStmt().
    */
+  cached
   abstract AstNode getAChildNode();
 
   /**
@@ -113,7 +114,10 @@ class Comprehension extends Comprehension_, AstNode {
     result = Comprehension_.super.getLocation()
   }
 
-  override AstNode getAChildNode() { result = this.getASubExpression() }
+  override AstNode getAChildNode() {
+    Stages::SSA::ref() and
+    result = this.getASubExpression()
+  }
 
   Expr getASubExpression() {
     result = this.getIter() or

python/ql/lib/semmle/python/filters/GeneratedCode.qll

@@ -32,18 +32,26 @@ class GenericGeneratedFile extends GeneratedFile {
   override string getTool() { lax_generated_by(this, result) or strict_generated_by(this, result) }
 }
 
+pragma[nomagic]
+int minStmtLine(File file) {
+  result =
+    min(int line |
+      line = any(Stmt s | s.getLocation().getFile() = file).getLocation().getStartLine()
+    )
+}
+
+pragma[nomagic]
+predicate isBeforeCode(Comment c, File f) {
+  f = c.getLocation().getFile() and
+  minStmtLine(f) < c.getLocation().getStartLine()
+}
+
 private string comment_or_docstring(File f, boolean before_code) {
   exists(Comment c |
     c.getLocation().getFile() = f and
     result = c.getText()
   |
-    if
-      exists(Stmt s |
-        s.getEnclosingModule().getFile() = f and
-        s.getLocation().getStartLine() < c.getLocation().getStartLine()
-      )
-    then before_code = false
-    else before_code = true
+    if isBeforeCode(c, f) then before_code = true else before_code = false
   )
   or
   exists(Module m | m.getFile() = f |

python/ql/lib/semmle/python/internal/CachedStages.qll

@@ -73,6 +73,8 @@ module Stages {
       or
       exists(any(AstExtended::AstNode n).getLocation())
       or
+      exists(any(AstExtended::AstNode n).getAChildNode())
+      or
       exists(any(PyFlow::BasicBlock b).getImmediateDominator())
       or
       any(PyFlow::BasicBlock b).strictlyDominates(_)