Allow inlining all the drops

Author: scottmcm

compiler/rustc_mir_transform/src/inline.rs

@@ -13,10 +13,10 @@ 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;
@@ -75,6 +75,87 @@ impl<'tcx> MirPass<'tcx> for Inline {
     }
 }
 
+/// Change all the `Drop(place)` terminators to `drop_in_place(&place)` calls.
+///
+/// This needs to be done before starting the inlining analysis or else the cycle
+/// detection logic will miss things. But it's a waste of time to do the work of
+/// adding locals unless we're going to try to inline things, so might as well
+/// leave the drop terminators alone in lower optimization levels.
+fn lower_drops_to_calls<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>, param_env: ParamEnv<'tcx>) {
+    // I wish this could just be `require_lang_item`, but there are so many
+    // obnoxious `no_core` tests that I'd have to update to do that.
+    let Some(drop_in_place_fn) = tcx.lang_items().drop_in_place_fn() else {
+        error!("No `drop_in_place` function; not even trying to simplify drops!");
+        return;
+    };
+
+    // We need a unit local to store the "result" of the `drop_in_place` call
+    let mut unit_local = None;
+
+    // Changing `Drop` to `Call` actually preserves the CFG because
+    // it keeps the same `target` and the same `unwind` action.
+    for block in body.basic_blocks.as_mut_preserves_cfg().iter_mut() {
+        let terminator = block.terminator.as_mut().unwrap();
+        let TerminatorKind::Drop { place: dropped_place, target, unwind, replace: _ } =
+            terminator.kind
+        else {
+            continue;
+        };
+
+        if block.is_cleanup {
+            // We don't inline into cleanup blocks, so don't rewrite drops in them either.
+            continue;
+        }
+
+        let dropped_ty = dropped_place.ty(&body.local_decls, tcx).ty;
+        if !dropped_ty.needs_drop(tcx, param_env) {
+            // Leave it for RemoveUnneededDrops to remove, since that's
+            // cheaper than doing the work to inline an empty shim.
+            continue;
+        }
+
+        if !tcx.consider_optimizing(|| format!("Drop -> drop_in_place on type {dropped_ty:?}")) {
+            return;
+        }
+
+        let dropped_operand = if let [PlaceElem::Deref] = **dropped_place.projection
+            && body.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(tcx, dropped_ty);
+            let ptr_local =
+                body.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 = *unit_local.get_or_insert_with(|| {
+            let unit_ty = tcx.types.unit;
+            body.local_decls.push(LocalDecl::new(unit_ty, terminator.source_info.span))
+        });
+        terminator.kind = TerminatorKind::Call {
+            func: Operand::function_handle(
+                tcx,
+                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 inline<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) -> bool {
     let def_id = body.source.def_id().expect_local();
 
@@ -94,6 +175,7 @@ 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);
+    lower_drops_to_calls(tcx, body, param_env);
 
     let mut this = Inliner {
         tcx,
@@ -509,7 +591,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)

compiler/rustc_mir_transform/src/remove_unneeded_drops.rs

@@ -4,6 +4,7 @@
 //! useful because (unlike MIR building) it runs after type checking, so it can make use of
 //! `Reveal::All` to provide more precise type information.
 
+use rustc_hir::LangItem;
 use rustc_middle::mir::*;
 use rustc_middle::ty::TyCtxt;
 
@@ -21,18 +22,28 @@ impl<'tcx> MirPass<'tcx> for RemoveUnneededDrops {
 
         for block in body.basic_blocks.as_mut() {
             let terminator = block.terminator_mut();
-            if let TerminatorKind::Drop { place, target, .. } = terminator.kind {
-                let ty = place.ty(&body.local_decls, tcx);
-                if ty.ty.needs_drop(tcx, param_env) {
-                    continue;
+            let (dropped_ty, target) = match terminator.kind {
+                TerminatorKind::Drop { place, target, .. } => {
+                    (place.ty(&body.local_decls, tcx).ty, target)
                 }
-                if !tcx.consider_optimizing(|| format!("RemoveUnneededDrops {did:?} ")) {
-                    continue;
+                TerminatorKind::Call { ref func, target, .. }
+                    if let Some((def_id, generics)) = func.const_fn_def()
+                        && tcx.is_lang_item(def_id, LangItem::DropInPlace) =>
+                {
+                    (generics.type_at(0), target.unwrap())
                 }
-                debug!("SUCCESS: replacing `drop` with goto({:?})", target);
-                terminator.kind = TerminatorKind::Goto { target };
-                should_simplify = true;
+                _ => continue,
+            };
+
+            if dropped_ty.needs_drop(tcx, param_env) {
+                continue;
+            }
+            if !tcx.consider_optimizing(|| format!("RemoveUnneededDrops {did:?} ")) {
+                continue;
             }
+            debug!("SUCCESS: replacing `drop` with goto({:?})", target);
+            terminator.kind = TerminatorKind::Goto { target };
+            should_simplify = true;
         }
 
         // if we applied optimizations, we potentially have some cfg to cleanup to

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/codegen/issues/issue-112509-slice-get-andthen-get.rs

@@ -3,8 +3,8 @@
 
 // CHECK-LABEL: @write_u8_variant_a
 // CHECK-NEXT: {{.*}}:
-// CHECK-NEXT: icmp ugt
 // CHECK-NEXT: getelementptr
+// CHECK-NEXT: icmp ugt
 // CHECK-NEXT: select i1 {{.+}} null
 // CHECK-NEXT: insertvalue
 // CHECK-NEXT: insertvalue

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

@@ -5,32 +5,34 @@
       let mut _0: ();
       let _1: A;
       let mut _2: std::boxed::Box<[bool]>;
+      let mut _3: *mut A;
+      let mut _4: ();
       scope 1 {
           debug a => _1;
       }
       scope 2 (inlined <Box<[bool]> as Default>::default) {
-          let _3: std::ptr::Unique<[bool]>;
-          let mut _4: std::ptr::Unique<[bool; 0]>;
+          let _5: std::ptr::Unique<[bool]>;
+          let mut _6: std::ptr::Unique<[bool; 0]>;
           scope 3 {
           }
           scope 4 (inlined Unique::<[bool; 0]>::dangling) {
-              let mut _5: std::ptr::NonNull<[bool; 0]>;
+              let mut _7: std::ptr::NonNull<[bool; 0]>;
               scope 5 (inlined NonNull::<[bool; 0]>::dangling) {
-                  let _6: *mut [bool; 0];
+                  let _8: *mut [bool; 0];
                   scope 6 {
                       scope 10 (inlined NonNull::<[bool; 0]>::new_unchecked) {
-                          let mut _8: bool;
-                          let _9: ();
-                          let mut _10: *mut ();
-                          let mut _11: *const [bool; 0];
+                          let mut _10: bool;
+                          let _11: ();
+                          let mut _12: *mut ();
+                          let mut _13: *const [bool; 0];
                           scope 11 (inlined core::ub_checks::check_language_ub) {
                               scope 12 (inlined core::ub_checks::check_language_ub::runtime) {
                               }
                           }
                       }
                   }
                   scope 7 (inlined dangling_mut::<[bool; 0]>) {
-                      let mut _7: usize;
+                      let mut _9: usize;
                       scope 8 (inlined align_of::<[bool; 0]>) {
                       }
                       scope 9 (inlined without_provenance_mut::<[bool; 0]>) {
@@ -39,58 +41,61 @@
               }
           }
       }
+      scope 13 (inlined std::ptr::drop_in_place::<A> - shim(Some(A))) {
+      }
 
       bb0: {
           StorageLive(_1);
           StorageLive(_2);
-          StorageLive(_3);
-          StorageLive(_9);
-          StorageLive(_4);
           StorageLive(_5);
+          StorageLive(_11);
           StorageLive(_6);
           StorageLive(_7);
-          _7 = const 1_usize;
-          _6 = const {0x1 as *mut [bool; 0]};
-          StorageDead(_7);
-          StorageLive(_11);
           StorageLive(_8);
-          _8 = UbChecks();
-          switchInt(move _8) -> [0: bb4, otherwise: bb2];
+          StorageLive(_9);
+          _9 = const 1_usize;
+          _8 = const {0x1 as *mut [bool; 0]};
+          StorageDead(_9);
+          StorageLive(_13);
+          StorageLive(_10);
+          _10 = UbChecks();
+          switchInt(move _10) -> [0: bb3, otherwise: bb1];
       }
 
       bb1: {
-          StorageDead(_1);
-          return;
+          StorageLive(_12);
+          _12 = const {0x1 as *mut ()};
+          _11 = NonNull::<T>::new_unchecked::precondition_check(const {0x1 as *mut ()}) -> [return: bb2, unwind unreachable];
       }
 
       bb2: {
-          StorageLive(_10);
-          _10 = const {0x1 as *mut ()};
-          _9 = NonNull::<T>::new_unchecked::precondition_check(const {0x1 as *mut ()}) -> [return: bb3, unwind unreachable];
+          StorageDead(_12);
+          goto -> bb3;
       }
 
       bb3: {
           StorageDead(_10);
-          goto -> bb4;
-      }
-  
-      bb4: {
+          _13 = const {0x1 as *const [bool; 0]};
+          _7 = const NonNull::<[bool; 0]> {{ pointer: {0x1 as *const [bool; 0]} }};
+          StorageDead(_13);
           StorageDead(_8);
-          _11 = const {0x1 as *const [bool; 0]};
-          _5 = const NonNull::<[bool; 0]> {{ pointer: {0x1 as *const [bool; 0]} }};
-          StorageDead(_11);
+          _6 = const Unique::<[bool; 0]> {{ pointer: NonNull::<[bool; 0]> {{ pointer: {0x1 as *const [bool; 0]} }}, _marker: PhantomData::<[bool; 0]> }};
+          StorageDead(_7);
+          _5 = const Unique::<[bool]> {{ pointer: NonNull::<[bool]> {{ pointer: Indirect { alloc_id: ALLOC0, offset: Size(0 bytes) }: *const [bool] }}, _marker: PhantomData::<[bool]> }};
           StorageDead(_6);
-          _4 = const Unique::<[bool; 0]> {{ pointer: NonNull::<[bool; 0]> {{ pointer: {0x1 as *const [bool; 0]} }}, _marker: PhantomData::<[bool; 0]> }};
-          StorageDead(_5);
-          _3 = const Unique::<[bool]> {{ pointer: NonNull::<[bool]> {{ pointer: Indirect { alloc_id: ALLOC0, offset: Size(0 bytes) }: *const [bool] }}, _marker: PhantomData::<[bool]> }};
-          StorageDead(_4);
           _2 = const Box::<[bool]>(Unique::<[bool]> {{ pointer: NonNull::<[bool]> {{ pointer: Indirect { alloc_id: ALLOC1, offset: Size(0 bytes) }: *const [bool] }}, _marker: PhantomData::<[bool]> }}, std::alloc::Global);
-          StorageDead(_9);
-          StorageDead(_3);
+          StorageDead(_11);
+          StorageDead(_5);
           _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];
+          _3 = &raw mut _1;
+          drop(((*_3).0: std::boxed::Box<[bool]>)) -> [return: bb4, unwind unreachable];
+      }
+  
+      bb4: {
+          StorageDead(_1);
+          return;
       }
   }