POC: Move drop elaboration before borrowck

POC to show what's needed to move drop elaboration before borrow
checking. Ideally we could have it right aftr MIR build, but this
tries to be as minimal as possible.

There are a few hacks, mostly around drop-elaboration code that
does not pass borrowck / move analysis.

I plan to work on the remaining failing tests in the near future.

Author: zeegomo

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

@@ -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(_)) {
+                    self.consume_rvalue(location, (rhs, span), flow_state);
+                }
 
                 self.mutate_place(location, (*lhs, span), Shallow(None), flow_state);
             }
@@ -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,11 +681,22 @@ impl<'cx, 'tcx> rustc_mir_dataflow::ResultsVisitor<'cx, 'tcx> for MirBorrowckCtx
                 from_hir_call: _,
                 fn_span: _,
             } => {
+                self.mutate_place(loc, (*destination, span), Deep, flow_state);
                 self.consume_operand(loc, (func, span), flow_state);
+
+                // drop glue for box does not pass borrowck
+                let func_ty = func.ty(self.body, self.infcx.tcx);
+                use rustc_hir::lang_items::LangItem;
+                if let ty::FnDef(func_id, _) = func_ty.kind() {
+                    if Some(func_id) == self.infcx.tcx.lang_items().get(LangItem::BoxFree).as_ref()
+                    {
+                        return;
+                    }
+                }
+
                 for arg in args {
                     self.consume_operand(loc, (arg, span), flow_state);
                 }
-                self.mutate_place(loc, (*destination, span), Deep, flow_state);
             }
             TerminatorKind::Assert { cond, expected: _, msg, target: _, cleanup: _ } => {
                 self.consume_operand(loc, (cond, span), flow_state);
@@ -1100,13 +1114,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)
                         }

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

compiler/rustc_middle/src/mir/visit.rs

@@ -370,7 +370,7 @@ macro_rules! make_mir_visitor {
                     StatementKind::FakeRead(box (_, place)) => {
                         self.visit_place(
                             place,
-                            PlaceContext::NonMutatingUse(NonMutatingUseContext::Inspect),
+                            PlaceContext::NonMutatingUse(NonMutatingUseContext::Inspect(InspectKind::Other)),
                             location
                         );
                     }
@@ -652,7 +652,7 @@ macro_rules! make_mir_visitor {
                     Rvalue::CopyForDeref(place) => {
                         self.visit_place(
                             place,
-                            PlaceContext::NonMutatingUse(NonMutatingUseContext::Inspect),
+                            PlaceContext::NonMutatingUse(NonMutatingUseContext::Inspect(InspectKind::Other)),
                             location
                         );
                     }
@@ -672,7 +672,7 @@ macro_rules! make_mir_visitor {
                     Rvalue::Len(path) => {
                         self.visit_place(
                             path,
-                            PlaceContext::NonMutatingUse(NonMutatingUseContext::Inspect),
+                            PlaceContext::NonMutatingUse(NonMutatingUseContext::Inspect(InspectKind::Other)),
                             location
                         );
                     }
@@ -695,7 +695,7 @@ macro_rules! make_mir_visitor {
                     Rvalue::Discriminant(place) => {
                         self.visit_place(
                             place,
-                            PlaceContext::NonMutatingUse(NonMutatingUseContext::Inspect),
+                            PlaceContext::NonMutatingUse(NonMutatingUseContext::Inspect(InspectKind::Discriminant)),
                             location
                         );
                     }
@@ -1132,7 +1132,7 @@ macro_rules! visit_place_fns {
             let mut context = context;
 
             if !place.projection.is_empty() {
-                if context.is_use() {
+                if context.is_use() && !context.is_drop() {
                     // ^ Only change the context if it is a real use, not a "use" in debuginfo.
                     context = if context.is_mutating_use() {
                         PlaceContext::MutatingUse(MutatingUseContext::Projection)
@@ -1236,10 +1236,16 @@ pub enum TyContext {
     Location(Location),
 }
 
+#[derive(Copy, Clone, Debug, PartialEq, Eq)]
+pub enum InspectKind {
+    Other,
+    Discriminant,
+}
+
 #[derive(Copy, Clone, Debug, PartialEq, Eq)]
 pub enum NonMutatingUseContext {
     /// Being inspected in some way, like loading a len.
-    Inspect,
+    Inspect(InspectKind),
     /// Consumed as part of an operand.
     Copy,
     /// Consumed as part of an operand.

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

@@ -39,10 +39,19 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
 
                 // Generate better code for things that don't need to be
                 // dropped.
+                use rustc_span::DesugaringKind;
                 if lhs.ty.needs_drop(this.tcx, this.param_env) {
                     let rhs = unpack!(block = this.as_local_operand(block, rhs));
                     let lhs = unpack!(block = this.as_place(block, lhs));
-                    unpack!(block = this.build_drop_and_replace(block, lhs_span, lhs, rhs));
+                    let span = this.tcx.with_stable_hashing_context(|hcx| {
+                        lhs_span.mark_with_reason(
+                            None,
+                            DesugaringKind::Replace,
+                            this.tcx.sess.edition(),
+                            hcx,
+                        )
+                    });
+                    unpack!(block = this.build_drop_and_replace(block, span, lhs, rhs));
                 } else {
                     let rhs = unpack!(block = this.as_local_rvalue(block, rhs));
                     let lhs = unpack!(block = this.as_place(block, lhs));

compiler/rustc_mir_dataflow/src/drop_flag_effects.rs

@@ -29,6 +29,35 @@ where
     None
 }
 
+// Drop elaboration of a Box makes the inner fields (pointer and allocator) explicit in MIR,
+// which causes those path to be tracked by dataflow and checked by the borrowck.
+// However, those are not considered to be children of *_box_place (and only of _box_place),
+// which causes issues in initialization analysis when accessing _box_place.
+// This code basically check if the parent of the path is a box and if so, apply the callback
+// to its children to fix such cases.
+fn maybe_recurse_box_parent<'tcx, F>(
+    tcx: TyCtxt<'tcx>,
+    body: &Body<'tcx>,
+    move_data: &MoveData<'tcx>,
+    move_path_index: MovePathIndex,
+    mut each_child: F,
+) where
+    F: FnMut(MovePathIndex),
+{
+    if let Some(path) = move_data.move_paths[move_path_index].parent {
+        let place = move_data.move_paths[path].place;
+        let ty = place.ty(body, tcx).ty;
+        if let ty::Adt(def, _) = ty.kind() {
+            if def.is_box() {
+                each_child(move_path_index);
+                on_all_children_bits(tcx, body, move_data, path, each_child);
+                return;
+            }
+        }
+    }
+    on_all_children_bits(tcx, body, move_data, move_path_index, each_child);
+}
+
 pub fn on_lookup_result_bits<'tcx, F>(
     tcx: TyCtxt<'tcx>,
     body: &Body<'tcx>,
@@ -42,7 +71,7 @@ pub fn on_lookup_result_bits<'tcx, F>(
         LookupResult::Parent(..) => {
             // access to untracked value - do not touch children
         }
-        LookupResult::Exact(e) => on_all_children_bits(tcx, body, move_data, e, each_child),
+        LookupResult::Exact(e) => maybe_recurse_box_parent(tcx, body, move_data, e, each_child),
     }
 }
 
@@ -194,6 +223,19 @@ pub fn drop_flag_effects_for_location<'tcx, F>(
         on_all_children_bits(tcx, body, move_data, path, |mpi| callback(mpi, DropFlagState::Absent))
     }
 
+    use rustc_middle::mir::{Terminator, TerminatorKind};
+
+    // Drop does not count as a move but we should still consider the variable uninitialized.
+    if let Some(Terminator { kind: TerminatorKind::Drop { place, .. }, .. }) =
+        body.stmt_at(loc).right()
+    {
+        if let LookupResult::Exact(mpi) = move_data.rev_lookup.find(place.as_ref()) {
+            on_all_children_bits(tcx, body, move_data, mpi, |mpi| {
+                callback(mpi, DropFlagState::Absent)
+            })
+        }
+    }
+
     debug!("drop_flag_effects: assignment for location({:?})", loc);
 
     for_location_inits(tcx, body, move_data, loc, |mpi| callback(mpi, DropFlagState::Present));
@@ -212,9 +254,7 @@ pub fn for_location_inits<'tcx, F>(
         let init = move_data.inits[*ii];
         match init.kind {
             InitKind::Deep => {
-                let path = init.path;
-
-                on_all_children_bits(tcx, body, move_data, path, &mut callback)
+                maybe_recurse_box_parent(tcx, body, move_data, init.path, &mut callback)
             }
             InitKind::Shallow => {
                 let mpi = init.path;