Rust: Cache Locatable.getLocation and Location

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

@@ -47,11 +47,9 @@ module Impl {
     ) {
       // TODO: handle locations in multi-line comments
       // TODO: handle the case where the template is from a nested macro call
-      Synth::convertFormatArgsExprFromRaw(parent)
-          .(FormatArgsExpr)
-          .getTemplate()
-          .getLocation()
-          .hasLocationInfo(file.getAbsolutePath(), startline, startcolumn - offset, _, _) and
+      LocatableImpl::getLocationDefault(Synth::convertFormatArgsExprFromRaw(parent)
+            .(FormatArgsExpr)
+            .getTemplate()).hasLocationFileInfo(file, startline, startcolumn - offset, _, _) and
       endline = startline and
       endcolumn = startcolumn + name.length() - 1
     }

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

@@ -5,6 +5,7 @@
  */
 
 private import codeql.rust.elements.internal.generated.Raw
+private import codeql.rust.internal.CachedStages
 
 /**
  *  The characteristic predicate of `Format` synthesized instances.
@@ -22,6 +23,7 @@ predicate constructFormat(Raw::FormatArgsExpr parent, int index, string text, in
  * Match an element of a format string, either text (`Hello`) or a format placeholder (`{}`).
  */
 string formatElement(Raw::FormatArgsExpr parent, int occurrenceIndex, int occurrenceOffset) {
+  Stages::AstStage::ref() and
   result =
     parent
         .getTemplate()

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

@@ -79,10 +79,8 @@ module Impl {
     override predicate hasSynthLocationInfo(
       File file, int startline, int startcolumn, int endline, int endcolumn
     ) {
-      this.getParent()
-          .getTemplate()
-          .getLocation()
-          .hasLocationInfo(file.getAbsolutePath(), startline, startcolumn - offset, _, _) and
+      LocatableImpl::getLocationDefault(this.getParent().getTemplate())
+          .hasLocationFileInfo(file, startline, startcolumn - offset, _, _) and
       endline = startline and
       endcolumn = startcolumn + text.length() - 1
     }

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

@@ -9,44 +9,45 @@ private import codeql.rust.elements.internal.LocationImpl
 private import codeql.rust.elements.internal.generated.Locatable
 private import codeql.rust.elements.internal.generated.Synth
 private import codeql.rust.elements.internal.generated.Raw
+private import codeql.rust.internal.CachedStages
 
 /**
  * INTERNAL: This module contains the customizable definition of `Locatable` and should not
  * be referenced directly.
  */
 module Impl {
   abstract class SynthLocatable extends Locatable {
+    pragma[nomagic]
     abstract predicate hasSynthLocationInfo(
       File file, int startline, int startcolumn, int endline, int endcolumn
     );
 
     final override Location getLocation() {
-      not locatable_locations(Synth::convertLocatableToRaw(this), _) and
-      exists(File file, int beginLine, int beginColumn, int endLine, int endColumn |
-        this.hasSynthLocationInfo(file, beginLine, beginColumn, endLine, endColumn)
-      |
-        result = LocationImpl::TLocationSynth(file, beginLine, beginColumn, endLine, endColumn)
-        or
-        exists(@location_default location |
-          result = LocationImpl::TLocationDefault(location) and
-          locations_default(location, file, beginLine, beginColumn, endLine, endColumn)
-        )
+      exists(File file, int startline, int startcolumn, int endline, int endcolumn |
+        this.hasSynthLocationInfo(file, startline, startcolumn, endline, endcolumn) and
+        result.hasLocationFileInfo(file, startline, startcolumn, endline, endcolumn)
       )
     }
   }
 
   class Locatable extends Generated::Locatable {
-    pragma[nomagic]
+    cached
     Location getLocation() {
-      exists(@location_default location |
-        result = LocationImpl::TLocationDefault(location) and
-        locatable_locations(Synth::convertLocatableToRaw(this), location)
-      )
+      Stages::AstStage::ref() and
+      result = getLocationDefault(this)
     }
 
     /**
      * Gets the primary file where this element occurs.
      */
     File getFile() { result = this.getLocation().getFile() }
   }
+
+  /** Gets the non-synthesized location of `l`, if any. */
+  LocationImpl::LocationDefault getLocationDefault(Locatable l) {
+    exists(@location_default location |
+      result = LocationImpl::TLocationDefault(location) and
+      locatable_locations(Synth::convertLocatableToRaw(l), location)
+    )
+  }
 }

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

@@ -3,10 +3,12 @@ private import codeql.rust.elements.internal.LocatableImpl::Impl as LocatableImp
 private import codeql.rust.elements.Locatable
 private import codeql.rust.elements.Format
 private import codeql.rust.elements.FormatArgument
+private import codeql.rust.internal.CachedStages
 
 module LocationImpl {
+  cached
   newtype TLocation =
-    TLocationDefault(@location_default location) or
+    TLocationDefault(@location_default location) { Stages::AstStage::ref() } or
     TLocationSynth(File file, int beginLine, int beginColumn, int endLine, int endColumn) {
       not locations_default(_, file, beginLine, beginColumn, endLine, endColumn) and
       any(LocatableImpl::SynthLocatable l)
@@ -55,10 +57,26 @@ module LocationImpl {
      * For more information, see
      * [Providing locations in CodeQL queries](https://codeql.github.com/docs/writing-codeql-queries/providing-locations-in-codeql-queries/).
      */
-    abstract predicate hasLocationInfo(
-      string filepath, int startline, int startcolumn, int endline, int endcolumn
+    abstract predicate hasLocationFileInfo(
+      File file, int startline, int startcolumn, int endline, int endcolumn
     );
 
+    /**
+     * Holds if this element is at the specified location.
+     * The location spans column `startcolumn` of line `startline` to
+     * column `endcolumn` of line `endline` in file `filepath`.
+     * For more information, see
+     * [Providing locations in CodeQL queries](https://codeql.github.com/docs/writing-codeql-queries/providing-locations-in-codeql-queries/).
+     */
+    final predicate hasLocationInfo(
+      string filepath, int startline, int startcolumn, int endline, int endcolumn
+    ) {
+      exists(File file |
+        this.hasLocationFileInfo(file, startline, startcolumn, endline, endcolumn) and
+        filepath = file.getAbsolutePath()
+      )
+    }
+
     /** Holds if this location starts strictly before the specified location. */
     pragma[inline]
     predicate strictlyBefore(Location other) {
@@ -68,18 +86,15 @@ module LocationImpl {
     }
   }
 
-  private class LocationDefault extends Location, TLocationDefault {
+  class LocationDefault extends Location, TLocationDefault {
     @location_default self;
 
     LocationDefault() { this = TLocationDefault(self) }
 
-    override predicate hasLocationInfo(
-      string filepath, int startline, int startcolumn, int endline, int endcolumn
+    override predicate hasLocationFileInfo(
+      File file, int startline, int startcolumn, int endline, int endcolumn
     ) {
-      exists(File f |
-        locations_default(self, f, startline, startcolumn, endline, endcolumn) and
-        filepath = f.getAbsolutePath()
-      )
+      locations_default(self, file, startline, startcolumn, endline, endcolumn)
     }
   }
 
@@ -88,13 +103,11 @@ module LocationImpl {
     EmptyLocation() { empty_location(self) }
   }
 
-  private class LocationSynth extends Location, TLocationSynth {
-    override predicate hasLocationInfo(
-      string filepath, int startline, int startcolumn, int endline, int endcolumn
+  class LocationSynth extends Location, TLocationSynth {
+    override predicate hasLocationFileInfo(
+      File file, int startline, int startcolumn, int endline, int endcolumn
     ) {
-      this =
-        TLocationSynth(any(File f | f.getAbsolutePath() = filepath), startline, startcolumn,
-          endline, endcolumn)
+      this = TLocationSynth(file, startline, startcolumn, endline, endcolumn)
     }
   }
 }

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

@@ -229,41 +229,41 @@ module Impl {
     name = v.getName() and
     (
       parameterDeclInScope(_, v, scope) and
-      scope.getLocation().hasLocationInfo(_, line, column, _, _)
+      scope.getLocation().hasLocationFileInfo(_, line, column, _, _)
       or
       exists(Pat pat | pat = getAVariablePatAncestor(v) |
         scope =
           any(MatchArmScope arm |
             arm.getPat() = pat and
-            arm.getLocation().hasLocationInfo(_, line, column, _, _)
+            arm.getLocation().hasLocationFileInfo(_, line, column, _, _)
           )
         or
         exists(LetStmt let |
           let.getPat() = pat and
           scope = getEnclosingScope(let) and
           // for `let` statements, variables are bound _after_ the statement, i.e.
           // not in the RHS
-          let.getLocation().hasLocationInfo(_, _, _, line, column)
+          let.getLocation().hasLocationFileInfo(_, _, _, line, column)
         )
         or
         exists(IfExpr ie, LetExpr let |
           let.getPat() = pat and
           ie.getCondition() = let and
           scope = ie.getThen() and
-          scope.getLocation().hasLocationInfo(_, line, column, _, _)
+          scope.getLocation().hasLocationFileInfo(_, line, column, _, _)
         )
         or
         exists(ForExpr fe |
           fe.getPat() = pat and
           scope = fe.getLoopBody() and
-          scope.getLocation().hasLocationInfo(_, line, column, _, _)
+          scope.getLocation().hasLocationFileInfo(_, line, column, _, _)
         )
         or
         exists(WhileExpr we, LetExpr let |
           let.getPat() = pat and
           we.getCondition() = let and
           scope = we.getLoopBody() and
-          scope.getLocation().hasLocationInfo(_, line, column, _, _)
+          scope.getLocation().hasLocationFileInfo(_, line, column, _, _)
         )
       )
     )
@@ -283,7 +283,7 @@ module Impl {
   ) {
     name = cand.getName() and
     scope = [cand.(VariableScope), getEnclosingScope(cand)] and
-    cand.getLocation().hasLocationInfo(_, startline, startcolumn, endline, endcolumn) and
+    cand.getLocation().hasLocationFileInfo(_, startline, startcolumn, endline, endcolumn) and
     nestLevel = 0
     or
     exists(VariableScope inner |
@@ -292,7 +292,7 @@ module Impl {
       // Use the location of the inner scope as the location of the access, instead of the
       // actual access location. This allows us to collapse multiple accesses in inner
       // scopes to a single entity
-      inner.getLocation().hasLocationInfo(_, startline, startcolumn, endline, endcolumn)
+      inner.getLocation().hasLocationFileInfo(_, startline, startcolumn, endline, endcolumn)
     )
   }
 

rust/ql/lib/codeql/rust/internal/CachedStages.qll

@@ -30,6 +30,39 @@ import rust
  * The `backref` predicate starts with `1 = 1 or` to ensure that the predicate will be optimized down to a constant by the optimizer.
  */
 module Stages {
+  /**
+   * The abstract syntex tree (AST) stage.
+   */
+  cached
+  module AstStage {
+    private import codeql.rust.controlflow.internal.Splitting
+    private import codeql.rust.controlflow.internal.SuccessorType
+    private import codeql.rust.controlflow.internal.ControlFlowGraphImpl
+
+    /**
+     * Always holds.
+     * Ensures that a predicate is evaluated as part of the AST stage.
+     */
+    cached
+    predicate ref() { 1 = 1 }
+
+    /**
+     * DO NOT USE!
+     *
+     * Contains references to each predicate that use the above `ref` predicate.
+     */
+    cached
+    predicate backref() {
+      1 = 1
+      or
+      exists(Location l)
+      or
+      exists(any(Locatable l).getLocation())
+      or
+      exists(Format f)
+    }
+  }
+
   /**
    * The control flow graph (CFG) stage.
    */
@@ -41,7 +74,7 @@ module Stages {
 
     /**
      * Always holds.
-     * Ensures that a predicate is evaluated as part of the BasicBlocks stage.
+     * Ensures that a predicate is evaluated as part of the CFG stage.
      */
     cached
     predicate ref() { 1 = 1 }