Add some standard traversal iterators for MIR

Adds Preorder, Postorder and Reverse Postorder traversal iterators.

Also makes trans/mir use Reverse Postorder traversal for blocks.

4 files changed, 284 insertions, 2 deletions

diff --git a/src/librustc_data_structures/bitvec.rs b/src/librustc_data_structures/bitvec.rs
index 3677c8c5e59..e45c6cfc6dc 100644
--- a/src/librustc_data_structures/bitvec.rs
+++ b/src/librustc_data_structures/bitvec.rs
@@ -9,6 +9,7 @@
 // except according to those terms.
 
 /// A very simple BitVector type.
+#[derive(Clone)]
 pub struct BitVector {
     data: Vec<u64>,
 }
diff --git a/src/librustc_mir/lib.rs b/src/librustc_mir/lib.rs
index 67bac196f48..cbbad057cf9 100644
--- a/src/librustc_mir/lib.rs
+++ b/src/librustc_mir/lib.rs
@@ -40,3 +40,4 @@ mod hair;
 pub mod mir_map;
 pub mod pretty;
 pub mod transform;
+pub mod traversal;
diff --git a/src/librustc_mir/traversal.rs b/src/librustc_mir/traversal.rs
new file mode 100644
index 00000000000..fc53881bd81
--- /dev/null
+++ b/src/librustc_mir/traversal.rs
@@ -0,0 +1,276 @@
+// Copyright 2016 The Rust Project Developers. See the COPYRIGHT
+// file at the top-level directory of this distribution and at
+// http://rust-lang.org/COPYRIGHT.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+use std::vec;
+
+use rustc_data_structures::bitvec::BitVector;
+
+use rustc::mir::repr::*;
+
+/// Preorder traversal of a graph.
+///
+/// Preorder traversal is when each node is visited before an of it's
+/// successors
+///
+///         A
+///        / \
+///       /   \
+///      B     C
+///       \   /
+///        \ /
+///         D
+///
+/// A preorder traversal of this graph is either `A B D C` or `A C D B`
+#[derive(Clone)]
+pub struct Preorder<'a, 'tcx: 'a> {
+    mir: &'a Mir<'tcx>,
+    visited: BitVector,
+    worklist: Vec<BasicBlock>,
+}
+
+impl<'a, 'tcx> Preorder<'a, 'tcx> {
+    pub fn new(mir: &'a Mir<'tcx>, root: BasicBlock) -> Preorder<'a, 'tcx> {
+        let worklist = vec![root];
+
+        Preorder {
+            mir: mir,
+            visited: BitVector::new(mir.basic_blocks.len()),
+            worklist: worklist
+        }
+    }
+}
+
+pub fn preorder<'a, 'tcx>(mir: &'a Mir<'tcx>) -> Preorder<'a, 'tcx> {
+    Preorder::new(mir, START_BLOCK)
+}
+
+impl<'a, 'tcx> Iterator for Preorder<'a, 'tcx> {
+    type Item = (BasicBlock, &'a BasicBlockData<'tcx>);
+
+    fn next(&mut self) -> Option<(BasicBlock, &'a BasicBlockData<'tcx>)> {
+        while let Some(idx) = self.worklist.pop() {
+            if !self.visited.insert(idx.index()) {
+                continue;
+            }
+
+            let data = self.mir.basic_block_data(idx);
+
+            if let Some(ref term) = data.terminator {
+                for &succ in term.successors().iter() {
+                    self.worklist.push(succ);
+                }
+            }
+
+            return Some((idx, data));
+        }
+
+        None
+    }
+}
+
+/// Postorder traversal of a graph.
+///
+/// Postorder traversal is when each node is visited after all of it's
+/// successors, except when the successor is only reachable by a back-edge
+///
+///         A
+///        / \
+///       /   \
+///      B     C
+///       \   /
+///        \ /
+///         D
+///
+/// A Postorder traversal of this graph is `D B C A` or `D C B A`
+pub struct Postorder<'a, 'tcx: 'a> {
+    mir: &'a Mir<'tcx>,
+    visited: BitVector,
+    visit_stack: Vec<(BasicBlock, vec::IntoIter<BasicBlock>)>
+}
+
+impl<'a, 'tcx> Postorder<'a, 'tcx> {
+    pub fn new(mir: &'a Mir<'tcx>, root: BasicBlock) -> Postorder<'a, 'tcx> {
+        let mut po = Postorder {
+            mir: mir,
+            visited: BitVector::new(mir.basic_blocks.len()),
+            visit_stack: Vec::new()
+        };
+
+
+        let data = po.mir.basic_block_data(root);
+
+        if let Some(ref term) = data.terminator {
+            po.visited.insert(root.index());
+
+            let succs = term.successors().into_owned().into_iter();
+
+            po.visit_stack.push((root, succs));
+            po.traverse_successor();
+        }
+
+        po
+    }
+
+    fn traverse_successor(&mut self) {
+        // This is quite a complex loop due to 1. the borrow checker not liking it much
+        // and 2. what exactly is going on is not clear
+        //
+        // It does the actual traversal of the graph, while the `next` method on the iterator
+        // just pops off of the stack. `visit_stack` is a stack containing pairs of nodes and
+        // iterators over the sucessors of those nodes. Each iteration attempts to get the next
+        // node from the top of the stack, then pushes that node and an iterator over the
+        // successors to the top of the stack. This loop only grows `visit_stack`, stopping when
+        // we reach a child that has no children that we haven't already visited.
+        //
+        // For a graph that looks like this:
+        //
+        //         A
+        //        / \
+        //       /   \
+        //      B     C
+        //      |     |
+        //      |     |
+        //      D     |
+        //       \   /
+        //        \ /
+        //         E
+        //
+        // The state of the stack starts out with just the root node (`A` in this case);
+        //     [(A, [B, C])]
+        //
+        // When the first call to `traverse_sucessor` happens, the following happens:
+        //
+        //     [(B, [D]),  // `B` taken from the successors of `A`, pushed to the
+        //                 // top of the stack along with the successors of `B`
+        //      (A, [C])]
+        //
+        //     [(D, [E]),  // `D` taken from successors of `B`, pushed to stack
+        //      (B, []),
+        //      (A, [C])]
+        //
+        //     [(E, []),   // `E` taken from successors of `D`, pushed to stack
+        //      (D, []),
+        //      (B, []),
+        //      (A, [C])]
+        //
+        // Now that the top of the stack has no successors we can traverse, each item will
+        // be popped off during iteration until we get back to `A`. This yeilds [E, D, B].
+        //
+        // When we yeild `B` and call `traverse_successor`, We push `C` to the stack, but
+        // since we've already visited `E`, that child isn't added to the stack. The last
+        // two iterations yield `C` and finally `A` for a final traversal of [E, D, B, C, A]
+        loop {
+            let bb = if let Some(&mut (_, ref mut iter)) = self.visit_stack.last_mut() {
+                if let Some(bb) = iter.next() {
+                    bb
+                } else {
+                    break;
+                }
+            } else {
+                break;
+            };
+
+            if self.visited.insert(bb.index()) {
+                let data = self.mir.basic_block_data(bb);
+
+                if let Some(ref term) = data.terminator {
+                    let succs = term.successors().into_owned().into_iter();
+                    self.visit_stack.push((bb, succs));
+                }
+            }
+        }
+    }
+}
+
+pub fn postorder<'a, 'tcx>(mir: &'a Mir<'tcx>) -> Postorder<'a, 'tcx> {
+    Postorder::new(mir, START_BLOCK)
+}
+
+impl<'a, 'tcx> Iterator for Postorder<'a, 'tcx> {
+    type Item = (BasicBlock, &'a BasicBlockData<'tcx>);
+
+    fn next(&mut self) -> Option<(BasicBlock, &'a BasicBlockData<'tcx>)> {
+        let next = self.visit_stack.pop();
+        if next.is_some() {
+            self.traverse_successor();
+        }
+
+        next.map(|(bb, _)| {
+            let data = self.mir.basic_block_data(bb);
+            (bb, data)
+        })
+    }
+}
+
+/// Reverse postorder traversal of a graph
+///
+/// Reverse postorder is the reverse order of a postorder traversal.
+/// This is different to a preorder traversal and represents a natural
+/// linearisation of control-flow.
+///
+///         A
+///        / \
+///       /   \
+///      B     C
+///       \   /
+///        \ /
+///         D
+///
+/// A reverse postorder traversal of this graph is either `A B C D` or `A C B D`
+/// Note that for a graph containing no loops (i.e. A DAG), this is equivalent to
+/// a topological sort.
+///
+/// Construction of a `ReversePostorder` traversal requires doing a full
+/// postorder traversal of the graph, therefore this traversal should be
+/// constructed as few times as possible. Use the `reset` method to be able
+/// to re-use the traversal
+#[derive(Clone)]
+pub struct ReversePostorder<'a, 'tcx: 'a> {
+    mir: &'a Mir<'tcx>,
+    blocks: Vec<BasicBlock>,
+    idx: usize
+}
+
+impl<'a, 'tcx> ReversePostorder<'a, 'tcx> {
+    pub fn new(mir: &'a Mir<'tcx>, root: BasicBlock) -> ReversePostorder<'a, 'tcx> {
+        let blocks : Vec<_> = Postorder::new(mir, root).map(|(bb, _)| bb).collect();
+
+        let len = blocks.len();
+
+        ReversePostorder {
+            mir: mir,
+            blocks: blocks,
+            idx: len
+        }
+    }
+
+    pub fn reset(&mut self) {
+        self.idx = self.blocks.len();
+    }
+}
+
+
+pub fn reverse_postorder<'a, 'tcx>(mir: &'a Mir<'tcx>) -> ReversePostorder<'a, 'tcx> {
+    ReversePostorder::new(mir, START_BLOCK)
+}
+
+impl<'a, 'tcx> Iterator for ReversePostorder<'a, 'tcx> {
+    type Item = (BasicBlock, &'a BasicBlockData<'tcx>);
+
+    fn next(&mut self) -> Option<(BasicBlock, &'a BasicBlockData<'tcx>)> {
+        if self.idx == 0 { return None; }
+        self.idx -= 1;
+
+        self.blocks.get(self.idx).map(|&bb| {
+            let data = self.mir.basic_block_data(bb);
+            (bb, data)
+        })
+    }
+}
diff --git a/src/librustc_trans/mir/mod.rs b/src/librustc_trans/mir/mod.rs
index 61536f3ecae..152834a07ba 100644
--- a/src/librustc_trans/mir/mod.rs
+++ b/src/librustc_trans/mir/mod.rs
@@ -20,6 +20,9 @@ use std::ops::Deref;
 use std::rc::Rc;
 
 use self::lvalue::{LvalueRef, get_dataptr, get_meta};
+use rustc_mir::traversal;
+
+use self::lvalue::LvalueRef;
 use self::operand::OperandRef;
 
 #[derive(Clone)]
@@ -152,8 +155,9 @@ pub fn trans_mir<'blk, 'tcx>(fcx: &'blk FunctionContext<'blk, 'tcx>) {
         args: args,
     };
 
-    // Translate the body of each block
-    for &bb in &mir_blocks {
+    let rpo = traversal::reverse_postorder(mir);
+    // Translate the body of each block using reverse postorder
+    for (bb, _) in rpo {
         mircx.trans_block(bb);
     }