Drop elaboration: Only create a reset block if there are flags to reset.

compiler/rustc_mir_transform/src/elaborate_drop.rs

@@ -11,6 +11,7 @@ use rustc_middle::ty::adjustment::PointerCoercion;
 use rustc_middle::ty::util::{Discr, IntTypeExt};
 use rustc_middle::ty::{self, GenericArgsRef, Ty, TyCtxt};
 use rustc_middle::{bug, span_bug};
+use rustc_mir_dataflow::DropFlagState;
 use rustc_span::{DUMMY_SP, dummy_spanned};
 use tracing::{debug, instrument};
 
@@ -107,11 +108,10 @@ pub(crate) trait DropElaborator<'a, 'tcx>: fmt::Debug {
     /// Returns the drop flag of `path` as a MIR `Operand` (or `None` if `path` has no drop flag).
     fn get_drop_flag(&mut self, path: Self::Path) -> Option<Operand<'tcx>>;
 
-    /// Modifies the MIR patch so that the drop flag of `path` (if any) is cleared at `location`.
+    /// Return the drop flag of `path`, if any.
     ///
-    /// If `mode` is deep, drop flags of all child paths should also be cleared by inserting
-    /// additional statements.
-    fn clear_drop_flag(&mut self, location: Location, path: Self::Path, mode: DropFlagMode);
+    /// If `mode` is deep, drop flags of all child paths should be returned.
+    fn drop_flags_for(&mut self, path: Self::Path, mode: DropFlagMode) -> Vec<Place<'tcx>>;
 
     // Subpaths
 
@@ -1568,10 +1568,20 @@ where
             // bother setting it.
             return succ;
         }
-        let block = self.new_block(unwind, TerminatorKind::Goto { target: succ });
-        let block_start = Location { block, statement_index: 0 };
-        self.elaborator.clear_drop_flag(block_start, self.path, mode);
-        block
+        let flags = self.elaborator.drop_flags_for(self.path, mode);
+        let statements: Vec<_> = flags
+            .into_iter()
+            .map(|flag| {
+                self.assign(
+                    flag,
+                    Rvalue::Use(self.constant_bool(DropFlagState::Absent.value()), WithRetag::Yes),
+                )
+            })
+            .collect();
+        if statements.is_empty() {
+            return succ;
+        }
+        self.new_block_with_statements(unwind, statements, TerminatorKind::Goto { target: succ })
     }
 
     #[instrument(level = "debug", skip(self), ret)]
@@ -1666,6 +1676,14 @@ where
         }))
     }
 
+    fn constant_bool(&self, val: bool) -> Operand<'tcx> {
+        Operand::Constant(Box::new(ConstOperand {
+            span: self.source_info.span,
+            user_ty: None,
+            const_: Const::from_bool(self.tcx(), val),
+        }))
+    }
+
     fn assign(&self, lhs: Place<'tcx>, rhs: Rvalue<'tcx>) -> Statement<'tcx> {
         Statement::new(self.source_info, StatementKind::Assign(Box::new((lhs, rhs))))
     }

compiler/rustc_mir_transform/src/elaborate_drops.rs

@@ -189,17 +189,23 @@ impl<'a, 'tcx> DropElaborator<'a, 'tcx> for ElaborateDropsCtxt<'a, 'tcx> {
         }
     }
 
-    fn clear_drop_flag(&mut self, loc: Location, path: Self::Path, mode: DropFlagMode) {
+    fn drop_flags_for(&mut self, path: Self::Path, mode: DropFlagMode) -> Vec<Place<'tcx>> {
+        let mut flags = vec![];
         match mode {
             DropFlagMode::Shallow => {
-                self.set_drop_flag(loc, path, DropFlagState::Absent);
+                if let Some(flag) = self.drop_flags[path] {
+                    flags.push(flag.into());
+                }
             }
             DropFlagMode::Deep => {
                 on_all_children_bits(self.move_data(), path, |child| {
-                    self.set_drop_flag(loc, child, DropFlagState::Absent)
+                    if let Some(flag) = self.drop_flags[child] {
+                        flags.push(flag.into());
+                    }
                 });
             }
         }
+        flags
     }
 
     fn field_subpath(&self, path: Self::Path, field: FieldIdx) -> Option<Self::Path> {

compiler/rustc_mir_transform/src/shim.rs

@@ -460,7 +460,9 @@ impl<'a, 'tcx> DropElaborator<'a, 'tcx> for DropShimElaborator<'a, 'tcx> {
         None
     }
 
-    fn clear_drop_flag(&mut self, _location: Location, _path: Self::Path, _mode: DropFlagMode) {}
+    fn drop_flags_for(&mut self, _path: Self::Path, _mode: DropFlagMode) -> Vec<Place<'tcx>> {
+        Vec::new()
+    }
 
     fn field_subpath(&self, _path: Self::Path, _field: FieldIdx) -> Option<Self::Path> {
         None

tests/mir-opt/box_partial_move.maybe_move.ElaborateDrops.diff

@@ -44,7 +44,7 @@
       bb4: {
           StorageDead(_3);
 -         drop(_2) -> [return: bb5, unwind continue];
-+         goto -> bb14;
++         goto -> bb13;
       }
 
       bb5: {
@@ -61,34 +61,30 @@
 +     }
 + 
 +     bb8: {
-+         goto -> bb5;
-+     }
-+ 
-+     bb9: {
 +         _6 = &mut _2;
-+         _7 = <Box<String> as Drop>::drop(move _6) -> [return: bb8, unwind: bb7];
++         _7 = <Box<String> as Drop>::drop(move _6) -> [return: bb5, unwind: bb7];
 +     }
 + 
-+     bb10 (cleanup): {
++     bb9 (cleanup): {
 +         _8 = &mut _2;
 +         _9 = <Box<String> as Drop>::drop(move _8) -> [return: bb7, unwind terminate(cleanup)];
 +     }
 + 
++     bb10: {
++         goto -> bb12;
++     }
++ 
 +     bb11: {
-+         goto -> bb13;
++         drop((*_10)) -> [return: bb8, unwind: bb9];
 +     }
 + 
 +     bb12: {
-+         drop((*_10)) -> [return: bb9, unwind: bb10];
++         switchInt(copy _5) -> [0: bb8, otherwise: bb11];
 +     }
 + 
 +     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;
++         goto -> bb10;
       }
   }
 

tests/mir-opt/building/fallible_struct_drop.build.ElaborateDrops.diff

@@ -404,34 +404,34 @@
       bb37: {
           StorageDead(_2);
 -         drop(_44) -> [return: bb38, unwind: bb78];
-+         goto -> bb93;
++         goto -> bb38;
       }
 
       bb38: {
           StorageDead(_44);
 -         drop(_34) -> [return: bb39, unwind: bb83];
-+         goto -> bb94;
++         goto -> bb39;
       }
 
       bb39: {
 +         _53 = const false;
           StorageDead(_34);
 -         drop(_24) -> [return: bb40, unwind: bb87];
-+         goto -> bb95;
++         goto -> bb40;
       }
 
       bb40: {
 +         _54 = const false;
           StorageDead(_24);
 -         drop(_14) -> [return: bb41, unwind: bb90];
-+         goto -> bb96;
++         goto -> bb41;
       }
 
       bb41: {
 +         _55 = const false;
           StorageDead(_14);
 -         drop(_4) -> [return: bb42, unwind continue];
-+         goto -> bb97;
++         goto -> bb42;
       }
 
       bb42: {
@@ -459,7 +459,7 @@
           StorageDead(_3);
           StorageDead(_2);
 -         drop(_44) -> [return: bb47, unwind: bb78];
-+         goto -> bb98;
++         goto -> bb47;
       }
 
       bb47: {
@@ -485,7 +485,7 @@
 
       bb51: {
 -         drop(_34) -> [return: bb52, unwind: bb83];
-+         goto -> bb99;
++         goto -> bb52;
       }
 
       bb52: {
@@ -507,7 +507,7 @@
 
       bb55: {
 -         drop(_24) -> [return: bb56, unwind: bb87];
-+         goto -> bb100;
++         goto -> bb56;
       }
 
       bb56: {
@@ -524,7 +524,7 @@
 
       bb58: {
 -         drop(_14) -> [return: bb59, unwind: bb90];
-+         goto -> bb101;
++         goto -> bb59;
       }
 
       bb59: {
@@ -535,7 +535,7 @@
 
       bb60: {
 -         drop(_4) -> [return: bb61, unwind continue];
-+         goto -> bb102;
++         goto -> bb61;
       }
 
       bb61: {
@@ -681,46 +681,6 @@
 
       bb92 (cleanup): {
           resume;
-+     }
-+ 
-+     bb93: {
-+         goto -> bb38;
-+     }
-+ 
-+     bb94: {
-+         goto -> bb39;
-+     }
-+ 
-+     bb95: {
-+         goto -> bb40;
-+     }
-+ 
-+     bb96: {
-+         goto -> bb41;
-+     }
-+ 
-+     bb97: {
-+         goto -> bb42;
-+     }
-+ 
-+     bb98: {
-+         goto -> bb47;
-+     }
-+ 
-+     bb99: {
-+         goto -> bb52;
-+     }
-+ 
-+     bb100: {
-+         goto -> bb56;
-+     }
-+ 
-+     bb101: {
-+         goto -> bb59;
-+     }
-+ 
-+     bb102: {
-+         goto -> bb61;
       }
   }
 

tests/mir-opt/coroutine/async_await.a-{closure#0}.StateTransform.diff

@@ -30,7 +30,7 @@
 +         _6 = std::future::ResumeTy(move _7);
 +         _5 = copy (_1.0: &mut {async fn body of a()});
 +         _4 = discriminant((*_5));
-+         switchInt(move _4) -> [0: bb5, 1: bb3, otherwise: bb4];
++         switchInt(move _4) -> [0: bb4, 1: bb2, otherwise: bb3];
       }
 
       bb1: {
@@ -42,20 +42,16 @@
 -     bb2 (cleanup): {
 -         resume;
 +     bb2: {
-+         goto -> bb1;
++         assert(const false, "`async fn` resumed after completion") -> [success: bb2, unwind continue];
 +     }
 + 
 +     bb3: {
-+         assert(const false, "`async fn` resumed after completion") -> [success: bb3, unwind continue];
-+     }
-+ 
-+     bb4: {
 +         unreachable;
 +     }
 + 
-+     bb5: {
++     bb4: {
 +         _3 = const ();
-+         goto -> bb2;
++         goto -> bb1;
       }
   }
 

tests/mir-opt/coroutine/async_await.b-{closure#0}.StateTransform.diff

@@ -86,7 +86,7 @@
 +         _40 = std::future::ResumeTy(move _41);
 +         _39 = copy (_1.0: &mut {async fn body of b()});
 +         _38 = discriminant((*_39));
-+         switchInt(move _38) -> [0: bb47, 1: bb46, 2: bb45, 3: bb43, 4: bb44, otherwise: bb7];
++         switchInt(move _38) -> [0: bb46, 1: bb45, 2: bb44, 3: bb42, 4: bb43, otherwise: bb7];
       }
 
       bb1: {
@@ -309,7 +309,7 @@
       bb23: {
           StorageDead(_21);
 -         drop(_1) -> [return: bb24, unwind: bb50];
-+         goto -> bb41;
++         goto -> bb24;
       }
 
       bb24: {
@@ -486,22 +486,18 @@
 
 -     bb50 (cleanup): {
 +     bb40 (cleanup): {
-+         goto -> bb42;
-+     }
-+ 
-+     bb41: {
-+         goto -> bb24;
++         goto -> bb41;
 +     }
 + 
-+     bb42 (cleanup): {
++     bb41 (cleanup): {
 +         discriminant((*_39)) = 2;
           resume;
       }
 
 -     bb51 (cleanup): {
 -         StorageDead(_23);
 -         goto -> bb52;
-+     bb43: {
++     bb42: {
 +         StorageLive(_3);
 +         StorageLive(_4);
 +         StorageLive(_19);
@@ -513,7 +509,7 @@
 -     bb52 (cleanup): {
 -         StorageDead(_21);
 -         goto -> bb55;
-+     bb44: {
++     bb43: {
 +         StorageLive(_21);
 +         StorageLive(_35);
 +         StorageLive(_36);
@@ -524,21 +520,21 @@
 -     bb53 (cleanup): {
 -         StorageDead(_6);
 -         goto -> bb54;
-+     bb45: {
-+         assert(const false, "`async fn` resumed after panicking") -> [success: bb45, unwind continue];
++     bb44: {
++         assert(const false, "`async fn` resumed after panicking") -> [success: bb44, unwind continue];
       }
 
 -     bb54 (cleanup): {
 -         StorageDead(_4);
 -         StorageDead(_3);
 -         goto -> bb55;
-+     bb46: {
-+         assert(const false, "`async fn` resumed after completion") -> [success: bb46, unwind continue];
++     bb45: {
++         assert(const false, "`async fn` resumed after completion") -> [success: bb45, unwind continue];
       }
 
 -     bb55 (cleanup): {
 -         drop(_1) -> [return: bb50, unwind terminate(cleanup)];
-+     bb47: {
++     bb46: {
 +         StorageLive(_3);
 +         StorageLive(_4);
 +         StorageLive(_5);