1 use rustc_index::vec::{Idx, IndexVec};
2 use rustc_middle::mir::*;
3 use rustc_middle::ty::Ty;
6 /// This struct represents a patch to MIR, which can add
7 /// new statements and basic blocks and patch over block
9 pub struct MirPatch<'tcx> {
10 patch_map: IndexVec<BasicBlock, Option<TerminatorKind<'tcx>>>,
11 new_blocks: Vec<BasicBlockData<'tcx>>,
12 new_statements: Vec<(Location, StatementKind<'tcx>)>,
13 new_locals: Vec<LocalDecl<'tcx>>,
14 resume_block: Option<BasicBlock>,
19 impl<'tcx> MirPatch<'tcx> {
20 pub fn new(body: &Body<'tcx>) -> Self {
21 let mut result = MirPatch {
22 patch_map: IndexVec::from_elem(None, &body.basic_blocks),
24 new_statements: vec![],
26 next_local: body.local_decls.len(),
31 // Check if we already have a resume block
32 for (bb, block) in body.basic_blocks.iter_enumerated() {
33 if let TerminatorKind::Resume = block.terminator().kind && block.statements.is_empty() {
34 result.resume_block = Some(bb);
42 pub fn resume_block(&mut self) -> BasicBlock {
43 if let Some(bb) = self.resume_block {
47 let bb = self.new_block(BasicBlockData {
49 terminator: Some(Terminator {
50 source_info: SourceInfo::outermost(self.body_span),
51 kind: TerminatorKind::Resume,
55 self.resume_block = Some(bb);
59 pub fn is_patched(&self, bb: BasicBlock) -> bool {
60 self.patch_map[bb].is_some()
63 pub fn terminator_loc(&self, body: &Body<'tcx>, bb: BasicBlock) -> Location {
64 let offset = match bb.index().checked_sub(body.basic_blocks.len()) {
65 Some(index) => self.new_blocks[index].statements.len(),
66 None => body[bb].statements.len(),
68 Location { block: bb, statement_index: offset }
71 pub fn new_internal_with_info(
75 local_info: Option<Box<LocalInfo<'tcx>>>,
77 let index = self.next_local;
79 let mut new_decl = LocalDecl::new(ty, span).internal();
80 new_decl.local_info = local_info;
81 self.new_locals.push(new_decl);
82 Local::new(index as usize)
85 pub fn new_temp(&mut self, ty: Ty<'tcx>, span: Span) -> Local {
86 let index = self.next_local;
88 self.new_locals.push(LocalDecl::new(ty, span));
89 Local::new(index as usize)
92 pub fn new_internal(&mut self, ty: Ty<'tcx>, span: Span) -> Local {
93 let index = self.next_local;
95 self.new_locals.push(LocalDecl::new(ty, span).internal());
96 Local::new(index as usize)
99 pub fn new_block(&mut self, data: BasicBlockData<'tcx>) -> BasicBlock {
100 let block = BasicBlock::new(self.patch_map.len());
101 debug!("MirPatch: new_block: {:?}: {:?}", block, data);
102 self.new_blocks.push(data);
103 self.patch_map.push(None);
107 pub fn patch_terminator(&mut self, block: BasicBlock, new: TerminatorKind<'tcx>) {
108 assert!(self.patch_map[block].is_none());
109 debug!("MirPatch: patch_terminator({:?}, {:?})", block, new);
110 self.patch_map[block] = Some(new);
113 pub fn add_statement(&mut self, loc: Location, stmt: StatementKind<'tcx>) {
114 debug!("MirPatch: add_statement({:?}, {:?})", loc, stmt);
115 self.new_statements.push((loc, stmt));
118 pub fn add_assign(&mut self, loc: Location, place: Place<'tcx>, rv: Rvalue<'tcx>) {
119 self.add_statement(loc, StatementKind::Assign(Box::new((place, rv))));
122 pub fn apply(self, body: &mut Body<'tcx>) {
124 "MirPatch: {:?} new temps, starting from index {}: {:?}",
125 self.new_locals.len(),
126 body.local_decls.len(),
130 "MirPatch: {} new blocks, starting from index {}",
131 self.new_blocks.len(),
132 body.basic_blocks.len()
134 let bbs = if self.patch_map.is_empty() && self.new_blocks.is_empty() {
135 body.basic_blocks.as_mut_preserves_cfg()
137 body.basic_blocks.as_mut()
139 bbs.extend(self.new_blocks);
140 body.local_decls.extend(self.new_locals);
141 for (src, patch) in self.patch_map.into_iter_enumerated() {
142 if let Some(patch) = patch {
143 debug!("MirPatch: patching block {:?}", src);
144 bbs[src].terminator_mut().kind = patch;
148 let mut new_statements = self.new_statements;
149 new_statements.sort_by_key(|s| s.0);
152 let mut last_bb = START_BLOCK;
153 for (mut loc, stmt) in new_statements {
154 if loc.block != last_bb {
158 debug!("MirPatch: adding statement {:?} at loc {:?}+{}", stmt, loc, delta);
159 loc.statement_index += delta;
160 let source_info = Self::source_info_for_index(&body[loc.block], loc);
163 .insert(loc.statement_index, Statement { source_info, kind: stmt });
168 pub fn source_info_for_index(data: &BasicBlockData<'_>, loc: Location) -> SourceInfo {
169 match data.statements.get(loc.statement_index) {
170 Some(stmt) => stmt.source_info,
171 None => data.terminator().source_info,
175 pub fn source_info_for_location(&self, body: &Body<'tcx>, loc: Location) -> SourceInfo {
176 let data = match loc.block.index().checked_sub(body.basic_blocks.len()) {
177 Some(new) => &self.new_blocks[new],
178 None => &body[loc.block],
180 Self::source_info_for_index(data, loc)