Stop clearing box's drop flags early

compiler/rustc_mir_dataflow/src/elaborate_drops.rs

@@ -460,13 +460,19 @@ where
 
         if adt.is_box() {
             // we need to drop the inside of the box before running the destructor
-            let succ = self.destructor_call_block(contents_drop);
+            let destructor = self.destructor_call_block(contents_drop);
             let unwind = contents_drop
                 .1
                 .map(|unwind| self.destructor_call_block((unwind, Unwind::InCleanup)));
 
-            self.open_drop_for_box_contents(adt, args, succ, unwind)
+            let boxed_drop = self.open_drop_for_box_contents(adt, args, destructor, unwind);
+
+            // the drop flag will be at the end of contents_drop
+            self.drop_flag_test_block(boxed_drop, self.succ, unwind)
         } else if adt.has_dtor(self.tcx()) {
+            // We don't need to test drop flags here because
+            // this path is only taken with DropShimElaborator
+            // where testing drop flags is a noop
             self.destructor_call_block(contents_drop)
         } else {
             contents_drop.0
@@ -659,13 +665,7 @@ where
             }),
             is_cleanup: unwind.is_cleanup(),
         };
-
-        let destructor_block = self.elaborator.patch().new_block(result);
-
-        let block_start = Location { block: destructor_block, statement_index: 0 };
-        self.elaborator.clear_drop_flag(block_start, self.path, DropFlagMode::Shallow);
-
-        self.drop_flag_test_block(destructor_block, succ, unwind)
+        self.elaborator.patch().new_block(result)
     }
 
     /// Create a loop that drops an array:

compiler/rustc_mir_transform/src/elaborate_drops.rs

@@ -241,8 +241,8 @@ struct ElaborateDropsCtxt<'a, 'tcx> {
 }
 
 impl fmt::Debug for ElaborateDropsCtxt<'_, '_> {
-    fn fmt(&self, _f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        Ok(())
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("ElaborateDropsCtxt").finish_non_exhaustive()
     }
 }
 

compiler/rustc_mir_transform/src/shim.rs

@@ -339,8 +339,8 @@ pub(super) struct DropShimElaborator<'a, 'tcx> {
 }
 
 impl fmt::Debug for DropShimElaborator<'_, '_> {
-    fn fmt(&self, _f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
-        Ok(())
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
+        f.debug_struct("DropShimElaborator").finish_non_exhaustive()
     }
 }
 

tests/mir-opt/box_conditional_drop_allocator.main.ElaborateDrops.diff

@@ -0,0 +1,174 @@
+- // MIR for `main` before ElaborateDrops
++ // MIR for `main` after ElaborateDrops
+
+  fn main() -> () {
+      let mut _0: ();
+      let _1: std::boxed::Box<HasDrop, DropAllocator>;
+      let mut _2: HasDrop;
+      let mut _3: DropAllocator;
+      let mut _4: bool;
+      let _5: ();
+      let mut _6: HasDrop;
+      let _7: ();
+      let mut _8: std::boxed::Box<HasDrop, DropAllocator>;
++     let mut _9: bool;
++     let mut _10: &mut std::boxed::Box<HasDrop, DropAllocator>;
++     let mut _11: ();
++     let mut _12: &mut std::boxed::Box<HasDrop, DropAllocator>;
++     let mut _13: ();
++     let mut _14: &mut std::boxed::Box<HasDrop, DropAllocator>;
++     let mut _15: ();
+      scope 1 {
+          debug b => _1;
+      }
+
+      bb0: {
++         _9 = const false;
+          StorageLive(_1);
+          StorageLive(_2);
+          _2 = HasDrop;
+          StorageLive(_3);
+          _3 = DropAllocator;
+          _1 = Box::<HasDrop, DropAllocator>::new_in(move _2, move _3) -> [return: bb1, unwind: bb11];
+      }
+
+      bb1: {
++         _9 = const true;
+          StorageDead(_3);
+          StorageDead(_2);
+          StorageLive(_4);
+          _4 = const true;
+          switchInt(move _4) -> [0: bb4, otherwise: bb2];
+      }
+
+      bb2: {
+          StorageLive(_5);
+          StorageLive(_6);
+          _6 = move (*_1);
+          _5 = std::mem::drop::<HasDrop>(move _6) -> [return: bb3, unwind: bb9];
+      }
+
+      bb3: {
+          StorageDead(_6);
+          StorageDead(_5);
+          _0 = const ();
+          goto -> bb6;
+      }
+
+      bb4: {
+          StorageLive(_7);
+          StorageLive(_8);
++         _9 = const false;
+          _8 = move _1;
+          _7 = std::mem::drop::<Box<HasDrop, DropAllocator>>(move _8) -> [return: bb5, unwind: bb8];
+      }
+
+      bb5: {
+          StorageDead(_8);
+          StorageDead(_7);
+          _0 = const ();
+          goto -> bb6;
+      }
+
+      bb6: {
+          StorageDead(_4);
+-         drop(_1) -> [return: bb7, unwind continue];
++         goto -> bb22;
+      }
+
+      bb7: {
++         _9 = const false;
+          StorageDead(_1);
+          return;
+      }
+
+      bb8 (cleanup): {
+-         drop(_8) -> [return: bb10, unwind terminate(cleanup)];
++         goto -> bb10;
+      }
+
+      bb9 (cleanup): {
+-         drop(_6) -> [return: bb10, unwind terminate(cleanup)];
++         goto -> bb10;
+      }
+
+      bb10 (cleanup): {
+-         drop(_1) -> [return: bb13, unwind terminate(cleanup)];
++         goto -> bb27;
+      }
+
+      bb11 (cleanup): {
+-         drop(_3) -> [return: bb12, unwind terminate(cleanup)];
++         goto -> bb12;
+      }
+
+      bb12 (cleanup): {
+-         drop(_2) -> [return: bb13, unwind terminate(cleanup)];
++         goto -> bb13;
+      }
+
+      bb13 (cleanup): {
+          resume;
++     }
++ 
++     bb14: {
++         _9 = const false;
++         goto -> bb7;
++     }
++ 
++     bb15 (cleanup): {
++         drop((_1.1: DropAllocator)) -> [return: bb13, unwind terminate(cleanup)];
++     }
++ 
++     bb16 (cleanup): {
++         switchInt(copy _9) -> [0: bb13, otherwise: bb15];
++     }
++ 
++     bb17: {
++         drop((_1.1: DropAllocator)) -> [return: bb14, unwind: bb13];
++     }
++ 
++     bb18: {
++         switchInt(copy _9) -> [0: bb14, otherwise: bb17];
++     }
++ 
++     bb19: {
++         _10 = &mut _1;
++         _11 = <Box<HasDrop, DropAllocator> as Drop>::drop(move _10) -> [return: bb18, unwind: bb16];
++     }
++ 
++     bb20 (cleanup): {
++         _12 = &mut _1;
++         _13 = <Box<HasDrop, DropAllocator> as Drop>::drop(move _12) -> [return: bb16, unwind terminate(cleanup)];
++     }
++ 
++     bb21: {
++         goto -> bb19;
++     }
++ 
++     bb22: {
++         switchInt(copy _9) -> [0: bb7, otherwise: bb21];
++     }
++ 
++     bb23 (cleanup): {
++         drop((_1.1: DropAllocator)) -> [return: bb13, unwind terminate(cleanup)];
++     }
++ 
++     bb24 (cleanup): {
++         switchInt(copy _9) -> [0: bb13, otherwise: bb23];
++     }
++ 
++     bb25 (cleanup): {
++         _14 = &mut _1;
++         _15 = <Box<HasDrop, DropAllocator> as Drop>::drop(move _14) -> [return: bb24, unwind terminate(cleanup)];
++     }
++ 
++     bb26 (cleanup): {
++         goto -> bb25;
++     }
++ 
++     bb27 (cleanup): {
++         switchInt(copy _9) -> [0: bb13, otherwise: bb26];
+      }
+  }
+

tests/mir-opt/box_conditional_drop_allocator.rs

@@ -0,0 +1,38 @@
+// skip-filecheck
+//@ test-mir-pass: ElaborateDrops
+#![feature(allocator_api)]
+
+// Regression test for #131082.
+// Testing that the allocator of a Box is dropped in conditional drops
+
+use std::alloc::{AllocError, Allocator, Global, Layout};
+use std::ptr::NonNull;
+
+struct DropAllocator;
+
+unsafe impl Allocator for DropAllocator {
+    fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
+        Global.allocate(layout)
+    }
+    unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
+        Global.deallocate(ptr, layout);
+    }
+}
+impl Drop for DropAllocator {
+    fn drop(&mut self) {}
+}
+
+struct HasDrop;
+impl Drop for HasDrop {
+    fn drop(&mut self) {}
+}
+
+// EMIT_MIR box_conditional_drop_allocator.main.ElaborateDrops.diff
+fn main() {
+    let b = Box::new_in(HasDrop, DropAllocator);
+    if true {
+        drop(*b);
+    } else {
+        drop(b);
+    }
+}

tests/ui/drop/box-conditional-drop-allocator.rs

@@ -0,0 +1,43 @@
+//@ run-pass
+#![feature(allocator_api)]
+
+// Regression test for #131082.
+// Testing that the allocator of a Box is dropped in conditional drops
+
+use std::alloc::{AllocError, Allocator, Global, Layout};
+use std::cell::Cell;
+use std::ptr::NonNull;
+
+struct DropCheckingAllocator<'a>(&'a Cell<bool>);
+
+unsafe impl Allocator for DropCheckingAllocator<'_> {
+    fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
+        Global.allocate(layout)
+    }
+    unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
+        Global.deallocate(ptr, layout);
+    }
+}
+impl Drop for DropCheckingAllocator<'_> {
+    fn drop(&mut self) {
+        self.0.set(true);
+    }
+}
+
+struct HasDrop;
+impl Drop for HasDrop {
+    fn drop(&mut self) {}
+}
+
+fn main() {
+    let dropped = Cell::new(false);
+    {
+        let b = Box::new_in(HasDrop, DropCheckingAllocator(&dropped));
+        if true {
+            drop(*b);
+        } else {
+            drop(b);
+        }
+    }
+    assert!(dropped.get());
+}