1
Fork 0

Improve MirPatch documentation and naming.

It's currently lacking comments. This commit adds some, which is useful
because there are some methods with non-obvious behaviour.

The commit also renames two things:
- `patch_map` becomes `term_patch_map`, because it's only about
  terminators.
- `is_patched` becomes `is_term_patched`, for the same reason.

(I would guess that originally `MirPatch` only handled terminators, and
then over time it expanded to allow other modifications, but these names
weren't updated.)
This commit is contained in:
Nicholas Nethercote 2025-02-18 12:48:27 +11:00
parent a1daa34ad0
commit e3316ae453
2 changed files with 41 additions and 16 deletions

View file

@ -417,7 +417,7 @@ impl<'a, 'tcx> ElaborateDropsCtxt<'a, 'tcx> {
.. ..
} = data.terminator().kind } = data.terminator().kind
{ {
assert!(!self.patch.is_patched(bb)); assert!(!self.patch.is_term_patched(bb));
let loc = Location { block: tgt, statement_index: 0 }; let loc = Location { block: tgt, statement_index: 0 };
let path = self.move_data().rev_lookup.find(destination.as_ref()); let path = self.move_data().rev_lookup.find(destination.as_ref());
@ -462,7 +462,7 @@ impl<'a, 'tcx> ElaborateDropsCtxt<'a, 'tcx> {
// a Goto; see `MirPatch::new`). // a Goto; see `MirPatch::new`).
} }
_ => { _ => {
assert!(!self.patch.is_patched(bb)); assert!(!self.patch.is_term_patched(bb));
} }
} }
} }
@ -486,7 +486,7 @@ impl<'a, 'tcx> ElaborateDropsCtxt<'a, 'tcx> {
.. ..
} = data.terminator().kind } = data.terminator().kind
{ {
assert!(!self.patch.is_patched(bb)); assert!(!self.patch.is_term_patched(bb));
let loc = Location { block: bb, statement_index: data.statements.len() }; let loc = Location { block: bb, statement_index: data.statements.len() };
let path = self.move_data().rev_lookup.find(destination.as_ref()); let path = self.move_data().rev_lookup.find(destination.as_ref());

View file

@ -4,11 +4,12 @@ use rustc_middle::ty::Ty;
use rustc_span::Span; use rustc_span::Span;
use tracing::debug; use tracing::debug;
/// This struct represents a patch to MIR, which can add /// This struct lets you "patch" a MIR body, i.e. modify it. You can queue up
/// new statements and basic blocks and patch over block /// various changes, such as the addition of new statements and basic blocks
/// terminators. /// 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> { pub(crate) struct MirPatch<'tcx> {
patch_map: IndexVec<BasicBlock, Option<TerminatorKind<'tcx>>>, term_patch_map: IndexVec<BasicBlock, Option<TerminatorKind<'tcx>>>,
new_blocks: Vec<BasicBlockData<'tcx>>, new_blocks: Vec<BasicBlockData<'tcx>>,
new_statements: Vec<(Location, StatementKind<'tcx>)>, new_statements: Vec<(Location, StatementKind<'tcx>)>,
new_locals: Vec<LocalDecl<'tcx>>, new_locals: Vec<LocalDecl<'tcx>>,
@ -24,9 +25,10 @@ pub(crate) struct MirPatch<'tcx> {
} }
impl<'tcx> MirPatch<'tcx> { impl<'tcx> MirPatch<'tcx> {
/// Creates a new, empty patch.
pub(crate) fn new(body: &Body<'tcx>) -> Self { pub(crate) fn new(body: &Body<'tcx>) -> Self {
let mut result = MirPatch { let mut result = MirPatch {
patch_map: IndexVec::from_elem(None, &body.basic_blocks), term_patch_map: IndexVec::from_elem(None, &body.basic_blocks),
new_blocks: vec![], new_blocks: vec![],
new_statements: vec![], new_statements: vec![],
new_locals: vec![], new_locals: vec![],
@ -141,10 +143,12 @@ impl<'tcx> MirPatch<'tcx> {
bb bb
} }
pub(crate) fn is_patched(&self, bb: BasicBlock) -> bool { /// Has a replacement of this block's terminator been queued in this patch?
self.patch_map[bb].is_some() pub(crate) fn is_term_patched(&self, bb: BasicBlock) -> bool {
self.term_patch_map[bb].is_some()
} }
/// Queues the addition of a new temporary with additional local info.
pub(crate) fn new_local_with_info( pub(crate) fn new_local_with_info(
&mut self, &mut self,
ty: Ty<'tcx>, ty: Ty<'tcx>,
@ -159,6 +163,7 @@ impl<'tcx> MirPatch<'tcx> {
Local::new(index) Local::new(index)
} }
/// Queues the addition of a new temporary.
pub(crate) fn new_temp(&mut self, ty: Ty<'tcx>, span: Span) -> Local { pub(crate) fn new_temp(&mut self, ty: Ty<'tcx>, span: Span) -> Local {
let index = self.next_local; let index = self.next_local;
self.next_local += 1; self.next_local += 1;
@ -174,29 +179,46 @@ impl<'tcx> MirPatch<'tcx> {
self.new_locals[new_local_idx].ty self.new_locals[new_local_idx].ty
} }
/// Queues the addition of a new basic block.
pub(crate) fn new_block(&mut self, data: BasicBlockData<'tcx>) -> BasicBlock { pub(crate) fn new_block(&mut self, data: BasicBlockData<'tcx>) -> BasicBlock {
let block = BasicBlock::new(self.patch_map.len()); let block = BasicBlock::new(self.term_patch_map.len());
debug!("MirPatch: new_block: {:?}: {:?}", block, data); debug!("MirPatch: new_block: {:?}: {:?}", block, data);
self.new_blocks.push(data); self.new_blocks.push(data);
self.patch_map.push(None); self.term_patch_map.push(None);
block block
} }
/// Queues the replacement of a block's terminator.
pub(crate) fn patch_terminator(&mut self, block: BasicBlock, new: TerminatorKind<'tcx>) { pub(crate) fn patch_terminator(&mut self, block: BasicBlock, new: TerminatorKind<'tcx>) {
assert!(self.patch_map[block].is_none()); assert!(self.term_patch_map[block].is_none());
debug!("MirPatch: patch_terminator({:?}, {:?})", block, new); debug!("MirPatch: patch_terminator({:?}, {:?})", block, new);
self.patch_map[block] = Some(new); self.term_patch_map[block] = Some(new);
} }
/// Queues the insertion of a statement at a given location. The statement
/// currently at that location, and all statements that follow, are shifted
/// down. If multiple statements are queued for addition at the same
/// location, the final statement order after calling `apply` will match
/// the queue insertion order.
///
/// E.g. if we have `s0` at location `loc` and do these calls:
///
/// p.add_statement(loc, s1);
/// p.add_statement(loc, s2);
/// p.apply(body);
///
/// then the final order will be `s1, s2, s0`, with `s1` at `loc`.
pub(crate) fn add_statement(&mut self, loc: Location, stmt: StatementKind<'tcx>) { pub(crate) fn add_statement(&mut self, loc: Location, stmt: StatementKind<'tcx>) {
debug!("MirPatch: add_statement({:?}, {:?})", loc, stmt); debug!("MirPatch: add_statement({:?}, {:?})", loc, stmt);
self.new_statements.push((loc, stmt)); self.new_statements.push((loc, stmt));
} }
/// Like `add_statement`, but specialized for assignments.
pub(crate) fn add_assign(&mut self, loc: Location, place: Place<'tcx>, rv: Rvalue<'tcx>) { pub(crate) fn add_assign(&mut self, loc: Location, place: Place<'tcx>, rv: Rvalue<'tcx>) {
self.add_statement(loc, StatementKind::Assign(Box::new((place, rv)))); self.add_statement(loc, StatementKind::Assign(Box::new((place, rv))));
} }
/// Applies the queued changes.
pub(crate) fn apply(self, body: &mut Body<'tcx>) { pub(crate) fn apply(self, body: &mut Body<'tcx>) {
debug!( debug!(
"MirPatch: {:?} new temps, starting from index {}: {:?}", "MirPatch: {:?} new temps, starting from index {}: {:?}",
@ -209,14 +231,14 @@ impl<'tcx> MirPatch<'tcx> {
self.new_blocks.len(), self.new_blocks.len(),
body.basic_blocks.len() body.basic_blocks.len()
); );
let bbs = if self.patch_map.is_empty() && self.new_blocks.is_empty() { let bbs = if self.term_patch_map.is_empty() && self.new_blocks.is_empty() {
body.basic_blocks.as_mut_preserves_cfg() body.basic_blocks.as_mut_preserves_cfg()
} else { } else {
body.basic_blocks.as_mut() body.basic_blocks.as_mut()
}; };
bbs.extend(self.new_blocks); bbs.extend(self.new_blocks);
body.local_decls.extend(self.new_locals); body.local_decls.extend(self.new_locals);
for (src, patch) in self.patch_map.into_iter_enumerated() { for (src, patch) in self.term_patch_map.into_iter_enumerated() {
if let Some(patch) = patch { if let Some(patch) = patch {
debug!("MirPatch: patching block {:?}", src); debug!("MirPatch: patching block {:?}", src);
bbs[src].terminator_mut().kind = patch; bbs[src].terminator_mut().kind = patch;
@ -224,6 +246,9 @@ impl<'tcx> MirPatch<'tcx> {
} }
let mut new_statements = self.new_statements; let mut new_statements = self.new_statements;
// This must be a stable sort to provide the ordering described in the
// comment for `add_statement`.
new_statements.sort_by_key(|s| s.0); new_statements.sort_by_key(|s| s.0);
let mut delta = 0; let mut delta = 0;