compiler/rustc_middle/src/mir/traversal.rs

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