Auto merge of #146697 - cjgillot:invalidate-patch, r=lcnr

Avoid invalidating CFG caches from MirPatch::apply.

Small effort to reduce invalidating CFG caches.
This commit is contained in:
bors
2025-09-25 17:19:29 +00:00
+21 -15
View File
@@ -1,4 +1,5 @@
use rustc_index::{Idx, IndexVec};
use rustc_data_structures::fx::FxHashMap;
use rustc_index::Idx;
use rustc_middle::mir::*;
use rustc_middle::ty::Ty;
use rustc_span::Span;
@@ -9,7 +10,7 @@
/// and replacement of terminators, and then apply the queued changes all at
/// once with `apply`. This is useful for MIR transformation passes.
pub(crate) struct MirPatch<'tcx> {
term_patch_map: IndexVec<BasicBlock, Option<TerminatorKind<'tcx>>>,
term_patch_map: FxHashMap<BasicBlock, TerminatorKind<'tcx>>,
new_blocks: Vec<BasicBlockData<'tcx>>,
new_statements: Vec<(Location, StatementKind<'tcx>)>,
new_locals: Vec<LocalDecl<'tcx>>,
@@ -22,17 +23,21 @@ pub(crate) struct MirPatch<'tcx> {
terminate_block: Option<(BasicBlock, UnwindTerminateReason)>,
body_span: Span,
next_local: usize,
/// The number of blocks at the start of the transformation. New blocks
/// get appended at the end.
next_block: usize,
}
impl<'tcx> MirPatch<'tcx> {
/// Creates a new, empty patch.
pub(crate) fn new(body: &Body<'tcx>) -> Self {
let mut result = MirPatch {
term_patch_map: IndexVec::from_elem(None, &body.basic_blocks),
term_patch_map: Default::default(),
new_blocks: vec![],
new_statements: vec![],
new_locals: vec![],
next_local: body.local_decls.len(),
next_block: body.basic_blocks.len(),
resume_block: None,
unreachable_cleanup_block: None,
unreachable_no_cleanup_block: None,
@@ -141,7 +146,7 @@ pub(crate) fn terminate_block(&mut self, reason: UnwindTerminateReason) -> Basic
/// Has a replacement of this block's terminator been queued in this patch?
pub(crate) fn is_term_patched(&self, bb: BasicBlock) -> bool {
self.term_patch_map[bb].is_some()
self.term_patch_map.contains_key(&bb)
}
/// Universal getter for block data, either it is in 'old' blocks or in patched ones
@@ -194,18 +199,17 @@ pub(crate) fn local_ty(&self, local: Local) -> Ty<'tcx> {
/// Queues the addition of a new basic block.
pub(crate) fn new_block(&mut self, data: BasicBlockData<'tcx>) -> BasicBlock {
let block = self.term_patch_map.next_index();
let block = BasicBlock::from_usize(self.next_block + self.new_blocks.len());
debug!("MirPatch: new_block: {:?}: {:?}", block, data);
self.new_blocks.push(data);
self.term_patch_map.push(None);
block
}
/// Queues the replacement of a block's terminator.
pub(crate) fn patch_terminator(&mut self, block: BasicBlock, new: TerminatorKind<'tcx>) {
assert!(self.term_patch_map[block].is_none());
assert!(!self.term_patch_map.contains_key(&block));
debug!("MirPatch: patch_terminator({:?}, {:?})", block, new);
self.term_patch_map[block] = Some(new);
self.term_patch_map.insert(block, new);
}
/// Queues the insertion of a statement at a given location. The statement
@@ -244,6 +248,7 @@ pub(crate) fn apply(self, body: &mut Body<'tcx>) {
self.new_blocks.len(),
body.basic_blocks.len()
);
debug_assert_eq!(self.next_block, body.basic_blocks.len());
let bbs = if self.term_patch_map.is_empty() && self.new_blocks.is_empty() {
body.basic_blocks.as_mut_preserves_cfg()
} else {
@@ -251,11 +256,12 @@ pub(crate) fn apply(self, body: &mut Body<'tcx>) {
};
bbs.extend(self.new_blocks);
body.local_decls.extend(self.new_locals);
for (src, patch) in self.term_patch_map.into_iter_enumerated() {
if let Some(patch) = patch {
debug!("MirPatch: patching block {:?}", src);
bbs[src].terminator_mut().kind = patch;
}
// The order in which we patch terminators does not change the result.
#[allow(rustc::potential_query_instability)]
for (src, patch) in self.term_patch_map {
debug!("MirPatch: patching block {:?}", src);
bbs[src].terminator_mut().kind = patch;
}
let mut new_statements = self.new_statements;
@@ -273,8 +279,8 @@ pub(crate) fn apply(self, body: &mut Body<'tcx>) {
}
debug!("MirPatch: adding statement {:?} at loc {:?}+{}", stmt, loc, delta);
loc.statement_index += delta;
let source_info = Self::source_info_for_index(&body[loc.block], loc);
body[loc.block]
let source_info = Self::source_info_for_index(&bbs[loc.block], loc);
bbs[loc.block]
.statements
.insert(loc.statement_index, Statement::new(source_info, stmt));
delta += 1;