Skip to content

Commit

Permalink
coverage: Replace ShortCircuitPreorder with a single function
Browse files Browse the repository at this point in the history
Instead of defining a named struct, we can use `std::iter::from_fn` and store
intermediate state in a closure.
  • Loading branch information
Zalathar committed Oct 10, 2023
1 parent fa6d1e7 commit f214497
Showing 1 changed file with 17 additions and 55 deletions.
72 changes: 17 additions & 55 deletions compiler/rustc_mir_transform/src/coverage/graph.rs
Original file line number Diff line number Diff line change
@@ -1,8 +1,9 @@
use rustc_data_structures::captures::Captures;
use rustc_data_structures::graph::dominators::{self, Dominators};
use rustc_data_structures::graph::{self, GraphSuccessors, WithNumNodes, WithStartNode};
use rustc_index::bit_set::BitSet;
use rustc_index::{IndexSlice, IndexVec};
use rustc_middle::mir::{self, BasicBlock, BasicBlockData, Terminator, TerminatorKind};
use rustc_middle::mir::{self, BasicBlock, Terminator, TerminatorKind};

use std::cmp::Ordering;
use std::ops::{Index, IndexMut};
Expand Down Expand Up @@ -80,10 +81,9 @@ impl CoverageGraph {
// intentionally omits unwind paths.
// FIXME(#78544): MIR InstrumentCoverage: Improve coverage of `#[should_panic]` tests and
// `catch_unwind()` handlers.
let mir_cfg_without_unwind = ShortCircuitPreorder::new(&mir_body, bcb_filtered_successors);

let mut basic_blocks = Vec::new();
for (bb, data) in mir_cfg_without_unwind {
for bb in short_circuit_preorder(mir_body, bcb_filtered_successors) {
if let Some(last) = basic_blocks.last() {
let predecessors = &mir_body.basic_blocks.predecessors()[bb];
if predecessors.len() > 1 || !predecessors.contains(last) {
Expand All @@ -109,7 +109,7 @@ impl CoverageGraph {
}
basic_blocks.push(bb);

let term = data.terminator();
let term = mir_body[bb].terminator();

match term.kind {
TerminatorKind::Return { .. }
Expand Down Expand Up @@ -553,66 +553,28 @@ pub(super) fn find_loop_backedges(
backedges
}

pub struct ShortCircuitPreorder<
'a,
'tcx,
F: Fn(&'a mir::Body<'tcx>, &'a TerminatorKind<'tcx>) -> Box<dyn Iterator<Item = BasicBlock> + 'a>,
> {
fn short_circuit_preorder<'a, 'tcx, F, Iter>(
body: &'a mir::Body<'tcx>,
visited: BitSet<BasicBlock>,
worklist: Vec<BasicBlock>,
filtered_successors: F,
}

impl<
'a,
'tcx,
F: Fn(&'a mir::Body<'tcx>, &'a TerminatorKind<'tcx>) -> Box<dyn Iterator<Item = BasicBlock> + 'a>,
> ShortCircuitPreorder<'a, 'tcx, F>
{
pub fn new(
body: &'a mir::Body<'tcx>,
filtered_successors: F,
) -> ShortCircuitPreorder<'a, 'tcx, F> {
let worklist = vec![mir::START_BLOCK];

ShortCircuitPreorder {
body,
visited: BitSet::new_empty(body.basic_blocks.len()),
worklist,
filtered_successors,
}
}
}

impl<
'a,
'tcx,
F: Fn(&'a mir::Body<'tcx>, &'a TerminatorKind<'tcx>) -> Box<dyn Iterator<Item = BasicBlock> + 'a>,
> Iterator for ShortCircuitPreorder<'a, 'tcx, F>
) -> impl Iterator<Item = BasicBlock> + Captures<'a> + Captures<'tcx>
where
F: Fn(&'a mir::Body<'tcx>, &'a TerminatorKind<'tcx>) -> Iter,
Iter: Iterator<Item = BasicBlock>,
{
type Item = (BasicBlock, &'a BasicBlockData<'tcx>);
let mut visited = BitSet::new_empty(body.basic_blocks.len());
let mut worklist = vec![mir::START_BLOCK];

fn next(&mut self) -> Option<(BasicBlock, &'a BasicBlockData<'tcx>)> {
while let Some(idx) = self.worklist.pop() {
if !self.visited.insert(idx) {
std::iter::from_fn(move || {
while let Some(bb) = worklist.pop() {
if !visited.insert(bb) {
continue;
}

let data = &self.body[idx];

if let Some(ref term) = data.terminator {
self.worklist.extend((self.filtered_successors)(&self.body, &term.kind));
}
worklist.extend(filtered_successors(body, &body[bb].terminator().kind));

return Some((idx, data));
return Some(bb);
}

None
}

fn size_hint(&self) -> (usize, Option<usize>) {
let size = self.body.basic_blocks.len() - self.visited.count();
(size, Some(size))
}
})
}

0 comments on commit f214497

Please sign in to comment.