Rust: Evaluate mayInvokeCallback in type inference stage

Author: hvitved

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

@@ -58,6 +58,7 @@ private import codeql.rust.dataflow.FlowBarrier
 private import codeql.rust.dataflow.FlowSummary
 private import codeql.rust.dataflow.FlowSource
 private import codeql.rust.dataflow.FlowSink
+private import codeql.rust.internal.CachedStages
 private import codeql.rust.internal.typeinference.FunctionType
 private import codeql.rust.internal.typeinference.TypeMention
 private import codeql.rust.frameworks.stdlib.Stdlib
@@ -215,7 +216,9 @@ private class SummarizedCallableFromModel extends SummarizedCallable::Range {
  * add a flow model that achieves the effect of simulating that the callback is
  * invoked, which is needed for flow through captured variables to work.
  */
-private predicate mayInvokeCallback(Function f, int n) {
+cached
+predicate mayInvokeCallback(Function f, int n) {
+  Stages::TypeInferenceStage::ref() and
   exists(TypeMention tm, Trait trait |
     tm = f.getParam(n).getTypeRepr() and
     trait = getALookupTrait(f, tm.getType()) and

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

@@ -130,6 +130,7 @@ module Stages {
   module TypeInferenceStage {
     private import codeql.rust.internal.typeinference.Type
     private import codeql.rust.internal.typeinference.TypeInference
+    private import codeql.rust.dataflow.internal.ModelsAsData
 
     /**
      * Always holds.
@@ -150,6 +151,8 @@ module Stages {
       exists(Type t)
       or
       exists(inferType(_))
+      or
+      mayInvokeCallback(_, _)
     }
   }