From f214497d22588c95b94f8ef89ca2db7d71424c36 Mon Sep 17 00:00:00 2001 From: Zalathar Date: Tue, 10 Oct 2023 15:44:33 +1100 Subject: [PATCH 1/2] coverage: Replace `ShortCircuitPreorder` with a single function Instead of defining a named struct, we can use `std::iter::from_fn` and store intermediate state in a closure. --- .../rustc_mir_transform/src/coverage/graph.rs | 72 +++++-------------- 1 file changed, 17 insertions(+), 55 deletions(-) diff --git a/compiler/rustc_mir_transform/src/coverage/graph.rs b/compiler/rustc_mir_transform/src/coverage/graph.rs index 812633348e3bd..c94e0baf8be81 100644 --- a/compiler/rustc_mir_transform/src/coverage/graph.rs +++ b/compiler/rustc_mir_transform/src/coverage/graph.rs @@ -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}; @@ -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) { @@ -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 { .. } @@ -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 + 'a>, -> { +fn short_circuit_preorder<'a, 'tcx, F, Iter>( body: &'a mir::Body<'tcx>, - visited: BitSet, - worklist: Vec, filtered_successors: F, -} - -impl< - 'a, - 'tcx, - F: Fn(&'a mir::Body<'tcx>, &'a TerminatorKind<'tcx>) -> Box + '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 + 'a>, -> Iterator for ShortCircuitPreorder<'a, 'tcx, F> +) -> impl Iterator + Captures<'a> + Captures<'tcx> +where + F: Fn(&'a mir::Body<'tcx>, &'a TerminatorKind<'tcx>) -> Iter, + Iter: Iterator, { - 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) { - let size = self.body.basic_blocks.len() - self.visited.count(); - (size, Some(size)) - } + }) } From 5d629457fd3bb580b87800f7c2858ddfe3f1c0ac Mon Sep 17 00:00:00 2001 From: Zalathar Date: Tue, 10 Oct 2023 15:09:04 +1100 Subject: [PATCH 2/2] coverage: Unbox and simplify `bcb_filtered_successors` This function already has access to the MIR body, so instead of taking a reference to a terminator, it's simpler and easier to pass in a basic block index. There is no need to box the returned iterator if we instead add appropriate lifetime captures, since `short_circuit_preorder` is now generic over the type of iterator it expects. We can also greatly simplify the function's implementation by observing that the only difference between its two cases is whether we take all of a BB's successors, or just the first one. --- .../rustc_mir_transform/src/coverage/graph.rs | 50 +++++++++---------- .../rustc_mir_transform/src/coverage/tests.rs | 2 +- 2 files changed, 24 insertions(+), 28 deletions(-) diff --git a/compiler/rustc_mir_transform/src/coverage/graph.rs b/compiler/rustc_mir_transform/src/coverage/graph.rs index c94e0baf8be81..39027d21f066a 100644 --- a/compiler/rustc_mir_transform/src/coverage/graph.rs +++ b/compiler/rustc_mir_transform/src/coverage/graph.rs @@ -3,7 +3,7 @@ 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, Terminator, TerminatorKind}; +use rustc_middle::mir::{self, BasicBlock, TerminatorKind}; use std::cmp::Ordering; use std::ops::{Index, IndexMut}; @@ -37,9 +37,8 @@ impl CoverageGraph { } let bcb_data = &bcbs[bcb]; let mut bcb_successors = Vec::new(); - for successor in - bcb_filtered_successors(&mir_body, &bcb_data.terminator(mir_body).kind) - .filter_map(|successor_bb| bb_to_bcb[successor_bb]) + for successor in bcb_filtered_successors(&mir_body, bcb_data.last_bb()) + .filter_map(|successor_bb| bb_to_bcb[successor_bb]) { if !seen[successor] { seen[successor] = true; @@ -316,11 +315,6 @@ impl BasicCoverageBlockData { pub fn last_bb(&self) -> BasicBlock { *self.basic_blocks.last().unwrap() } - - #[inline(always)] - pub fn terminator<'a, 'tcx>(&self, mir_body: &'a mir::Body<'tcx>) -> &'a Terminator<'tcx> { - &mir_body[self.last_bb()].terminator() - } } /// Represents a successor from a branching BasicCoverageBlock (such as the arms of a `SwitchInt`) @@ -362,26 +356,28 @@ impl std::fmt::Debug for BcbBranch { } } -// Returns the `Terminator`s non-unwind successors. +// Returns the subset of a block's successors that are relevant to the coverage +// graph, i.e. those that do not represent unwinds or unreachable branches. // FIXME(#78544): MIR InstrumentCoverage: Improve coverage of `#[should_panic]` tests and // `catch_unwind()` handlers. fn bcb_filtered_successors<'a, 'tcx>( body: &'a mir::Body<'tcx>, - term_kind: &'a TerminatorKind<'tcx>, -) -> Box + 'a> { - Box::new( - match &term_kind { - // SwitchInt successors are never unwind, and all of them should be traversed. - TerminatorKind::SwitchInt { ref targets, .. } => { - None.into_iter().chain(targets.all_targets().into_iter().copied()) - } - // For all other kinds, return only the first successor, if any, and ignore unwinds. - // NOTE: `chain(&[])` is required to coerce the `option::iter` (from - // `next().into_iter()`) into the `mir::Successors` aliased type. - _ => term_kind.successors().next().into_iter().chain((&[]).into_iter().copied()), - } - .filter(move |&successor| body[successor].terminator().kind != TerminatorKind::Unreachable), - ) + bb: BasicBlock, +) -> impl Iterator + Captures<'a> + Captures<'tcx> { + let terminator = body[bb].terminator(); + + let take_n_successors = match terminator.kind { + // SwitchInt successors are never unwinds, so all of them should be traversed. + TerminatorKind::SwitchInt { .. } => usize::MAX, + // For all other kinds, return only the first successor (if any), ignoring any + // unwind successors. + _ => 1, + }; + + terminator + .successors() + .take(take_n_successors) + .filter(move |&successor| body[successor].terminator().kind != TerminatorKind::Unreachable) } /// Maintains separate worklists for each loop in the BasicCoverageBlock CFG, plus one for the @@ -558,7 +554,7 @@ fn short_circuit_preorder<'a, 'tcx, F, Iter>( filtered_successors: F, ) -> impl Iterator + Captures<'a> + Captures<'tcx> where - F: Fn(&'a mir::Body<'tcx>, &'a TerminatorKind<'tcx>) -> Iter, + F: Fn(&'a mir::Body<'tcx>, BasicBlock) -> Iter, Iter: Iterator, { let mut visited = BitSet::new_empty(body.basic_blocks.len()); @@ -570,7 +566,7 @@ where continue; } - worklist.extend(filtered_successors(body, &body[bb].terminator().kind)); + worklist.extend(filtered_successors(body, bb)); return Some(bb); } diff --git a/compiler/rustc_mir_transform/src/coverage/tests.rs b/compiler/rustc_mir_transform/src/coverage/tests.rs index 7476d3ce9272a..ee7cb791c8180 100644 --- a/compiler/rustc_mir_transform/src/coverage/tests.rs +++ b/compiler/rustc_mir_transform/src/coverage/tests.rs @@ -241,7 +241,7 @@ fn print_coverage_graphviz( " {:?} [label=\"{:?}: {}\"];\n{}", bcb, bcb, - bcb_data.terminator(mir_body).kind.name(), + mir_body[bcb_data.last_bb()].terminator().kind.name(), basic_coverage_blocks .successors(bcb) .map(|successor| { format!(" {:?} -> {:?};", bcb, successor) })