Merge pull request github#17829 from hvitved/rust/cfg-stage

Rust: Collapse cached CFG logic into one stage

Author: hvitved

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

@@ -5,6 +5,7 @@ private import Scope as Scope
 private import codeql.rust.controlflow.ControlFlowGraph as Cfg
 
 private module CfgInput implements InputSig<Location> {
+  private import codeql.rust.internal.CachedStages
   private import rust as Rust
   private import Completion as C
 
@@ -21,7 +22,10 @@ private module CfgInput implements InputSig<Location> {
   /** An AST node with an associated control-flow graph. */
   class CfgScope = Scope::CfgScope;
 
-  CfgScope getCfgScope(AstNode n) { result = n.getEnclosingCallable() }
+  CfgScope getCfgScope(AstNode n) {
+    result = n.getEnclosingCallable() and
+    Stages::CfgStage::ref()
+  }
 
   class SuccessorType = Cfg::SuccessorType;
 

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

@@ -4,8 +4,10 @@ private import Scope
 
 cached
 private module Cached {
+  private import codeql.rust.internal.CachedStages
+
   cached
-  newtype TSplitKind = TConditionalCompletionSplitKind()
+  newtype TSplitKind = TConditionalCompletionSplitKind() { Stages::CfgStage::ref() }
 
   cached
   newtype TSplit = TConditionalCompletionSplit(ConditionalCompletion c)

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

@@ -1,10 +1,11 @@
 private import rust
 private import codeql.util.Boolean
 private import Completion
+private import codeql.rust.internal.CachedStages
 
 cached
 newtype TSuccessorType =
-  TSuccessorSuccessor() or
+  TNormalSuccessor() { Stages::CfgStage::ref() } or
   TBooleanSuccessor(Boolean b) or
   TMatchSuccessor(Boolean b) or
   TBreakSuccessor() or
@@ -18,7 +19,7 @@ abstract class SuccessorTypeImpl extends TSuccessorType {
 }
 
 /** A normal control flow successor. */
-class NormalSuccessorImpl extends SuccessorTypeImpl, TSuccessorSuccessor {
+class NormalSuccessorImpl extends SuccessorTypeImpl, TNormalSuccessor {
   override string toString() { result = "successor" }
 }
 

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

@@ -0,0 +1,65 @@
+/**
+ * The purpose of this file is to control which cached predicates belong to the same stage.
+ *
+ * Combining stages can improve performance as we are more likely to reuse shared, non-cached predicates.
+ *
+ * To make a predicate `p` belong to a stage `A`:
+ * - make `p` depend on `A::ref()`, and
+ * - make `A::backref()` depend on `p`.
+ *
+ * Since `A` is a cached module, `ref` and `backref` must be in the same stage, and the dependency
+ * chain above thus forces `p` to be in that stage as well.
+ *
+ * With these two predicates in a `cached module` we ensure that all the cached predicates will be in a single stage at runtime.
+ *
+ * Grouping stages can cause unnecessary computation, as a concrete query might not depend on
+ * all the cached predicates in a stage.
+ * Care should therefore be taken not to combine two stages, if it is likely that a query only depend
+ * on some but not all the cached predicates in the combined stage.
+ */
+
+import rust
+
+/**
+ * Contains a `cached module` for each stage.
+ * Each `cached module` ensures that predicates that are supposed to be in the same stage, are in the same stage.
+ *
+ * Each `cached module` contain two predicates:
+ * The first, `ref`, always holds, and is referenced from `cached` predicates.
+ * The second, `backref`, contains references to the same `cached` predicates.
+ * 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 control flow graph (CFG) stage.
+   */
+  cached
+  module CfgStage {
+    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 BasicBlocks 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(TConditionalCompletionSplitKind())
+      or
+      exists(TNormalSuccessor())
+      or
+      exists(AstCfgNode n)
+    }
+  }
+}