diff --git a/src/librustc/mir/traversal.rs b/src/librustc/mir/traversal.rs index db1bc3e7519c2..a1e2b7a06468d 100644 --- a/src/librustc/mir/traversal.rs +++ b/src/librustc/mir/traversal.rs @@ -34,6 +34,7 @@ pub struct Preorder<'a, 'tcx: 'a> { mir: &'a Mir<'tcx>, visited: BitSet, worklist: Vec, + root_is_start_block: bool, } impl<'a, 'tcx> Preorder<'a, 'tcx> { @@ -44,6 +45,7 @@ impl<'a, 'tcx> Preorder<'a, 'tcx> { mir, visited: BitSet::new_empty(mir.basic_blocks().len()), worklist, + root_is_start_block: root == START_BLOCK, } } } @@ -75,15 +77,19 @@ impl<'a, 'tcx> Iterator for Preorder<'a, 'tcx> { fn size_hint(&self) -> (usize, Option) { // All the blocks, minus the number of blocks we've visited. - let remaining = self.mir.basic_blocks().len() - self.visited.count(); + let upper = self.mir.basic_blocks().len() - self.visited.count(); - // We will visit all remaining blocks exactly once. - (remaining, Some(remaining)) + let lower = if self.root_is_start_block { + // We will visit all remaining blocks exactly once. + upper + } else { + self.worklist.len() + }; + + (lower, Some(upper)) } } -impl<'a, 'tcx> ExactSizeIterator for Preorder<'a, 'tcx> {} - /// Postorder traversal of a graph. /// /// Postorder traversal is when each node is visited after all of it's @@ -105,7 +111,8 @@ impl<'a, 'tcx> ExactSizeIterator for Preorder<'a, 'tcx> {} pub struct Postorder<'a, 'tcx: 'a> { mir: &'a Mir<'tcx>, visited: BitSet, - visit_stack: Vec<(BasicBlock, Successors<'a>)> + visit_stack: Vec<(BasicBlock, Successors<'a>)>, + root_is_start_block: bool, } impl<'a, 'tcx> Postorder<'a, 'tcx> { @@ -113,7 +120,8 @@ impl<'a, 'tcx> Postorder<'a, 'tcx> { let mut po = Postorder { mir, visited: BitSet::new_empty(mir.basic_blocks().len()), - visit_stack: Vec::new() + visit_stack: Vec::new(), + root_is_start_block: root == START_BLOCK, }; @@ -214,15 +222,19 @@ impl<'a, 'tcx> Iterator for Postorder<'a, 'tcx> { fn size_hint(&self) -> (usize, Option) { // All the blocks, minus the number of blocks we've visited. - let remaining = self.mir.basic_blocks().len() - self.visited.count(); + let upper = self.mir.basic_blocks().len() - self.visited.count(); - // We will visit all remaining blocks exactly once. - (remaining, Some(remaining)) + let lower = if self.root_is_start_block { + // We will visit all remaining blocks exactly once. + upper + } else { + self.visit_stack.len() + }; + + (lower, Some(upper)) } } -impl<'a, 'tcx> ExactSizeIterator for Postorder<'a, 'tcx> {} - /// Reverse postorder traversal of a graph /// /// Reverse postorder is the reverse order of a postorder traversal.