Inline trivial leaf functions

Co-authored-by: Tomasz Miąsko <[email protected]>

Author: lcnr

compiler/rustc_metadata/src/rmeta/decoder.rs

@@ -1174,6 +1174,10 @@ impl<'a, 'tcx> CrateMetadataRef<'a> {
         }
     }
 
+    fn get_is_trivial_mir(&self, id: DefIndex) -> bool {
+        self.root.tables.is_trivial_mir.get(self, id).is_some()
+    }
+
     fn get_optimized_mir(&self, tcx: TyCtxt<'tcx>, id: DefIndex) -> Body<'tcx> {
         self.root
             .tables

compiler/rustc_metadata/src/rmeta/decoder/cstore_impl.rs

@@ -117,6 +117,7 @@ provide! { <'tcx> tcx, def_id, other, cdata,
             bug!("coerce_unsized_info: `{:?}` is missing its info", def_id);
         })
     }
+    is_trivial_mir => { cdata.get_is_trivial_mir(def_id.index) }
     optimized_mir => { tcx.arena.alloc(cdata.get_optimized_mir(tcx, def_id.index)) }
     mir_for_ctfe => { tcx.arena.alloc(cdata.get_mir_for_ctfe(tcx, def_id.index)) }
     promoted_mir => { tcx.arena.alloc(cdata.get_promoted_mir(tcx, def_id.index)) }

compiler/rustc_metadata/src/rmeta/encoder.rs

@@ -1164,6 +1164,12 @@ impl EncodeContext<'a, 'tcx> {
         debug!("EntryBuilder::encode_optimized_mir({:?})", def_id);
         record!(self.tables.mir[def_id.to_def_id()] <- self.tcx.optimized_mir(def_id));
 
+        if self.tcx.is_trivial_mir(def_id) {
+            // We don't store anything if `is_trivial_mir` is `false`
+            // so we can use a unit type here.
+            self.tables.is_trivial_mir.set(def_id.local_def_index, ());
+        }
+
         let unused = self.tcx.unused_generic_params(def_id);
         if !unused.is_empty() {
             record!(self.tables.unused_generic_params[def_id.to_def_id()] <- unused);

compiler/rustc_metadata/src/rmeta/mod.rs

@@ -301,6 +301,7 @@ define_tables! {
     super_predicates: Table<DefIndex, Lazy!(ty::GenericPredicates<'tcx>)>,
     // As an optimization, a missing entry indicates an empty `&[]`.
     explicit_item_bounds: Table<DefIndex, Lazy!([(ty::Predicate<'tcx>, Span)])>,
+    is_trivial_mir: Table<DefIndex, ()>,
     mir: Table<DefIndex, Lazy!(mir::Body<'tcx>)>,
     mir_for_ctfe: Table<DefIndex, Lazy!(mir::Body<'tcx>)>,
     promoted_mir: Table<DefIndex, Lazy!(IndexVec<mir::Promoted, mir::Body<'tcx>>)>,

compiler/rustc_metadata/src/rmeta/table.rs

@@ -4,7 +4,7 @@ use rustc_index::vec::Idx;
 use rustc_serialize::opaque::Encoder;
 use std::convert::TryInto;
 use std::marker::PhantomData;
-use std::num::NonZeroUsize;
+use std::num::{NonZeroU8, NonZeroUsize};
 use tracing::debug;
 
 /// Helper trait, for encoding to, and decoding from, a fixed number of bytes.
@@ -75,6 +75,30 @@ impl FixedSizeEncoding for u32 {
     }
 }
 
+impl FixedSizeEncoding for Option<NonZeroU8> {
+    fixed_size_encoding_byte_len_and_defaults!(1);
+
+    fn from_bytes(b: &[u8]) -> Self {
+        NonZeroU8::new(b[0])
+    }
+
+    fn write_to_bytes(self, b: &mut [u8]) {
+        b[0] = self.map_or(0, |x| x.get());
+    }
+}
+
+impl FixedSizeEncoding for Option<()> {
+    fixed_size_encoding_byte_len_and_defaults!(Option::<NonZeroU8>::BYTE_LEN);
+
+    fn from_bytes(b: &[u8]) -> Self {
+        Option::<NonZeroU8>::from_bytes(b).map(|_| ())
+    }
+
+    fn write_to_bytes(self, b: &mut [u8]) {
+        self.map(|()| NonZeroU8::new(1).unwrap()).write_to_bytes(b)
+    }
+}
+
 // NOTE(eddyb) there could be an impl for `usize`, which would enable a more
 // generic `Lazy<T>` impl, but in the general case we might not need / want to
 // fit every `usize` in `u32`.

compiler/rustc_middle/src/query/mod.rs

@@ -339,6 +339,11 @@ rustc_queries! {
             }
         }
 
+        query is_trivial_mir(key: DefId) -> bool {
+            desc { |tcx| "checking if MIR for `{}` is trivial", tcx.def_path_str(key) }
+            cache_on_disk_if { key.is_local() }
+        }
+
         /// MIR after our optimization passes have run. This is MIR that is ready
         /// for codegen. This is also the only query that can fetch non-local MIR, at present.
         query optimized_mir(key: DefId) -> &'tcx mir::Body<'tcx> {

compiler/rustc_mir/src/transform/inline.rs

@@ -2,6 +2,7 @@
 
 use rustc_attr as attr;
 use rustc_hir as hir;
+use rustc_hir::def_id::DefId;
 use rustc_index::bit_set::BitSet;
 use rustc_index::vec::Idx;
 use rustc_middle::middle::codegen_fn_attrs::{CodegenFnAttrFlags, CodegenFnAttrs};
@@ -37,10 +38,6 @@ struct CallSite<'tcx> {
 
 impl<'tcx> MirPass<'tcx> for Inline {
     fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
-        if tcx.sess.opts.debugging_opts.mir_opt_level < 2 {
-            return;
-        }
-
         if tcx.sess.opts.debugging_opts.instrument_coverage {
             // Since `Inline` happens after `InstrumentCoverage`, the function-specific coverage
             // counters can be invalidated, such as by merging coverage counter statements from
@@ -50,6 +47,13 @@ impl<'tcx> MirPass<'tcx> for Inline {
             return;
         }
 
+        if body.generator_kind.is_some() {
+            // 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.
+            return;
+        }
+
         if inline(tcx, body) {
             debug!("running simplify cfg on {:?}", body.source);
             CfgSimplifier::new(body).simplify();
@@ -104,7 +108,7 @@ impl Inliner<'tcx> {
                 Some(it) => it,
             };
 
-            if !self.is_mir_available(&callsite.callee, caller_body) {
+            if !self.is_mir_available(&callsite.callee) {
                 debug!("MIR unavailable {}", callsite.callee);
                 continue;
             }
@@ -137,11 +141,27 @@ 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;
-            }
+    fn is_mir_available(&self, callee: &Instance<'tcx>) -> bool {
+        match callee.def {
+            InstanceDef::Virtual(..) | InstanceDef::Intrinsic(..) => return false,
+
+            InstanceDef::VtableShim(..)
+            | InstanceDef::ReifyShim(..)
+            | InstanceDef::FnPtrShim(..)
+            | InstanceDef::ClosureOnceShim { .. }
+            | InstanceDef::DropGlue(..)
+            | InstanceDef::CloneShim(..) => return true,
+
+            InstanceDef::Item(_) => {}
+        };
+
+        if !self.tcx.is_mir_available(callee.def_id()) {
+            return false;
+        }
+
+        if self.tcx.sess.opts.debugging_opts.mir_opt_level <= 1 {
+            // Only inline trivial functions by default.
+            return self.tcx.is_trivial_mir(callee.def_id());
         }
 
         if let Some(callee_def_id) = callee.def_id().as_local() {
@@ -153,9 +173,7 @@ impl Inliner<'tcx> {
             // 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()
-                && self.hir_id < callee_hir_id
-                && caller_body.generator_kind.is_none()
+            !self.tcx.dep_graph.is_fully_enabled() && self.hir_id < callee_hir_id
         } else {
             // This cannot result in a cycle since the callee MIR is from another crate
             // and is already optimized.
@@ -885,3 +903,48 @@ impl<'a, 'tcx> MutVisitor<'tcx> for Integrator<'a, 'tcx> {
         }
     }
 }
+
+pub fn is_trivial_mir(tcx: TyCtxt<'tcx>, did: DefId) -> bool {
+    debug!("is_trivial_mir({:?})", did);
+    if tcx.is_constructor(did) {
+        debug!("is_trivial_mir = true (constructor)");
+        return true;
+    }
+
+    use rustc_hir::def::DefKind;
+    if !matches!(tcx.def_kind(did), DefKind::Fn | DefKind::AssocFn) {
+        debug!("is_trivial_mir = false (not a function)");
+        // Only inline functions, don't look at constants here.
+        return false;
+    }
+
+    if !did.is_local() {
+        // This branch is only taken if no `optimized_mir` is available for
+        // an extern crate, as `is_trivial_mir` has otherwise been encoded.
+        debug!("is_trivial_mir = false (no MIR available)");
+        return false;
+    };
+
+    let body = tcx
+        .mir_drops_elaborated_and_const_checked(ty::WithOptConstParam::unknown(did.expect_local()))
+        .borrow();
+
+    for bb in body.basic_blocks() {
+        let terminator = bb.terminator();
+        if let TerminatorKind::Call { func, .. } = &terminator.kind {
+            let func_ty = func.ty(body.local_decls(), tcx);
+            if let ty::FnDef(..) = *func_ty.kind() {
+                let fn_sig = func_ty.fn_sig(tcx);
+                if fn_sig.abi() == Abi::RustIntrinsic || fn_sig.abi() == Abi::PlatformIntrinsic {
+                    continue;
+                }
+            }
+
+            debug!("is_trivial_mir = false (function call)");
+            return false;
+        }
+    }
+
+    debug!("is_trivial_mir = true");
+    true
+}

compiler/rustc_mir/src/transform/mod.rs

@@ -71,6 +71,7 @@ pub(crate) fn provide(providers: &mut Providers) {
         },
         mir_promoted,
         mir_drops_elaborated_and_const_checked,
+        is_trivial_mir: inline::is_trivial_mir,
         mir_for_ctfe,
         mir_for_ctfe_of_const_arg,
         optimized_mir,
@@ -555,7 +556,8 @@ fn inner_optimized_mir(tcx: TyCtxt<'_>, did: LocalDefId) -> Body<'_> {
         // constructors.
         return shim::build_adt_ctor(tcx, did.to_def_id());
     }
-
+    // `is_trivial_mir` uses `mir_drops_elaborated_and_const_checked` so run that first.
+    tcx.ensure().is_trivial_mir(did.to_def_id());
     match tcx.hir().body_const_context(did) {
         // Run the `mir_for_ctfe` query, which depends on `mir_drops_elaborated_and_const_checked`
         // which we are going to steal below. Thus we need to run `mir_for_ctfe` first, so it