]>
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>>, |
064997fb | 14 | resume_block: Option<BasicBlock>, |
353b0b11 FG |
15 | // Only for unreachable in cleanup path. |
16 | unreachable_cleanup_block: Option<BasicBlock>, | |
17 | terminate_block: Option<BasicBlock>, | |
064997fb | 18 | body_span: Span, |
c30ab7b3 | 19 | next_local: usize, |
3157f602 XL |
20 | } |
21 | ||
22 | impl<'tcx> MirPatch<'tcx> { | |
dc9dc135 | 23 | pub fn new(body: &Body<'tcx>) -> Self { |
3157f602 | 24 | let mut result = MirPatch { |
f2b60f7d | 25 | patch_map: IndexVec::from_elem(None, &body.basic_blocks), |
3157f602 | 26 | new_blocks: vec![], |
3157f602 | 27 | new_statements: vec![], |
c30ab7b3 | 28 | new_locals: vec![], |
dc9dc135 | 29 | next_local: body.local_decls.len(), |
064997fb | 30 | resume_block: None, |
353b0b11 FG |
31 | unreachable_cleanup_block: None, |
32 | terminate_block: None, | |
064997fb | 33 | body_span: body.span, |
3157f602 XL |
34 | }; |
35 | ||
f2b60f7d | 36 | for (bb, block) in body.basic_blocks.iter_enumerated() { |
353b0b11 | 37 | // Check if we already have a resume block |
064997fb FG |
38 | if let TerminatorKind::Resume = block.terminator().kind && block.statements.is_empty() { |
39 | result.resume_block = Some(bb); | |
353b0b11 FG |
40 | continue; |
41 | } | |
42 | ||
43 | // Check if we already have an unreachable block | |
44 | if let TerminatorKind::Unreachable = block.terminator().kind | |
45 | && block.statements.is_empty() | |
46 | && block.is_cleanup | |
47 | { | |
48 | result.unreachable_cleanup_block = Some(bb); | |
49 | continue; | |
50 | } | |
51 | ||
52 | // Check if we already have a terminate block | |
53 | if let TerminatorKind::Terminate = block.terminator().kind && block.statements.is_empty() { | |
54 | result.terminate_block = Some(bb); | |
55 | continue; | |
3157f602 XL |
56 | } |
57 | } | |
064997fb | 58 | |
3157f602 XL |
59 | result |
60 | } | |
61 | ||
064997fb FG |
62 | pub fn resume_block(&mut self) -> BasicBlock { |
63 | if let Some(bb) = self.resume_block { | |
64 | return bb; | |
65 | } | |
66 | ||
67 | let bb = self.new_block(BasicBlockData { | |
68 | statements: vec![], | |
69 | terminator: Some(Terminator { | |
70 | source_info: SourceInfo::outermost(self.body_span), | |
71 | kind: TerminatorKind::Resume, | |
72 | }), | |
73 | is_cleanup: true, | |
74 | }); | |
75 | self.resume_block = Some(bb); | |
76 | bb | |
3157f602 XL |
77 | } |
78 | ||
353b0b11 FG |
79 | pub fn unreachable_cleanup_block(&mut self) -> BasicBlock { |
80 | if let Some(bb) = self.unreachable_cleanup_block { | |
81 | return bb; | |
82 | } | |
83 | ||
84 | let bb = self.new_block(BasicBlockData { | |
85 | statements: vec![], | |
86 | terminator: Some(Terminator { | |
87 | source_info: SourceInfo::outermost(self.body_span), | |
88 | kind: TerminatorKind::Unreachable, | |
89 | }), | |
90 | is_cleanup: true, | |
91 | }); | |
92 | self.unreachable_cleanup_block = Some(bb); | |
93 | bb | |
94 | } | |
95 | ||
96 | pub fn terminate_block(&mut self) -> BasicBlock { | |
97 | if let Some(bb) = self.terminate_block { | |
98 | return bb; | |
99 | } | |
100 | ||
101 | let bb = self.new_block(BasicBlockData { | |
102 | statements: vec![], | |
103 | terminator: Some(Terminator { | |
104 | source_info: SourceInfo::outermost(self.body_span), | |
105 | kind: TerminatorKind::Terminate, | |
106 | }), | |
107 | is_cleanup: true, | |
108 | }); | |
109 | self.terminate_block = Some(bb); | |
110 | bb | |
111 | } | |
112 | ||
3157f602 XL |
113 | pub fn is_patched(&self, bb: BasicBlock) -> bool { |
114 | self.patch_map[bb].is_some() | |
115 | } | |
116 | ||
dc9dc135 | 117 | pub fn terminator_loc(&self, body: &Body<'tcx>, bb: BasicBlock) -> Location { |
f2b60f7d | 118 | let offset = match bb.index().checked_sub(body.basic_blocks.len()) { |
3157f602 | 119 | Some(index) => self.new_blocks[index].statements.len(), |
dfeec247 | 120 | None => body[bb].statements.len(), |
3157f602 | 121 | }; |
dfeec247 | 122 | Location { block: bb, statement_index: offset } |
3157f602 XL |
123 | } |
124 | ||
f2b60f7d | 125 | pub fn new_internal_with_info( |
04454e1e FG |
126 | &mut self, |
127 | ty: Ty<'tcx>, | |
128 | span: Span, | |
353b0b11 | 129 | local_info: LocalInfo<'tcx>, |
04454e1e | 130 | ) -> Local { |
c30ab7b3 SL |
131 | let index = self.next_local; |
132 | self.next_local += 1; | |
f2b60f7d | 133 | let mut new_decl = LocalDecl::new(ty, span).internal(); |
353b0b11 | 134 | **new_decl.local_info.as_mut().assert_crate_local() = local_info; |
04454e1e | 135 | self.new_locals.push(new_decl); |
353b0b11 | 136 | Local::new(index) |
3157f602 XL |
137 | } |
138 | ||
04454e1e | 139 | pub fn new_temp(&mut self, ty: Ty<'tcx>, span: Span) -> Local { |
f2b60f7d FG |
140 | let index = self.next_local; |
141 | self.next_local += 1; | |
142 | self.new_locals.push(LocalDecl::new(ty, span)); | |
353b0b11 | 143 | Local::new(index) |
04454e1e FG |
144 | } |
145 | ||
ea8adc8c XL |
146 | pub fn new_internal(&mut self, ty: Ty<'tcx>, span: Span) -> Local { |
147 | let index = self.next_local; | |
148 | self.next_local += 1; | |
f9f354fc | 149 | self.new_locals.push(LocalDecl::new(ty, span).internal()); |
353b0b11 | 150 | Local::new(index) |
ea8adc8c XL |
151 | } |
152 | ||
3157f602 XL |
153 | pub fn new_block(&mut self, data: BasicBlockData<'tcx>) -> BasicBlock { |
154 | let block = BasicBlock::new(self.patch_map.len()); | |
155 | debug!("MirPatch: new_block: {:?}: {:?}", block, data); | |
156 | self.new_blocks.push(data); | |
157 | self.patch_map.push(None); | |
158 | block | |
159 | } | |
160 | ||
161 | pub fn patch_terminator(&mut self, block: BasicBlock, new: TerminatorKind<'tcx>) { | |
162 | assert!(self.patch_map[block].is_none()); | |
163 | debug!("MirPatch: patch_terminator({:?}, {:?})", block, new); | |
164 | self.patch_map[block] = Some(new); | |
165 | } | |
166 | ||
167 | pub fn add_statement(&mut self, loc: Location, stmt: StatementKind<'tcx>) { | |
168 | debug!("MirPatch: add_statement({:?}, {:?})", loc, stmt); | |
169 | self.new_statements.push((loc, stmt)); | |
170 | } | |
171 | ||
ff7c6d11 | 172 | pub fn add_assign(&mut self, loc: Location, place: Place<'tcx>, rv: Rvalue<'tcx>) { |
94222f64 | 173 | self.add_statement(loc, StatementKind::Assign(Box::new((place, rv)))); |
3157f602 XL |
174 | } |
175 | ||
f9f354fc | 176 | pub fn apply(self, body: &mut Body<'tcx>) { |
dfeec247 XL |
177 | debug!( |
178 | "MirPatch: {:?} new temps, starting from index {}: {:?}", | |
179 | self.new_locals.len(), | |
180 | body.local_decls.len(), | |
181 | self.new_locals | |
182 | ); | |
183 | debug!( | |
184 | "MirPatch: {} new blocks, starting from index {}", | |
185 | self.new_blocks.len(), | |
f2b60f7d | 186 | body.basic_blocks.len() |
dfeec247 | 187 | ); |
064997fb FG |
188 | let bbs = if self.patch_map.is_empty() && self.new_blocks.is_empty() { |
189 | body.basic_blocks.as_mut_preserves_cfg() | |
190 | } else { | |
191 | body.basic_blocks.as_mut() | |
192 | }; | |
193 | bbs.extend(self.new_blocks); | |
dc9dc135 | 194 | body.local_decls.extend(self.new_locals); |
3157f602 XL |
195 | for (src, patch) in self.patch_map.into_iter_enumerated() { |
196 | if let Some(patch) = patch { | |
197 | debug!("MirPatch: patching block {:?}", src); | |
064997fb | 198 | bbs[src].terminator_mut().kind = patch; |
3157f602 XL |
199 | } |
200 | } | |
201 | ||
202 | let mut new_statements = self.new_statements; | |
8faf50e0 | 203 | new_statements.sort_by_key(|s| s.0); |
3157f602 XL |
204 | |
205 | let mut delta = 0; | |
206 | let mut last_bb = START_BLOCK; | |
207 | for (mut loc, stmt) in new_statements { | |
208 | if loc.block != last_bb { | |
209 | delta = 0; | |
210 | last_bb = loc.block; | |
211 | } | |
dfeec247 | 212 | debug!("MirPatch: adding statement {:?} at loc {:?}+{}", stmt, loc, delta); |
9e0c209e | 213 | loc.statement_index += delta; |
dfeec247 XL |
214 | let source_info = Self::source_info_for_index(&body[loc.block], loc); |
215 | body[loc.block] | |
216 | .statements | |
217 | .insert(loc.statement_index, Statement { source_info, kind: stmt }); | |
3157f602 XL |
218 | delta += 1; |
219 | } | |
220 | } | |
221 | ||
9fa01778 | 222 | pub fn source_info_for_index(data: &BasicBlockData<'_>, loc: Location) -> SourceInfo { |
9e0c209e | 223 | match data.statements.get(loc.statement_index) { |
3157f602 | 224 | Some(stmt) => stmt.source_info, |
dfeec247 | 225 | None => data.terminator().source_info, |
3157f602 XL |
226 | } |
227 | } | |
228 | ||
923072b8 | 229 | pub fn source_info_for_location(&self, body: &Body<'tcx>, loc: Location) -> SourceInfo { |
f2b60f7d | 230 | let data = match loc.block.index().checked_sub(body.basic_blocks.len()) { |
3157f602 | 231 | Some(new) => &self.new_blocks[new], |
dfeec247 | 232 | None => &body[loc.block], |
3157f602 XL |
233 | }; |
234 | Self::source_info_for_index(data, loc) | |
235 | } | |
236 | } |