src/librustc/mir/traversal.rs

   1 // Copyright 2016 The Rust Project Developers. See the COPYRIGHT
   2 // file at the top-level directory of this distribution and at
   3 // http://rust-lang.org/COPYRIGHT.
   4 //
   5 // Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
   6 // http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
   7 // <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
   8 // option. This file may not be copied, modified, or distributed
   9 // except according to those terms.
  10
  11 use rustc_data_structures::bit_set::BitSet;
  12
  13 use super::*;
  14
  15 /// Preorder traversal of a graph.
  16 ///
  17 /// Preorder traversal is when each node is visited before an of it's
  18 /// successors
  19 ///
  20 /// ```text
  21 ///
  22 ///         A
  23 ///        / \
  24 ///       /   \
  25 ///      B     C
  26 ///       \   /
  27 ///        \ /
  28 ///         D
  29 /// ```
  30 ///
  31 /// A preorder traversal of this graph is either `A B D C` or `A C D B`
  32 #[derive(Clone)]
  33 pub struct Preorder<'a, 'tcx: 'a> {
  34     mir: &'a Mir<'tcx>,
  35     visited: BitSet<BasicBlock>,
  36     worklist: Vec<BasicBlock>,
  37     root_is_start_block: bool,
  38 }
  39
  40 impl<'a, 'tcx> Preorder<'a, 'tcx> {
  41     pub fn new(mir: &'a Mir<'tcx>, root: BasicBlock) -> Preorder<'a, 'tcx> {
  42         let worklist = vec![root];
  43
  44         Preorder {
  45             mir,
  46             visited: BitSet::new_empty(mir.basic_blocks().len()),
  47             worklist,
  48             root_is_start_block: root == START_BLOCK,
  49         }
  50     }
  51 }
  52
  53 pub fn preorder<'a, 'tcx>(mir: &'a Mir<'tcx>) -> Preorder<'a, 'tcx> {
  54     Preorder::new(mir, START_BLOCK)
  55 }
  56
  57 impl<'a, 'tcx> Iterator for Preorder<'a, 'tcx> {
  58     type Item = (BasicBlock, &'a BasicBlockData<'tcx>);
  59
  60     fn next(&mut self) -> Option<(BasicBlock, &'a BasicBlockData<'tcx>)> {
  61         while let Some(idx) = self.worklist.pop() {
  62             if !self.visited.insert(idx) {
  63                 continue;
  64             }
  65
  66             let data = &self.mir[idx];
  67
  68             if let Some(ref term) = data.terminator {
  69                 self.worklist.extend(term.successors());
  70             }
  71
  72             return Some((idx, data));
  73         }
  74
  75         None
  76     }
  77
  78     fn size_hint(&self) -> (usize, Option<usize>) {
  79         // All the blocks, minus the number of blocks we've visited.
  80         let upper = self.mir.basic_blocks().len() - self.visited.count();
  81
  82         let lower = if self.root_is_start_block {
  83             // We will visit all remaining blocks exactly once.
  84             upper
  85         } else {
  86             self.worklist.len()
  87         };
  88
  89         (lower, Some(upper))
  90     }
  91 }
  92
  93 /// Postorder traversal of a graph.
  94 ///
  95 /// Postorder traversal is when each node is visited after all of it's
  96 /// successors, except when the successor is only reachable by a back-edge
  97 ///
  98 ///
  99 /// ```text
 100 ///
 101 ///         A
 102 ///        / \
 103 ///       /   \
 104 ///      B     C
 105 ///       \   /
 106 ///        \ /
 107 ///         D
 108 /// ```
 109 ///
 110 /// A Postorder traversal of this graph is `D B C A` or `D C B A`
 111 pub struct Postorder<'a, 'tcx: 'a> {
 112     mir: &'a Mir<'tcx>,
 113     visited: BitSet<BasicBlock>,
 114     visit_stack: Vec<(BasicBlock, Successors<'a>)>,
 115     root_is_start_block: bool,
 116 }
 117
 118 impl<'a, 'tcx> Postorder<'a, 'tcx> {
 119     pub fn new(mir: &'a Mir<'tcx>, root: BasicBlock) -> Postorder<'a, 'tcx> {
 120         let mut po = Postorder {
 121             mir,
 122             visited: BitSet::new_empty(mir.basic_blocks().len()),
 123             visit_stack: Vec::new(),
 124             root_is_start_block: root == START_BLOCK,
 125         };
 126
 127
 128         let data = &po.mir[root];
 129
 130         if let Some(ref term) = data.terminator {
 131             po.visited.insert(root);
 132             po.visit_stack.push((root, term.successors()));
 133             po.traverse_successor();
 134         }
 135
 136         po
 137     }
 138
 139     fn traverse_successor(&mut self) {
 140         // This is quite a complex loop due to 1. the borrow checker not liking it much
 141         // and 2. what exactly is going on is not clear
 142         //
 143         // It does the actual traversal of the graph, while the `next` method on the iterator
 144         // just pops off of the stack. `visit_stack` is a stack containing pairs of nodes and
 145         // iterators over the successors of those nodes. Each iteration attempts to get the next
 146         // node from the top of the stack, then pushes that node and an iterator over the
 147         // successors to the top of the stack. This loop only grows `visit_stack`, stopping when
 148         // we reach a child that has no children that we haven't already visited.
 149         //
 150         // For a graph that looks like this:
 151         //
 152         //         A
 153         //        / \
 154         //       /   \
 155         //      B     C
 156         //      |     |
 157         //      |     |
 158         //      D     |
 159         //       \   /
 160         //        \ /
 161         //         E
 162         //
 163         // The state of the stack starts out with just the root node (`A` in this case);
 164         //     [(A, [B, C])]
 165         //
 166         // When the first call to `traverse_successor` happens, the following happens:
 167         //
 168         //     [(B, [D]),  // `B` taken from the successors of `A`, pushed to the
 169         //                 // top of the stack along with the successors of `B`
 170         //      (A, [C])]
 171         //
 172         //     [(D, [E]),  // `D` taken from successors of `B`, pushed to stack
 173         //      (B, []),
 174         //      (A, [C])]
 175         //
 176         //     [(E, []),   // `E` taken from successors of `D`, pushed to stack
 177         //      (D, []),
 178         //      (B, []),
 179         //      (A, [C])]
 180         //
 181         // Now that the top of the stack has no successors we can traverse, each item will
 182         // be popped off during iteration until we get back to `A`. This yields [E, D, B].
 183         //
 184         // When we yield `B` and call `traverse_successor`, we push `C` to the stack, but
 185         // since we've already visited `E`, that child isn't added to the stack. The last
 186         // two iterations yield `C` and finally `A` for a final traversal of [E, D, B, C, A]
 187         loop {
 188             let bb = if let Some(&mut (_, ref mut iter)) = self.visit_stack.last_mut() {
 189                 if let Some(&bb) = iter.next() {
 190                     bb
 191                 } else {
 192                     break;
 193                 }
 194             } else {
 195                 break;
 196             };
 197
 198             if self.visited.insert(bb) {
 199                 if let Some(term) = &self.mir[bb].terminator {
 200                     self.visit_stack.push((bb, term.successors()));
 201                 }
 202             }
 203         }
 204     }
 205 }
 206
 207 pub fn postorder<'a, 'tcx>(mir: &'a Mir<'tcx>) -> Postorder<'a, 'tcx> {
 208     Postorder::new(mir, START_BLOCK)
 209 }
 210
 211 impl<'a, 'tcx> Iterator for Postorder<'a, 'tcx> {
 212     type Item = (BasicBlock, &'a BasicBlockData<'tcx>);
 213
 214     fn next(&mut self) -> Option<(BasicBlock, &'a BasicBlockData<'tcx>)> {
 215         let next = self.visit_stack.pop();
 216         if next.is_some() {
 217             self.traverse_successor();
 218         }
 219
 220         next.map(|(bb, _)| (bb, &self.mir[bb]))
 221     }
 222
 223     fn size_hint(&self) -> (usize, Option<usize>) {
 224         // All the blocks, minus the number of blocks we've visited.
 225         let upper = self.mir.basic_blocks().len() - self.visited.count();
 226
 227         let lower = if self.root_is_start_block {
 228             // We will visit all remaining blocks exactly once.
 229             upper
 230         } else {
 231             self.visit_stack.len()
 232         };
 233
 234         (lower, Some(upper))
 235     }
 236 }
 237
 238 /// Reverse postorder traversal of a graph
 239 ///
 240 /// Reverse postorder is the reverse order of a postorder traversal.
 241 /// This is different to a preorder traversal and represents a natural
 242 /// linearization of control-flow.
 243 ///
 244 /// ```text
 245 ///
 246 ///         A
 247 ///        / \
 248 ///       /   \
 249 ///      B     C
 250 ///       \   /
 251 ///        \ /
 252 ///         D
 253 /// ```
 254 ///
 255 /// A reverse postorder traversal of this graph is either `A B C D` or `A C B D`
 256 /// Note that for a graph containing no loops (i.e., A DAG), this is equivalent to
 257 /// a topological sort.
 258 ///
 259 /// Construction of a `ReversePostorder` traversal requires doing a full
 260 /// postorder traversal of the graph, therefore this traversal should be
 261 /// constructed as few times as possible. Use the `reset` method to be able
 262 /// to re-use the traversal
 263 #[derive(Clone)]
 264 pub struct ReversePostorder<'a, 'tcx: 'a> {
 265     mir: &'a Mir<'tcx>,
 266     blocks: Vec<BasicBlock>,
 267     idx: usize
 268 }
 269
 270 impl<'a, 'tcx> ReversePostorder<'a, 'tcx> {
 271     pub fn new(mir: &'a Mir<'tcx>, root: BasicBlock) -> ReversePostorder<'a, 'tcx> {
 272         let blocks : Vec<_> = Postorder::new(mir, root).map(|(bb, _)| bb).collect();
 273
 274         let len = blocks.len();
 275
 276         ReversePostorder {
 277             mir,
 278             blocks,
 279             idx: len
 280         }
 281     }
 282
 283     pub fn reset(&mut self) {
 284         self.idx = self.blocks.len();
 285     }
 286 }
 287
 288
 289 pub fn reverse_postorder<'a, 'tcx>(mir: &'a Mir<'tcx>) -> ReversePostorder<'a, 'tcx> {
 290     ReversePostorder::new(mir, START_BLOCK)
 291 }
 292
 293 impl<'a, 'tcx> Iterator for ReversePostorder<'a, 'tcx> {
 294     type Item = (BasicBlock, &'a BasicBlockData<'tcx>);
 295
 296     fn next(&mut self) -> Option<(BasicBlock, &'a BasicBlockData<'tcx>)> {
 297         if self.idx == 0 { return None; }
 298         self.idx -= 1;
 299
 300         self.blocks.get(self.idx).map(|&bb| (bb, &self.mir[bb]))
 301     }
 302
 303     fn size_hint(&self) -> (usize, Option<usize>) {
 304         (self.idx, Some(self.idx))
 305     }
 306 }
 307
 308 impl<'a, 'tcx> ExactSizeIterator for ReversePostorder<'a, 'tcx> {}