From b216cf34b1d4dfce1a44f99f7c58e64724439a8e Mon Sep 17 00:00:00 2001 From: Camille Gillot Date: Thu, 18 Sep 2025 00:46:27 +0000 Subject: [PATCH 1/3] Avoid invalidating from MirPatch::apply. --- compiler/rustc_mir_transform/src/patch.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/compiler/rustc_mir_transform/src/patch.rs b/compiler/rustc_mir_transform/src/patch.rs index c781d1a5324b..f38881baf0f2 100644 --- a/compiler/rustc_mir_transform/src/patch.rs +++ b/compiler/rustc_mir_transform/src/patch.rs @@ -244,7 +244,7 @@ pub(crate) fn apply(self, body: &mut Body<'tcx>) { self.new_blocks.len(), body.basic_blocks.len() ); - let bbs = if self.term_patch_map.is_empty() && self.new_blocks.is_empty() { + let bbs = if self.term_patch_map.iter().all(Option::is_none) && self.new_blocks.is_empty() { body.basic_blocks.as_mut_preserves_cfg() } else { body.basic_blocks.as_mut() @@ -273,8 +273,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; From 3c232fe38f1629ae469ee97455ce0be07a1ae8ae Mon Sep 17 00:00:00 2001 From: Camille Gillot Date: Sat, 20 Sep 2025 13:53:58 +0000 Subject: [PATCH 2/3] Make term_patch_map sparse. --- compiler/rustc_mir_transform/src/patch.rs | 32 +++++++++++++---------- 1 file changed, 18 insertions(+), 14 deletions(-) diff --git a/compiler/rustc_mir_transform/src/patch.rs b/compiler/rustc_mir_transform/src/patch.rs index f38881baf0f2..8f2124405650 100644 --- a/compiler/rustc_mir_transform/src/patch.rs +++ b/compiler/rustc_mir_transform/src/patch.rs @@ -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>>, + term_patch_map: FxHashMap>, new_blocks: Vec>, new_statements: Vec<(Location, StatementKind<'tcx>)>, new_locals: Vec>, @@ -22,17 +23,19 @@ pub(crate) struct MirPatch<'tcx> { terminate_block: Option<(BasicBlock, UnwindTerminateReason)>, body_span: Span, next_local: usize, + 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 +144,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 +197,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,18 +246,20 @@ pub(crate) fn apply(self, body: &mut Body<'tcx>) { self.new_blocks.len(), body.basic_blocks.len() ); - let bbs = if self.term_patch_map.iter().all(Option::is_none) && self.new_blocks.is_empty() { + 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 { body.basic_blocks.as_mut() }; 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; From ce677c7db80a245e1eee5b948dd6b20b4bba0423 Mon Sep 17 00:00:00 2001 From: Camille Gillot Date: Tue, 23 Sep 2025 20:38:38 -0300 Subject: [PATCH 3/3] Update compiler/rustc_mir_transform/src/patch.rs Co-authored-by: lcnr --- compiler/rustc_mir_transform/src/patch.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/compiler/rustc_mir_transform/src/patch.rs b/compiler/rustc_mir_transform/src/patch.rs index 8f2124405650..d831ab50b1ac 100644 --- a/compiler/rustc_mir_transform/src/patch.rs +++ b/compiler/rustc_mir_transform/src/patch.rs @@ -23,6 +23,8 @@ 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, }