Rust: Introduce AstNode.getEnclosingCallable

rust/ql/lib/codeql/rust/controlflow/internal/ControlFlowGraphImpl.qll

@@ -21,7 +21,7 @@ private module CfgInput implements InputSig<Location> {
   /** An AST node with an associated control-flow graph. */
   class CfgScope = Scope::CfgScope;
 
-  CfgScope getCfgScope(AstNode n) { result = Scope::scopeOfAst(n) }
+  CfgScope getCfgScope(AstNode n) { result = n.getEnclosingCallable() }
 
   class SuccessorType = Cfg::SuccessorType;
 

rust/ql/lib/codeql/rust/controlflow/internal/Scope.qll

@@ -24,28 +24,3 @@ class CfgScope extends Callable {
     result = this.(ClosureExpr).getBody()
   }
 }
-
-/**
- * Gets the immediate parent of a non-`AstNode` element `e`.
- *
- * We restrict `e` to be a non-`AstNode` to skip past non-`AstNode` in
- * the transitive closure computation in `getParentOfAst`. This is
- * necessary because the parent of an `AstNode` is not necessarily an `AstNode`.
- */
-private Element getParentOfAstStep(Element e) {
-  not e instanceof AstNode and
-  result = getImmediateParent(e)
-}
-
-/** Gets the nearest enclosing parent of `ast` that is an `AstNode`. */
-private AstNode getParentOfAst(AstNode ast) {
-  result = getParentOfAstStep*(getImmediateParent(ast))
-}
-
-/** Gets the enclosing scope of a node */
-cached
-AstNode scopeOfAst(AstNode n) {
-  exists(AstNode p | p = getParentOfAst(n) |
-    if p instanceof CfgScope then p = result else result = scopeOfAst(p)
-  )
-}

rust/ql/lib/codeql/rust/elements/internal/AstNodeImpl.qll

@@ -1,4 +1,3 @@
-// generated by codegen, remove this comment if you wish to edit this file
 /**
  * This module provides a hand-modifiable wrapper around the generated class `AstNode`.
  *
@@ -12,5 +11,36 @@ private import codeql.rust.elements.internal.generated.AstNode
  * be referenced directly.
  */
 module Impl {
-  class AstNode extends Generated::AstNode { }
+  private import rust
+  private import codeql.rust.elements.internal.generated.ParentChild
+
+  /**
+   * Gets the immediate parent of a non-`AstNode` element `e`.
+   *
+   * We restrict `e` to be a non-`AstNode` to skip past non-`AstNode` in
+   * the transitive closure computation in `getParentOfAst`. This is
+   * necessary because the parent of an `AstNode` is not necessarily an `AstNode`.
+   */
+  private Element getParentOfAstStep(Element e) {
+    not e instanceof AstNode and
+    result = getImmediateParent(e)
+  }
+
+  /** Gets the nearest enclosing parent of `ast` that is an `AstNode`. */
+  private AstNode getParentOfAst(AstNode ast) {
+    result = getParentOfAstStep*(getImmediateParent(ast))
+  }
+
+  class AstNode extends Generated::AstNode {
+    /** Gets the immediately enclosing callable of this node, if any. */
+    cached
+    Callable getEnclosingCallable() {
+      exists(AstNode p | p = getParentOfAst(this) |
+        result = p
+        or
+        not p instanceof Callable and
+        result = p.getEnclosingCallable()
+      )
+    }
+  }
 }

rust/ql/lib/codeql/rust/elements/internal/VariableImpl.qll

@@ -421,8 +421,6 @@ module Impl {
     )
   }
 
-  private import codeql.rust.controlflow.internal.Scope
-
   /** A variable access. */
   class VariableAccess extends PathExprImpl::PathExpr instanceof VariableAccessCand {
     private string name;
@@ -434,7 +432,7 @@ module Impl {
     Variable getVariable() { result = v }
 
     /** Holds if this access is a capture. */
-    predicate isCapture() { scopeOfAst(this) != scopeOfAst(v.getPat()) }
+    predicate isCapture() { this.getEnclosingCallable() != v.getPat().getEnclosingCallable() }
 
     override string toString() { result = name }