Desugars drop and replace at MIR build

This commit desugars the drop and replace deriving from an
assignment at MIR build, avoiding the construction of the
DropAndReplace terminator (which will be removed in a followign PR)

In order to retain the same error messages for replaces a new
DesugaringKind::Replace variant is introduced.

7 files changed, 80 insertions, 23 deletions

diff --git a/compiler/rustc_borrowck/src/diagnostics/mod.rs b/compiler/rustc_borrowck/src/diagnostics/mod.rs
index a99fd594a07..a1817f85483 100644
--- a/compiler/rustc_borrowck/src/diagnostics/mod.rs
+++ b/compiler/rustc_borrowck/src/diagnostics/mod.rs
@@ -925,7 +925,20 @@ 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 TerminatorKind::Drop { target: drop_target, .. } =
+                self.body[location.block].terminator().kind
+            {
+                self.body[drop_target].statements.first()
+            } else {
+                None
+            };
+
+        let statements =
+            self.body[location.block].statements[location.statement_index + 1..].iter();
+
+        for stmt in statements.chain(maybe_additional_statement) {
             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),
diff --git a/compiler/rustc_borrowck/src/diagnostics/mutability_errors.rs b/compiler/rustc_borrowck/src/diagnostics/mutability_errors.rs
index 328ac880dd4..eded913ae96 100644
--- a/compiler/rustc_borrowck/src/diagnostics/mutability_errors.rs
+++ b/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!(
diff --git a/compiler/rustc_borrowck/src/lib.rs b/compiler/rustc_borrowck/src/lib.rs
index 0f591460e9d..11a746ffb0a 100644
--- a/compiler/rustc_borrowck/src/lib.rs
+++ b/compiler/rustc_borrowck/src/lib.rs
@@ -40,7 +40,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;
@@ -1184,13 +1184,22 @@ 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 => {
+                            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)
                         }
diff --git a/compiler/rustc_mir_build/src/build/expr/stmt.rs b/compiler/rustc_mir_build/src/build/expr/stmt.rs
index 78083685193..ea5aeb67d85 100644
--- a/compiler/rustc_mir_build/src/build/expr/stmt.rs
+++ b/compiler/rustc_mir_build/src/build/expr/stmt.rs
@@ -40,7 +40,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                 // Generate better code for things that don't need to be
                 // dropped.
                 if lhs.ty.needs_drop(this.tcx, this.param_env) {
-                    let rhs = unpack!(block = this.as_local_operand(block, rhs));
+                    let rhs = unpack!(block = this.as_local_rvalue(block, rhs));
                     let lhs = unpack!(block = this.as_place(block, lhs));
                     unpack!(block = this.build_drop_and_replace(block, lhs_span, lhs, rhs));
                 } else {
diff --git a/compiler/rustc_mir_build/src/build/scope.rs b/compiler/rustc_mir_build/src/build/scope.rs
index 591b416337b..770701b3750 100644
--- a/compiler/rustc_mir_build/src/build/scope.rs
+++ b/compiler/rustc_mir_build/src/build/scope.rs
@@ -91,7 +91,7 @@ use rustc_middle::middle::region;
 use rustc_middle::mir::*;
 use rustc_middle::thir::{Expr, LintLevel};
 
-use rustc_span::{Span, DUMMY_SP};
+use rustc_span::{DesugaringKind, Span, DUMMY_SP};
 
 #[derive(Debug)]
 pub struct Scopes<'tcx> {
@@ -1118,24 +1118,35 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
     }
 
     /// Utility function for *non*-scope code to build their own drops
+    /// Force a drop at this point in the MIR by creating a new block.
     pub(crate) fn build_drop_and_replace(
         &mut self,
         block: BasicBlock,
         span: Span,
         place: Place<'tcx>,
-        value: Operand<'tcx>,
+        value: Rvalue<'tcx>,
     ) -> BlockAnd<()> {
+        let span = self.tcx.with_stable_hashing_context(|hcx| {
+            span.mark_with_reason(None, DesugaringKind::Replace, self.tcx.sess.edition(), hcx)
+        });
         let source_info = self.source_info(span);
-        let next_target = self.cfg.start_new_block();
+
+        // create the new block for the assignment
+        let assign = self.cfg.start_new_block();
+        self.cfg.push_assign(assign, source_info, place, value.clone());
+
+        // create the new block for the assignment in the case of unwinding
+        let assign_unwind = self.cfg.start_new_cleanup_block();
+        self.cfg.push_assign(assign_unwind, source_info, place, value.clone());
 
         self.cfg.terminate(
             block,
             source_info,
-            TerminatorKind::DropAndReplace { place, value, target: next_target, unwind: None },
+            TerminatorKind::Drop { place, target: assign, unwind: Some(assign_unwind) },
         );
         self.diverge_from(block);
 
-        next_target.unit()
+        assign.unit()
     }
 
     /// Creates an `Assert` terminator and return the success block.
@@ -1413,8 +1424,15 @@ impl<'tcx> DropTreeBuilder<'tcx> for Unwind {
     fn add_entry(cfg: &mut CFG<'tcx>, from: BasicBlock, to: BasicBlock) {
         let term = &mut cfg.block_data_mut(from).terminator_mut();
         match &mut term.kind {
-            TerminatorKind::Drop { unwind, .. }
-            | TerminatorKind::DropAndReplace { unwind, .. }
+            TerminatorKind::Drop { unwind, .. } => {
+                if let Some(unwind) = *unwind {
+                    let source_info = term.source_info;
+                    cfg.terminate(unwind, source_info, TerminatorKind::Goto { target: to });
+                } else {
+                    *unwind = Some(to);
+                }
+            }
+            TerminatorKind::DropAndReplace { unwind, .. }
             | TerminatorKind::FalseUnwind { unwind, .. }
             | TerminatorKind::Call { cleanup: unwind, .. }
             | TerminatorKind::Assert { cleanup: unwind, .. }
diff --git a/compiler/rustc_mir_transform/src/elaborate_drops.rs b/compiler/rustc_mir_transform/src/elaborate_drops.rs
index bdfd8dc6e99..29424f09695 100644
--- a/compiler/rustc_mir_transform/src/elaborate_drops.rs
+++ b/compiler/rustc_mir_transform/src/elaborate_drops.rs
@@ -14,7 +14,7 @@ use rustc_mir_dataflow::un_derefer::UnDerefer;
 use rustc_mir_dataflow::MoveDataParamEnv;
 use rustc_mir_dataflow::{on_all_children_bits, on_all_drop_children_bits};
 use rustc_mir_dataflow::{Analysis, ResultsCursor};
-use rustc_span::Span;
+use rustc_span::{DesugaringKind, Span};
 use rustc_target::abi::VariantIdx;
 use std::fmt;
 
@@ -425,10 +425,19 @@ impl<'b, 'tcx> ElaborateDropsCtxt<'b, 'tcx> {
                             bb,
                         ),
                         LookupResult::Parent(..) => {
-                            self.tcx.sess.delay_span_bug(
-                                terminator.source_info.span,
-                                &format!("drop of untracked value {:?}", bb),
-                            );
+                            if !matches!(
+                                terminator.source_info.span.desugaring_kind(),
+                                Some(DesugaringKind::Replace),
+                            ) {
+                                self.tcx.sess.delay_span_bug(
+                                    terminator.source_info.span,
+                                    &format!("drop of untracked value {:?}", bb),
+                                );
+                            }
+                            // A drop and replace behind a pointer/array/whatever.
+                            // The borrow checker requires that these locations are initialized before the assignment,
+                            // so we just leave an unconditional drop.
+                            assert!(!data.is_cleanup);
                         }
                     }
                 }
diff --git a/compiler/rustc_span/src/hygiene.rs b/compiler/rustc_span/src/hygiene.rs
index dee823eefde..9f22e9776d4 100644
--- a/compiler/rustc_span/src/hygiene.rs
+++ b/compiler/rustc_span/src/hygiene.rs
@@ -1151,6 +1151,7 @@ pub enum DesugaringKind {
     Await,
     ForLoop,
     WhileLoop,
+    Replace,
 }
 
 impl DesugaringKind {
@@ -1166,6 +1167,7 @@ impl DesugaringKind {
             DesugaringKind::OpaqueTy => "`impl Trait`",
             DesugaringKind::ForLoop => "`for` loop",
             DesugaringKind::WhileLoop => "`while` loop",
+            DesugaringKind::Replace => "drop and replace",
         }
     }
 }