[WIP] Move drop elaboration before borrowck

compiler/rustc_borrowck/src/def_use.rs

@@ -54,7 +54,7 @@ pub fn categorize(context: PlaceContext) -> Option<DefUse> {
 
         PlaceContext::MutatingUse(MutatingUseContext::AddressOf) |
         PlaceContext::NonMutatingUse(NonMutatingUseContext::AddressOf) |
-        PlaceContext::NonMutatingUse(NonMutatingUseContext::Inspect) |
+        PlaceContext::NonMutatingUse(NonMutatingUseContext::Inspect(_)) |
         PlaceContext::NonMutatingUse(NonMutatingUseContext::Copy) |
         PlaceContext::NonMutatingUse(NonMutatingUseContext::Move) |
         PlaceContext::NonUse(NonUseContext::AscribeUserTy) |

compiler/rustc_borrowck/src/diagnostics/mod.rs

@@ -939,7 +939,22 @@ impl<'cx, 'tcx> MirBorrowckCtxt<'cx, 'tcx> {
             return OtherUse(use_span);
         }
 
-        for stmt in &self.body[location.block].statements[location.statement_index + 1..] {
+        // drop and replace might have moved the assignment to the next block
+        let maybe_additional_statement = if let Some(Terminator {
+            kind: TerminatorKind::Drop { target: drop_target, .. },
+            ..
+        }) = self.body[location.block].terminator
+        {
+            self.body[drop_target].statements.iter().take(1)
+        } else {
+            [].iter().take(0)
+        };
+
+        let statements = self.body[location.block].statements[location.statement_index + 1..]
+            .iter()
+            .chain(maybe_additional_statement);
+
+        for stmt in statements {
             if let StatementKind::Assign(box (_, Rvalue::Aggregate(kind, places))) = &stmt.kind {
                 let (&def_id, is_generator) = match kind {
                     box AggregateKind::Closure(def_id, _) => (def_id, false),

compiler/rustc_borrowck/src/diagnostics/mutability_errors.rs

@@ -828,7 +828,13 @@ impl<'a, 'tcx> MirBorrowckCtxt<'a, 'tcx> {
             let Some(hir::Node::Item(item)) = node else { return; };
             let hir::ItemKind::Fn(.., body_id) = item.kind else { return; };
             let body = self.infcx.tcx.hir().body(body_id);
-            let mut v = V { assign_span: span, err, ty, suggested: false };
+            let mut assign_span = span;
+            // Drop desugaring is done at MIR build so it's not in the HIR
+            if let Some(DesugaringKind::Replace) = span.desugaring_kind() {
+                assign_span.remove_mark();
+            }
+
+            let mut v = V { assign_span, err, ty, suggested: false };
             v.visit_body(body);
             if !v.suggested {
                 err.help(&format!(

compiler/rustc_borrowck/src/invalidation.rs

@@ -58,7 +58,9 @@ impl<'cx, 'tcx> Visitor<'tcx> for InvalidationGenerator<'cx, 'tcx> {
 
         match &statement.kind {
             StatementKind::Assign(box (lhs, rhs)) => {
-                self.consume_rvalue(location, rhs);
+                if !matches!(rhs, Rvalue::Discriminant(_)) {
+                    self.consume_rvalue(location, rhs);
+                }
 
                 self.mutate_place(location, *lhs, Shallow(None));
             }
@@ -119,13 +121,12 @@ impl<'cx, 'tcx> Visitor<'tcx> for InvalidationGenerator<'cx, 'tcx> {
                 );
             }
             TerminatorKind::DropAndReplace {
-                place: drop_place,
-                value: new_value,
+                place: _drop_place,
+                value: _new_value,
                 target: _,
                 unwind: _,
             } => {
-                self.mutate_place(location, *drop_place, Deep);
-                self.consume_operand(location, new_value);
+                bug!("undesugared drop and replace in borrowck")
             }
             TerminatorKind::Call {
                 func,

compiler/rustc_borrowck/src/lib.rs

@@ -36,7 +36,7 @@ use rustc_middle::mir::{InlineAsmOperand, Terminator, TerminatorKind};
 use rustc_middle::ty::query::Providers;
 use rustc_middle::ty::{self, CapturedPlace, ParamEnv, RegionVid, TyCtxt};
 use rustc_session::lint::builtin::UNUSED_MUT;
-use rustc_span::{Span, Symbol};
+use rustc_span::{DesugaringKind, Span, Symbol};
 
 use either::Either;
 use smallvec::SmallVec;
@@ -576,7 +576,11 @@ impl<'cx, 'tcx> rustc_mir_dataflow::ResultsVisitor<'cx, 'tcx> for MirBorrowckCtx
 
         match &stmt.kind {
             StatementKind::Assign(box (lhs, rhs)) => {
-                self.consume_rvalue(location, (rhs, span), flow_state);
+                // FIXME: drop-elaboration checks the discriminant of an enum after it has been
+                // moved out. This is a known isses and should be fixed in the future.
+                if !(matches!(rhs, Rvalue::Discriminant(_)) && span.desugaring_kind() == Some(DesugaringKind::Drop)) {
+                    self.consume_rvalue(location, (rhs, span), flow_state);
+                }
 
                 self.mutate_place(location, (*lhs, span), Shallow(None), flow_state);
             }
@@ -599,7 +603,7 @@ impl<'cx, 'tcx> rustc_mir_dataflow::ResultsVisitor<'cx, 'tcx> for MirBorrowckCtx
                 );
             }
             StatementKind::Intrinsic(box kind) => match kind {
-                NonDivergingIntrinsic::Assume(op) => self.consume_operand(location, (op, span), flow_state),
+                NonDivergingIntrinsic::Assume(op) => self.consume_operand(location, (op, span), flow_state, RequireInit::Yes),
                 NonDivergingIntrinsic::CopyNonOverlapping(..) => span_bug!(
                     span,
                     "Unexpected CopyNonOverlapping, should only appear after lower_intrinsics",
@@ -643,7 +647,7 @@ impl<'cx, 'tcx> rustc_mir_dataflow::ResultsVisitor<'cx, 'tcx> for MirBorrowckCtx
 
         match &term.kind {
             TerminatorKind::SwitchInt { discr, targets: _ } => {
-                self.consume_operand(loc, (discr, span), flow_state);
+                self.consume_operand(loc, (discr, span), flow_state, RequireInit::Yes);
             }
             TerminatorKind::Drop { place, target: _, unwind: _ } => {
                 debug!(
@@ -661,13 +665,12 @@ impl<'cx, 'tcx> rustc_mir_dataflow::ResultsVisitor<'cx, 'tcx> for MirBorrowckCtx
                 );
             }
             TerminatorKind::DropAndReplace {
-                place: drop_place,
-                value: new_value,
+                place: _drop_place,
+                value: _new_value,
                 target: _,
                 unwind: _,
             } => {
-                self.mutate_place(loc, (*drop_place, span), Deep, flow_state);
-                self.consume_operand(loc, (new_value, span), flow_state);
+                bug!("undesugared drop and replace in borrowck")
             }
             TerminatorKind::Call {
                 func,
@@ -678,23 +681,44 @@ impl<'cx, 'tcx> rustc_mir_dataflow::ResultsVisitor<'cx, 'tcx> for MirBorrowckCtx
                 from_hir_call: _,
                 fn_span: _,
             } => {
-                self.consume_operand(loc, (func, span), flow_state);
+                self.mutate_place(loc, (*destination, span), Deep, flow_state);
+                self.consume_operand(loc, (func, span), flow_state, RequireInit::Yes);
+
+                // the initialization dataflow analysis can't understand that drop elaboration
+                // inserts checks to determine whether a target is initialized before calling
+                // box_free. In other words, it can't correlate the value of the drop flag with
+                // the initialization status of the target.
+                // As a temporary fix, just trust that it does the right thing.
+                // Note that this is the same behavior of Drop actually, where the borrow checker
+                // doesn't check that the target is initialized and relies on drop elab to
+                // remove uninitialized ones.
+                let func_ty = func.ty(self.body, self.infcx.tcx);
+                use rustc_hir::lang_items::LangItem;
+                let require_init = match func_ty.kind() {
+                    ty::FnDef(func_id, _)
+                        if Some(func_id)
+                            == self.infcx.tcx.lang_items().get(LangItem::BoxFree).as_ref() =>
+                    {
+                        RequireInit::No
+                    }
+                    _ => RequireInit::Yes,
+                };
+
                 for arg in args {
-                    self.consume_operand(loc, (arg, span), flow_state);
+                    self.consume_operand(loc, (arg, span), flow_state, require_init);
                 }
-                self.mutate_place(loc, (*destination, span), Deep, flow_state);
             }
             TerminatorKind::Assert { cond, expected: _, msg, target: _, cleanup: _ } => {
-                self.consume_operand(loc, (cond, span), flow_state);
+                self.consume_operand(loc, (cond, span), flow_state, RequireInit::Yes);
                 use rustc_middle::mir::AssertKind;
                 if let AssertKind::BoundsCheck { len, index } = msg {
-                    self.consume_operand(loc, (len, span), flow_state);
-                    self.consume_operand(loc, (index, span), flow_state);
+                    self.consume_operand(loc, (len, span), flow_state, RequireInit::Yes);
+                    self.consume_operand(loc, (index, span), flow_state, RequireInit::Yes);
                 }
             }
 
             TerminatorKind::Yield { value, resume: _, resume_arg, drop: _ } => {
-                self.consume_operand(loc, (value, span), flow_state);
+                self.consume_operand(loc, (value, span), flow_state, RequireInit::Yes);
                 self.mutate_place(loc, (*resume_arg, span), Deep, flow_state);
             }
 
@@ -709,15 +733,20 @@ impl<'cx, 'tcx> rustc_mir_dataflow::ResultsVisitor<'cx, 'tcx> for MirBorrowckCtx
                 for op in operands {
                     match op {
                         InlineAsmOperand::In { reg: _, value } => {
-                            self.consume_operand(loc, (value, span), flow_state);
+                            self.consume_operand(loc, (value, span), flow_state, RequireInit::Yes);
                         }
                         InlineAsmOperand::Out { reg: _, late: _, place, .. } => {
                             if let Some(place) = place {
                                 self.mutate_place(loc, (*place, span), Shallow(None), flow_state);
                             }
                         }
                         InlineAsmOperand::InOut { reg: _, late: _, in_value, out_place } => {
-                            self.consume_operand(loc, (in_value, span), flow_state);
+                            self.consume_operand(
+                                loc,
+                                (in_value, span),
+                                flow_state,
+                                RequireInit::Yes,
+                            );
                             if let &Some(out_place) = out_place {
                                 self.mutate_place(
                                     loc,
@@ -860,6 +889,14 @@ enum WriteKind {
     Move,
 }
 
+/// Whether to require a place to be initialized
+/// when checking permissions.
+#[derive(Copy, Clone, PartialEq, Eq, Debug)]
+enum RequireInit {
+    Yes,
+    No,
+}
+
 /// When checking permissions for a place access, this flag is used to indicate that an immutable
 /// local place can be mutated.
 //
@@ -1100,13 +1137,23 @@ impl<'cx, 'tcx> MirBorrowckCtxt<'cx, 'tcx> {
                                 this.report_conflicting_borrow(location, place_span, bk, borrow);
                             this.buffer_error(err);
                         }
-                        WriteKind::StorageDeadOrDrop => this
-                            .report_borrowed_value_does_not_live_long_enough(
-                                location,
-                                borrow,
-                                place_span,
-                                Some(kind),
-                            ),
+                        WriteKind::StorageDeadOrDrop => {
+                            use rustc_span::DesugaringKind;
+                            if let Some(DesugaringKind::Replace) = place_span.1.desugaring_kind() {
+                                // If this is a drop triggered by a reassignment, it's more user friendly
+                                // to report a problem with the explicit assignment than the implicit drop.
+                                this.report_illegal_mutation_of_borrowed(
+                                    location, place_span, borrow,
+                                )
+                            } else {
+                                this.report_borrowed_value_does_not_live_long_enough(
+                                    location,
+                                    borrow,
+                                    place_span,
+                                    Some(kind),
+                                )
+                            }
+                        }
                         WriteKind::Mutate => {
                             this.report_illegal_mutation_of_borrowed(location, place_span, borrow)
                         }
@@ -1220,7 +1267,7 @@ impl<'cx, 'tcx> MirBorrowckCtxt<'cx, 'tcx> {
             | Rvalue::UnaryOp(_ /*un_op*/, operand)
             | Rvalue::Cast(_ /*cast_kind*/, operand, _ /*ty*/)
             | Rvalue::ShallowInitBox(operand, _ /*ty*/) => {
-                self.consume_operand(location, (operand, span), flow_state)
+                self.consume_operand(location, (operand, span), flow_state, RequireInit::Yes)
             }
 
             &Rvalue::CopyForDeref(place) => {
@@ -1264,8 +1311,8 @@ impl<'cx, 'tcx> MirBorrowckCtxt<'cx, 'tcx> {
 
             Rvalue::BinaryOp(_bin_op, box (operand1, operand2))
             | Rvalue::CheckedBinaryOp(_bin_op, box (operand1, operand2)) => {
-                self.consume_operand(location, (operand1, span), flow_state);
-                self.consume_operand(location, (operand2, span), flow_state);
+                self.consume_operand(location, (operand1, span), flow_state, RequireInit::Yes);
+                self.consume_operand(location, (operand2, span), flow_state, RequireInit::Yes);
             }
 
             Rvalue::NullaryOp(_op, _ty) => {
@@ -1292,7 +1339,7 @@ impl<'cx, 'tcx> MirBorrowckCtxt<'cx, 'tcx> {
                 }
 
                 for operand in operands {
-                    self.consume_operand(location, (operand, span), flow_state);
+                    self.consume_operand(location, (operand, span), flow_state, RequireInit::Yes);
                 }
             }
         }
@@ -1398,6 +1445,7 @@ impl<'cx, 'tcx> MirBorrowckCtxt<'cx, 'tcx> {
         location: Location,
         (operand, span): (&'cx Operand<'tcx>, Span),
         flow_state: &Flows<'cx, 'tcx>,
+        require_init: RequireInit,
     ) {
         match *operand {
             Operand::Copy(place) => {
@@ -1411,13 +1459,15 @@ impl<'cx, 'tcx> MirBorrowckCtxt<'cx, 'tcx> {
                     flow_state,
                 );
 
-                // Finally, check if path was already moved.
-                self.check_if_path_or_subpath_is_moved(
-                    location,
-                    InitializationRequiringAction::Use,
-                    (place.as_ref(), span),
-                    flow_state,
-                );
+                if let RequireInit::Yes = require_init {
+                    // Finally, check if path was already moved.
+                    self.check_if_path_or_subpath_is_moved(
+                        location,
+                        InitializationRequiringAction::Use,
+                        (place.as_ref(), span),
+                        flow_state,
+                    );
+                }
             }
             Operand::Move(place) => {
                 // move of place: check if this is move of already borrowed path
@@ -1429,13 +1479,15 @@ impl<'cx, 'tcx> MirBorrowckCtxt<'cx, 'tcx> {
                     flow_state,
                 );
 
-                // Finally, check if path was already moved.
-                self.check_if_path_or_subpath_is_moved(
-                    location,
-                    InitializationRequiringAction::Use,
-                    (place.as_ref(), span),
-                    flow_state,
-                );
+                if let RequireInit::Yes = require_init {
+                    // Finally, check if path was already moved.
+                    self.check_if_path_or_subpath_is_moved(
+                        location,
+                        InitializationRequiringAction::Use,
+                        (place.as_ref(), span),
+                        flow_state,
+                    );
+                }
             }
             Operand::Constant(_) => {}
         }

compiler/rustc_borrowck/src/type_check/liveness/local_use_map.rs

@@ -1,6 +1,6 @@
 use rustc_data_structures::vec_linked_list as vll;
 use rustc_index::vec::IndexVec;
-use rustc_middle::mir::visit::{PlaceContext, Visitor};
+use rustc_middle::mir::visit::{NonMutatingUseContext, PlaceContext, Visitor};
 use rustc_middle::mir::{Body, Local, Location};
 
 use crate::def_use::{self, DefUse};
@@ -158,6 +158,20 @@ impl LocalUseMapBuild<'_> {
 
 impl Visitor<'_> for LocalUseMapBuild<'_> {
     fn visit_local(&mut self, local: Local, context: PlaceContext, location: Location) {
+        if matches!(
+            context,
+            PlaceContext::NonMutatingUse(NonMutatingUseContext::Inspect(
+                rustc_middle::mir::visit::InspectKind::Discriminant
+            ))
+        ) {
+            // drop elaboration inserts discriminant reads for enums, but such reads should
+            // really be counted as a drop access for liveness computation.
+            // However, doing so messes with the way we calculate drop points, so for now just ignore
+            // those, as discriminant reads are ignored anyway by the borrowck.
+            // FIXME: found a better solution for drop-elab discriminant reads
+            return;
+        }
+
         if self.locals_with_use_data[local] {
             match def_use::categorize(context) {
                 Some(DefUse::Def) => self.insert_def(local, location),

compiler/rustc_codegen_ssa/src/mir/analyze.rs

@@ -230,7 +230,7 @@ impl<'mir, 'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> Visitor<'tcx>
                 | MutatingUseContext::Projection,
             )
             | PlaceContext::NonMutatingUse(
-                NonMutatingUseContext::Inspect
+                NonMutatingUseContext::Inspect(_)
                 | NonMutatingUseContext::SharedBorrow
                 | NonMutatingUseContext::UniqueBorrow
                 | NonMutatingUseContext::ShallowBorrow