Unified: Add caching to name binding

unified/ql/lib/codeql/unified/internal/LocalNameBinding.qll

@@ -9,6 +9,12 @@ private import codeql.unified.internal.NameBindingPlugin
 private import codeql.unified.internal.StaticNameBinding
 
 private module LocalNameBindingInput implements LocalNameBindingInputSig<Location> {
+  predicate cacheRevRef() {
+    (bindingContext(_, _, _) implies any())
+    or
+    (implicitDeclInScope(_, _, _) implies any())
+  }
+
   class AstNode = U::AstNode;
 
   private class LogicalAndRoot extends LogicalAndExpr {
@@ -201,7 +207,9 @@ private module LocalNameBindingInput implements LocalNameBindingInputSig<Locatio
     any(NameBindingPlugin p).isNonPattern(e)
   }
 
+  cached
   additional predicate bindingContext(AstNode pattern, AstNode scope, AstNode declaration) {
+    LocalNameBindingOutput::CachedStage::ref() and
     not isNonPattern(pattern) and
     (
       exists(SiblingShadowingDecl decl |
@@ -326,8 +334,9 @@ private module LocalNameBindingInput implements LocalNameBindingInputSig<Locatio
     )
   }
 
-  pragma[nomagic]
+  cached
   additional predicate implicitDeclInScope(string name, AstNode scope, boolean isLocalVariable) {
+    LocalNameBindingOutput::CachedStage::ref() and
     exists(Callable callable |
       isLocalVariable = true and
       name = any(NameBindingPlugin p).getImplicitReceiverParameterName(callable) and
@@ -356,7 +365,7 @@ private module LocalNameBindingInput implements LocalNameBindingInputSig<Locatio
   }
 }
 
-import LocalNameBindingInput
+predicate bindingContext = LocalNameBindingInput::bindingContext/3;
 
 module LocalNameBindingOutput = LocalNameBinding<Location, LocalNameBindingInput>;
 
@@ -383,10 +392,10 @@ module Public {
 
   /** An identifier appearing in a name-binding position, such as the `x` in `let x = 123`. */
   class NameBinding extends Identifier {
-    NameBinding() { LocalNameBindingInput::bindingContext(this, _, _) }
+    NameBinding() { bindingContext(this, _, _) }
 
     /** Gets the statement-like node declaring this name, such as a `VariableDeclaration` or `CatchClause`. */
-    AstNode getDeclaration() { LocalNameBindingInput::bindingContext(this, _, result) }
+    AstNode getDeclaration() { bindingContext(this, _, result) }
 
     /** Gets the name being declared. */
     string getName() { result = this.getValue() }

unified/ql/lib/codeql/unified/internal/StaticNameBinding.qll

@@ -6,8 +6,9 @@ private import unified
 private import codeql.unified.internal.LocalNameBinding
 private import codeql.unified.internal.NameBindingPlugin
 
+cached
 private newtype TNameBindingNode =
-  TIdentifier(Identifier n) or
+  TIdentifier(Identifier n) { CachedStage::ref() } or
   TBulkImport(BulkImportingPattern p) or
   TLocalName(LocalName local) or
   TStaticMemberNamespace(ClassLikeDeclaration cls) or
@@ -357,10 +358,15 @@ private predicate isInheritableMemberNode(NameBindingNode node) {
 
 /** A name-binding node that can have members. */
 class NamespaceNode extends NameBindingNode {
+  cached
   NamespaceNode() {
-    storeStep(_, _, this) or
-    inheritanceStep(_, this) or
-    this.isInstanceMemberNamespace(_) or
+    CachedStage::ref() and
+    storeStep(_, _, this)
+    or
+    inheritanceStep(_, this)
+    or
+    this.isInstanceMemberNamespace(_)
+    or
     this.isStaticMemberNamespace(_)
   }
 
@@ -410,8 +416,9 @@ class NamespaceNode extends NameBindingNode {
   NamespaceNode getAnInheritanceChild() { result.getAnInheritanceParent() = this }
 
   /** Gets a member of this namespace of the given name. */
-  pragma[nomagic]
+  cached
   NameBindingNode getMember(string name) {
+    CachedStage::ref() and
     result = this.getOwnMember(name)
     or
     not this.hasOwnMember(name) and
@@ -468,7 +475,7 @@ private module TrackNameBindingInput implements TrackInputSig {
 private module TrackNameBinding = Track<TrackNameBindingInput>;
 
 /** Gets a name-binding node that may refer to the given declaration. */
-NameBindingNode trackNameBinding(NameBinding decl) {
+private NameBindingNode trackNameBinding(NameBinding decl) {
   exists(NameBindingNode start |
     start.isIdentifier(decl) and
     result = TrackNameBinding::track(start)
@@ -677,10 +684,12 @@ private int unqualifiedMemberAccessDepth(PotentialLocalNameAccess access) {
  * `accessingClass` is the enclosing class in which the member was found, and
  * `instanceAccess` indicates if it is an instance member or static member.
  */
+cached
 predicate unqualifiedMemberAccess(
   PotentialLocalNameAccess access, boolean instanceAccess, NameBinding target,
   ClassLikeDeclaration accessingClass
 ) {
+  CachedStage::ref() and
   unqualifiedMemberAccessCand(access, instanceAccess, target, accessingClass) and
   accessingClass.getDepth() = unqualifiedMemberAccessDepth(access)
 }
@@ -708,6 +717,7 @@ module Public {
     predicate isInstanceAccess() { instanceAccess = true }
 
     /** Gets the local variable implicitly referenced as the base of this access. */
+    cached
     LocalVariable getImplicitQualifierVariable() {
       ResolveImplicitReceiverAccess::access(this, result)
     }
@@ -718,7 +728,9 @@ module Public {
 }
 
 /** Gets the declaration being accessed by identifier `i`, as determined by static name binding. */
+cached
 NameBinding getStaticBindingTargetFromIdentifier(Identifier i) {
+  CachedStage::ref() and
   // For unqualified accesses, use the shadowing-aware lookup
   result = i.(UnqualifiedMemberAccess).getTarget()
   or
@@ -758,3 +770,26 @@ private predicate implicitReceiverAccess(AstNode access, string name) {
 
 private module ResolveImplicitReceiverAccess =
   LocalNameBindingOutput::ResolveAccesses<implicitReceiverAccess/2>;
+
+cached
+private module CachedStage {
+  /** Reference to the cached stage of this module. */
+  cached
+  predicate ref() { any() }
+
+  /** Reverse references to the predicates that reference `ref()`. */
+  cached
+  predicate revRef() {
+    any()
+    or
+    (exists(NameBindingNode n) implies any())
+    or
+    (exists(any(NamespaceNode n).getMember(_)) implies any())
+    or
+    (unqualifiedMemberAccess(_, _, _, _) implies any())
+    or
+    (exists(any(UnqualifiedMemberAccess u).getImplicitQualifierVariable()) implies any())
+    or
+    (exists(getStaticBindingTargetFromIdentifier(_)) implies any())
+  }
+}