Allow inlining all the drops

Author: scottmcm

compiler/rustc_mir_transform/src/inline.rs

@@ -7,16 +7,16 @@ use rustc_attr::InlineAttr;
 use rustc_hir::def::DefKind;
 use rustc_hir::def_id::DefId;
 use rustc_index::bit_set::BitSet;
-use rustc_index::Idx;
+use rustc_index::{Idx, IndexVec};
 use rustc_middle::bug;
 use rustc_middle::middle::codegen_fn_attrs::{CodegenFnAttrFlags, CodegenFnAttrs};
 use rustc_middle::mir::visit::*;
 use rustc_middle::mir::*;
 use rustc_middle::ty::{
-    self, Instance, InstanceKind, ParamEnv, Ty, TyCtxt, TypeFlags, TypeVisitableExt,
+    self, GenericArg, Instance, InstanceKind, ParamEnv, Ty, TyCtxt, TypeFlags, TypeVisitableExt,
 };
 use rustc_session::config::{DebugInfo, OptLevel};
-use rustc_span::source_map::Spanned;
+use rustc_span::source_map::{dummy_spanned, Spanned};
 use rustc_span::sym;
 use rustc_target::abi::FieldIdx;
 use rustc_target::spec::abi::Abi;
@@ -94,6 +94,10 @@ fn inline<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) -> bool {
 
     let param_env = tcx.param_env_reveal_all_normalized(def_id);
     let codegen_fn_attrs = tcx.codegen_fn_attrs(def_id);
+    let Some(drop_in_place_fn) = tcx.lang_items().drop_in_place_fn() else {
+        error!("No `drop_in_place` function; not even trying to inline!");
+        return false;
+    };
 
     let mut this = Inliner {
         tcx,
@@ -105,6 +109,8 @@ fn inline<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) -> bool {
             codegen_fn_attrs.inline,
             InlineAttr::Hint | InlineAttr::Always
         ) && body_is_forwarder(body),
+        drop_in_place_fn,
+        unit_local: None,
     };
     let blocks = START_BLOCK..body.basic_blocks.next_index();
     this.process_blocks(body, blocks);
@@ -127,9 +133,78 @@ struct Inliner<'tcx> {
     /// Indicates that the caller is #[inline] and just calls another function,
     /// and thus we can inline less into it as it'll be inlined itself.
     caller_is_inline_forwarder: bool,
+    /// The compiler-magic function that actually drops something.
+    drop_in_place_fn: DefId,
+    /// When lowering `Drop(place)` to `drop_in_place(&place)`, we need a unit
+    /// local to use as the target of the call, but don't want multiple.
+    unit_local: Option<Local>,
 }
 
 impl<'tcx> Inliner<'tcx> {
+    fn lower_drop_to_call(
+        &mut self,
+        block: &mut BasicBlockData<'tcx>,
+        local_decls: &mut IndexVec<Local, LocalDecl<'tcx>>,
+    ) {
+        let terminator = block.terminator.as_mut().unwrap();
+        let TerminatorKind::Drop { place: dropped_place, target, unwind, replace: _ } =
+            terminator.kind
+        else {
+            return;
+        };
+
+        let dropped_ty = dropped_place.ty(local_decls, self.tcx).ty;
+        if matches!(dropped_ty.kind(), ty::Alias(..) | ty::Param(..)) {
+            // Not worth the extra locals, since we'll probably not be able to
+            // get MIR for it. If something non-generic happens to inline this
+            // block later, we'll have the opportunity to handle it then.
+            return;
+        }
+
+        if !dropped_ty.needs_drop(self.tcx, self.param_env) {
+            // Leave it for other passes to remove, which is cheaper than
+            // doing the work to inline an empty shim.
+            return;
+        }
+
+        self.changed = true;
+
+        let dropped_operand = if let [PlaceElem::Deref] = **dropped_place.projection
+            && local_decls[dropped_place.local].ty.is_mutable_ptr()
+        {
+            Operand::Copy(Place::from(dropped_place.local))
+        } else {
+            let dropped_ty_ptr = Ty::new_mut_ptr(self.tcx, dropped_ty);
+            let ptr_local =
+                local_decls.push(LocalDecl::new(dropped_ty_ptr, terminator.source_info.span));
+            block.statements.push(Statement {
+                source_info: terminator.source_info,
+                kind: StatementKind::Assign(Box::new((
+                    Place::from(ptr_local),
+                    Rvalue::AddressOf(Mutability::Mut, dropped_place),
+                ))),
+            });
+            Operand::Move(Place::from(ptr_local))
+        };
+        let unit_local = *self.unit_local.get_or_insert_with(|| {
+            local_decls.push(LocalDecl::new(self.tcx.types.unit, terminator.source_info.span))
+        });
+        terminator.kind = TerminatorKind::Call {
+            func: Operand::function_handle(
+                self.tcx,
+                self.drop_in_place_fn,
+                [GenericArg::from(dropped_ty)],
+                terminator.source_info.span,
+            ),
+            args: Box::new([dummy_spanned(dropped_operand)]),
+            destination: Place::from(unit_local),
+            target: Some(target),
+            unwind,
+            call_source: CallSource::Misc,
+            fn_span: terminator.source_info.span,
+        };
+    }
+
     fn process_blocks(&mut self, caller_body: &mut Body<'tcx>, blocks: Range<BasicBlock>) {
         // How many callsites in this body are we allowed to inline? We need to limit this in order
         // to prevent super-linear growth in MIR size
@@ -140,12 +215,18 @@ impl<'tcx> Inliner<'tcx> {
         };
         let mut inlined_count = 0;
         for bb in blocks {
-            let bb_data = &caller_body[bb];
-            if bb_data.is_cleanup {
+            if caller_body[bb].is_cleanup {
                 continue;
             }
 
-            let Some(callsite) = self.resolve_callsite(caller_body, bb, bb_data) else {
+            // Changing `Drop` to `Call` actually preserves the CFG because it
+            // keeps the same `target` and `unwind` action.
+            self.lower_drop_to_call(
+                &mut caller_body.basic_blocks.as_mut_preserves_cfg()[bb],
+                &mut caller_body.local_decls,
+            );
+
+            let Some(callsite) = self.resolve_callsite(caller_body, bb, &caller_body[bb]) else {
                 continue;
             };
 
@@ -509,7 +590,9 @@ impl<'tcx> Inliner<'tcx> {
             return Err("Body is tainted");
         }
 
-        let mut threshold = if self.caller_is_inline_forwarder {
+        let mut threshold = if self.caller_is_inline_forwarder
+            || matches!(callee_body.source.instance, InstanceKind::DropGlue(..))
+        {
             self.tcx.sess.opts.unstable_opts.inline_mir_forwarder_threshold.unwrap_or(30)
         } else if cross_crate_inlinable {
             self.tcx.sess.opts.unstable_opts.inline_mir_hint_threshold.unwrap_or(100)

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

@@ -1,8 +1,10 @@
-//@ compile-flags: -O -C no-prepopulate-passes
+//@ compile-flags: -O -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.
 
+// Needs to not run MIR inlining or else everything inlines away.
+
 #![crate_type = "lib"]
 
 use std::marker::PhantomPinned;

tests/codegen/drop.rs

@@ -1,5 +1,8 @@
 //@ 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
+
+// Needs to disable MIR inlining or else some of the calls are directly to
+// `<SomeUniqueName as Drop>::drop` (rather than via `drop_in_place`).
 
 #![crate_type = "lib"]
 

tests/mir-opt/dataflow-const-prop/default_boxed_slice.main.DataflowConstProp.32bit.panic-abort.diff

@@ -5,6 +5,9 @@
       let mut _0: ();
       let _1: A;
       let mut _2: std::boxed::Box<[bool]>;
+      let mut _12: *mut A;
+      let mut _13: ();
+      let mut _14: *mut std::boxed::Box<[bool]>;
       scope 1 {
           debug a => _1;
       }
@@ -39,6 +42,8 @@
               }
           }
       }
+      scope 13 (inlined std::ptr::drop_in_place::<A> - shim(Some(A))) {
+      }
 
       bb0: {
           StorageLive(_1);
@@ -55,26 +60,21 @@
           StorageLive(_11);
           StorageLive(_8);
           _8 = UbChecks();
-          switchInt(move _8) -> [0: bb4, otherwise: bb2];
+          switchInt(move _8) -> [0: bb3, otherwise: bb1];
       }
 
       bb1: {
-          StorageDead(_1);
-          return;
-      }
-  
-      bb2: {
           StorageLive(_10);
           _10 = const {0x1 as *mut ()};
-          _9 = NonNull::<T>::new_unchecked::precondition_check(const {0x1 as *mut ()}) -> [return: bb3, unwind unreachable];
+          _9 = NonNull::<T>::new_unchecked::precondition_check(const {0x1 as *mut ()}) -> [return: bb2, unwind unreachable];
       }
 
-      bb3: {
+      bb2: {
           StorageDead(_10);
-          goto -> bb4;
+          goto -> bb3;
       }
 
-      bb4: {
+      bb3: {
           StorageDead(_8);
           _11 = const {0x1 as *const [bool; 0]};
           _5 = const NonNull::<[bool; 0]> {{ pointer: {0x1 as *const [bool; 0]} }};
@@ -90,7 +90,14 @@
           _1 = const A {{ foo: Box::<[bool]>(Unique::<[bool]> {{ pointer: NonNull::<[bool]> {{ pointer: Indirect { alloc_id: ALLOC2, offset: Size(0 bytes) }: *const [bool] }}, _marker: PhantomData::<[bool]> }}, std::alloc::Global) }};
           StorageDead(_2);
           _0 = const ();
-          drop(_1) -> [return: bb1, unwind unreachable];
+          _12 = &raw mut _1;
+          _14 = &raw mut ((*_12).0: std::boxed::Box<[bool]>);
+          _13 = std::ptr::drop_in_place::<Box<[bool]>>(move _14) -> [return: bb4, unwind unreachable];
+      }
+  
+      bb4: {
+          StorageDead(_1);
+          return;
       }
   }
 

tests/mir-opt/dataflow-const-prop/default_boxed_slice.main.DataflowConstProp.32bit.panic-unwind.diff

@@ -5,6 +5,9 @@
       let mut _0: ();
       let _1: A;
       let mut _2: std::boxed::Box<[bool]>;
+      let mut _12: *mut A;
+      let mut _13: ();
+      let mut _14: *mut std::boxed::Box<[bool]>;
       scope 1 {
           debug a => _1;
       }
@@ -39,6 +42,8 @@
               }
           }
       }
+      scope 13 (inlined std::ptr::drop_in_place::<A> - shim(Some(A))) {
+      }
 
       bb0: {
           StorageLive(_1);
@@ -55,30 +60,25 @@
           StorageLive(_11);
           StorageLive(_8);
           _8 = UbChecks();
-          switchInt(move _8) -> [0: bb5, otherwise: bb3];
-      }
-  
-      bb1: {
-          StorageDead(_1);
-          return;
+          switchInt(move _8) -> [0: bb4, otherwise: bb2];
       }
 
-      bb2 (cleanup): {
+      bb1 (cleanup): {
           resume;
       }
 
-      bb3: {
+      bb2: {
           StorageLive(_10);
           _10 = const {0x1 as *mut ()};
-          _9 = NonNull::<T>::new_unchecked::precondition_check(const {0x1 as *mut ()}) -> [return: bb4, unwind unreachable];
+          _9 = NonNull::<T>::new_unchecked::precondition_check(const {0x1 as *mut ()}) -> [return: bb3, unwind unreachable];
       }
 
-      bb4: {
+      bb3: {
           StorageDead(_10);
-          goto -> bb5;
+          goto -> bb4;
       }
 
-      bb5: {
+      bb4: {
           StorageDead(_8);
           _11 = const {0x1 as *const [bool; 0]};
           _5 = const NonNull::<[bool; 0]> {{ pointer: {0x1 as *const [bool; 0]} }};
@@ -94,7 +94,14 @@
           _1 = const A {{ foo: Box::<[bool]>(Unique::<[bool]> {{ pointer: NonNull::<[bool]> {{ pointer: Indirect { alloc_id: ALLOC2, offset: Size(0 bytes) }: *const [bool] }}, _marker: PhantomData::<[bool]> }}, std::alloc::Global) }};
           StorageDead(_2);
           _0 = const ();
-          drop(_1) -> [return: bb1, unwind: bb2];
+          _12 = &raw mut _1;
+          _14 = &raw mut ((*_12).0: std::boxed::Box<[bool]>);
+          _13 = std::ptr::drop_in_place::<Box<[bool]>>(move _14) -> [return: bb5, unwind: bb1];
+      }
+  
+      bb5: {
+          StorageDead(_1);
+          return;
       }
   }
 

tests/mir-opt/dataflow-const-prop/default_boxed_slice.main.DataflowConstProp.64bit.panic-abort.diff

@@ -5,6 +5,9 @@
       let mut _0: ();
       let _1: A;
       let mut _2: std::boxed::Box<[bool]>;
+      let mut _12: *mut A;
+      let mut _13: ();
+      let mut _14: *mut std::boxed::Box<[bool]>;
       scope 1 {
           debug a => _1;
       }
@@ -39,6 +42,8 @@
               }
           }
       }
+      scope 13 (inlined std::ptr::drop_in_place::<A> - shim(Some(A))) {
+      }
 
       bb0: {
           StorageLive(_1);
@@ -55,26 +60,21 @@
           StorageLive(_11);
           StorageLive(_8);
           _8 = UbChecks();
-          switchInt(move _8) -> [0: bb4, otherwise: bb2];
+          switchInt(move _8) -> [0: bb3, otherwise: bb1];
       }
 
       bb1: {
-          StorageDead(_1);
-          return;
-      }
-  
-      bb2: {
           StorageLive(_10);
           _10 = const {0x1 as *mut ()};
-          _9 = NonNull::<T>::new_unchecked::precondition_check(const {0x1 as *mut ()}) -> [return: bb3, unwind unreachable];
+          _9 = NonNull::<T>::new_unchecked::precondition_check(const {0x1 as *mut ()}) -> [return: bb2, unwind unreachable];
       }
 
-      bb3: {
+      bb2: {
           StorageDead(_10);
-          goto -> bb4;
+          goto -> bb3;
       }
 
-      bb4: {
+      bb3: {
           StorageDead(_8);
           _11 = const {0x1 as *const [bool; 0]};
           _5 = const NonNull::<[bool; 0]> {{ pointer: {0x1 as *const [bool; 0]} }};
@@ -90,7 +90,14 @@
           _1 = const A {{ foo: Box::<[bool]>(Unique::<[bool]> {{ pointer: NonNull::<[bool]> {{ pointer: Indirect { alloc_id: ALLOC2, offset: Size(0 bytes) }: *const [bool] }}, _marker: PhantomData::<[bool]> }}, std::alloc::Global) }};
           StorageDead(_2);
           _0 = const ();
-          drop(_1) -> [return: bb1, unwind unreachable];
+          _12 = &raw mut _1;
+          _14 = &raw mut ((*_12).0: std::boxed::Box<[bool]>);
+          _13 = std::ptr::drop_in_place::<Box<[bool]>>(move _14) -> [return: bb4, unwind unreachable];
+      }
+  
+      bb4: {
+          StorageDead(_1);
+          return;
       }
   }