Mark more locals as moved to avoid building drops for them.

compiler/rustc_middle/src/mir/syntax.rs

@@ -732,42 +732,41 @@ pub enum TerminatorKind<'tcx> {
 
     /// The behavior of this statement differs significantly before and after drop elaboration.
     ///
-    /// After drop elaboration: `Drop` terminators are a complete nop for types that have no drop
+    /// After drop elaboration, `Drop` terminators are a complete nop for types that have no drop
     /// glue. For other types, `Drop` terminators behave exactly like a call to
     /// `core::mem::drop_glue` with a reference to the given place.
     ///
-    /// `Drop` before drop elaboration is a *conditional* execution of the drop glue. Specifically,
-    /// the `Drop` will be executed if...
+    /// Before drop elaboration, `Drop` behave as a *conditional* execution of the drop glue.
+    /// Specifically, the drop glue is executed if, among all statements executed within this
+    /// `Body`, an assignment to the place occurred more recently than a move out of it.
+    /// If a place is partially assigned-to or partially moved-from, the drop glue is only executed
+    /// on the assigned-to part.
     ///
-    /// **Needs clarification**: End of that sentence. This in effect should document the exact
-    /// behavior of drop elaboration. The following sounds vaguely right, but I'm not quite sure:
+    /// This considers the contents of a `Box` to be a sub-place, but does not consider indirect
+    /// assignments through references or pointers.
     ///
-    /// > The drop glue is executed if, among all statements executed within this `Body`, an assignment to
-    /// > the place or one of its "parents" occurred more recently than a move out of it. This does not
-    /// > consider indirect assignments.
+    /// **Async drop processing**:
+    /// MIR building detects possible async drops, and constructs a complete CFG. Async drops await
+    /// a coroutine, which means that the current coroutine may be itself dropped while dropping
+    /// `place`. That's the 'drop' field, similar to `Yield` terminator.
     ///
-    /// The `replace` flag indicates whether this terminator was created as part of an assignment.
-    /// This should only be used for diagnostic purposes, and does not have any operational
-    /// meaning.
+    /// Drop elaboration later refines the set of useful async drops. If there is no need for an
+    /// async drop, it is downgraded to a sync drop by setting `drop` to `None` If this is an
+    /// actual async drop, it is expanded to an `await` loop over the `async_drop_in_place` or
+    /// `AsyncDrop::drop` coroutine.
     ///
-    /// Async drop processing:
-    ///   MIR building detects possible async drops, and constructs a complete CFG. To correctly
-    ///   handle the coroutine being dropped while itself drops, we need a 'drop' target
-    ///   similar to `Yield` terminator (see `drops.build_mir::<CoroutineDrop>`).
-    ///
-    ///   Drop elaboration later refines the set of useful async drops. If there is no need for an
-    ///   async drop, it is downgraded to a sync drop by setting `drop` to `None` If this is an
-    ///   actual async drop, it is expanded to an `await` loop over the `async_drop_in_place` or
-    ///   `AsyncDrop::drop` coroutine.
-    ///
-    ///   When a coroutine has any internal async drop, the coroutine drop function will be async
-    ///   (generated by `create_coroutine_drop_shim_async`, not `create_coroutine_drop_shim`).
+    /// When a coroutine has any internal async drop, the coroutine drop function will be async
+    /// (generated by `create_coroutine_drop_shim_async`, not `create_coroutine_drop_shim`).
     Drop {
         place: Place<'tcx>,
         target: BasicBlock,
         unwind: UnwindAction,
+        /// The `replace` flag indicates whether this terminator was created as part of an
+        /// assignment. This should only be used for diagnostic purposes, and does not have any
+        /// operational meaning.
         replace: bool,
-        /// Cleanup to be done if the coroutine is dropped at this suspend point (for async drop).
+        /// Cleanup to be done if the coroutine is dropped at this suspend point, for async drop.
+        /// It is set to `None` after drop elaboration.
         drop: Option<BasicBlock>,
     },
 
@@ -1288,6 +1287,8 @@ pub enum Operand<'tcx> {
 
     /// Creates a value by performing loading the place, just like the `Copy` operand.
     ///
+    /// Before drop elaboration, this unschedules drops on the given place.
+    ///
     /// This *may* additionally overwrite the place with `uninit` bytes, depending on how we decide
     /// in [UCG#188]. You should not emit MIR that may attempt a subsequent second load of this
     /// place without first re-initializing it.

compiler/rustc_mir_build/src/builder/expr/as_rvalue.rs

@@ -11,7 +11,7 @@ use rustc_middle::ty::adjustment::PointerCoercion;
 use rustc_middle::ty::cast::{CastTy, mir_cast_kind};
 use rustc_middle::ty::util::IntTypeExt;
 use rustc_middle::ty::{self, Ty, UpvarArgs};
-use rustc_span::{DUMMY_SP, Span, Spanned};
+use rustc_span::Span;
 use tracing::debug;
 
 use crate::builder::expr::as_place::PlaceBase;
@@ -74,6 +74,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                             NeedsTemporary::No
                         )
                     );
+                    this.record_operand_moved(&value_operand);
                     block.and(Rvalue::Repeat(value_operand, count))
                 }
             }
@@ -219,6 +220,9 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                     })
                     .collect();
 
+                for operand in fields.iter() {
+                    this.record_operand_moved(operand);
+                }
                 block.and(Rvalue::Aggregate(Box::new(AggregateKind::Array(el_ty)), fields))
             }
             ExprKind::Tuple { ref fields } => {
@@ -240,6 +244,9 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                     })
                     .collect();
 
+                for operand in fields.iter() {
+                    this.record_operand_moved(operand);
+                }
                 block.and(Rvalue::Aggregate(Box::new(AggregateKind::Tuple), fields))
             }
             ExprKind::Closure(ClosureExpr {
@@ -342,6 +349,9 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                         Box::new(AggregateKind::CoroutineClosure(closure_id.to_def_id(), args))
                     }
                 };
+                for operand in operands.iter() {
+                    this.record_operand_moved(operand);
+                }
                 block.and(Rvalue::Aggregate(result, operands))
             }
             ExprKind::Assign { .. } | ExprKind::AssignOp { .. } => {
@@ -424,6 +434,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                         NeedsTemporary::No,
                     )
                 );
+                this.record_operand_moved(&operand);
                 block.and(Rvalue::Use(operand, WithRetag::Yes))
             }
 
@@ -647,7 +658,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                 this.diverge_from(block);
                 block = success;
             }
-            this.record_operands_moved(&[Spanned { node: value_operand, span: DUMMY_SP }]);
+            this.record_operand_moved(&value_operand);
         }
         block.and(Rvalue::Aggregate(Box::new(AggregateKind::Array(elem_ty)), IndexVec::new()))
     }

compiler/rustc_mir_build/src/builder/expr/into.rs

@@ -5,6 +5,7 @@ use rustc_ast::{AsmMacro, InlineAsmOptions};
 use rustc_data_structures::fx::FxHashMap;
 use rustc_hir as hir;
 use rustc_hir::attrs::lang_items::LangItem;
+use rustc_index::IndexVec;
 use rustc_middle::mir::*;
 use rustc_middle::span_bug;
 use rustc_middle::thir::*;
@@ -488,7 +489,9 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
 
                 let success = this.cfg.start_new_block();
 
-                this.record_operands_moved(&args);
+                for operand in args.iter() {
+                    this.record_operand_moved(&operand.node);
+                }
 
                 debug!("expr_into_dest: fn_span={:?}", fn_span);
 
@@ -628,7 +631,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                 let variant = adt_def.variant(variant_index);
                 let field_names = variant.fields.indices();
 
-                let fields = match base {
+                let fields: IndexVec<_, _> = match base {
                     AdtExprBase::None => {
                         field_names.filter_map(|n| fields_map.get(&n).cloned()).collect()
                     }
@@ -693,6 +696,9 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                     user_ty,
                     active_field_index,
                 ));
+                for operand in fields.iter() {
+                    this.record_operand_moved(operand);
+                }
                 this.cfg.push_assign(
                     block,
                     source_info,
@@ -841,7 +847,9 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                 debug_assert!(Category::of(&expr.kind) == Some(Category::Place));
 
                 let place = unpack!(block = this.as_place(block, expr_id));
-                let rvalue = Rvalue::Use(this.consume_by_copy_or_move(place), WithRetag::Yes);
+                let operand = this.consume_by_copy_or_move(place);
+                this.record_operand_moved(&operand);
+                let rvalue = Rvalue::Use(operand, WithRetag::Yes);
                 this.cfg.push_assign(block, source_info, destination, rvalue);
                 block.unit()
             }
@@ -867,6 +875,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                     block =
                         this.as_operand(block, scope, value, LocalInfo::Boring, NeedsTemporary::No)
                 );
+                this.record_operand_moved(&value);
                 let resume = this.cfg.start_new_block();
                 this.cfg.terminate(
                     block,

compiler/rustc_mir_build/src/builder/expr/stmt.rs

@@ -126,7 +126,9 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                         })
                         .collect();
 
-                    this.record_operands_moved(&args);
+                    for operand in args.iter() {
+                        this.record_operand_moved(&operand.node);
+                    }
 
                     debug!("expr_into_dest: fn_span={:?}", fn_span);
 

compiler/rustc_mir_build/src/builder/scope.rs

@@ -86,6 +86,7 @@ use std::mem;
 use interpret::ErrorHandled;
 use rustc_data_structures::fx::FxHashMap;
 use rustc_hir::HirId;
+use rustc_index::bit_set::GrowableBitSet;
 use rustc_index::{IndexSlice, IndexVec};
 use rustc_middle::middle::region;
 use rustc_middle::mir::{self, *};
@@ -137,7 +138,7 @@ struct Scope {
     /// end of the vector (top of the stack) first.
     drops: Vec<DropData>,
 
-    moved_locals: Vec<Local>,
+    moved_locals: GrowableBitSet<Local>,
 
     /// The drop index that will drop everything in and below this scope on an
     /// unwind path.
@@ -494,7 +495,7 @@ impl<'tcx> Scopes<'tcx> {
             source_scope: vis_scope,
             region_scope,
             drops: vec![],
-            moved_locals: vec![],
+            moved_locals: GrowableBitSet::new_empty(),
             cached_unwind_block: None,
             cached_coroutine_drop_block: None,
         });
@@ -1522,7 +1523,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
         self.schedule_drop(span, region_scope, local, DropKind::ForLint);
     }
 
-    /// Indicates that the "local operand" stored in `local` is
+    /// Indicates that the "local operand" stored in `operand` is
     /// *moved* at some point during execution (see `local_scope` for
     /// more information about what a "local operand" is -- in short,
     /// it's an intermediate operand created as part of preparing some
@@ -1558,26 +1559,22 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
     /// spurious borrow-check errors -- the problem, ironically, is
     /// not the `DROP(_X)` itself, but the (spurious) unwind pathways
     /// that it creates. See #64391 for an example.
-    pub(crate) fn record_operands_moved(&mut self, operands: &[Spanned<Operand<'tcx>>]) {
+    #[instrument(level = "debug", skip(self))]
+    pub(crate) fn record_operand_moved(&mut self, operand: &Operand<'tcx>) {
         let local_scope = self.local_scope();
         let scope = self.scopes.scopes.last_mut().unwrap();
-
-        assert_eq!(scope.region_scope, local_scope, "local scope is not the topmost scope!",);
+        assert_eq!(scope.region_scope, local_scope, "local scope is not the topmost scope!");
 
         // look for moves of a local variable, like `MOVE(_X)`
-        let locals_moved = operands.iter().flat_map(|operand| match operand.node {
-            Operand::Copy(_) | Operand::Constant(_) | Operand::RuntimeChecks(_) => None,
+        let local_moved = match operand {
+            Operand::Copy(_) | Operand::Constant(_) | Operand::RuntimeChecks(_) => return,
             Operand::Move(place) => place.as_local(),
-        });
+        };
 
-        for local in locals_moved {
-            // check if we have a Drop for this operand and -- if so
-            // -- add it to the list of moved operands. Note that this
-            // local might not have been an operand created for this
-            // call, it could come from other places too.
-            if scope.drops.iter().any(|drop| drop.local == local && drop.kind == DropKind::Value) {
-                scope.moved_locals.push(local);
-            }
+        // We have a move of a local. Mark its drop to be skipped when leaving top scope.
+        // If `local` is not dropped by the topmost scope, this is a no-op.
+        if let Some(local) = local_moved {
+            scope.moved_locals.insert(local);
         }
     }
 
@@ -1878,7 +1875,7 @@ where
                 // path, then don't generate the drop. (We only take this into
                 // account for non-unwind paths so as not to disturb the
                 // caching mechanism.)
-                if scope.moved_locals.contains(&local) {
+                if scope.moved_locals.contains(local) {
                     continue;
                 }
 
@@ -1922,7 +1919,7 @@ where
                 // path, then don't generate the drop. (We only take this into
                 // account for non-unwind paths so as not to disturb the
                 // caching mechanism.)
-                if scope.moved_locals.contains(&local) {
+                if scope.moved_locals.contains(local) {
                     continue;
                 }
 

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

@@ -19,72 +19,62 @@
 +         _5 = const true;
           StorageLive(_3);
           _3 = copy _1;
-          switchInt(move _3) -> [0: bb3, otherwise: bb1];
+          switchInt(move _3) -> [0: bb2, 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;
+          goto -> bb3;
       }
 
-      bb3: {
+      bb2: {
           _0 = Option::<String>::None;
-          goto -> bb4;
+          goto -> bb3;
       }
 
-      bb4: {
+      bb3: {
           StorageDead(_3);
--         drop(_2) -> [return: bb5, unwind continue];
-+         goto -> bb13;
+-         drop(_2) -> [return: bb4, unwind continue];
++         goto -> bb11;
       }
 
-      bb5: {
+      bb4: {
           return;
-      }
-
-      bb6 (cleanup): {
--         drop(_2) -> [return: bb7, unwind terminate(cleanup)];
-+         goto -> bb7;
-      }
-
-      bb7 (cleanup): {
-          resume;
 +     }
 + 
-+     bb8: {
++     bb5 (cleanup): {
++         resume;
++     }
++ 
++     bb6: {
 +         _6 = &mut _2;
-+         _7 = <Box<String> as Drop>::drop(move _6) -> [return: bb5, unwind: bb7];
++         _7 = <Box<String> as Drop>::drop(move _6) -> [return: bb4, unwind: bb5];
 +     }
 + 
-+     bb9 (cleanup): {
++     bb7 (cleanup): {
 +         _8 = &mut _2;
-+         _9 = <Box<String> as Drop>::drop(move _8) -> [return: bb7, unwind terminate(cleanup)];
++         _9 = <Box<String> as Drop>::drop(move _8) -> [return: bb5, unwind terminate(cleanup)];
 +     }
 + 
-+     bb10: {
-+         goto -> bb12;
++     bb8: {
++         goto -> bb10;
 +     }
 + 
-+     bb11: {
-+         drop((*_10)) -> [return: bb8, unwind: bb9];
++     bb9: {
++         drop((*_10)) -> [return: bb6, unwind: bb7];
 +     }
 + 
-+     bb12: {
-+         switchInt(copy _5) -> [0: bb8, otherwise: bb11];
++     bb10: {
++         switchInt(copy _5) -> [0: bb6, otherwise: bb9];
 +     }
 + 
-+     bb13: {
++     bb11: {
 +         _10 = copy ((_2.0: std::ptr::Unique<std::string::String>).0: std::ptr::NonNull<std::string::String>) as *const std::string::String (Transmute);
-+         goto -> bb10;
++         goto -> bb8;
       }
   }