Don't project into NonNull when dropping a Box

Author: scottmcm

Diff for: compiler/rustc_middle/src/mir/tcx.rs

@@ -86,6 +86,14 @@ impl<'tcx> PlaceTy<'tcx> {
         }
     }
 
+    pub fn multi_projection_ty(
+        self,
+        tcx: TyCtxt<'tcx>,
+        elems: &[PlaceElem<'tcx>],
+    ) -> PlaceTy<'tcx> {
+        elems.iter().fold(self, |place_ty, &elem| place_ty.projection_ty(tcx, elem))
+    }
+
     /// Convenience wrapper around `projection_ty_core` for
     /// `PlaceElem`, where we can just use the `Ty` that is already
     /// stored inline on field projection elems.
@@ -167,11 +175,7 @@ impl<'tcx> Place<'tcx> {
     where
         D: HasLocalDecls<'tcx>,
     {
-        projection
-            .iter()
-            .fold(PlaceTy::from_ty(local_decls.local_decls()[local].ty), |place_ty, &elem| {
-                place_ty.projection_ty(tcx, elem)
-            })
+        PlaceTy::from_ty(local_decls.local_decls()[local].ty).multi_projection_ty(tcx, projection)
     }
 
     pub fn ty<D: ?Sized>(&self, local_decls: &D, tcx: TyCtxt<'tcx>) -> PlaceTy<'tcx>

Diff for: compiler/rustc_mir_transform/src/elaborate_drop.rs

@@ -89,6 +89,7 @@ pub(crate) trait DropElaborator<'a, 'tcx>: fmt::Debug {
 
     // Accessors
 
+    fn patch_ref(&self) -> &MirPatch<'tcx>;
     fn patch(&mut self) -> &mut MirPatch<'tcx>;
     fn body(&self) -> &'a Body<'tcx>;
     fn tcx(&self) -> TyCtxt<'tcx>;
@@ -180,7 +181,14 @@ where
 {
     #[instrument(level = "trace", skip(self), ret)]
     fn place_ty(&self, place: Place<'tcx>) -> Ty<'tcx> {
-        place.ty(self.elaborator.body(), self.tcx()).ty
+        if place.local < self.elaborator.body().local_decls.next_index() {
+            place.ty(self.elaborator.body(), self.tcx()).ty
+        } else {
+            // We don't have a slice with all the locals, since some are in the patch.
+            tcx::PlaceTy::from_ty(self.elaborator.patch_ref().local_ty(place.local))
+                .multi_projection_ty(self.elaborator.tcx(), place.projection)
+                .ty
+        }
     }
 
     fn tcx(&self) -> TyCtxt<'tcx> {
@@ -410,12 +418,26 @@ where
 
         let unique_place = self.tcx().mk_place_field(self.place, FieldIdx::ZERO, unique_ty);
         let nonnull_place = self.tcx().mk_place_field(unique_place, FieldIdx::ZERO, nonnull_ty);
-        let ptr_place = self.tcx().mk_place_field(nonnull_place, FieldIdx::ZERO, ptr_ty);
-        let interior = self.tcx().mk_place_deref(ptr_place);
 
+        let ptr_local = self.new_temp(ptr_ty);
+
+        let interior = self.tcx().mk_place_deref(Place::from(ptr_local));
         let interior_path = self.elaborator.deref_subpath(self.path);
 
-        self.drop_subpath(interior, interior_path, succ, unwind)
+        let do_drop_bb = self.drop_subpath(interior, interior_path, succ, unwind);
+
+        let setup_bbd = BasicBlockData {
+            statements: vec![self.assign(
+                Place::from(ptr_local),
+                Rvalue::Cast(CastKind::Transmute, Operand::Copy(nonnull_place), ptr_ty),
+            )],
+            terminator: Some(Terminator {
+                kind: TerminatorKind::Goto { target: do_drop_bb },
+                source_info: self.source_info,
+            }),
+            is_cleanup: unwind.is_cleanup(),
+        };
+        self.elaborator.patch().new_block(setup_bbd)
     }
 
     #[instrument(level = "debug", ret)]

Diff for: compiler/rustc_mir_transform/src/elaborate_drops.rs

@@ -138,6 +138,10 @@ impl InitializationData<'_, '_> {
 impl<'a, 'tcx> DropElaborator<'a, 'tcx> for ElaborateDropsCtxt<'a, 'tcx> {
     type Path = MovePathIndex;
 
+    fn patch_ref(&self) -> &MirPatch<'tcx> {
+        &self.patch
+    }
+
     fn patch(&mut self) -> &mut MirPatch<'tcx> {
         &mut self.patch
     }

Diff for: compiler/rustc_mir_transform/src/patch.rs

@@ -166,6 +166,14 @@ impl<'tcx> MirPatch<'tcx> {
         Local::new(index)
     }
 
+    /// Returns the type of a local that's newly-added in the patch.
+    pub(crate) fn local_ty(&self, local: Local) -> Ty<'tcx> {
+        let local = local.as_usize();
+        assert!(local < self.next_local);
+        let new_local_idx = self.new_locals.len() - (self.next_local - local);
+        self.new_locals[new_local_idx].ty
+    }
+
     pub(crate) fn new_block(&mut self, data: BasicBlockData<'tcx>) -> BasicBlock {
         let block = BasicBlock::new(self.patch_map.len());
         debug!("MirPatch: new_block: {:?}: {:?}", block, data);

Diff for: compiler/rustc_mir_transform/src/shim.rs

@@ -350,6 +350,9 @@ impl fmt::Debug for DropShimElaborator<'_, '_> {
 impl<'a, 'tcx> DropElaborator<'a, 'tcx> for DropShimElaborator<'a, 'tcx> {
     type Path = ();
 
+    fn patch_ref(&self) -> &MirPatch<'tcx> {
+        &self.patch
+    }
     fn patch(&mut self) -> &mut MirPatch<'tcx> {
         &mut self.patch
     }

Diff for: tests/mir-opt/box_expr.main.ElaborateDrops.diff

@@ -12,6 +12,7 @@
       let mut _7: std::boxed::Box<S>;
 +     let mut _8: &mut std::boxed::Box<S>;
 +     let mut _9: ();
++     let mut _10: *const S;
       scope 1 {
           debug x => _1;
       }
@@ -68,7 +69,7 @@
 
       bb8 (cleanup): {
 -         drop(_5) -> [return: bb9, unwind terminate(cleanup)];
-+         goto -> bb11;
++         goto -> bb12;
       }
 
       bb9 (cleanup): {
@@ -82,6 +83,11 @@
 + 
 +     bb11 (cleanup): {
 +         goto -> bb10;
++     }
++ 
++     bb12 (cleanup): {
++         _10 = copy ((_5.0: std::ptr::Unique<S>).0: std::ptr::NonNull<S>) as *const S (Transmute);
++         goto -> bb11;
       }
   }
 

Diff for: tests/mir-opt/box_partial_move.maybe_move.ElaborateDrops.diff

@@ -0,0 +1,94 @@
+- // MIR for `maybe_move` before ElaborateDrops
++ // MIR for `maybe_move` after ElaborateDrops
+  
+  fn maybe_move(_1: bool, _2: Box<String>) -> Option<String> {
+      debug cond => _1;
+      debug thing => _2;
+      let mut _0: std::option::Option<std::string::String>;
+      let mut _3: bool;
+      let mut _4: std::string::String;
++     let mut _5: bool;
++     let mut _6: &mut std::boxed::Box<std::string::String>;
++     let mut _7: ();
++     let mut _8: &mut std::boxed::Box<std::string::String>;
++     let mut _9: ();
++     let mut _10: *const std::string::String;
+  
+      bb0: {
++         _5 = const false;
++         _5 = const true;
+          StorageLive(_3);
+          _3 = copy _1;
+          switchInt(move _3) -> [0: bb3, otherwise: bb1];
+      }
+  
+      bb1: {
+          StorageLive(_4);
++         _5 = const false;
+          _4 = move (*_2);
+          _0 = Option::<String>::Some(move _4);
+-         drop(_4) -> [return: bb2, unwind: bb6];
++         goto -> bb2;
+      }
+  
+      bb2: {
+          StorageDead(_4);
+          goto -> bb4;
+      }
+  
+      bb3: {
+          _0 = Option::<String>::None;
+          goto -> bb4;
+      }
+  
+      bb4: {
+          StorageDead(_3);
+-         drop(_2) -> [return: bb5, unwind continue];
++         goto -> bb14;
+      }
+  
+      bb5: {
+          return;
+      }
+  
+      bb6 (cleanup): {
+-         drop(_2) -> [return: bb7, unwind terminate(cleanup)];
++         goto -> bb7;
+      }
+  
+      bb7 (cleanup): {
+          resume;
++     }
++ 
++     bb8: {
++         goto -> bb5;
++     }
++ 
++     bb9: {
++         _6 = &mut _2;
++         _7 = <Box<String> as Drop>::drop(move _6) -> [return: bb8, unwind: bb7];
++     }
++ 
++     bb10 (cleanup): {
++         _8 = &mut _2;
++         _9 = <Box<String> as Drop>::drop(move _8) -> [return: bb7, unwind terminate(cleanup)];
++     }
++ 
++     bb11: {
++         goto -> bb13;
++     }
++ 
++     bb12: {
++         drop((*_10)) -> [return: bb9, unwind: bb10];
++     }
++ 
++     bb13: {
++         switchInt(copy _5) -> [0: bb9, otherwise: bb12];
++     }
++ 
++     bb14: {
++         _10 = copy ((_2.0: std::ptr::Unique<std::string::String>).0: std::ptr::NonNull<std::string::String>) as *const std::string::String (Transmute);
++         goto -> bb11;
+      }
+  }
+  

Diff for: tests/mir-opt/box_partial_move.rs

@@ -0,0 +1,17 @@
+//@ test-mir-pass: ElaborateDrops
+//@ needs-unwind
+
+#![feature(rustc_attrs, liballoc_internals)]
+
+// EMIT_MIR box_partial_move.maybe_move.ElaborateDrops.diff
+fn maybe_move(cond: bool, thing: Box<String>) -> Option<String> {
+    // CHECK-LABEL: fn maybe_move(
+    // CHECK: let mut [[PTR:_[0-9]+]]: *const std::string::String;
+    // CHECK: [[PTR]] = copy ((_2.0: std::ptr::Unique<std::string::String>).0: std::ptr::NonNull<std::string::String>) as *const std::string::String (Transmute);
+    // CHECK: drop((*[[PTR]]))
+    if cond { Some(*thing) } else { None }
+}
+
+fn main() {
+    maybe_move(false, Box::new("hello".to_string()));
+}