C++: Speed up Ssa::getAnUltimateDefinition

cpp/ql/lib/semmle/code/cpp/dataflow/internal/FlowSummaryImpl.qll

@@ -171,6 +171,28 @@ private module Input2 implements Impl::Private::InputSig2 {
     result = e.(ConversionCall).getQualifier().(LambdaExpression).getLambdaFunction()
   }
 
+  private predicate isRelevantUltimateDefinition(Ssa::DirectExplicitDefinition def, Function f) {
+    f =
+      getFunctionFromExpr(def.getAssignedInstruction()
+            .(StoreInstruction)
+            .getSourceValue()
+            .getUnconvertedResultExpression())
+  }
+
+  private module GetAnUltimateDefinitionInput implements Ssa::GetAnUltimateDefinitionSig {
+    predicate isRelevantUltimateDefinition(Ssa::Definition def) {
+      isRelevantUltimateDefinition(def, _)
+    }
+  }
+
+  private predicate hasAnUltimateFunctionAccessDefinition(Ssa::Definition def, Function f) {
+    exists(Ssa::Definition ultimate |
+      ultimate =
+        Ssa::GetAnUltimateDefinition<GetAnUltimateDefinitionInput>::getAnUltimateDefinition(def) and
+      isRelevantUltimateDefinition(ultimate, f)
+    )
+  }
+
   class SourceSinkReportingElement extends Element {
     SourceSinkReportingElement() { this instanceof Expr or this instanceof Parameter }
 
@@ -190,13 +212,7 @@ private module Input2 implements Impl::Private::InputSig2 {
       // The expression is an SSA read of an assignment of a callable
       exists(Ssa::Definition def |
         def.getAUse().getDef().getUnconvertedResultExpression() = this and
-        result =
-          getFunctionFromExpr(def.getAnUltimateDefinition()
-                .(Ssa::DirectExplicitDefinition)
-                .getAssignedInstruction()
-                .(StoreInstruction)
-                .getSourceValue()
-                .getUnconvertedResultExpression())
+        hasAnUltimateFunctionAccessDefinition(def, result)
       )
     }
 

cpp/ql/lib/semmle/code/cpp/ir/dataflow/internal/DataFlowPrivate.qll

@@ -1977,13 +1977,23 @@ module IteratorFlow {
     }
 
     /**
-     * Gets an ultimate definition of `def`.
-     *
-     * Note: Unlike `def.getAnUltimateDefinition()` this predicate also
-     * traverses back through iterator increment and decrement operations.
+     * Holds if `write` is an instruction that writes to address `address`
      */
-    private Ssa::Definition getAnUltimateDefinition(Ssa::Definition def) {
-      result = def.getAnUltimateDefinition()
+    private predicate isIteratorWrite(Instruction write, Operand address) {
+      exists(Ssa::DefImpl writeDef, IRBlock bb, int i |
+        writeDef.hasIndexInBlock(_, bb, i) and
+        bb.getInstruction(i) = write and
+        address = writeDef.getAddressOperand()
+      )
+    }
+
+    private module GetAnUltimateDefinitionInput implements Ssa::GetAnUltimateDefinitionSig {
+      predicate isRelevantUltimateDefinition(Ssa::Definition def) { fwd(_, def) }
+    }
+
+    private Ssa::Definition getAnUltimateDefinitionStep(Ssa::Definition def) {
+      result =
+        Ssa::GetAnUltimateDefinition<GetAnUltimateDefinitionInput>::getAnUltimateDefinition(def)
       or
       exists(IRBlock bb, int i, IteratorCrementCall crementCall, Ssa::SourceVariable sv |
         crementCall = def.getValue().asInstruction().(StoreInstruction).getSourceValue() and
@@ -1993,14 +2003,28 @@ module IteratorFlow {
       )
     }
 
-    /**
-     * Holds if `write` is an instruction that writes to address `address`
-     */
-    private predicate isIteratorWrite(Instruction write, Operand address) {
-      exists(Ssa::DefImpl writeDef, IRBlock bb, int i |
-        writeDef.hasIndexInBlock(_, bb, i) and
-        bb.getInstruction(i) = write and
-        address = writeDef.getAddressOperand()
+    private predicate isSource(GetsIteratorCall beginCall, Ssa::Definition def) {
+      exists(StoreInstruction beginStore |
+        beginStore = def.getValue().asInstruction() and
+        operandForFullyConvertedCall(beginStore.getSourceValueOperand(), beginCall)
+      )
+    }
+
+    private predicate isSink(Instruction writeToDeref, Ssa::Definition def) {
+      exists(IteratorPointerDereferenceCall starCall, Operand address, IRBlock bbStar, int iStar |
+        isIteratorWrite(writeToDeref, address) and
+        operandForFullyConvertedCall(address, starCall) and
+        bbStar.getInstruction(iStar) = starCall and
+        Ssa::ssaDefReachesRead(_, def, bbStar, iStar)
+      )
+    }
+
+    private predicate fwd(GetsIteratorCall beginCall, Ssa::Definition def) {
+      isSource(beginCall, def)
+      or
+      exists(Ssa::Definition def0 |
+        fwd(beginCall, def0) and
+        def0 = getAnUltimateDefinitionStep(def)
       )
     }
 
@@ -2016,17 +2040,9 @@ module IteratorFlow {
     private predicate isIteratorStoreInstruction(
       GetsIteratorCall beginCall, Instruction writeToDeref
     ) {
-      exists(
-        StoreInstruction beginStore, IRBlock bbStar, int iStar, Ssa::Definition def,
-        IteratorPointerDereferenceCall starCall, Ssa::Definition ultimate, Operand address
-      |
-        isIteratorWrite(writeToDeref, address) and
-        operandForFullyConvertedCall(address, starCall) and
-        bbStar.getInstruction(iStar) = starCall and
-        Ssa::ssaDefReachesRead(_, def, bbStar, iStar) and
-        ultimate = getAnUltimateDefinition*(def) and
-        beginStore = ultimate.getValue().asInstruction() and
-        operandForFullyConvertedCall(beginStore.getSourceValueOperand(), beginCall)
+      exists(Ssa::Definition def |
+        fwd(beginCall, def) and
+        isSink(writeToDeref, def)
       )
     }
 

cpp/ql/lib/semmle/code/cpp/ir/dataflow/internal/DataFlowUtil.qll

@@ -1018,4 +1018,6 @@ module Ssa {
   class IndirectExplicitDefinition = SsaImpl::IndirectExplicitDefinition;
 
   class PhiNode = SsaImpl::PhiNode;
+
+  import SsaImpl::Public
 }

cpp/ql/lib/semmle/code/cpp/ir/dataflow/internal/SsaImpl.qll

@@ -1405,14 +1405,80 @@ private class PhiCycle extends PhiCycleEquivalence::EquivalenceClass {
   }
 }
 
-/** An static single assignment (SSA) definition. */
-class Definition extends SsaImpl::Definition {
-  private Definition getAPhiInputOrPriorDefinition() {
-    result = this.(PhiNode).getAnInput()
-    or
-    uncertainWriteDefinitionInput(this, result)
+private Definition getAPhiInputOrPriorDefinition(Definition def) {
+  result = def.(PhiNode).getAnInput()
+  or
+  uncertainWriteDefinitionInput(def, result)
+}
+
+module Public {
+  /**
+   * A module signature to define relevant ultimate definitions for an
+   * optimized version of `Definition.getAnUltimateDefinition`.
+   */
+  signature module GetAnUltimateDefinitionSig {
+    /**
+     * Holds if `def` is a relevant definition. This defines the set
+     * of `Definition`s which may be returned by
+     * `GetAnUltimateDefinition::getAnUltimateDefinition`.
+     */
+    predicate isRelevantUltimateDefinition(Definition def);
+  }
+
+  /**
+   * A module which constructs an optimized version of
+   * ```
+   * Definition.getAnUltimateDefinition
+   * ```
+   * by restricting the set of possible ultimate definitions.
+   *
+   * Use this module by defining a module `M` which implements
+   * `GetAnUltimateDefinitionSig` and then call:
+   * ```
+   * GetAnUltimateDefinition<M>::getAnUltimateDefinition
+   * ```
+   */
+  module GetAnUltimateDefinition<GetAnUltimateDefinitionSig Sig> {
+    private import Sig
+
+    private predicate relevantUltimateDefinition(Definition def) {
+      isRelevantUltimateDefinition(def) and
+      not def instanceof PhiNode
+    }
+
+    private predicate fwd(Definition def) {
+      relevantUltimateDefinition(def)
+      or
+      exists(Definition def0 |
+        fwd(def0) and
+        def0 = getAPhiInputOrPriorDefinition(def)
+      )
+    }
+
+    private predicate step(Definition def1, Definition def2) {
+      fwd(def1) and
+      fwd(def2) and
+      def1 = getAPhiInputOrPriorDefinition(def2)
+    }
+
+    /**
+     * Gets a definition that ultimately defines this SSA definition and is
+     * not itself a phi node.
+     *
+     * This predicate is restricted to ultimate definitions which
+     * satisfy `isRelevantUltimateDefinition`.
+     */
+    Definition getAnUltimateDefinition(Definition def) {
+      step*(result, def) and
+      relevantUltimateDefinition(result)
+    }
   }
+}
+
+import Public
 
+/** A static single assignment (SSA) definition. */
+class Definition extends SsaImpl::Definition {
   /**
    * Holds if this SSA definition is live at the end of basic block `bb`.
    * That is, this definition reaches the end of basic block `bb`, at which
@@ -1424,9 +1490,13 @@ class Definition extends SsaImpl::Definition {
   /**
    * Gets a definition that ultimately defines this SSA definition and is
    * not itself a phi node.
+   *
+   * Note: A more efficient implementation of this predicate exists. See the
+   * `GetAnUltimateDefinition` module for a description of how to access
+   * the more efficient implementation.
    */
   final Definition getAnUltimateDefinition() {
-    result = this.getAPhiInputOrPriorDefinition*() and
+    result = getAPhiInputOrPriorDefinition*(this) and
     not result instanceof PhiNode
   }