Only compute recursive callees once.

4 files changed, 152 insertions, 134 deletions

diff --git a/compiler/rustc_middle/src/query/mod.rs b/compiler/rustc_middle/src/query/mod.rs
index 3668f4e12f5..de3b15c5a6f 100644
--- a/compiler/rustc_middle/src/query/mod.rs
+++ b/compiler/rustc_middle/src/query/mod.rs
@@ -1283,14 +1283,13 @@ rustc_queries! {
         return_result_from_ensure_ok
     }
 
-    /// Check whether the function has any recursion that could cause the inliner to trigger
-    /// a cycle.
-    query mir_callgraph_reachable(key: (ty::Instance<'tcx>, LocalDefId)) -> bool {
+    /// Return the set of (transitive) callees that may result in a recursive call to `key`.
+    query mir_callgraph_cyclic(key: LocalDefId) -> &'tcx UnordSet<ty::Instance<'tcx>> {
         fatal_cycle
+        arena_cache
         desc { |tcx|
-            "computing if `{}` (transitively) calls `{}`",
-            key.0,
-            tcx.def_path_str(key.1),
+            "computing (transitive) callees of `{}` that may recurse",
+            tcx.def_path_str(key),
         }
     }
 
diff --git a/compiler/rustc_mir_transform/src/inline.rs b/compiler/rustc_mir_transform/src/inline.rs
index f48dba9663a..0972e95a238 100644
--- a/compiler/rustc_mir_transform/src/inline.rs
+++ b/compiler/rustc_mir_transform/src/inline.rs
@@ -777,7 +777,7 @@ fn check_mir_is_available<'tcx, I: Inliner<'tcx>>(
     {
         // If we know for sure that the function we're calling will itself try to
         // call us, then we avoid inlining that function.
-        if inliner.tcx().mir_callgraph_reachable((callee, caller_def_id.expect_local())) {
+        if inliner.tcx().mir_callgraph_cyclic(caller_def_id.expect_local()).contains(&callee) {
             debug!("query cycle avoidance");
             return Err("caller might be reachable from callee");
         }
diff --git a/compiler/rustc_mir_transform/src/inline/cycle.rs b/compiler/rustc_mir_transform/src/inline/cycle.rs
index a944960ce4a..4ef41b9ba44 100644
--- a/compiler/rustc_mir_transform/src/inline/cycle.rs
+++ b/compiler/rustc_mir_transform/src/inline/cycle.rs
@@ -1,5 +1,6 @@
 use rustc_data_structures::fx::{FxHashMap, FxHashSet, FxIndexSet};
 use rustc_data_structures::stack::ensure_sufficient_stack;
+use rustc_data_structures::unord::UnordSet;
 use rustc_hir::def_id::{DefId, LocalDefId};
 use rustc_middle::mir::TerminatorKind;
 use rustc_middle::ty::{self, GenericArgsRef, InstanceKind, TyCtxt, TypeVisitableExt};
@@ -7,137 +8,139 @@ use rustc_session::Limit;
 use rustc_span::sym;
 use tracing::{instrument, trace};
 
-// FIXME: check whether it is cheaper to precompute the entire call graph instead of invoking
-// this query ridiculously often.
-#[instrument(level = "debug", skip(tcx, root, target))]
-pub(crate) fn mir_callgraph_reachable<'tcx>(
+fn should_recurse<'tcx>(tcx: TyCtxt<'tcx>, callee: ty::Instance<'tcx>) -> bool {
+    match callee.def {
+        // If there is no MIR available (either because it was not in metadata or
+        // because it has no MIR because it's an extern function), then the inliner
+        // won't cause cycles on this.
+        InstanceKind::Item(_) => {
+            if !tcx.is_mir_available(callee.def_id()) {
+                return false;
+            }
+        }
+
+        // These have no own callable MIR.
+        InstanceKind::Intrinsic(_) | InstanceKind::Virtual(..) => return false,
+
+        // These have MIR and if that MIR is inlined, instantiated and then inlining is run
+        // again, a function item can end up getting inlined. Thus we'll be able to cause
+        // a cycle that way
+        InstanceKind::VTableShim(_)
+        | InstanceKind::ReifyShim(..)
+        | InstanceKind::FnPtrShim(..)
+        | InstanceKind::ClosureOnceShim { .. }
+        | InstanceKind::ConstructCoroutineInClosureShim { .. }
+        | InstanceKind::ThreadLocalShim { .. }
+        | InstanceKind::CloneShim(..) => {}
+
+        // This shim does not call any other functions, thus there can be no recursion.
+        InstanceKind::FnPtrAddrShim(..) => return false,
+
+        // FIXME: A not fully instantiated drop shim can cause ICEs if one attempts to
+        // have its MIR built. Likely oli-obk just screwed up the `ParamEnv`s, so this
+        // needs some more analysis.
+        InstanceKind::DropGlue(..)
+        | InstanceKind::FutureDropPollShim(..)
+        | InstanceKind::AsyncDropGlue(..)
+        | InstanceKind::AsyncDropGlueCtorShim(..) => {
+            if callee.has_param() {
+                return false;
+            }
+        }
+    }
+
+    crate::pm::should_run_pass(tcx, &crate::inline::Inline, crate::pm::Optimizations::Allowed)
+        || crate::inline::ForceInline::should_run_pass_for_callee(tcx, callee.def.def_id())
+}
+
+#[instrument(
+    level = "debug",
+    skip(tcx, typing_env, seen, involved, recursion_limiter, recursion_limit),
+    ret
+)]
+fn process<'tcx>(
     tcx: TyCtxt<'tcx>,
-    (root, target): (ty::Instance<'tcx>, LocalDefId),
+    typing_env: ty::TypingEnv<'tcx>,
+    caller: ty::Instance<'tcx>,
+    target: LocalDefId,
+    seen: &mut FxHashSet<ty::Instance<'tcx>>,
+    involved: &mut FxHashSet<ty::Instance<'tcx>>,
+    recursion_limiter: &mut FxHashMap<DefId, usize>,
+    recursion_limit: Limit,
 ) -> bool {
-    trace!(%root, target = %tcx.def_path_str(target));
-    assert_ne!(
-        root.def_id().expect_local(),
-        target,
-        "you should not call `mir_callgraph_reachable` on immediate self recursion"
-    );
-    assert!(
-        matches!(root.def, InstanceKind::Item(_)),
-        "you should not call `mir_callgraph_reachable` on shims"
-    );
-    assert!(
-        !tcx.is_constructor(root.def_id()),
-        "you should not call `mir_callgraph_reachable` on enum/struct constructor functions"
-    );
-    #[instrument(
-        level = "debug",
-        skip(tcx, typing_env, target, stack, seen, recursion_limiter, caller, recursion_limit)
-    )]
-    fn process<'tcx>(
-        tcx: TyCtxt<'tcx>,
-        typing_env: ty::TypingEnv<'tcx>,
-        caller: ty::Instance<'tcx>,
-        target: LocalDefId,
-        stack: &mut Vec<ty::Instance<'tcx>>,
-        seen: &mut FxHashSet<ty::Instance<'tcx>>,
-        recursion_limiter: &mut FxHashMap<DefId, usize>,
-        recursion_limit: Limit,
-    ) -> bool {
-        trace!(%caller);
-        for &(callee, args) in tcx.mir_inliner_callees(caller.def) {
-            let Ok(args) = caller.try_instantiate_mir_and_normalize_erasing_regions(
-                tcx,
-                typing_env,
-                ty::EarlyBinder::bind(args),
-            ) else {
-                trace!(?caller, ?typing_env, ?args, "cannot normalize, skipping");
-                continue;
-            };
-            let Ok(Some(callee)) = ty::Instance::try_resolve(tcx, typing_env, callee, args) else {
-                trace!(?callee, "cannot resolve, skipping");
-                continue;
-            };
+    trace!(%caller);
+    let mut cycle_found = false;
 
-            // Found a path.
-            if callee.def_id() == target.to_def_id() {
-                return true;
-            }
+    for &(callee, args) in tcx.mir_inliner_callees(caller.def) {
+        let Ok(args) = caller.try_instantiate_mir_and_normalize_erasing_regions(
+            tcx,
+            typing_env,
+            ty::EarlyBinder::bind(args),
+        ) else {
+            trace!(?caller, ?typing_env, ?args, "cannot normalize, skipping");
+            continue;
+        };
+        let Ok(Some(callee)) = ty::Instance::try_resolve(tcx, typing_env, callee, args) else {
+            trace!(?callee, "cannot resolve, skipping");
+            continue;
+        };
 
-            if tcx.is_constructor(callee.def_id()) {
-                trace!("constructors always have MIR");
-                // Constructor functions cannot cause a query cycle.
-                continue;
-            }
+        // Found a path.
+        if callee.def_id() == target.to_def_id() {
+            cycle_found = true;
+        }
 
-            match callee.def {
-                InstanceKind::Item(_) => {
-                    // If there is no MIR available (either because it was not in metadata or
-                    // because it has no MIR because it's an extern function), then the inliner
-                    // won't cause cycles on this.
-                    if !tcx.is_mir_available(callee.def_id()) {
-                        trace!(?callee, "no mir available, skipping");
-                        continue;
-                    }
-                }
-                // These have no own callable MIR.
-                InstanceKind::Intrinsic(_) | InstanceKind::Virtual(..) => continue,
-                // These have MIR and if that MIR is inlined, instantiated and then inlining is run
-                // again, a function item can end up getting inlined. Thus we'll be able to cause
-                // a cycle that way
-                InstanceKind::VTableShim(_)
-                | InstanceKind::ReifyShim(..)
-                | InstanceKind::FnPtrShim(..)
-                | InstanceKind::ClosureOnceShim { .. }
-                | InstanceKind::ConstructCoroutineInClosureShim { .. }
-                | InstanceKind::ThreadLocalShim { .. }
-                | InstanceKind::CloneShim(..) => {}
-
-                // This shim does not call any other functions, thus there can be no recursion.
-                InstanceKind::FnPtrAddrShim(..) => {
-                    continue;
-                }
-                InstanceKind::DropGlue(..)
-                | InstanceKind::FutureDropPollShim(..)
-                | InstanceKind::AsyncDropGlue(..)
-                | InstanceKind::AsyncDropGlueCtorShim(..) => {
-                    // FIXME: A not fully instantiated drop shim can cause ICEs if one attempts to
-                    // have its MIR built. Likely oli-obk just screwed up the `ParamEnv`s, so this
-                    // needs some more analysis.
-                    if callee.has_param() {
-                        continue;
-                    }
-                }
-            }
+        if tcx.is_constructor(callee.def_id()) {
+            trace!("constructors always have MIR");
+            // Constructor functions cannot cause a query cycle.
+            continue;
+        }
 
-            if seen.insert(callee) {
-                let recursion = recursion_limiter.entry(callee.def_id()).or_default();
-                trace!(?callee, recursion = *recursion);
-                if recursion_limit.value_within_limit(*recursion) {
-                    *recursion += 1;
-                    stack.push(callee);
-                    let found_recursion = ensure_sufficient_stack(|| {
-                        process(
-                            tcx,
-                            typing_env,
-                            callee,
-                            target,
-                            stack,
-                            seen,
-                            recursion_limiter,
-                            recursion_limit,
-                        )
-                    });
-                    if found_recursion {
-                        return true;
-                    }
-                    stack.pop();
-                } else {
-                    // Pessimistically assume that there could be recursion.
-                    return true;
-                }
+        if !should_recurse(tcx, callee) {
+            continue;
+        }
+
+        if seen.insert(callee) {
+            let recursion = recursion_limiter.entry(callee.def_id()).or_default();
+            trace!(?callee, recursion = *recursion);
+            let found_recursion = if recursion_limit.value_within_limit(*recursion) {
+                *recursion += 1;
+                ensure_sufficient_stack(|| {
+                    process(
+                        tcx,
+                        typing_env,
+                        callee,
+                        target,
+                        seen,
+                        involved,
+                        recursion_limiter,
+                        recursion_limit,
+                    )
+                })
+            } else {
+                // Pessimistically assume that there could be recursion.
+                true
+            };
+            if found_recursion {
+                involved.insert(callee);
+                cycle_found = true;
             }
         }
-        false
     }
+
+    cycle_found
+}
+
+#[instrument(level = "debug", skip(tcx), ret)]
+pub(crate) fn mir_callgraph_cyclic<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    root: LocalDefId,
+) -> UnordSet<ty::Instance<'tcx>> {
+    assert!(
+        !tcx.is_constructor(root.to_def_id()),
+        "you should not call `mir_callgraph_reachable` on enum/struct constructor functions"
+    );
+
     // FIXME(-Znext-solver=no): Remove this hack when trait solver overflow can return an error.
     // In code like that pointed out in #128887, the type complexity we ask the solver to deal with
     // grows as we recurse into the call graph. If we use the same recursion limit here and in the
@@ -146,16 +149,32 @@ pub(crate) fn mir_callgraph_reachable<'tcx>(
     // the default recursion limits are quite generous for us. If we need to recurse 64 times
     // into the call graph, we're probably not going to find any useful MIR inlining.
     let recursion_limit = tcx.recursion_limit() / 2;
+    let mut involved = FxHashSet::default();
+    let typing_env = ty::TypingEnv::post_analysis(tcx, root);
+    let Ok(Some(root_instance)) = ty::Instance::try_resolve(
+        tcx,
+        typing_env,
+        root.to_def_id(),
+        ty::GenericArgs::identity_for_item(tcx, root.to_def_id()),
+    ) else {
+        trace!("cannot resolve, skipping");
+        return involved.into();
+    };
+    if !should_recurse(tcx, root_instance) {
+        trace!("cannot walk, skipping");
+        return involved.into();
+    }
     process(
         tcx,
-        ty::TypingEnv::post_analysis(tcx, target),
+        typing_env,
+        root_instance,
         root,
-        target,
-        &mut Vec::new(),
         &mut FxHashSet::default(),
+        &mut involved,
         &mut FxHashMap::default(),
         recursion_limit,
-    )
+    );
+    involved.into()
 }
 
 pub(crate) fn mir_inliner_callees<'tcx>(
diff --git a/compiler/rustc_mir_transform/src/lib.rs b/compiler/rustc_mir_transform/src/lib.rs
index 572ad585c8c..48a4fed3993 100644
--- a/compiler/rustc_mir_transform/src/lib.rs
+++ b/compiler/rustc_mir_transform/src/lib.rs
@@ -215,7 +215,7 @@ pub fn provide(providers: &mut Providers) {
         optimized_mir,
         is_mir_available,
         is_ctfe_mir_available: is_mir_available,
-        mir_callgraph_reachable: inline::cycle::mir_callgraph_reachable,
+        mir_callgraph_cyclic: inline::cycle::mir_callgraph_cyclic,
         mir_inliner_callees: inline::cycle::mir_inliner_callees,
         promoted_mir,
         deduced_param_attrs: deduce_param_attrs::deduced_param_attrs,