Prevent query cycles during inlining

4 files changed, 239 insertions, 17 deletions

diff --git a/compiler/rustc_mir/src/lib.rs b/compiler/rustc_mir/src/lib.rs
index e6d822086f5..8b3881ef9de 100644
--- a/compiler/rustc_mir/src/lib.rs
+++ b/compiler/rustc_mir/src/lib.rs
@@ -57,6 +57,8 @@ pub fn provide(providers: &mut Providers) {
     providers.eval_to_const_value_raw = const_eval::eval_to_const_value_raw_provider;
     providers.eval_to_allocation_raw = const_eval::eval_to_allocation_raw_provider;
     providers.const_caller_location = const_eval::const_caller_location;
+    providers.mir_callgraph_reachable = transform::inline::cycle::mir_callgraph_reachable;
+    providers.mir_inliner_callees = transform::inline::cycle::mir_inliner_callees;
     providers.destructure_const = |tcx, param_env_and_value| {
         let (param_env, value) = param_env_and_value.into_parts();
         const_eval::destructure_const(tcx, param_env, value)
diff --git a/compiler/rustc_mir/src/transform/inline.rs b/compiler/rustc_mir/src/transform/inline.rs
index 07e637b88f9..e157ef9c682 100644
--- a/compiler/rustc_mir/src/transform/inline.rs
+++ b/compiler/rustc_mir/src/transform/inline.rs
@@ -17,6 +17,8 @@ use crate::transform::MirPass;
 use std::iter;
 use std::ops::{Range, RangeFrom};
 
+crate mod cycle;
+
 const INSTR_COST: usize = 5;
 const CALL_PENALTY: usize = 25;
 const LANDINGPAD_PENALTY: usize = 50;
@@ -50,6 +52,8 @@ impl<'tcx> MirPass<'tcx> for Inline {
             return;
         }
 
+        let span = trace_span!("inline", body = %tcx.def_path_str(body.source.def_id()));
+        let _guard = span.enter();
         if inline(tcx, body) {
             debug!("running simplify cfg on {:?}", body.source);
             CfgSimplifier::new(body).simplify();
@@ -90,8 +94,8 @@ struct Inliner<'tcx> {
     codegen_fn_attrs: &'tcx CodegenFnAttrs,
     /// Caller HirID.
     hir_id: hir::HirId,
-    /// Stack of inlined instances.
-    history: Vec<Instance<'tcx>>,
+    /// Stack of inlined Instances.
+    history: Vec<ty::Instance<'tcx>>,
     /// Indicates that the caller body has been modified.
     changed: bool,
 }
@@ -103,13 +107,28 @@ impl Inliner<'tcx> {
                 None => continue,
                 Some(it) => it,
             };
+            let span = trace_span!("process_blocks", %callsite.callee, ?bb);
+            let _guard = span.enter();
+
+            trace!(
+                "checking for self recursion ({:?} vs body_source: {:?})",
+                callsite.callee.def_id(),
+                caller_body.source.def_id()
+            );
+            if callsite.callee.def_id() == caller_body.source.def_id() {
+                debug!("Not inlining a function into itself");
+                continue;
+            }
 
-            if !self.is_mir_available(&callsite.callee, caller_body) {
+            if !self.is_mir_available(callsite.callee, caller_body) {
                 debug!("MIR unavailable {}", callsite.callee);
                 continue;
             }
 
+            let span = trace_span!("instance_mir", %callsite.callee);
+            let instance_mir_guard = span.enter();
             let callee_body = self.tcx.instance_mir(callsite.callee.def);
+            drop(instance_mir_guard);
             if !self.should_inline(callsite, callee_body) {
                 continue;
             }
@@ -137,28 +156,61 @@ impl Inliner<'tcx> {
         }
     }
 
-    fn is_mir_available(&self, callee: &Instance<'tcx>, caller_body: &Body<'tcx>) -> bool {
-        if let InstanceDef::Item(_) = callee.def {
-            if !self.tcx.is_mir_available(callee.def_id()) {
-                return false;
+    #[instrument(skip(self, caller_body))]
+    fn is_mir_available(&self, callee: Instance<'tcx>, caller_body: &Body<'tcx>) -> bool {
+        match callee.def {
+            InstanceDef::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 !self.tcx.is_mir_available(callee.def_id()) {
+                    return false;
+                }
             }
+            // These have no own callable MIR.
+            InstanceDef::Intrinsic(_) | InstanceDef::Virtual(..) => return false,
+            // This cannot result in an immediate cycle since the callee MIR is a shim, which does
+            // not get any optimizations run on it. Any subsequent inlining may cause cycles, but we
+            // do not need to catch this here, we can wait until the inliner decides to continue
+            // inlining a second time.
+            InstanceDef::VtableShim(_)
+            | InstanceDef::ReifyShim(_)
+            | InstanceDef::FnPtrShim(..)
+            | InstanceDef::ClosureOnceShim { .. }
+            | InstanceDef::DropGlue(..)
+            | InstanceDef::CloneShim(..) => return true,
+        }
+
+        if self.tcx.is_constructor(callee.def_id()) {
+            trace!("constructors always have MIR");
+            // Constructor functions cannot cause a query cycle.
+            return true;
         }
 
         if let Some(callee_def_id) = callee.def_id().as_local() {
             let callee_hir_id = self.tcx.hir().local_def_id_to_hir_id(callee_def_id);
-            // Avoid a cycle here by only using `instance_mir` only if we have
-            // a lower `HirId` than the callee. This ensures that the callee will
-            // not inline us. This trick only works without incremental compilation.
-            // So don't do it if that is enabled. Also avoid inlining into generators,
+            // Avoid inlining into generators,
             // since their `optimized_mir` is used for layout computation, which can
             // create a cycle, even when no attempt is made to inline the function
             // in the other direction.
-            !self.tcx.dep_graph.is_fully_enabled()
+            caller_body.generator_kind.is_none()
+                && (
+                    // Avoid a cycle here by only using `instance_mir` only if we have
+                    // a lower `HirId` than the callee. This ensures that the callee will
+                    // not inline us. This trick only works without incremental compilation.
+                    // So don't do it if that is enabled.
+                    !self.tcx.dep_graph.is_fully_enabled()
                 && self.hir_id < callee_hir_id
-                && caller_body.generator_kind.is_none()
+                // If we know for sure that the function we're calling will itself try to
+                // call us, then we avoid inlining that function.
+                || !self.tcx.mir_callgraph_reachable((callee, caller_body.source.def_id().expect_local()))
+                )
         } else {
-            // This cannot result in a cycle since the callee MIR is from another crate
-            // and is already optimized.
+            // This cannot result in an immediate cycle since the callee MIR is from another crate
+            // and is already optimized. Any subsequent inlining may cause cycles, but we do
+            // not need to catch this here, we can wait until the inliner decides to continue
+            // inlining a second time.
+            trace!("functions from other crates always have MIR");
             true
         }
     }
@@ -203,8 +255,8 @@ impl Inliner<'tcx> {
         None
     }
 
+    #[instrument(skip(self, callee_body))]
     fn should_inline(&self, callsite: CallSite<'tcx>, callee_body: &Body<'tcx>) -> bool {
-        debug!("should_inline({:?})", callsite);
         let tcx = self.tcx;
 
         if callsite.fn_sig.c_variadic() {
@@ -333,7 +385,9 @@ impl Inliner<'tcx> {
                         if let Ok(Some(instance)) =
                             Instance::resolve(self.tcx, self.param_env, def_id, substs)
                         {
-                            if callsite.callee == instance || self.history.contains(&instance) {
+                            if callsite.callee.def_id() == instance.def_id()
+                                || self.history.contains(&instance)
+                            {
                                 debug!("`callee is recursive - not inlining");
                                 return false;
                             }
diff --git a/compiler/rustc_mir/src/transform/inline/cycle.rs b/compiler/rustc_mir/src/transform/inline/cycle.rs
new file mode 100644
index 00000000000..e4d403fbf60
--- /dev/null
+++ b/compiler/rustc_mir/src/transform/inline/cycle.rs
@@ -0,0 +1,157 @@
+use rustc_data_structures::fx::{FxHashMap, FxHashSet};
+use rustc_data_structures::stack::ensure_sufficient_stack;
+use rustc_hir::def_id::{DefId, LocalDefId};
+use rustc_middle::mir::TerminatorKind;
+use rustc_middle::ty::TypeFoldable;
+use rustc_middle::ty::{self, subst::SubstsRef, InstanceDef, TyCtxt};
+
+// FIXME: check whether it is cheaper to precompute the entire call graph instead of invoking
+// this query riddiculously often.
+#[instrument(skip(tcx, root, target))]
+crate fn mir_callgraph_reachable(
+    tcx: TyCtxt<'tcx>,
+    (root, target): (ty::Instance<'tcx>, LocalDefId),
+) -> bool {
+    trace!(%root, target = %tcx.def_path_str(target.to_def_id()));
+    let param_env = tcx.param_env_reveal_all_normalized(target);
+    assert_ne!(
+        root.def_id().expect_local(),
+        target,
+        "you should not call `mir_callgraph_reachable` on immediate self recursion"
+    );
+    assert!(
+        matches!(root.def, InstanceDef::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(skip(tcx, param_env, target, stack, seen, recursion_limiter, caller))]
+    fn process(
+        tcx: TyCtxt<'tcx>,
+        param_env: ty::ParamEnv<'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>,
+    ) -> bool {
+        trace!(%caller);
+        for &(callee, substs) in tcx.mir_inliner_callees(caller.def) {
+            let substs = caller.subst_mir_and_normalize_erasing_regions(tcx, param_env, substs);
+            let callee = match ty::Instance::resolve(tcx, param_env, callee, substs).unwrap() {
+                Some(callee) => callee,
+                None => {
+                    trace!(?callee, "cannot resolve, skipping");
+                    continue;
+                }
+            };
+
+            // Found a path.
+            if callee.def_id() == target.to_def_id() {
+                return true;
+            }
+
+            if tcx.is_constructor(callee.def_id()) {
+                trace!("constructors always have MIR");
+                // Constructor functions cannot cause a query cycle.
+                continue;
+            }
+
+            match callee.def {
+                InstanceDef::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.
+                InstanceDef::Intrinsic(_) | InstanceDef::Virtual(..) => continue,
+                // These have MIR and if that MIR is inlined, substituted 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
+                InstanceDef::VtableShim(_)
+                | InstanceDef::ReifyShim(_)
+                | InstanceDef::FnPtrShim(..)
+                | InstanceDef::ClosureOnceShim { .. }
+                | InstanceDef::CloneShim(..) => {}
+                InstanceDef::DropGlue(..) => {
+                    // FIXME: A not fully substituted 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.needs_subst() {
+                        continue;
+                    }
+                }
+            }
+
+            if seen.insert(callee) {
+                let recursion = recursion_limiter.entry(callee.def_id()).or_default();
+                trace!(?callee, recursion = *recursion);
+                if tcx.sess.recursion_limit().value_within_limit(*recursion) {
+                    *recursion += 1;
+                    stack.push(callee);
+                    let found_recursion = ensure_sufficient_stack(|| {
+                        process(tcx, param_env, callee, target, stack, seen, recursion_limiter)
+                    });
+                    if found_recursion {
+                        return true;
+                    }
+                    stack.pop();
+                } else {
+                    // Pessimistically assume that there could be recursion.
+                    return true;
+                }
+            }
+        }
+        false
+    }
+    process(
+        tcx,
+        param_env,
+        root,
+        target,
+        &mut Vec::new(),
+        &mut FxHashSet::default(),
+        &mut FxHashMap::default(),
+    )
+}
+
+crate fn mir_inliner_callees<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    instance: ty::InstanceDef<'tcx>,
+) -> &'tcx [(DefId, SubstsRef<'tcx>)] {
+    let steal;
+    let guard;
+    let body = match (instance, instance.def_id().as_local()) {
+        (InstanceDef::Item(_), Some(def_id)) => {
+            let def = ty::WithOptConstParam::unknown(def_id);
+            steal = tcx.mir_promoted(def).0;
+            guard = steal.borrow();
+            &*guard
+        }
+        // Functions from other crates and MIR shims
+        _ => tcx.instance_mir(instance),
+    };
+    let mut calls = Vec::new();
+    for bb_data in body.basic_blocks() {
+        let terminator = bb_data.terminator();
+        if let TerminatorKind::Call { func, .. } = &terminator.kind {
+            let ty = func.ty(&body.local_decls, tcx);
+            let call = match ty.kind() {
+                ty::FnDef(def_id, substs) => (*def_id, *substs),
+                _ => continue,
+            };
+            // We've seen this before
+            if calls.contains(&call) {
+                continue;
+            }
+            calls.push(call);
+        }
+    }
+    tcx.arena.alloc_slice(&calls)
+}
diff --git a/compiler/rustc_mir/src/transform/mod.rs b/compiler/rustc_mir/src/transform/mod.rs
index e509c35de40..ea62f9a8f95 100644
--- a/compiler/rustc_mir/src/transform/mod.rs
+++ b/compiler/rustc_mir/src/transform/mod.rs
@@ -419,6 +419,15 @@ fn mir_drops_elaborated_and_const_checked<'tcx>(
         tcx.ensure().mir_borrowck(def.did);
     }
 
+    let hir_id = tcx.hir().local_def_id_to_hir_id(def.did);
+    use rustc_middle::hir::map::blocks::FnLikeNode;
+    let is_fn_like = FnLikeNode::from_node(tcx.hir().get(hir_id)).is_some();
+    if is_fn_like {
+        let did = def.did.to_def_id();
+        let def = ty::WithOptConstParam::unknown(did);
+        let _ = tcx.mir_inliner_callees(ty::InstanceDef::Item(def));
+    }
+
     let (body, _) = tcx.mir_promoted(def);
     let mut body = body.steal();