From b8727e2d604e37a1510afaa9bd8777fecdcb5da4 Mon Sep 17 00:00:00 2001 From: oli Date: Tue, 29 Dec 2020 16:21:52 +0000 Subject: [PATCH] Prevent query cycles during inlining --- compiler/rustc_middle/src/query/mod.rs | 21 +++ compiler/rustc_middle/src/ty/query/keys.rs | 11 ++ compiler/rustc_mir/src/lib.rs | 2 + compiler/rustc_mir/src/transform/inline.rs | 88 ++++++++-- .../rustc_mir/src/transform/inline/cycle.rs | 157 ++++++++++++++++++ compiler/rustc_mir/src/transform/mod.rs | 9 + .../mir-opt/inline/inline-cycle-generic.rs | 40 +++++ .../inline_cycle_generic.main.Inline.diff | 29 ++++ src/test/ui/inline_cycle.rs | 17 ++ 9 files changed, 357 insertions(+), 17 deletions(-) create mode 100644 compiler/rustc_mir/src/transform/inline/cycle.rs create mode 100644 src/test/mir-opt/inline/inline-cycle-generic.rs create mode 100644 src/test/mir-opt/inline/inline_cycle_generic.main.Inline.diff create mode 100644 src/test/ui/inline_cycle.rs diff --git a/compiler/rustc_middle/src/query/mod.rs b/compiler/rustc_middle/src/query/mod.rs index 00ee7b8ec77..f088e05b49e 100644 --- a/compiler/rustc_middle/src/query/mod.rs +++ b/compiler/rustc_middle/src/query/mod.rs @@ -782,6 +782,27 @@ rustc_queries! { } Other { + /// Check whether the function has any recursion that could cause the inliner to trigger + /// a cycle. Returns the call stack causing the cycle. The call stack does not contain the + /// current function, just all intermediate functions. + query mir_callgraph_reachable(key: (ty::Instance<'tcx>, LocalDefId)) -> bool { + fatal_cycle + desc { |tcx| + "computing if `{}` (transitively) calls `{}`", + key.0, + tcx.def_path_str(key.1.to_def_id()), + } + } + + /// Obtain all the calls into other local functions + query mir_inliner_callees(key: ty::InstanceDef<'tcx>) -> &'tcx [(DefId, SubstsRef<'tcx>)] { + fatal_cycle + desc { |tcx| + "computing all local function calls in `{}`", + tcx.def_path_str(key.def_id()), + } + } + /// Evaluates a constant and returns the computed allocation. /// /// **Do not use this** directly, use the `tcx.eval_static_initializer` wrapper. diff --git a/compiler/rustc_middle/src/ty/query/keys.rs b/compiler/rustc_middle/src/ty/query/keys.rs index a005990264c..bfa1581aaae 100644 --- a/compiler/rustc_middle/src/ty/query/keys.rs +++ b/compiler/rustc_middle/src/ty/query/keys.rs @@ -127,6 +127,17 @@ impl Key for (DefId, DefId) { } } +impl Key for (ty::Instance<'tcx>, LocalDefId) { + type CacheSelector = DefaultCacheSelector; + + fn query_crate(&self) -> CrateNum { + self.0.query_crate() + } + fn default_span(&self, tcx: TyCtxt<'_>) -> Span { + self.0.default_span(tcx) + } +} + impl Key for (DefId, LocalDefId) { type CacheSelector = DefaultCacheSelector; 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>, + /// Stack of inlined Instances. + history: Vec>, /// 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(); - if !self.is_mir_available(&callsite.callee, caller_body) { + 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) { 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>, + seen: &mut FxHashSet>, + recursion_limiter: &mut FxHashMap, + ) -> 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(); diff --git a/src/test/mir-opt/inline/inline-cycle-generic.rs b/src/test/mir-opt/inline/inline-cycle-generic.rs new file mode 100644 index 00000000000..24b4f37939a --- /dev/null +++ b/src/test/mir-opt/inline/inline-cycle-generic.rs @@ -0,0 +1,40 @@ +// Check that inliner handles various forms of recursion and doesn't fall into +// an infinite inlining cycle. The particular outcome of inlining is not +// crucial otherwise. +// +// Regression test for issue #78573. + +// EMIT_MIR inline_cycle_generic.main.Inline.diff +fn main() { + ::call(); +} + +pub trait Call { + fn call(); +} + +pub struct A; +pub struct B(T); +pub struct C; + +impl Call for A { + #[inline] + fn call() { + as Call>::call() + } +} + + +impl Call for B { + #[inline] + fn call() { + ::call() + } +} + +impl Call for C { + #[inline] + fn call() { + as Call>::call() + } +} diff --git a/src/test/mir-opt/inline/inline_cycle_generic.main.Inline.diff b/src/test/mir-opt/inline/inline_cycle_generic.main.Inline.diff new file mode 100644 index 00000000000..9709f273779 --- /dev/null +++ b/src/test/mir-opt/inline/inline_cycle_generic.main.Inline.diff @@ -0,0 +1,29 @@ +- // MIR for `main` before Inline ++ // MIR for `main` after Inline + + fn main() -> () { + let mut _0: (); // return place in scope 0 at $DIR/inline-cycle-generic.rs:8:11: 8:11 + let _1: (); // in scope 0 at $DIR/inline-cycle-generic.rs:9:5: 9:24 ++ scope 1 (inlined ::call) { // at $DIR/inline-cycle-generic.rs:9:5: 9:24 ++ scope 2 (inlined as Call>::call) { // at $DIR/inline-cycle-generic.rs:9:5: 9:24 ++ } ++ } + + bb0: { + StorageLive(_1); // scope 0 at $DIR/inline-cycle-generic.rs:9:5: 9:24 +- _1 = ::call() -> bb1; // scope 0 at $DIR/inline-cycle-generic.rs:9:5: 9:24 ++ _1 = ::call() -> bb1; // scope 2 at $DIR/inline-cycle-generic.rs:9:5: 9:24 + // mir::Constant +- // + span: $DIR/inline-cycle-generic.rs:9:5: 9:22 +- // + literal: Const { ty: fn() {::call}, val: Value(Scalar()) } ++ // + span: $DIR/inline-cycle-generic.rs:9:5: 9:24 ++ // + literal: Const { ty: fn() {::call}, val: Value(Scalar()) } + } + + bb1: { + StorageDead(_1); // scope 0 at $DIR/inline-cycle-generic.rs:9:24: 9:25 + _0 = const (); // scope 0 at $DIR/inline-cycle-generic.rs:8:11: 10:2 + return; // scope 0 at $DIR/inline-cycle-generic.rs:10:2: 10:2 + } + } + diff --git a/src/test/ui/inline_cycle.rs b/src/test/ui/inline_cycle.rs new file mode 100644 index 00000000000..60767f10660 --- /dev/null +++ b/src/test/ui/inline_cycle.rs @@ -0,0 +1,17 @@ +// Needs build-pass to trigger `optimized_mir` on all mir bodies +// build-pass +// compile-flags: -Zmir-opt-level=2 + +#[inline(always)] +fn f(g: impl Fn()) { + g(); +} + +#[inline(always)] +fn g() { + f(main); +} + +fn main() { + f(g); +}