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: BasicBlock,
16 make_nop: Vec<Location>,
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(),
27 resume_block: START_BLOCK,
31 // make sure the MIR we create has a resume block. It is
32 // completely legal to convert jumps to the resume block
33 // to jumps to None, but we occasionally have to add
34 // instructions just before that.
36 let mut resume_block = None;
37 let mut resume_stmt_block = None;
38 for (bb, block) in body.basic_blocks().iter_enumerated() {
39 if let TerminatorKind::Resume = block.terminator().kind {
40 if !block.statements.is_empty() {
41 assert!(resume_stmt_block.is_none());
42 resume_stmt_block = Some(bb);
44 resume_block = Some(bb);
49 let resume_block = resume_block.unwrap_or_else(|| {
50 result.new_block(BasicBlockData {
52 terminator: Some(Terminator {
53 source_info: SourceInfo { span: body.span, scope: OUTERMOST_SOURCE_SCOPE },
54 kind: TerminatorKind::Resume,
59 result.resume_block = resume_block;
60 if let Some(resume_stmt_block) = resume_stmt_block {
62 .patch_terminator(resume_stmt_block, TerminatorKind::Goto { target: resume_block });
67 pub fn resume_block(&self) -> BasicBlock {
71 pub fn is_patched(&self, bb: BasicBlock) -> bool {
72 self.patch_map[bb].is_some()
75 pub fn terminator_loc(&self, body: &Body<'tcx>, bb: BasicBlock) -> Location {
76 let offset = match bb.index().checked_sub(body.basic_blocks().len()) {
77 Some(index) => self.new_blocks[index].statements.len(),
78 None => body[bb].statements.len(),
80 Location { block: bb, statement_index: offset }
83 pub fn new_temp(&mut self, ty: Ty<'tcx>, span: Span) -> Local {
84 let index = self.next_local;
86 self.new_locals.push(LocalDecl::new_temp(ty, span));
87 Local::new(index as usize)
90 pub fn new_internal(&mut self, ty: Ty<'tcx>, span: Span) -> Local {
91 let index = self.next_local;
93 self.new_locals.push(LocalDecl::new_internal(ty, span));
94 Local::new(index as usize)
97 pub fn new_block(&mut self, data: BasicBlockData<'tcx>) -> BasicBlock {
98 let block = BasicBlock::new(self.patch_map.len());
99 debug!("MirPatch: new_block: {:?}: {:?}", block, data);
100 self.new_blocks.push(data);
101 self.patch_map.push(None);
105 pub fn patch_terminator(&mut self, block: BasicBlock, new: TerminatorKind<'tcx>) {
106 assert!(self.patch_map[block].is_none());
107 debug!("MirPatch: patch_terminator({:?}, {:?})", block, new);
108 self.patch_map[block] = Some(new);
111 pub fn add_statement(&mut self, loc: Location, stmt: StatementKind<'tcx>) {
112 debug!("MirPatch: add_statement({:?}, {:?})", loc, stmt);
113 self.new_statements.push((loc, stmt));
116 pub fn add_assign(&mut self, loc: Location, place: Place<'tcx>, rv: Rvalue<'tcx>) {
117 self.add_statement(loc, StatementKind::Assign(box (place, rv)));
120 pub fn make_nop(&mut self, loc: Location) {
121 self.make_nop.push(loc);
124 pub fn apply(self, body: &mut BodyAndCache<'tcx>) {
125 debug!("MirPatch: make nops at: {:?}", self.make_nop);
126 for loc in self.make_nop {
127 body.make_statement_nop(loc);
130 "MirPatch: {:?} new temps, starting from index {}: {:?}",
131 self.new_locals.len(),
132 body.local_decls.len(),
136 "MirPatch: {} new blocks, starting from index {}",
137 self.new_blocks.len(),
138 body.basic_blocks().len()
140 body.basic_blocks_mut().extend(self.new_blocks);
141 body.local_decls.extend(self.new_locals);
142 for (src, patch) in self.patch_map.into_iter_enumerated() {
143 if let Some(patch) = patch {
144 debug!("MirPatch: patching block {:?}", src);
145 body[src].terminator_mut().kind = patch;
149 let mut new_statements = self.new_statements;
150 new_statements.sort_by_key(|s| s.0);
153 let mut last_bb = START_BLOCK;
154 for (mut loc, stmt) in new_statements {
155 if loc.block != last_bb {
159 debug!("MirPatch: adding statement {:?} at loc {:?}+{}", stmt, loc, delta);
160 loc.statement_index += delta;
161 let source_info = Self::source_info_for_index(&body[loc.block], loc);
164 .insert(loc.statement_index, Statement { source_info, kind: stmt });
169 pub fn source_info_for_index(data: &BasicBlockData<'_>, loc: Location) -> SourceInfo {
170 match data.statements.get(loc.statement_index) {
171 Some(stmt) => stmt.source_info,
172 None => data.terminator().source_info,
176 pub fn source_info_for_location(&self, body: &Body<'_>, loc: Location) -> SourceInfo {
177 let data = match loc.block.index().checked_sub(body.basic_blocks().len()) {
178 Some(new) => &self.new_blocks[new],
179 None => &body[loc.block],
181 Self::source_info_for_index(data, loc)