]>
Commit | Line | Data |
---|---|---|
dfeec247 | 1 | use rustc_index::vec::{Idx, IndexVec}; |
ba9703b0 XL |
2 | use rustc_middle::mir::*; |
3 | use rustc_middle::ty::Ty; | |
dfeec247 | 4 | use rustc_span::Span; |
3157f602 XL |
5 | |
6 | /// This struct represents a patch to MIR, which can add | |
7 | /// new statements and basic blocks and patch over block | |
8 | /// terminators. | |
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>)>, | |
c30ab7b3 | 13 | new_locals: Vec<LocalDecl<'tcx>>, |
3157f602 | 14 | resume_block: BasicBlock, |
c30ab7b3 | 15 | next_local: usize, |
3157f602 XL |
16 | } |
17 | ||
18 | impl<'tcx> MirPatch<'tcx> { | |
dc9dc135 | 19 | pub fn new(body: &Body<'tcx>) -> Self { |
3157f602 | 20 | let mut result = MirPatch { |
dc9dc135 | 21 | patch_map: IndexVec::from_elem(None, body.basic_blocks()), |
3157f602 | 22 | new_blocks: vec![], |
3157f602 | 23 | new_statements: vec![], |
c30ab7b3 | 24 | new_locals: vec![], |
dc9dc135 | 25 | next_local: body.local_decls.len(), |
0531ce1d | 26 | resume_block: START_BLOCK, |
3157f602 XL |
27 | }; |
28 | ||
29 | // make sure the MIR we create has a resume block. It is | |
30 | // completely legal to convert jumps to the resume block | |
31 | // to jumps to None, but we occasionally have to add | |
32 | // instructions just before that. | |
33 | ||
34 | let mut resume_block = None; | |
35 | let mut resume_stmt_block = None; | |
dc9dc135 | 36 | for (bb, block) in body.basic_blocks().iter_enumerated() { |
3157f602 | 37 | if let TerminatorKind::Resume = block.terminator().kind { |
74b04a01 | 38 | if !block.statements.is_empty() { |
041b39d2 | 39 | assert!(resume_stmt_block.is_none()); |
3157f602 XL |
40 | resume_stmt_block = Some(bb); |
41 | } else { | |
42 | resume_block = Some(bb); | |
43 | } | |
dfeec247 | 44 | break; |
3157f602 XL |
45 | } |
46 | } | |
47 | let resume_block = resume_block.unwrap_or_else(|| { | |
48 | result.new_block(BasicBlockData { | |
49 | statements: vec![], | |
50 | terminator: Some(Terminator { | |
f9f354fc | 51 | source_info: SourceInfo::outermost(body.span), |
dfeec247 | 52 | kind: TerminatorKind::Resume, |
3157f602 | 53 | }), |
dfeec247 XL |
54 | is_cleanup: true, |
55 | }) | |
56 | }); | |
3157f602 XL |
57 | result.resume_block = resume_block; |
58 | if let Some(resume_stmt_block) = resume_stmt_block { | |
dfeec247 XL |
59 | result |
60 | .patch_terminator(resume_stmt_block, TerminatorKind::Goto { target: resume_block }); | |
3157f602 XL |
61 | } |
62 | result | |
63 | } | |
64 | ||
65 | pub fn resume_block(&self) -> BasicBlock { | |
66 | self.resume_block | |
67 | } | |
68 | ||
69 | pub fn is_patched(&self, bb: BasicBlock) -> bool { | |
70 | self.patch_map[bb].is_some() | |
71 | } | |
72 | ||
dc9dc135 XL |
73 | pub fn terminator_loc(&self, body: &Body<'tcx>, bb: BasicBlock) -> Location { |
74 | let offset = match bb.index().checked_sub(body.basic_blocks().len()) { | |
3157f602 | 75 | Some(index) => self.new_blocks[index].statements.len(), |
dfeec247 | 76 | None => body[bb].statements.len(), |
3157f602 | 77 | }; |
dfeec247 | 78 | Location { block: bb, statement_index: offset } |
3157f602 XL |
79 | } |
80 | ||
cc61c64b | 81 | pub fn new_temp(&mut self, ty: Ty<'tcx>, span: Span) -> Local { |
c30ab7b3 SL |
82 | let index = self.next_local; |
83 | self.next_local += 1; | |
f9f354fc | 84 | self.new_locals.push(LocalDecl::new(ty, span)); |
c30ab7b3 | 85 | Local::new(index as usize) |
3157f602 XL |
86 | } |
87 | ||
ea8adc8c XL |
88 | pub fn new_internal(&mut self, ty: Ty<'tcx>, span: Span) -> Local { |
89 | let index = self.next_local; | |
90 | self.next_local += 1; | |
f9f354fc | 91 | self.new_locals.push(LocalDecl::new(ty, span).internal()); |
ea8adc8c XL |
92 | Local::new(index as usize) |
93 | } | |
94 | ||
3157f602 XL |
95 | pub fn new_block(&mut self, data: BasicBlockData<'tcx>) -> BasicBlock { |
96 | let block = BasicBlock::new(self.patch_map.len()); | |
97 | debug!("MirPatch: new_block: {:?}: {:?}", block, data); | |
98 | self.new_blocks.push(data); | |
99 | self.patch_map.push(None); | |
100 | block | |
101 | } | |
102 | ||
103 | pub fn patch_terminator(&mut self, block: BasicBlock, new: TerminatorKind<'tcx>) { | |
104 | assert!(self.patch_map[block].is_none()); | |
105 | debug!("MirPatch: patch_terminator({:?}, {:?})", block, new); | |
106 | self.patch_map[block] = Some(new); | |
107 | } | |
108 | ||
109 | pub fn add_statement(&mut self, loc: Location, stmt: StatementKind<'tcx>) { | |
110 | debug!("MirPatch: add_statement({:?}, {:?})", loc, stmt); | |
111 | self.new_statements.push((loc, stmt)); | |
112 | } | |
113 | ||
ff7c6d11 | 114 | pub fn add_assign(&mut self, loc: Location, place: Place<'tcx>, rv: Rvalue<'tcx>) { |
94222f64 | 115 | self.add_statement(loc, StatementKind::Assign(Box::new((place, rv)))); |
3157f602 XL |
116 | } |
117 | ||
f9f354fc | 118 | pub fn apply(self, body: &mut Body<'tcx>) { |
dfeec247 XL |
119 | debug!( |
120 | "MirPatch: {:?} new temps, starting from index {}: {:?}", | |
121 | self.new_locals.len(), | |
122 | body.local_decls.len(), | |
123 | self.new_locals | |
124 | ); | |
125 | debug!( | |
126 | "MirPatch: {} new blocks, starting from index {}", | |
127 | self.new_blocks.len(), | |
128 | body.basic_blocks().len() | |
129 | ); | |
dc9dc135 XL |
130 | body.basic_blocks_mut().extend(self.new_blocks); |
131 | body.local_decls.extend(self.new_locals); | |
3157f602 XL |
132 | for (src, patch) in self.patch_map.into_iter_enumerated() { |
133 | if let Some(patch) = patch { | |
134 | debug!("MirPatch: patching block {:?}", src); | |
dc9dc135 | 135 | body[src].terminator_mut().kind = patch; |
3157f602 XL |
136 | } |
137 | } | |
138 | ||
139 | let mut new_statements = self.new_statements; | |
8faf50e0 | 140 | new_statements.sort_by_key(|s| s.0); |
3157f602 XL |
141 | |
142 | let mut delta = 0; | |
143 | let mut last_bb = START_BLOCK; | |
144 | for (mut loc, stmt) in new_statements { | |
145 | if loc.block != last_bb { | |
146 | delta = 0; | |
147 | last_bb = loc.block; | |
148 | } | |
dfeec247 | 149 | debug!("MirPatch: adding statement {:?} at loc {:?}+{}", stmt, loc, delta); |
9e0c209e | 150 | loc.statement_index += delta; |
dfeec247 XL |
151 | let source_info = Self::source_info_for_index(&body[loc.block], loc); |
152 | body[loc.block] | |
153 | .statements | |
154 | .insert(loc.statement_index, Statement { source_info, kind: stmt }); | |
3157f602 XL |
155 | delta += 1; |
156 | } | |
157 | } | |
158 | ||
9fa01778 | 159 | pub fn source_info_for_index(data: &BasicBlockData<'_>, loc: Location) -> SourceInfo { |
9e0c209e | 160 | match data.statements.get(loc.statement_index) { |
3157f602 | 161 | Some(stmt) => stmt.source_info, |
dfeec247 | 162 | None => data.terminator().source_info, |
3157f602 XL |
163 | } |
164 | } | |
165 | ||
dc9dc135 XL |
166 | pub fn source_info_for_location(&self, body: &Body<'_>, loc: Location) -> SourceInfo { |
167 | let data = match loc.block.index().checked_sub(body.basic_blocks().len()) { | |
3157f602 | 168 | Some(new) => &self.new_blocks[new], |
dfeec247 | 169 | None => &body[loc.block], |
3157f602 XL |
170 | }; |
171 | Self::source_info_for_index(data, loc) | |
172 | } | |
173 | } |