Track drop points by MIR Place, splitting partially-moved locals (fix #121, #122)

src/analyze/basic_block.rs

@@ -1,5 +1,5 @@
 use std::borrow::Cow;
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 
 use rustc_hir::def::DefKind;
 use rustc_index::IndexVec;
@@ -136,7 +136,7 @@ pub struct Analyzer<'tcx, 'ctx> {
     tcx: TyCtxt<'tcx>,
 
     local_def_id: LocalDefId,
-    drop_points: DropPoints,
+    drop_points: DropPoints<'tcx>,
     basic_block: BasicBlock,
     body: Cow<'tcx, Body<'tcx>>,
 
@@ -951,14 +951,17 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
         }
     }
 
-    fn drop_local(&mut self, local: Local) {
-        self.env.drop_local(local);
+    fn drop_places(&mut self, places: HashSet<mir::Place<'tcx>>) {
+        for place in places {
+            tracing::info!(?place, "implicitly dropped");
+            self.env.drop_place(place);
+        }
     }
 
-    /// Schedules `local` to be implicitly dropped after this block's terminator,
+    /// Schedules `place` to be implicitly dropped after this block's terminator,
     /// in addition to the liveness-derived drop points.
-    fn drop_after_terminator(&mut self, local: Local) {
-        self.drop_points.insert_after_terminator(local);
+    fn drop_after_terminator(&mut self, place: mir::Place<'tcx>) {
+        self.drop_points.insert_after_terminator(place);
     }
 
     fn add_prophecy_var(&mut self, statement_index: usize, ty: mir_ty::Ty<'tcx>) {
@@ -1050,10 +1053,8 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
     }
 
     fn analyze_statements(&mut self) {
-        for local in self.drop_points.before_statements.clone() {
-            tracing::info!(?local, "implicitly dropped before statements");
-            self.drop_local(local);
-        }
+        let before_statements = self.drop_points.before_statements.clone();
+        self.drop_places(before_statements);
         let statements = self.body.basic_blocks[self.basic_block].statements.clone();
         for (stmt_idx, mut stmt) in statements.iter().cloned().enumerate() {
             if stmt_idx == statements.len() - 1 && self.terminator_is_drop_call().is_some() {
@@ -1072,10 +1073,8 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
                 | StatementKind::StorageDead(_) => {}
                 _ => unimplemented!("stmt={:?}", stmt.kind),
             }
-            for local in self.drop_points.after_statement(stmt_idx).iter() {
-                tracing::info!(?local, ?stmt_idx, "implicitly dropped after statement");
-                self.drop_local(local);
-            }
+            let after_statement = self.drop_points.after_statement(stmt_idx);
+            self.drop_places(after_statement);
         }
     }
 
@@ -1155,27 +1154,21 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
                     targets.clone(),
                     outer_fn_param_vars,
                     |a, target| {
-                        for local in a.drop_points.after_terminator(&target) {
-                            tracing::info!(?local, ?target, "implicitly dropped for target");
-                            a.drop_local(local);
-                        }
+                        let set = a.drop_points.after_terminator(&target);
+                        a.drop_places(set);
                     },
                 );
             }
             TerminatorKind::Call { target, .. } => {
                 if let Some(target) = target {
-                    for local in self.drop_points.after_terminator(target) {
-                        tracing::info!(?local, "implicitly dropped after call");
-                        self.drop_local(local);
-                    }
+                    let set = self.drop_points.after_terminator(target);
+                    self.drop_places(set);
                     self.type_goto(*target, outer_fn_param_vars);
                 }
             }
             TerminatorKind::Drop { target, .. } => {
-                for local in self.drop_points.after_terminator(target) {
-                    tracing::info!(?local, "dropped");
-                    self.drop_local(local);
-                }
+                let set = self.drop_points.after_terminator(target);
+                self.drop_places(set);
                 self.type_goto(*target, outer_fn_param_vars);
             }
             TerminatorKind::Assert {
@@ -1184,10 +1177,8 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
                 target,
                 ..
             } => {
-                for local in self.drop_points.after_terminator(target) {
-                    tracing::info!(?local, "dropped");
-                    self.drop_local(local);
-                }
+                let set = self.drop_points.after_terminator(target);
+                self.drop_places(set);
                 self.type_operand(
                     cond.clone(),
                     &rty::RefinedType::refined_with_term(
@@ -1368,7 +1359,7 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
         }
     }
 
-    pub fn drop_points(&mut self, drop_points: DropPoints) -> &mut Self {
+    pub fn drop_points(&mut self, drop_points: DropPoints<'tcx>) -> &mut Self {
         self.drop_points = drop_points;
         self
     }

src/analyze/basic_block/drop_point.rs

@@ -1,117 +1,116 @@
-use std::collections::{BTreeSet, HashMap};
+use std::collections::{HashMap, HashSet};
 
 use rustc_index::bit_set::DenseBitSet;
-use rustc_middle::mir::{self, BasicBlock, Body, Local};
+use rustc_middle::mir::{self, BasicBlock, Body, Local, Location};
+use rustc_middle::ty::TyCtxt;
 use rustc_mir_dataflow::{impls::MaybeLiveLocals, ResultsCursor};
 
+/// Implicit-drop targets. A drop is a `Place`; a whole-local drop is just a
+/// place with an empty projection (semantically a bare `Local`).
 #[derive(Debug, Clone, Default)]
-pub struct DropPoints {
-    // TODO: ad-hoc
-    pub before_statements: Vec<Local>,
-    after_statements: Vec<DenseBitSet<Local>>,
-    after_terminator: HashMap<BasicBlock, DenseBitSet<Local>>,
-    /// Locals dropped after the terminator regardless of the target, in
-    /// addition to the liveness-derived sets above. A set, since the same local
-    /// must not be dropped twice; ordered by index to keep drops deterministic.
-    after_terminator_extra: BTreeSet<Local>,
+pub struct DropPoints<'tcx> {
+    pub before_statements: HashSet<mir::Place<'tcx>>,
+    after_statements: Vec<HashSet<mir::Place<'tcx>>>,
+    after_terminator: HashMap<BasicBlock, HashSet<mir::Place<'tcx>>>,
+    /// Drops scheduled after the terminator regardless of the target, in
+    /// addition to the liveness-derived sets above.
+    after_terminator_extra: HashSet<mir::Place<'tcx>>,
 }
 
-impl DropPoints {
-    pub fn builder<'mir, 'tcx>(body: &'mir Body<'tcx>) -> DropPointsBuilder<'mir, 'tcx> {
+impl DropPoints<'_> {
+    pub fn builder<'mir, 'tcx>(
+        tcx: TyCtxt<'tcx>,
+        body: &'mir Body<'tcx>,
+    ) -> DropPointsBuilder<'mir, 'tcx> {
         DropPointsBuilder {
             body,
             bb_ins_cache: HashMap::new(),
+            moves: Moves::collect(tcx, body),
         }
     }
+}
 
-    pub fn position(&self, local: Local) -> Option<usize> {
-        self.after_statements
-            .iter()
-            .position(|s| s.contains(local))
-            .or_else(|| {
-                self.is_after_terminator(local)
-                    .then_some(self.after_statements.len())
-            })
-    }
-
-    fn is_after_terminator(&self, local: Local) -> bool {
-        self.after_terminator.values().any(|s| s.contains(local))
-            || self.after_terminator_extra.contains(&local)
-    }
-
-    pub fn remove_after_statement(&mut self, statement_index: usize, local: Local) -> bool {
-        self.after_statements[statement_index].remove(local)
-    }
-
-    pub fn insert_after_statement(&mut self, statement_index: usize, local: Local) -> bool {
-        self.after_statements[statement_index].insert(local)
-    }
-
-    pub fn after_statement(&self, statement_index: usize) -> DenseBitSet<Local> {
+impl<'tcx> DropPoints<'tcx> {
+    pub fn after_statement(&self, statement_index: usize) -> HashSet<mir::Place<'tcx>> {
         self.after_statements[statement_index].clone()
     }
 
-    pub fn insert_after_terminator(&mut self, local: Local) {
-        self.after_terminator_extra.insert(local);
+    pub fn insert_after_terminator(&mut self, place: mir::Place<'tcx>) {
+        self.after_terminator_extra.insert(place);
     }
 
-    pub fn after_terminator(&self, target: &BasicBlock) -> Vec<Local> {
-        let mut t = self.after_terminator[target].clone();
-        t.union(self.after_statements.last().unwrap());
-        t.iter()
-            .chain(self.after_terminator_extra.iter().copied())
-            .collect()
+    pub fn after_terminator(&self, target: &BasicBlock) -> HashSet<mir::Place<'tcx>> {
+        let mut set = self.after_terminator[target].clone();
+        set.extend(self.after_statements.last().unwrap().iter().copied());
+        set.extend(self.after_terminator_extra.iter().copied());
+        set
     }
 }
 
-#[derive(Debug, Clone)]
+#[derive(Clone)]
 pub struct DropPointsBuilder<'mir, 'tcx> {
     body: &'mir Body<'tcx>,
     bb_ins_cache: HashMap<BasicBlock, DenseBitSet<Local>>,
+    moves: Moves,
 }
 
-/// Locals whose ownership is fully transferred away by the statement (or
-/// terminator) at `statement_index`. Such a local is left uninitialized, so its
-/// drop obligation (including resolving any mutable-borrow prophecies it owns)
-/// moves to the destination and it must not be dropped at the move site.
+/// The `move`d operands of a body, excluding reference-typed places.
 ///
-/// Only owned (non-reference) operands are reported: `move`d references are
-/// turned into reborrows by `ReborrowVisitor`/`RustCallVisitor`, so the source
-/// local remains live and must still be dropped.
-fn moved_locals<'tcx>(
-    body: &Body<'tcx>,
-    bb: BasicBlock,
-    statement_index: usize,
-) -> DenseBitSet<Local> {
-    struct Visitor<'a, 'tcx> {
-        body: &'a Body<'tcx>,
-        locals: DenseBitSet<Local>,
-    }
-    impl<'tcx> mir::visit::Visitor<'tcx> for Visitor<'_, 'tcx> {
-        fn visit_operand(&mut self, operand: &mir::Operand<'tcx>, _location: mir::Location) {
-            if let mir::Operand::Move(place) = operand {
-                if place.projection.is_empty() && !self.body.local_decls[place.local].ty.is_ref() {
-                    self.locals.insert(place.local);
+/// `move`d references are turned into reborrows by
+/// `ReborrowVisitor`/`RustCallVisitor`, so the source still owns its prophecy
+/// and must be dropped normally; they are therefore not recorded here.
+///
+/// A whole-local move leaves the local uninitialized, transferring its entire
+/// drop obligation (including resolving any mutable-borrow prophecies it owns)
+/// to the destination, so it must not be dropped at the move site — where the
+/// local also becomes dead, hence the per-location keying. A partial (projected)
+/// field move transfers only a sub-place, but the parent stays live until its
+/// remaining parts die later; dropping it wholesale at that point would walk
+/// into the moved-out sub-place and resolve its `&mut` prophecy a second time,
+/// so a partially-moved local is excluded from dropping entirely.
+#[derive(Clone, Default)]
+struct Moves {
+    /// Whole-local moves, keyed by the location performing the move.
+    whole: HashMap<Location, DenseBitSet<Local>>,
+    /// Locals that have a partial field move somewhere in the body.
+    partial: HashSet<Local>,
+}
+
+impl Moves {
+    fn collect<'tcx>(tcx: TyCtxt<'tcx>, body: &Body<'tcx>) -> Moves {
+        struct Visitor<'a, 'tcx> {
+            tcx: TyCtxt<'tcx>,
+            body: &'a Body<'tcx>,
+            moves: Moves,
+        }
+        impl<'tcx> mir::visit::Visitor<'tcx> for Visitor<'_, 'tcx> {
+            fn visit_operand(&mut self, operand: &mir::Operand<'tcx>, location: Location) {
+                let mir::Operand::Move(place) = operand else {
+                    return;
+                };
+                if place.ty(&self.body.local_decls, self.tcx).ty.is_ref() {
+                    return;
+                }
+                if place.projection.is_empty() {
+                    self.moves
+                        .whole
+                        .entry(location)
+                        .or_insert_with(|| DenseBitSet::new_empty(self.body.local_decls.len()))
+                        .insert(place.local);
+                } else {
+                    self.moves.partial.insert(place.local);
                 }
             }
         }
+        let mut visitor = Visitor {
+            tcx,
+            body,
+            moves: Moves::default(),
+        };
+        use mir::visit::Visitor as _;
+        visitor.visit_body(body);
+        visitor.moves
     }
-    let mut visitor = Visitor {
-        body,
-        locals: DenseBitSet::new_empty(body.local_decls.len()),
-    };
-    let loc = mir::Location {
-        statement_index,
-        block: bb,
-    };
-    let data = &body.basic_blocks[bb];
-    use mir::visit::Visitor as _;
-    if statement_index < data.statements.len() {
-        visitor.visit_statement(&data.statements[statement_index], loc);
-    } else if let Some(tmnt) = &data.terminator {
-        visitor.visit_terminator(tmnt, loc);
-    }
-    visitor.locals
 }
 
 fn def_local<'tcx>(data: &mir::BasicBlockData<'tcx>, statement_index: usize) -> Option<Local> {
@@ -143,16 +142,28 @@ fn def_local<'tcx>(data: &mir::BasicBlockData<'tcx>, statement_index: usize) ->
 }
 
 impl<'mir, 'tcx> DropPointsBuilder<'mir, 'tcx> {
+    /// Turn a set of locals that become dead into the drop targets. A local with
+    /// a partial field move is excluded: dropping it wholesale would resolve the
+    /// moved-out sub-place's `&mut` prophecy a second time (it is resolved at the
+    /// move destination instead).
+    fn drop_set(&self, locals: DenseBitSet<Local>) -> HashSet<mir::Place<'tcx>> {
+        locals
+            .iter()
+            .filter(|local| !self.moves.partial.contains(local))
+            .map(mir::Place::from)
+            .collect()
+    }
+
     pub fn build(
         &mut self,
         results: &mut ResultsCursor<'mir, 'tcx, MaybeLiveLocals>,
         bb: BasicBlock,
-    ) -> DropPoints {
+    ) -> DropPoints<'tcx> {
         let data = &self.body.basic_blocks[bb];
 
         let mut after_terminator = HashMap::new();
         let mut after_statements = Vec::new();
-        after_statements.resize_with(data.statements.len() + 1, || DenseBitSet::new_empty(0));
+        after_statements.resize_with(data.statements.len() + 1, HashSet::new);
 
         results.seek_to_block_end(bb);
         let live_locals_after_terminator = results.get().clone();
@@ -169,7 +180,7 @@ impl<'mir, 'tcx> DropPointsBuilder<'mir, 'tcx> {
                 t.subtract(&self.bb_ins_cache[&succ_bb]);
                 t
             };
-            after_terminator.insert(succ_bb, edge_drops);
+            after_terminator.insert(succ_bb, self.drop_set(edge_drops));
             ins.union(&self.bb_ins_cache[&succ_bb]);
         }
 
@@ -192,8 +203,10 @@ impl<'mir, 'tcx> DropPointsBuilder<'mir, 'tcx> {
                     t.insert(def);
                 }
                 t.subtract(&last_live_locals);
-                t.subtract(&moved_locals(self.body, bb, statement_index));
-                t
+                if let Some(moved) = self.moves.whole.get(&loc) {
+                    t.subtract(moved);
+                }
+                self.drop_set(t)
             };
             last_live_locals = live_locals;
         }
@@ -207,10 +220,10 @@ impl<'mir, 'tcx> DropPointsBuilder<'mir, 'tcx> {
             "analyzed implicit drop points"
         );
         DropPoints {
-            before_statements: Default::default(),
+            before_statements: HashSet::default(),
             after_statements,
             after_terminator,
-            after_terminator_extra: Default::default(),
+            after_terminator_extra: HashSet::default(),
         }
     }
 }