|
1 | 1 | use super::*; |
2 | | -use crate::mir::visit::Visitor; |
3 | 2 |
|
4 | 3 | /// Preorder traversal of a graph. |
5 | 4 | /// |
@@ -105,36 +104,46 @@ impl<'a, 'tcx> Iterator for Preorder<'a, 'tcx> { |
105 | 104 | /// ``` |
106 | 105 | /// |
107 | 106 | /// A Postorder traversal of this graph is `D B C A` or `D C B A` |
108 | | -pub struct Postorder<'a, 'tcx> { |
| 107 | +pub struct Postorder<'a, 'tcx, C> { |
109 | 108 | basic_blocks: &'a IndexSlice<BasicBlock, BasicBlockData<'tcx>>, |
110 | 109 | visited: BitSet<BasicBlock>, |
111 | 110 | visit_stack: Vec<(BasicBlock, Successors<'a>)>, |
112 | 111 | root_is_start_block: bool, |
| 112 | + extra: C, |
113 | 113 | } |
114 | 114 |
|
115 | | -impl<'a, 'tcx> Postorder<'a, 'tcx> { |
| 115 | +impl<'a, 'tcx, C> Postorder<'a, 'tcx, C> |
| 116 | +where |
| 117 | + C: Customization<'tcx>, |
| 118 | +{ |
116 | 119 | pub fn new( |
117 | 120 | basic_blocks: &'a IndexSlice<BasicBlock, BasicBlockData<'tcx>>, |
118 | 121 | root: BasicBlock, |
119 | | - ) -> Postorder<'a, 'tcx> { |
| 122 | + extra: C, |
| 123 | + ) -> Postorder<'a, 'tcx, C> { |
120 | 124 | let mut po = Postorder { |
121 | 125 | basic_blocks, |
122 | 126 | visited: BitSet::new_empty(basic_blocks.len()), |
123 | 127 | visit_stack: Vec::new(), |
124 | 128 | root_is_start_block: root == START_BLOCK, |
| 129 | + extra, |
125 | 130 | }; |
126 | 131 |
|
127 | | - let data = &po.basic_blocks[root]; |
128 | | - |
129 | | - if let Some(ref term) = data.terminator { |
130 | | - po.visited.insert(root); |
131 | | - po.visit_stack.push((root, term.successors())); |
132 | | - po.traverse_successor(); |
133 | | - } |
| 132 | + po.visit(root); |
| 133 | + po.traverse_successor(); |
134 | 134 |
|
135 | 135 | po |
136 | 136 | } |
137 | 137 |
|
| 138 | + fn visit(&mut self, bb: BasicBlock) { |
| 139 | + if !self.visited.insert(bb) { |
| 140 | + return; |
| 141 | + } |
| 142 | + let data = &self.basic_blocks[bb]; |
| 143 | + let successors = C::successors(data, self.extra); |
| 144 | + self.visit_stack.push((bb, successors)); |
| 145 | + } |
| 146 | + |
138 | 147 | fn traverse_successor(&mut self) { |
139 | 148 | // This is quite a complex loop due to 1. the borrow checker not liking it much |
140 | 149 | // and 2. what exactly is going on is not clear |
@@ -184,16 +193,15 @@ impl<'a, 'tcx> Postorder<'a, 'tcx> { |
184 | 193 | // since we've already visited `E`, that child isn't added to the stack. The last |
185 | 194 | // two iterations yield `B` and finally `A` for a final traversal of [E, D, C, B, A] |
186 | 195 | while let Some(bb) = self.visit_stack.last_mut().and_then(|(_, iter)| iter.next_back()) { |
187 | | - if self.visited.insert(bb) { |
188 | | - if let Some(term) = &self.basic_blocks[bb].terminator { |
189 | | - self.visit_stack.push((bb, term.successors())); |
190 | | - } |
191 | | - } |
| 196 | + self.visit(bb); |
192 | 197 | } |
193 | 198 | } |
194 | 199 | } |
195 | 200 |
|
196 | | -impl<'tcx> Iterator for Postorder<'_, 'tcx> { |
| 201 | +impl<'tcx, C> Iterator for Postorder<'_, 'tcx, C> |
| 202 | +where |
| 203 | + C: Customization<'tcx>, |
| 204 | +{ |
197 | 205 | type Item = BasicBlock; |
198 | 206 |
|
199 | 207 | fn next(&mut self) -> Option<BasicBlock> { |
@@ -233,88 +241,38 @@ pub fn postorder<'a, 'tcx>( |
233 | 241 | reverse_postorder(body).rev() |
234 | 242 | } |
235 | 243 |
|
236 | | -struct UsedLocals(BitSet<Local>); |
237 | | - |
238 | | -impl<'tcx> Visitor<'tcx> for UsedLocals { |
239 | | - fn visit_local( |
240 | | - &mut self, |
241 | | - local: Local, |
242 | | - _ctx: crate::mir::visit::PlaceContext, |
243 | | - _location: Location, |
244 | | - ) { |
245 | | - self.0.insert(local); |
246 | | - } |
247 | | -} |
248 | | - |
249 | | -struct MonoReachablePostorder<'a, 'tcx> { |
250 | | - basic_blocks: &'a IndexSlice<BasicBlock, BasicBlockData<'tcx>>, |
251 | | - visited: BitSet<BasicBlock>, |
252 | | - visit_stack: Vec<(BasicBlock, Successors<'a>)>, |
253 | | - locals: UsedLocals, |
254 | | - tcx: TyCtxt<'tcx>, |
255 | | - instance: Instance<'tcx>, |
| 244 | +/// Lets us plug in some additional logic and data into a Postorder traversal. Or not. |
| 245 | +pub trait Customization<'tcx>: Copy { |
| 246 | + fn successors<'a>(_: &'a BasicBlockData<'tcx>, _: Self) -> Successors<'a>; |
256 | 247 | } |
257 | 248 |
|
258 | | -impl<'a, 'tcx> MonoReachablePostorder<'a, 'tcx> { |
259 | | - fn new( |
260 | | - body: &'a Body<'tcx>, |
261 | | - tcx: TyCtxt<'tcx>, |
262 | | - instance: Instance<'tcx>, |
263 | | - ) -> MonoReachablePostorder<'a, 'tcx> { |
264 | | - let basic_blocks = &body.basic_blocks; |
265 | | - let mut po = MonoReachablePostorder { |
266 | | - basic_blocks, |
267 | | - visited: BitSet::new_empty(basic_blocks.len()), |
268 | | - visit_stack: Vec::new(), |
269 | | - locals: UsedLocals(BitSet::new_empty(body.local_decls.len())), |
270 | | - tcx, |
271 | | - instance, |
272 | | - }; |
273 | | - |
274 | | - po.visit(START_BLOCK); |
275 | | - po.traverse_successor(); |
276 | | - po |
277 | | - } |
278 | | - |
279 | | - fn visit(&mut self, bb: BasicBlock) { |
280 | | - if !self.visited.insert(bb) { |
281 | | - return; |
282 | | - } |
283 | | - let data = &self.basic_blocks[bb]; |
284 | | - self.locals.visit_basic_block_data(bb, data); |
285 | | - let successors = data.mono_successors(self.tcx, self.instance); |
286 | | - self.visit_stack.push((bb, successors)); |
287 | | - } |
288 | | - |
289 | | - fn traverse_successor(&mut self) { |
290 | | - while let Some(bb) = self.visit_stack.last_mut().and_then(|(_, iter)| iter.next_back()) { |
291 | | - self.visit(bb); |
292 | | - } |
| 249 | +impl<'tcx> Customization<'tcx> for () { |
| 250 | + fn successors<'a>(data: &'a BasicBlockData<'tcx>, _: ()) -> Successors<'a> { |
| 251 | + data.terminator().successors() |
293 | 252 | } |
294 | 253 | } |
295 | 254 |
|
296 | | -impl<'tcx> Iterator for MonoReachablePostorder<'_, 'tcx> { |
297 | | - type Item = BasicBlock; |
298 | | - |
299 | | - fn next(&mut self) -> Option<BasicBlock> { |
300 | | - let (bb, _) = self.visit_stack.pop()?; |
301 | | - self.traverse_successor(); |
302 | | - Some(bb) |
| 255 | +impl<'tcx> Customization<'tcx> for (TyCtxt<'tcx>, Instance<'tcx>) { |
| 256 | + fn successors<'a>( |
| 257 | + data: &'a BasicBlockData<'tcx>, |
| 258 | + (tcx, instance): (TyCtxt<'tcx>, Instance<'tcx>), |
| 259 | + ) -> Successors<'a> { |
| 260 | + data.mono_successors(tcx, instance) |
303 | 261 | } |
304 | 262 | } |
305 | 263 |
|
306 | 264 | pub fn mono_reachable_reverse_postorder<'a, 'tcx>( |
307 | 265 | body: &'a Body<'tcx>, |
308 | 266 | tcx: TyCtxt<'tcx>, |
309 | 267 | instance: Instance<'tcx>, |
310 | | -) -> (Vec<BasicBlock>, BitSet<Local>) { |
311 | | - let mut iter = MonoReachablePostorder::new(body, tcx, instance); |
| 268 | +) -> Vec<BasicBlock> { |
| 269 | + let mut iter = Postorder::new(&body.basic_blocks, START_BLOCK, (tcx, instance)); |
312 | 270 | let mut items = Vec::with_capacity(body.basic_blocks.len()); |
313 | 271 | while let Some(block) = iter.next() { |
314 | 272 | items.push(block); |
315 | 273 | } |
316 | 274 | items.reverse(); |
317 | | - (items, iter.locals.0) |
| 275 | + items |
318 | 276 | } |
319 | 277 |
|
320 | 278 | /// Returns an iterator over all basic blocks reachable from the `START_BLOCK` in no particular |
|
0 commit comments