Rust: Cache defaultAdditionalTaintStep

rust/ql/lib/codeql/rust/dataflow/internal/DataFlowImpl.qll

@@ -1101,9 +1101,11 @@ import MakeImpl<Location, RustDataFlow>
 /** A collection of cached types and predicates to be evaluated in the same stage. */
 cached
 private module Cached {
+  private import codeql.rust.internal.CachedStages
+
   cached
   newtype TNode =
-    TExprNode(ExprCfgNode n) or
+    TExprNode(ExprCfgNode n) { Stages::DataFlowStage::ref() } or
     TSourceParameterNode(ParamBaseCfgNode p) or
     TPatNode(PatCfgNode p) or
     TExprPostUpdateNode(ExprCfgNode e) {

rust/ql/lib/codeql/rust/dataflow/internal/TaintTrackingImpl.qll

@@ -4,6 +4,7 @@ private import codeql.rust.controlflow.CfgNodes
 private import codeql.rust.dataflow.FlowSummary
 private import DataFlowImpl
 private import FlowSummaryImpl as FlowSummaryImpl
+private import codeql.rust.internal.CachedStages
 
 module RustTaintTracking implements InputSig<Location, RustDataFlow> {
   predicate defaultTaintSanitizer(Node::Node node) { none() }
@@ -12,7 +13,9 @@ module RustTaintTracking implements InputSig<Location, RustDataFlow> {
    * Holds if the additional step from `pred` to `succ` should be included in all
    * global taint flow configurations.
    */
+  cached
   predicate defaultAdditionalTaintStep(Node::Node pred, Node::Node succ, string model) {
+    Stages::DataFlowStage::ref() and
     model = "" and
     (
       exists(BinaryExprCfgNode binary |

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

@@ -22,7 +22,7 @@ module Impl {
    * ```
    */
   class ArrayExpr extends Generated::ArrayExpr {
-    cached
+    pragma[nomagic]
     private Raw::ArrayExprInternal getUnderlyingEntity() {
       this = Synth::TArrayListExpr(result) or this = Synth::TArrayRepeatExpr(result)
     }

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

@@ -35,10 +35,6 @@ module Stages {
    */
   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.
@@ -98,4 +94,34 @@ module Stages {
       exists(CallExprCfgNode n | exists(n.getFunction()))
     }
   }
+
+  /**
+   * The data flow stage.
+   */
+  cached
+  module DataFlowStage {
+    private import codeql.rust.dataflow.internal.DataFlowImpl
+    private import codeql.rust.dataflow.internal.TaintTrackingImpl
+
+    /**
+     * Always holds.
+     * Ensures that a predicate is evaluated as part of the data flow 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(Node n)
+      or
+      RustTaintTracking::defaultAdditionalTaintStep(_, _, _)
+    }
+  }
 }