Track drop points by MIR Place

src/analyze/basic_block.rs

@@ -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,14 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
         }
     }
 
-    fn drop_local(&mut self, local: Local) {
-        self.env.drop_local(local);
+    fn drop_place(&mut self, place: mir::Place<'tcx>) {
+        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,9 +1050,9 @@ 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);
+        for place in self.drop_points.before_statements.clone() {
+            tracing::info!(?place, "implicitly dropped before statements");
+            self.drop_place(place);
         }
         let statements = self.body.basic_blocks[self.basic_block].statements.clone();
         for (stmt_idx, mut stmt) in statements.iter().cloned().enumerate() {
@@ -1072,9 +1072,9 @@ 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);
+            for place in self.drop_points.after_statement(stmt_idx) {
+                tracing::info!(?place, ?stmt_idx, "implicitly dropped after statement");
+                self.drop_place(place);
             }
         }
     }
@@ -1155,26 +1155,26 @@ 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);
+                        for place in a.drop_points.after_terminator(&target) {
+                            tracing::info!(?place, ?target, "implicitly dropped for target");
+                            a.drop_place(place);
                         }
                     },
                 );
             }
             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);
+                    for place in self.drop_points.after_terminator(target) {
+                        tracing::info!(?place, "implicitly dropped after call");
+                        self.drop_place(place);
                     }
                     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);
+                for place in self.drop_points.after_terminator(target) {
+                    tracing::info!(?place, "dropped");
+                    self.drop_place(place);
                 }
                 self.type_goto(*target, outer_fn_param_vars);
             }
@@ -1184,9 +1184,9 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
                 target,
                 ..
             } => {
-                for local in self.drop_points.after_terminator(target) {
-                    tracing::info!(?local, "dropped");
-                    self.drop_local(local);
+                for place in self.drop_points.after_terminator(target) {
+                    tracing::info!(?place, "dropped");
+                    self.drop_place(place);
                 }
                 self.type_operand(
                     cond.clone(),
@@ -1368,7 +1368,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,66 +1,97 @@
-use std::collections::{BTreeSet, HashMap};
+use std::collections::HashMap;
 
 use rustc_index::bit_set::DenseBitSet;
 use rustc_middle::mir::{self, BasicBlock, Body, Local};
 use rustc_mir_dataflow::{impls::MaybeLiveLocals, ResultsCursor};
 
 #[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: Vec<mir::Place<'tcx>>,
+    after_statements: Vec<Vec<mir::Place<'tcx>>>,
+    after_terminator: HashMap<BasicBlock, Vec<mir::Place<'tcx>>>,
+    /// Places dropped after the terminator regardless of the target, in
+    /// addition to the liveness-derived sets above.
+    after_terminator_extra: Vec<mir::Place<'tcx>>,
 }
 
-impl DropPoints {
+impl DropPoints<'_> {
     pub fn builder<'mir, 'tcx>(body: &'mir Body<'tcx>) -> DropPointsBuilder<'mir, 'tcx> {
         DropPointsBuilder {
             body,
             bb_ins_cache: HashMap::new(),
         }
     }
+}
 
-    pub fn position(&self, local: Local) -> Option<usize> {
+impl<'tcx> DropPoints<'tcx> {
+    pub fn position(&self, place: mir::Place<'tcx>) -> Option<usize> {
         self.after_statements
             .iter()
-            .position(|s| s.contains(local))
+            .position(|s| s.contains(&place))
             .or_else(|| {
-                self.is_after_terminator(local)
+                self.is_after_terminator(place)
                     .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)
+    fn is_after_terminator(&self, place: mir::Place<'tcx>) -> bool {
+        self.after_terminator.values().any(|s| s.contains(&place))
+            || self.after_terminator_extra.contains(&place)
     }
 
-    pub fn remove_after_statement(&mut self, statement_index: usize, local: Local) -> bool {
-        self.after_statements[statement_index].remove(local)
+    pub fn remove_after_statement(
+        &mut self,
+        statement_index: usize,
+        place: mir::Place<'tcx>,
+    ) -> bool {
+        if let Some(pos) = self.after_statements[statement_index]
+            .iter()
+            .position(|p| *p == place)
+        {
+            self.after_statements[statement_index].remove(pos);
+            true
+        } else {
+            false
+        }
     }
 
-    pub fn insert_after_statement(&mut self, statement_index: usize, local: Local) -> bool {
-        self.after_statements[statement_index].insert(local)
+    pub fn insert_after_statement(
+        &mut self,
+        statement_index: usize,
+        place: mir::Place<'tcx>,
+    ) -> bool {
+        let statements = &mut self.after_statements[statement_index];
+        if statements.contains(&place) {
+            false
+        } else {
+            statements.push(place);
+            true
+        }
     }
 
-    pub fn after_statement(&self, statement_index: usize) -> DenseBitSet<Local> {
+    pub fn after_statement(&self, statement_index: usize) -> Vec<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>) {
+        if !self.after_terminator_extra.contains(&place) {
+            self.after_terminator_extra.push(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) -> Vec<mir::Place<'tcx>> {
+        let mut places = self.after_terminator[target].clone();
+        for place in self.after_statements.last().unwrap() {
+            if !places.contains(place) {
+                places.push(*place);
+            }
+        }
+        for place in &self.after_terminator_extra {
+            if !places.contains(place) {
+                places.push(*place);
+            }
+        }
+        places
     }
 }
 
@@ -147,12 +178,12 @@ impl<'mir, 'tcx> DropPointsBuilder<'mir, 'tcx> {
         &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, Vec::new);
 
         results.seek_to_block_end(bb);
         let live_locals_after_terminator = results.get().clone();
@@ -169,7 +200,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, edge_drops.iter().map(Into::into).collect());
             ins.union(&self.bb_ins_cache[&succ_bb]);
         }
 
@@ -193,7 +224,7 @@ impl<'mir, 'tcx> DropPointsBuilder<'mir, 'tcx> {
                 }
                 t.subtract(&last_live_locals);
                 t.subtract(&moved_locals(self.body, bb, statement_index));
-                t
+                t.iter().map(Into::into).collect()
             };
             last_live_locals = live_locals;
         }

src/analyze/basic_block/visitor/rust_call.rs

@@ -122,17 +122,11 @@ impl<'a, 'tcx, 'ctx> mir::visit::MutVisitor<'tcx> for RustCallVisitor<'a, 'tcx,
                     // FnOnce::call_once consumes the closure, but the resolved function
                     // only borrows it: drop the borrow and the environment after the
                     // call to resolve the prophecies of the captured mutable borrows.
-                    self.analyzer.drop_after_terminator(borrowed_closure_local);
-                    // The original MIR moves the closure into the call, so `moved_locals`
-                    // dropped its drop obligation, expecting the callee to consume it; we
-                    // must restore it. `moved_locals` only steals whole-local moves, so we
-                    // only restore those: with a projection the obligation was never stolen
-                    // and the normal drop machinery still handles it (re-adding it would
-                    // double-drop). In practice a non-`Copy` closure (the only kind reaching
-                    // this case) is always moved through a projection-less temporary.
-                    if arg_closure_place.projection.is_empty() {
-                        self.analyzer.drop_after_terminator(arg_closure_place.local);
-                    }
+                    self.analyzer
+                        .drop_after_terminator(borrowed_closure_local.into());
+                    // The original MIR moves the closure into the call, so restore a
+                    // drop obligation for the moved closure environment at the precise place.
+                    self.analyzer.drop_after_terminator(arg_closure_place);
                     tracing::debug!("applied mut-borrow for closure argument");
                 }
             }

src/analyze/local_def.rs

@@ -49,7 +49,7 @@ pub struct Analyzer<'tcx, 'ctx> {
     body: Body<'tcx>,
     /// to substitute HIR types during translation in [`crate::analyze::annot_fn`]
     generic_args: mir_ty::GenericArgsRef<'tcx>,
-    drop_points: HashMap<BasicBlock, analyze::basic_block::DropPoints>,
+    drop_points: HashMap<BasicBlock, analyze::basic_block::DropPoints<'tcx>>,
     type_builder: TypeBuilder<'tcx>,
 }
 
@@ -1098,7 +1098,7 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
                         .get_mut(&bb)
                         .unwrap()
                         .before_statements
-                        .push(a);
+                        .push(a.into());
                 }
             }
             // function return type is basic block return type

src/refine/env.rs

@@ -1189,10 +1189,14 @@ where
         }
     }
 
-    pub fn drop_local(&mut self, local: Local) {
-        let assumption = self.dropping_assumption(&Path::Local(local));
+    pub fn drop_place(&mut self, place: Place<'_>) {
+        let assumption = self.dropping_assumption(&place.into());
         if !assumption.is_top() {
             self.assume(assumption);
         }
     }
+
+    pub fn drop_local(&mut self, local: Local) {
+        self.drop_place(local.into());
+    }
 }