compiler/rustc_middle/src/mir/traversal.rs

   1 use rustc_data_structures::stable_hasher::{HashStable, StableHasher};
   2 use rustc_data_structures::sync::OnceCell;
   3 use rustc_index::bit_set::BitSet;
   4 use rustc_serialize::{Decodable, Decoder, Encodable, Encoder};
   5
   6 use super::*;
   7
   8 /// Preorder traversal of a graph.
   9 ///
  10 /// Preorder traversal is when each node is visited after at least one of its predecessors. If you
  11 /// are familiar with some basic graph theory, then this performs a depth first search and returns
  12 /// nodes in order of discovery time.
  13 ///
  14 /// ```text
  15 ///
  16 ///         A
  17 ///        / \
  18 ///       /   \
  19 ///      B     C
  20 ///       \   /
  21 ///        \ /
  22 ///         D
  23 /// ```
  24 ///
  25 /// A preorder traversal of this graph is either `A B D C` or `A C D B`
  26 #[derive(Clone)]
  27 pub struct Preorder<'a, 'tcx> {
  28     body: &'a Body<'tcx>,
  29     visited: BitSet<BasicBlock>,
  30     worklist: Vec<BasicBlock>,
  31     root_is_start_block: bool,
  32 }
  33
  34 impl<'a, 'tcx> Preorder<'a, 'tcx> {
  35     pub fn new(body: &'a Body<'tcx>, root: BasicBlock) -> Preorder<'a, 'tcx> {
  36         let worklist = vec![root];
  37
  38         Preorder {
  39             body,
  40             visited: BitSet::new_empty(body.basic_blocks.len()),
  41             worklist,
  42             root_is_start_block: root == START_BLOCK,
  43         }
  44     }
  45 }
  46
  47 pub fn preorder<'a, 'tcx>(body: &'a Body<'tcx>) -> Preorder<'a, 'tcx> {
  48     Preorder::new(body, START_BLOCK)
  49 }
  50
  51 impl<'a, 'tcx> Iterator for Preorder<'a, 'tcx> {
  52     type Item = (BasicBlock, &'a BasicBlockData<'tcx>);
  53
  54     fn next(&mut self) -> Option<(BasicBlock, &'a BasicBlockData<'tcx>)> {
  55         while let Some(idx) = self.worklist.pop() {
  56             if !self.visited.insert(idx) {
  57                 continue;
  58             }
  59
  60             let data = &self.body[idx];
  61
  62             if let Some(ref term) = data.terminator {
  63                 self.worklist.extend(term.successors());
  64             }
  65
  66             return Some((idx, data));
  67         }
  68
  69         None
  70     }
  71
  72     fn size_hint(&self) -> (usize, Option<usize>) {
  73         // All the blocks, minus the number of blocks we've visited.
  74         let upper = self.body.basic_blocks.len() - self.visited.count();
  75
  76         let lower = if self.root_is_start_block {
  77             // We will visit all remaining blocks exactly once.
  78             upper
  79         } else {
  80             self.worklist.len()
  81         };
  82
  83         (lower, Some(upper))
  84     }
  85 }
  86
  87 /// Postorder traversal of a graph.
  88 ///
  89 /// Postorder traversal is when each node is visited after all of its successors, except when the
  90 /// successor is only reachable by a back-edge. If you are familiar with some basic graph theory,
  91 /// then this performs a depth first search and returns nodes in order of completion time.
  92 ///
  93 ///
  94 /// ```text
  95 ///
  96 ///         A
  97 ///        / \
  98 ///       /   \
  99 ///      B     C
 100 ///       \   /
 101 ///        \ /
 102 ///         D
 103 /// ```
 104 ///
 105 /// A Postorder traversal of this graph is `D B C A` or `D C B A`
 106 pub struct Postorder<'a, 'tcx> {
 107     basic_blocks: &'a IndexVec<BasicBlock, BasicBlockData<'tcx>>,
 108     visited: BitSet<BasicBlock>,
 109     visit_stack: Vec<(BasicBlock, Successors<'a>)>,
 110     root_is_start_block: bool,
 111 }
 112
 113 impl<'a, 'tcx> Postorder<'a, 'tcx> {
 114     pub fn new(
 115         basic_blocks: &'a IndexVec<BasicBlock, BasicBlockData<'tcx>>,
 116         root: BasicBlock,
 117     ) -> Postorder<'a, 'tcx> {
 118         let mut po = Postorder {
 119             basic_blocks,
 120             visited: BitSet::new_empty(basic_blocks.len()),
 121             visit_stack: Vec::new(),
 122             root_is_start_block: root == START_BLOCK,
 123         };
 124
 125         let data = &po.basic_blocks[root];
 126
 127         if let Some(ref term) = data.terminator {
 128             po.visited.insert(root);
 129             po.visit_stack.push((root, term.successors()));
 130             po.traverse_successor();
 131         }
 132
 133         po
 134     }
 135
 136     fn traverse_successor(&mut self) {
 137         // This is quite a complex loop due to 1. the borrow checker not liking it much
 138         // and 2. what exactly is going on is not clear
 139         //
 140         // It does the actual traversal of the graph, while the `next` method on the iterator
 141         // just pops off of the stack. `visit_stack` is a stack containing pairs of nodes and
 142         // iterators over the successors of those nodes. Each iteration attempts to get the next
 143         // node from the top of the stack, then pushes that node and an iterator over the
 144         // successors to the top of the stack. This loop only grows `visit_stack`, stopping when
 145         // we reach a child that has no children that we haven't already visited.
 146         //
 147         // For a graph that looks like this:
 148         //
 149         //         A
 150         //        / \
 151         //       /   \
 152         //      B     C
 153         //      |     |
 154         //      |     |
 155         //      D     |
 156         //       \   /
 157         //        \ /
 158         //         E
 159         //
 160         // The state of the stack starts out with just the root node (`A` in this case);
 161         //     [(A, [B, C])]
 162         //
 163         // When the first call to `traverse_successor` happens, the following happens:
 164         //
 165         //     [(B, [D]),  // `B` taken from the successors of `A`, pushed to the
 166         //                 // top of the stack along with the successors of `B`
 167         //      (A, [C])]
 168         //
 169         //     [(D, [E]),  // `D` taken from successors of `B`, pushed to stack
 170         //      (B, []),
 171         //      (A, [C])]
 172         //
 173         //     [(E, []),   // `E` taken from successors of `D`, pushed to stack
 174         //      (D, []),
 175         //      (B, []),
 176         //      (A, [C])]
 177         //
 178         // Now that the top of the stack has no successors we can traverse, each item will
 179         // be popped off during iteration until we get back to `A`. This yields [E, D, B].
 180         //
 181         // When we yield `B` and call `traverse_successor`, we push `C` to the stack, but
 182         // since we've already visited `E`, that child isn't added to the stack. The last
 183         // two iterations yield `C` and finally `A` for a final traversal of [E, D, B, C, A]
 184         loop {
 185             let bb = if let Some(&mut (_, ref mut iter)) = self.visit_stack.last_mut() {
 186                 if let Some(bb) = iter.next() {
 187                     bb
 188                 } else {
 189                     break;
 190                 }
 191             } else {
 192                 break;
 193             };
 194
 195             if self.visited.insert(bb) {
 196                 if let Some(term) = &self.basic_blocks[bb].terminator {
 197                     self.visit_stack.push((bb, term.successors()));
 198                 }
 199             }
 200         }
 201     }
 202 }
 203
 204 pub fn postorder<'a, 'tcx>(body: &'a Body<'tcx>) -> Postorder<'a, 'tcx> {
 205     Postorder::new(&body.basic_blocks, START_BLOCK)
 206 }
 207
 208 impl<'a, 'tcx> Iterator for Postorder<'a, 'tcx> {
 209     type Item = (BasicBlock, &'a BasicBlockData<'tcx>);
 210
 211     fn next(&mut self) -> Option<(BasicBlock, &'a BasicBlockData<'tcx>)> {
 212         let next = self.visit_stack.pop();
 213         if next.is_some() {
 214             self.traverse_successor();
 215         }
 216
 217         next.map(|(bb, _)| (bb, &self.basic_blocks[bb]))
 218     }
 219
 220     fn size_hint(&self) -> (usize, Option<usize>) {
 221         // All the blocks, minus the number of blocks we've visited.
 222         let upper = self.basic_blocks.len() - self.visited.count();
 223
 224         let lower = if self.root_is_start_block {
 225             // We will visit all remaining blocks exactly once.
 226             upper
 227         } else {
 228             self.visit_stack.len()
 229         };
 230
 231         (lower, Some(upper))
 232     }
 233 }
 234
 235 /// Reverse postorder traversal of a graph
 236 ///
 237 /// Reverse postorder is the reverse order of a postorder traversal.
 238 /// This is different to a preorder traversal and represents a natural
 239 /// linearization of control-flow.
 240 ///
 241 /// ```text
 242 ///
 243 ///         A
 244 ///        / \
 245 ///       /   \
 246 ///      B     C
 247 ///       \   /
 248 ///        \ /
 249 ///         D
 250 /// ```
 251 ///
 252 /// A reverse postorder traversal of this graph is either `A B C D` or `A C B D`
 253 /// Note that for a graph containing no loops (i.e., A DAG), this is equivalent to
 254 /// a topological sort.
 255 ///
 256 /// Construction of a `ReversePostorder` traversal requires doing a full
 257 /// postorder traversal of the graph, therefore this traversal should be
 258 /// constructed as few times as possible. Use the `reset` method to be able
 259 /// to re-use the traversal
 260 #[derive(Clone)]
 261 pub struct ReversePostorder<'a, 'tcx> {
 262     body: &'a Body<'tcx>,
 263     blocks: Vec<BasicBlock>,
 264     idx: usize,
 265 }
 266
 267 impl<'a, 'tcx> ReversePostorder<'a, 'tcx> {
 268     pub fn new(body: &'a Body<'tcx>, root: BasicBlock) -> ReversePostorder<'a, 'tcx> {
 269         let blocks: Vec<_> = Postorder::new(&body.basic_blocks, root).map(|(bb, _)| bb).collect();
 270         let len = blocks.len();
 271         ReversePostorder { body, blocks, idx: len }
 272     }
 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 }
 310
 311 #[derive(Clone)]
 312 pub struct ReversePostorderIter<'a, 'tcx> {
 313     body: &'a Body<'tcx>,
 314     blocks: &'a [BasicBlock],
 315     idx: usize,
 316 }
 317
 318 impl<'a, 'tcx> Iterator for ReversePostorderIter<'a, 'tcx> {
 319     type Item = (BasicBlock, &'a BasicBlockData<'tcx>);
 320
 321     fn next(&mut self) -> Option<(BasicBlock, &'a BasicBlockData<'tcx>)> {
 322         if self.idx == 0 {
 323             return None;
 324         }
 325         self.idx -= 1;
 326
 327         self.blocks.get(self.idx).map(|&bb| (bb, &self.body[bb]))
 328     }
 329
 330     fn size_hint(&self) -> (usize, Option<usize>) {
 331         (self.idx, Some(self.idx))
 332     }
 333 }
 334
 335 impl<'a, 'tcx> ExactSizeIterator for ReversePostorderIter<'a, 'tcx> {}
 336
 337 pub fn reverse_postorder<'a, 'tcx>(body: &'a Body<'tcx>) -> ReversePostorderIter<'a, 'tcx> {
 338     let blocks = body.basic_blocks.postorder();
 339     let len = blocks.len();
 340     ReversePostorderIter { body, blocks, idx: len }
 341 }
 342
 343 #[derive(Clone, Debug)]
 344 pub(super) struct PostorderCache {
 345     cache: OnceCell<Vec<BasicBlock>>,
 346 }
 347
 348 impl PostorderCache {
 349     #[inline]
 350     pub(super) fn new() -> Self {
 351         PostorderCache { cache: OnceCell::new() }
 352     }
 353
 354     /// Invalidates the postorder cache.
 355     #[inline]
 356     pub(super) fn invalidate(&mut self) {
 357         self.cache = OnceCell::new();
 358     }
 359
 360     /// Returns the `&[BasicBlocks]` represents the postorder graph for this MIR.
 361     #[inline]
 362     pub(super) fn compute(&self, body: &IndexVec<BasicBlock, BasicBlockData<'_>>) -> &[BasicBlock] {
 363         self.cache.get_or_init(|| Postorder::new(body, START_BLOCK).map(|(bb, _)| bb).collect())
 364     }
 365 }
 366
 367 impl<S: Encoder> Encodable<S> for PostorderCache {
 368     #[inline]
 369     fn encode(&self, _s: &mut S) {}
 370 }
 371
 372 impl<D: Decoder> Decodable<D> for PostorderCache {
 373     #[inline]
 374     fn decode(_: &mut D) -> Self {
 375         Self::new()
 376     }
 377 }
 378
 379 impl<CTX> HashStable<CTX> for PostorderCache {
 380     #[inline]
 381     fn hash_stable(&self, _: &mut CTX, _: &mut StableHasher) {
 382         // do nothing
 383     }
 384 }
 385
 386 TrivialTypeTraversalAndLiftImpls! {
 387     PostorderCache,
 388 }