stop clearing box's drop flags early

Author: beepster4096

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/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());
+}