Also allow inlining drop shims

Author: scottmcm

compiler/rustc_middle/src/mir/visit.rs

@@ -1328,9 +1328,9 @@ pub enum NonMutatingUseContext {
 pub enum MutatingUseContext {
     /// Appears as LHS of an assignment.
     Store,
-    /// Appears on `SetDiscriminant`
+    /// Appears on [`StatementKind::SetDiscriminant`]
     SetDiscriminant,
-    /// Appears on `Deinit`
+    /// Appears on [`StatementKind::Deinit`]
     Deinit,
     /// Output operand of an inline assembly block.
     AsmOutput,

compiler/rustc_mir_transform/src/inline.rs

@@ -6,6 +6,7 @@ use std::ops::{Range, RangeFrom};
 
 use rustc_abi::{ExternAbi, FieldIdx};
 use rustc_attr_data_structures::{InlineAttr, OptimizeAttr};
+use rustc_hir::LangItem;
 use rustc_hir::def::DefKind;
 use rustc_hir::def_id::DefId;
 use rustc_index::Idx;
@@ -116,6 +117,9 @@ trait Inliner<'tcx> {
     /// Has the caller body been changed?
     fn changed(self) -> bool;
 
+    /// Whether to also attempt to inline `Drop` terminators (not just `Call`s)
+    fn consider_drops(&self) -> bool;
+
     /// Should inlining happen for a given callee?
     fn should_inline_for_callee(&self, def_id: DefId) -> bool;
 
@@ -187,6 +191,10 @@ impl<'tcx> Inliner<'tcx> for ForceInliner<'tcx> {
         self.changed
     }
 
+    fn consider_drops(&self) -> bool {
+        false
+    }
+
     fn should_inline_for_callee(&self, def_id: DefId) -> bool {
         ForceInline::should_run_pass_for_callee(self.tcx(), def_id)
     }
@@ -272,6 +280,7 @@ struct NormalInliner<'tcx> {
     typing_env: ty::TypingEnv<'tcx>,
     /// `DefId` of caller.
     def_id: DefId,
+    caller_is_coroutine: bool,
     /// Stack of inlined instances.
     /// We only check the `DefId` and not the args because we want to
     /// avoid inlining cases of polymorphic recursion.
@@ -304,6 +313,7 @@ impl<'tcx> Inliner<'tcx> for NormalInliner<'tcx> {
             tcx,
             typing_env,
             def_id,
+            caller_is_coroutine: tcx.is_coroutine(def_id),
             history: Vec::new(),
             top_down_counter: 0,
             changed: false,
@@ -334,6 +344,10 @@ impl<'tcx> Inliner<'tcx> for NormalInliner<'tcx> {
         self.changed
     }
 
+    fn consider_drops(&self) -> bool {
+        !self.caller_is_coroutine
+    }
+
     fn should_inline_for_callee(&self, _: DefId) -> bool {
         true
     }
@@ -549,50 +563,77 @@ fn resolve_callsite<'tcx, I: Inliner<'tcx>>(
     bb_data: &BasicBlockData<'tcx>,
 ) -> Option<CallSite<'tcx>> {
     let tcx = inliner.tcx();
-    // Only consider direct calls to functions
     let terminator = bb_data.terminator();
 
     // FIXME(explicit_tail_calls): figure out if we can inline tail calls
-    if let TerminatorKind::Call { ref func, fn_span, .. } = terminator.kind {
-        let func_ty = func.ty(caller_body, tcx);
-        if let ty::FnDef(def_id, args) = *func_ty.kind() {
-            if !inliner.should_inline_for_callee(def_id) {
-                debug!("not enabled");
-                return None;
-            }
+    let (def_id, args, fn_span) = match &terminator.kind {
+        TerminatorKind::Call { func, fn_span, .. } => {
+            let func_ty = func.ty(caller_body, tcx);
+            if let ty::FnDef(def_id, args) = *func_ty.kind() {
+                if !inliner.should_inline_for_callee(def_id) {
+                    debug!("not enabled");
+                    return None;
+                }
 
-            // To resolve an instance its args have to be fully normalized.
-            let args = tcx.try_normalize_erasing_regions(inliner.typing_env(), args).ok()?;
-            let callee =
-                Instance::try_resolve(tcx, inliner.typing_env(), def_id, args).ok().flatten()?;
+                // Allow RemoveUnneededDrops to handle these, rather than inlining,
+                // since it doesn't add the extra locals nor the metadata.
+                if inliner.consider_drops()
+                    && tcx.is_lang_item(def_id, LangItem::DropInPlace)
+                    && let drop_ty = args.type_at(0)
+                    && !drop_ty.needs_drop(tcx, inliner.typing_env())
+                {
+                    return None;
+                }
 
-            if let InstanceKind::Virtual(..) | InstanceKind::Intrinsic(_) = callee.def {
+                (def_id, args, *fn_span)
+            } else {
                 return None;
             }
-
-            if inliner.history().contains(&callee.def_id()) {
+        }
+        TerminatorKind::Drop { place, .. } if inliner.consider_drops() => {
+            let drop_ty = place.ty(&caller_body.local_decls, tcx).ty;
+            if !drop_ty.needs_drop(tcx, inliner.typing_env())
+                // FIXME: Box's shim looks like it inlines fine, but it's quite
+                // a bit of MIR right now, so is currently skipped to avoid churn.
+                || drop_ty.is_box()
+            {
                 return None;
             }
 
-            let fn_sig = tcx.fn_sig(def_id).instantiate(tcx, args);
+            let drop_def_id =
+                tcx.require_lang_item(LangItem::DropInPlace, terminator.source_info.span);
+            let args = tcx.mk_args(&[drop_ty.into()]);
+            (drop_def_id, args, rustc_span::DUMMY_SP)
+        }
+        _ => return None,
+    };
 
-            // Additionally, check that the body that we're inlining actually agrees
-            // with the ABI of the trait that the item comes from.
-            if let InstanceKind::Item(instance_def_id) = callee.def
-                && tcx.def_kind(instance_def_id) == DefKind::AssocFn
-                && let instance_fn_sig = tcx.fn_sig(instance_def_id).skip_binder()
-                && instance_fn_sig.abi() != fn_sig.abi()
-            {
-                return None;
-            }
+    // To resolve an instance its args have to be fully normalized.
+    let args = tcx.try_normalize_erasing_regions(inliner.typing_env(), args).ok()?;
+    let callee = Instance::try_resolve(tcx, inliner.typing_env(), def_id, args).ok().flatten()?;
 
-            let source_info = SourceInfo { span: fn_span, ..terminator.source_info };
+    if let InstanceKind::Virtual(..) | InstanceKind::Intrinsic(_) = callee.def {
+        return None;
+    }
 
-            return Some(CallSite { callee, fn_sig, block: bb, source_info });
-        }
+    if inliner.history().contains(&callee.def_id()) {
+        return None;
+    }
+
+    let fn_sig = tcx.fn_sig(def_id).instantiate(tcx, args);
+
+    // Additionally, check that the body that we're inlining actually agrees
+    // with the ABI of the trait that the item comes from.
+    if let InstanceKind::Item(instance_def_id) = callee.def
+        && tcx.def_kind(instance_def_id) == DefKind::AssocFn
+        && let instance_fn_sig = tcx.fn_sig(instance_def_id).skip_binder()
+        && instance_fn_sig.abi() != fn_sig.abi()
+    {
+        return None;
     }
 
-    None
+    let source_info = SourceInfo { span: fn_span, ..terminator.source_info };
+    Some(CallSite { callee, fn_sig, block: bb, source_info })
 }
 
 /// Attempts to inline a callsite into the caller body. When successful returns basic blocks
@@ -604,14 +645,72 @@ fn try_inlining<'tcx, I: Inliner<'tcx>>(
     callsite: &CallSite<'tcx>,
 ) -> Result<std::ops::Range<BasicBlock>, &'static str> {
     let tcx = inliner.tcx();
+
     check_mir_is_available(inliner, caller_body, callsite.callee)?;
 
     let callee_attrs = tcx.codegen_fn_attrs(callsite.callee.def_id());
     check_inline::is_inline_valid_on_fn(tcx, callsite.callee.def_id())?;
     check_codegen_attributes(inliner, callsite, callee_attrs)?;
     inliner.check_codegen_attributes_extra(callee_attrs)?;
 
-    let terminator = caller_body[callsite.block].terminator.as_ref().unwrap();
+    let callee_body = try_instance_mir(tcx, callsite.callee.def)?;
+    check_inline::is_inline_valid_on_body(tcx, callee_body)?;
+    inliner.check_callee_mir_body(callsite, callee_body, callee_attrs)?;
+
+    if inliner.consider_drops() {
+        let block_mut = &mut caller_body.basic_blocks.as_mut()[callsite.block];
+        let terminator = block_mut.terminator.as_mut().unwrap();
+        if let TerminatorKind::Drop {
+            place,
+            target,
+            unwind,
+            replace: _,
+            drop: None,
+            async_fut: None,
+        } = terminator.kind
+        {
+            // Rather than updating everything after here to also handle `Drop`,
+            // just replace the terminator with a `Call`, since we'll need things
+            // like the local for the argument anyway.
+            let source_info = terminator.source_info;
+            let drop_ty = place.ty(&caller_body.local_decls, tcx).ty;
+            if !drop_ty.needs_drop(tcx, inliner.typing_env()) {
+                //return None;
+                bug!("Why does it think {drop_ty:?} needs a drop *now*?");
+            }
+
+            let drop_ty_mut = Ty::new_mut_ptr(tcx, drop_ty);
+            let arg = caller_body.local_decls.push(LocalDecl::new(drop_ty_mut, source_info.span));
+            block_mut.statements.push(Statement::new(
+                source_info,
+                StatementKind::Assign(Box::new((
+                    Place::from(arg),
+                    Rvalue::RawPtr(RawPtrKind::Mut, place),
+                ))),
+            ));
+            let unit_local =
+                caller_body.local_decls.push(LocalDecl::new(tcx.types.unit, source_info.span));
+            terminator.kind = TerminatorKind::Call {
+                func: Operand::function_handle(
+                    tcx,
+                    callsite.callee.def_id(),
+                    [drop_ty.into()],
+                    source_info.span,
+                ),
+                args: Box::new([Spanned {
+                    node: Operand::Move(Place::from(arg)),
+                    span: source_info.span,
+                }]),
+                destination: Place::from(unit_local),
+                target: Some(target),
+                unwind,
+                call_source: CallSource::Misc,
+                fn_span: source_info.span,
+            };
+        }
+    }
+
+    let terminator = caller_body[callsite.block].terminator();
     let TerminatorKind::Call { args, destination, .. } = &terminator.kind else { bug!() };
     let destination_ty = destination.ty(&caller_body.local_decls, tcx).ty;
     for arg in args {
@@ -622,10 +721,6 @@ fn try_inlining<'tcx, I: Inliner<'tcx>>(
         }
     }
 
-    let callee_body = try_instance_mir(tcx, callsite.callee.def)?;
-    check_inline::is_inline_valid_on_body(tcx, callee_body)?;
-    inliner.check_callee_mir_body(callsite, callee_body, callee_attrs)?;
-
     let Ok(callee_body) = callsite.callee.try_instantiate_mir_and_normalize_erasing_regions(
         tcx,
         inliner.typing_env(),

compiler/rustc_mir_transform/src/inline/cycle.rs

@@ -1,9 +1,10 @@
 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::LangItem;
 use rustc_hir::def_id::{DefId, LocalDefId};
 use rustc_middle::mir::TerminatorKind;
-use rustc_middle::ty::{self, GenericArgsRef, InstanceKind, TyCtxt, TypeVisitableExt};
+use rustc_middle::ty::{self, GenericArgsRef, InstanceKind, TyCtxt, TypeVisitableExt, TypingEnv};
 use rustc_session::Limit;
 use rustc_span::sym;
 use tracing::{instrument, trace};
@@ -204,23 +205,40 @@ pub(crate) fn mir_inliner_callees<'tcx>(
     let mut calls = FxIndexSet::default();
     for bb_data in body.basic_blocks.iter() {
         let terminator = bb_data.terminator();
-        if let TerminatorKind::Call { func, args: call_args, .. } = &terminator.kind {
-            let ty = func.ty(&body.local_decls, tcx);
-            let ty::FnDef(def_id, generic_args) = ty.kind() else {
-                continue;
-            };
-            let call = if tcx.is_intrinsic(*def_id, sym::const_eval_select) {
-                let func = &call_args[2].node;
+        let call = match &terminator.kind {
+            TerminatorKind::Call { func, args: call_args, .. } => {
                 let ty = func.ty(&body.local_decls, tcx);
                 let ty::FnDef(def_id, generic_args) = ty.kind() else {
                     continue;
                 };
-                (*def_id, *generic_args)
-            } else {
-                (*def_id, *generic_args)
-            };
-            calls.insert(call);
-        }
+                if tcx.is_intrinsic(*def_id, sym::const_eval_select) {
+                    let func = &call_args[2].node;
+                    let ty = func.ty(&body.local_decls, tcx);
+                    let ty::FnDef(def_id, generic_args) = ty.kind() else {
+                        continue;
+                    };
+                    (*def_id, *generic_args)
+                } else {
+                    (*def_id, *generic_args)
+                }
+            }
+            // Coroutines need optimized MIR for layout, so skip them to avoid cycles
+            TerminatorKind::Drop { place, drop: None, async_fut: None, .. }
+                if !tcx.is_coroutine(instance.def_id()) =>
+            {
+                let drop_ty = place.ty(&body.local_decls, tcx).ty;
+                if !drop_ty.needs_drop(tcx, TypingEnv::post_analysis(tcx, instance.def_id())) {
+                    continue;
+                }
+
+                let drop_def_id =
+                    tcx.require_lang_item(LangItem::DropInPlace, terminator.source_info.span);
+                let args = tcx.mk_args(&[drop_ty.into()]);
+                (drop_def_id, args)
+            }
+            _ => continue,
+        };
+        calls.insert(call);
     }
     tcx.arena.alloc_from_iter(calls.iter().copied())
 }

library/core/src/future/ready.rs

@@ -20,10 +20,20 @@ impl<T> Future for Ready<T> {
 
     #[inline]
     fn poll(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<T> {
-        Poll::Ready(self.0.take().expect("`Ready` polled after completion"))
+        match self.0.take() {
+            Some(v) => Poll::Ready(v),
+            None => polled_after_completion(),
+        }
     }
 }
 
+// Have only one of these, rather than mono'ing the panic for every `T`
+#[cold]
+#[inline]
+fn polled_after_completion() -> ! {
+    panic!("`Ready` polled after completion")
+}
+
 impl<T> Ready<T> {
     /// Consumes the `Ready`, returning the wrapped value.
     ///

tests/codegen-llvm/drop-in-place-noalias.rs

@@ -1,4 +1,4 @@
-//@ compile-flags: -Copt-level=3 -C no-prepopulate-passes
+//@ compile-flags: -Copt-level=3 -C no-prepopulate-passes -Z inline-mir=no
 
 // Tests that the compiler can apply `noalias` and other &mut attributes to `drop_in_place`.
 // Note that non-Unpin types should not get `noalias`, matching &mut behavior.

tests/codegen-llvm/drop.rs

@@ -1,5 +1,5 @@
 //@ needs-unwind - this test verifies the amount of drop calls when unwinding is used
-//@ compile-flags: -C no-prepopulate-passes
+//@ compile-flags: -C no-prepopulate-passes -Z inline-mir=no
 
 #![crate_type = "lib"]
 

tests/mir-opt/c_unwind_terminate.rs

@@ -2,6 +2,7 @@
 
 struct Noise;
 impl Drop for Noise {
+    #[inline(never)]
     fn drop(&mut self) {
         eprintln!("Noisy Drop");
     }
@@ -14,7 +15,9 @@ fn panic() {
 // EMIT_MIR c_unwind_terminate.test.AbortUnwindingCalls.after.mir
 extern "C" fn test() {
     // CHECK-LABEL: fn test(
-    // CHECK: drop
+    // CHECK: inlined drop_in_place::<Noise>
+    // CHECK-NOT: drop
+    // CHECK: <Noise as Drop>::drop
     // CHECK-SAME: unwind: [[unwind:bb.*]]]
     // CHECK: [[unwind]] (cleanup)
     // CHECK-NEXT: terminate(abi)