From 39dc672ff09b6a13d5493bff3b4104ec68728272 Mon Sep 17 00:00:00 2001 From: Hiromi Ogawa Date: Sun, 28 Jun 2026 22:06:02 +0900 Subject: [PATCH] Track drop points by MIR place --- src/analyze/basic_block.rs | 50 +++++----- src/analyze/basic_block/drop_point.rs | 99 +++++++++++++------- src/analyze/basic_block/visitor/rust_call.rs | 16 +--- src/analyze/local_def.rs | 4 +- src/refine/env.rs | 8 +- 5 files changed, 103 insertions(+), 74 deletions(-) diff --git a/src/analyze/basic_block.rs b/src/analyze/basic_block.rs index 6e5aecf2..8761240d 100644 --- a/src/analyze/basic_block.rs +++ b/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 } diff --git a/src/analyze/basic_block/drop_point.rs b/src/analyze/basic_block/drop_point.rs index e660e3c3..72fdfdcd 100644 --- a/src/analyze/basic_block/drop_point.rs +++ b/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, - after_statements: Vec>, - after_terminator: HashMap>, - /// 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, +pub struct DropPoints<'tcx> { + pub before_statements: Vec>, + after_statements: Vec>>, + after_terminator: HashMap>>, + /// Places dropped after the terminator regardless of the target, in + /// addition to the liveness-derived sets above. + after_terminator_extra: Vec>, } -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 { +impl<'tcx> DropPoints<'tcx> { + pub fn position(&self, place: mir::Place<'tcx>) -> Option { 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 { + pub fn after_statement(&self, statement_index: usize) -> Vec> { 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 { - 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> { + 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; } diff --git a/src/analyze/basic_block/visitor/rust_call.rs b/src/analyze/basic_block/visitor/rust_call.rs index b46c0467..1de02ea5 100644 --- a/src/analyze/basic_block/visitor/rust_call.rs +++ b/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"); } } diff --git a/src/analyze/local_def.rs b/src/analyze/local_def.rs index 3938e071..1c85b3ee 100644 --- a/src/analyze/local_def.rs +++ b/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, + drop_points: HashMap>, 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 diff --git a/src/refine/env.rs b/src/refine/env.rs index ba494acb..b700a9eb 100644 --- a/src/refine/env.rs +++ b/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()); + } }