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