bjoernager/rust
bjoernager/rust
1
Fork 0
Code Activity

coverage: Completely overhaul counter assignment, using node-flow graphs

Browse source
This commit is contained in:
Zalathar 2025-01-12 21:36:07 +11:00
parent e70112caf8
commit f1300c860e
51 changed files with 1930 additions and 1973 deletions
Download patch file Download diff file Expand all files Collapse all files

446
compiler/rustc_mir_transform/src/coverage/counters.rs
View file

@ -1,18 +1,24 @@
use std::cmp::Ordering;
use std::fmt::{self, Debug};
use either::Either;
use itertools::Itertools;
use rustc_data_structures::captures::Captures;
use rustc_data_structures::fx::FxHashMap;
use rustc_data_structures::graph::DirectedGraph;
use rustc_index::IndexVec;
use rustc_index::bit_set::DenseBitSet;
use rustc_middle::mir::coverage::{CounterId, CovTerm, Expression, ExpressionId, Op};
use tracing::{debug, debug_span, instrument};
use crate::coverage::graph::{BasicCoverageBlock, CoverageGraph, ReadyFirstTraversal};
use crate::coverage::counters::balanced_flow::BalancedFlowGraph;
use crate::coverage::counters::iter_nodes::IterNodes;
use crate::coverage::counters::node_flow::{CounterTerm, MergedNodeFlowGraph, NodeCounters};
use crate::coverage::graph::{BasicCoverageBlock, CoverageGraph};
#[cfg(test)]
mod tests;
mod balanced_flow;
mod iter_nodes;
mod node_flow;
mod union_find;
/// The coverage counter or counter expression associated with a particular
/// BCB node or BCB edge.
@ -48,10 +54,12 @@ struct BcbExpression {
}
/// Enum representing either a node or an edge in the coverage graph.
///
/// FIXME(#135481): This enum is no longer needed now that we only instrument
/// nodes and not edges. It can be removed in a subsequent PR.
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub(super) enum Site {
Node { bcb: BasicCoverageBlock },
Edge { from_bcb: BasicCoverageBlock, to_bcb: BasicCoverageBlock },
}
/// Generates and stores coverage counter and coverage expression information
@ -79,10 +87,38 @@ impl CoverageCounters {
graph: &CoverageGraph,
bcb_needs_counter: &DenseBitSet<BasicCoverageBlock>,
) -> Self {
let mut builder = CountersBuilder::new(graph, bcb_needs_counter);
builder.make_bcb_counters();
let balanced_graph = BalancedFlowGraph::for_graph(graph, |n| !graph[n].is_out_summable);
let merged_graph = MergedNodeFlowGraph::for_balanced_graph(&balanced_graph);
builder.into_coverage_counters()
// A "reloop" node has exactly one out-edge, which jumps back to the top
// of an enclosing loop. Reloop nodes are typically visited more times
// than loop-exit nodes, so try to avoid giving them physical counters.
let is_reloop_node = IndexVec::from_fn_n(
|node| match graph.successors[node].as_slice() {
&[succ] => graph.dominates(succ, node),
_ => false,
},
graph.num_nodes(),
);
let mut nodes = balanced_graph.iter_nodes().rev().collect::<Vec<_>>();
// The first node is the sink, which must not get a physical counter.
assert_eq!(nodes[0], balanced_graph.sink);
// Sort the real nodes, such that earlier (lesser) nodes take priority
// in being given a counter expression instead of a physical counter.
nodes[1..].sort_by(|&a, &b| {
// Start with a dummy `Equal` to make the actual tests line up nicely.
Ordering::Equal
// Prefer a physical counter for return/yield nodes.
.then_with(|| Ord::cmp(&graph[a].is_out_summable, &graph[b].is_out_summable))
// Prefer an expression for reloop nodes (see definition above).
.then_with(|| Ord::cmp(&is_reloop_node[a], &is_reloop_node[b]).reverse())
// Otherwise, prefer a physical counter for dominating nodes.
.then_with(|| graph.cmp_in_dominator_order(a, b).reverse())
});
let node_counters = merged_graph.make_node_counters(&nodes);
Transcriber::new(graph.num_nodes(), node_counters).transcribe_counters(bcb_needs_counter)
}
fn with_num_bcbs(num_bcbs: usize) -> Self {
@ -182,321 +218,51 @@ impl CoverageCounters {
}
}
/// Symbolic representation of the coverage counter to be used for a particular
/// node or edge in the coverage graph. The same site counter can be used for
/// multiple sites, if they have been determined to have the same count.
#[derive(Clone, Copy, Debug)]
enum SiteCounter {
/// A physical counter at some node/edge.
Phys { site: Site },
/// A counter expression for a node that takes the sum of all its in-edge
/// counters.
NodeSumExpr { bcb: BasicCoverageBlock },
/// A counter expression for an edge that takes the counter of its source
/// node, and subtracts the counters of all its sibling out-edges.
EdgeDiffExpr { from_bcb: BasicCoverageBlock, to_bcb: BasicCoverageBlock },
}
/// Yields the graph successors of `from_bcb` that aren't `to_bcb`. This is
/// used when creating a counter expression for [`SiteCounter::EdgeDiffExpr`].
///
/// For example, in this diagram the sibling out-edge targets of edge `AC` are
/// the nodes `B` and `D`.
///
/// ```text
/// A
/// / | \
/// B C D
/// ```
fn sibling_out_edge_targets(
graph: &CoverageGraph,
from_bcb: BasicCoverageBlock,
to_bcb: BasicCoverageBlock,
) -> impl Iterator<Item = BasicCoverageBlock> + Captures<'_> {
graph.successors[from_bcb].iter().copied().filter(move |&t| t != to_bcb)
}
/// Helper struct that allows counter creation to inspect the BCB graph, and
/// the set of nodes that need counters.
struct CountersBuilder<'a> {
graph: &'a CoverageGraph,
bcb_needs_counter: &'a DenseBitSet<BasicCoverageBlock>,
site_counters: FxHashMap<Site, SiteCounter>,
}
impl<'a> CountersBuilder<'a> {
fn new(
graph: &'a CoverageGraph,
bcb_needs_counter: &'a DenseBitSet<BasicCoverageBlock>,
) -> Self {
assert_eq!(graph.num_nodes(), bcb_needs_counter.domain_size());
Self { graph, bcb_needs_counter, site_counters: FxHashMap::default() }
}
fn make_bcb_counters(&mut self) {
debug!("make_bcb_counters(): adding a counter or expression to each BasicCoverageBlock");
// Traverse the coverage graph, ensuring that every node that needs a
// coverage counter has one.
for bcb in ReadyFirstTraversal::new(self.graph) {
let _span = debug_span!("traversal", ?bcb).entered();
if self.bcb_needs_counter.contains(bcb) {
self.make_node_counter_and_out_edge_counters(bcb);
}
}
}
/// Make sure the given node has a node counter, and then make sure each of
/// its out-edges has an edge counter (if appropriate).
#[instrument(level = "debug", skip(self))]
fn make_node_counter_and_out_edge_counters(&mut self, from_bcb: BasicCoverageBlock) {
// First, ensure that this node has a counter of some kind.
// We might also use that counter to compute one of the out-edge counters.
self.get_or_make_node_counter(from_bcb);
// If this node's out-edges won't sum to the node's counter,
// then there's no reason to create edge counters here.
if !self.graph[from_bcb].is_out_summable {
return;
}
// When choosing which out-edge should be given a counter expression, ignore edges that
// already have counters, or could use the existing counter of their target node.
let out_edge_has_counter = |to_bcb| {
if self.site_counters.contains_key(&Site::Edge { from_bcb, to_bcb }) {
return true;
}
self.graph.sole_predecessor(to_bcb) == Some(from_bcb)
&& self.site_counters.contains_key(&Site::Node { bcb: to_bcb })
};
// Determine the set of out-edges that could benefit from being given an expression.
let candidate_successors = self.graph.successors[from_bcb]
.iter()
.copied()
.filter(|&to_bcb| !out_edge_has_counter(to_bcb))
.collect::<Vec<_>>();
debug!(?candidate_successors);
// If there are out-edges without counters, choose one to be given an expression
// (computed from this node and the other out-edges) instead of a physical counter.
let Some(to_bcb) = self.choose_out_edge_for_expression(from_bcb, &candidate_successors)
else {
return;
};
// For each out-edge other than the one that was chosen to get an expression,
// ensure that it has a counter (existing counter/expression or a new counter).
for target in sibling_out_edge_targets(self.graph, from_bcb, to_bcb) {
self.get_or_make_edge_counter(from_bcb, target);
}
// Now create an expression for the chosen edge, by taking the counter
// for its source node and subtracting the sum of its sibling out-edges.
let counter = SiteCounter::EdgeDiffExpr { from_bcb, to_bcb };
self.site_counters.insert(Site::Edge { from_bcb, to_bcb }, counter);
}
#[instrument(level = "debug", skip(self))]
fn get_or_make_node_counter(&mut self, bcb: BasicCoverageBlock) -> SiteCounter {
// If the BCB already has a counter, return it.
if let Some(&counter) = self.site_counters.get(&Site::Node { bcb }) {
debug!("{bcb:?} already has a counter: {counter:?}");
return counter;
}
let counter = self.make_node_counter_inner(bcb);
self.site_counters.insert(Site::Node { bcb }, counter);
counter
}
fn make_node_counter_inner(&mut self, bcb: BasicCoverageBlock) -> SiteCounter {
// If the node's sole in-edge already has a counter, use that.
if let Some(sole_pred) = self.graph.sole_predecessor(bcb)
&& let Some(&edge_counter) =
self.site_counters.get(&Site::Edge { from_bcb: sole_pred, to_bcb: bcb })
{
return edge_counter;
}
let predecessors = self.graph.predecessors[bcb].as_slice();
// Handle cases where we can't compute a node's count from its in-edges:
// - START_BCB has no in-edges, so taking the sum would panic (or be wrong).
// - For nodes with one in-edge, or that directly loop to themselves,
// trying to get the in-edge counts would require this node's counter,
// leading to infinite recursion.
if predecessors.len() <= 1 || predecessors.contains(&bcb) {
debug!(?bcb, ?predecessors, "node has <=1 predecessors or is its own predecessor");
let counter = SiteCounter::Phys { site: Site::Node { bcb } };
debug!(?bcb, ?counter, "node gets a physical counter");
return counter;
}
// A BCB with multiple incoming edges can compute its count by ensuring that counters
// exist for each of those edges, and then adding them up to get a total count.
for &from_bcb in predecessors {
self.get_or_make_edge_counter(from_bcb, bcb);
}
let sum_of_in_edges = SiteCounter::NodeSumExpr { bcb };
debug!("{bcb:?} gets a new counter (sum of predecessor counters): {sum_of_in_edges:?}");
sum_of_in_edges
}
#[instrument(level = "debug", skip(self))]
fn get_or_make_edge_counter(
&mut self,
from_bcb: BasicCoverageBlock,
to_bcb: BasicCoverageBlock,
) -> SiteCounter {
// If the edge already has a counter, return it.
if let Some(&counter) = self.site_counters.get(&Site::Edge { from_bcb, to_bcb }) {
debug!("Edge {from_bcb:?}->{to_bcb:?} already has a counter: {counter:?}");
return counter;
}
let counter = self.make_edge_counter_inner(from_bcb, to_bcb);
self.site_counters.insert(Site::Edge { from_bcb, to_bcb }, counter);
counter
}
fn make_edge_counter_inner(
&mut self,
from_bcb: BasicCoverageBlock,
to_bcb: BasicCoverageBlock,
) -> SiteCounter {
// If the target node has exactly one in-edge (i.e. this one), then just
// use the node's counter, since it will have the same value.
if let Some(sole_pred) = self.graph.sole_predecessor(to_bcb) {
assert_eq!(sole_pred, from_bcb);
// This call must take care not to invoke `get_or_make_edge` for
// this edge, since that would result in infinite recursion!
return self.get_or_make_node_counter(to_bcb);
}
// If the source node has exactly one out-edge (i.e. this one) and would have
// the same execution count as that edge, then just use the node's counter.
if let Some(simple_succ) = self.graph.simple_successor(from_bcb) {
assert_eq!(simple_succ, to_bcb);
return self.get_or_make_node_counter(from_bcb);
}
// Make a new counter to count this edge.
let counter = SiteCounter::Phys { site: Site::Edge { from_bcb, to_bcb } };
debug!(?from_bcb, ?to_bcb, ?counter, "edge gets a physical counter");
counter
}
/// Given a set of candidate out-edges (represented by their successor node),
/// choose one to be given a counter expression instead of a physical counter.
fn choose_out_edge_for_expression(
&self,
from_bcb: BasicCoverageBlock,
candidate_successors: &[BasicCoverageBlock],
) -> Option<BasicCoverageBlock> {
// Try to find a candidate that leads back to the top of a loop,
// because reloop edges tend to be executed more times than loop-exit edges.
if let Some(reloop_target) = self.find_good_reloop_edge(from_bcb, &candidate_successors) {
debug!("Selecting reloop target {reloop_target:?} to get an expression");
return Some(reloop_target);
}
// We couldn't identify a "good" edge, so just choose an arbitrary one.
let arbitrary_target = candidate_successors.first().copied()?;
debug!(?arbitrary_target, "selecting arbitrary out-edge to get an expression");
Some(arbitrary_target)
}
/// Given a set of candidate out-edges (represented by their successor node),
/// tries to find one that leads back to the top of a loop.
///
/// Reloop edges are good candidates for counter expressions, because they
/// will tend to be executed more times than a loop-exit edge, so it's nice
/// for them to be able to avoid a physical counter increment.
fn find_good_reloop_edge(
&self,
from_bcb: BasicCoverageBlock,
candidate_successors: &[BasicCoverageBlock],
) -> Option<BasicCoverageBlock> {
// If there are no candidates, avoid iterating over the loop stack.
if candidate_successors.is_empty() {
return None;
}
// Consider each loop on the current traversal context stack, top-down.
for loop_header_node in self.graph.loop_headers_containing(from_bcb) {
// Try to find a candidate edge that doesn't exit this loop.
for &target_bcb in candidate_successors {
// An edge is a reloop edge if its target dominates any BCB that has
// an edge back to the loop header. (Otherwise it's an exit edge.)
let is_reloop_edge = self
.graph
.reloop_predecessors(loop_header_node)
.any(|reloop_bcb| self.graph.dominates(target_bcb, reloop_bcb));
if is_reloop_edge {
// We found a good out-edge to be given an expression.
return Some(target_bcb);
}
}
// All of the candidate edges exit this loop, so keep looking
// for a good reloop edge for one of the outer loops.
}
None
}
fn into_coverage_counters(self) -> CoverageCounters {
Transcriber::new(&self).transcribe_counters()
}
}
/// Helper struct for converting `CountersBuilder` into a final `CoverageCounters`.
struct Transcriber<'a> {
old: &'a CountersBuilder<'a>,
struct Transcriber {
old: NodeCounters<BasicCoverageBlock>,
new: CoverageCounters,
phys_counter_for_site: FxHashMap<Site, BcbCounter>,
}
impl<'a> Transcriber<'a> {
fn new(old: &'a CountersBuilder<'a>) -> Self {
impl Transcriber {
fn new(num_nodes: usize, old: NodeCounters<BasicCoverageBlock>) -> Self {
Self {
old,
new: CoverageCounters::with_num_bcbs(old.graph.num_nodes()),
new: CoverageCounters::with_num_bcbs(num_nodes),
phys_counter_for_site: FxHashMap::default(),
}
}
fn transcribe_counters(mut self) -> CoverageCounters {
for bcb in self.old.bcb_needs_counter.iter() {
fn transcribe_counters(
mut self,
bcb_needs_counter: &DenseBitSet<BasicCoverageBlock>,
) -> CoverageCounters {
for bcb in bcb_needs_counter.iter() {
let site = Site::Node { bcb };
let site_counter = self.site_counter(site);
// Resolve the site counter into flat lists of nodes/edges whose
// physical counts contribute to the counter for this node.
// Distinguish between counts that will be added vs subtracted.
let mut pos = vec![];
let mut neg = vec![];
self.push_resolved_sites(site_counter, &mut pos, &mut neg);
// Simplify by cancelling out sites that appear on both sides.
let (mut pos, mut neg) = sort_and_cancel(pos, neg);
let (mut pos, mut neg): (Vec<_>, Vec<_>) =
self.old.counter_expr(bcb).iter().partition_map(
|&CounterTerm { node, op }| match op {
Op::Add => Either::Left(node),
Op::Subtract => Either::Right(node),
},
);
if pos.is_empty() {
// If we somehow end up with no positive terms after cancellation,
// fall back to creating a physical counter. There's no known way
// for this to happen, but it's hard to confidently rule it out.
// If we somehow end up with no positive terms, fall back to
// creating a physical counter. There's no known way for this
// to happen, but we can avoid an ICE if it does.
debug_assert!(false, "{site:?} has no positive counter terms");
pos = vec![Some(site)];
pos = vec![bcb];
neg = vec![];
}
let mut new_counters_for_sites = |sites: Vec<Option<Site>>| {
pos.sort();
neg.sort();
let mut new_counters_for_sites = |sites: Vec<BasicCoverageBlock>| {
sites
.into_iter()
.filter_map(|id| try { self.ensure_phys_counter(id?) })
.map(|node| self.ensure_phys_counter(Site::Node { bcb: node }))
.collect::<Vec<_>>()
};
let mut pos = new_counters_for_sites(pos);
@ -513,79 +279,7 @@ impl<'a> Transcriber<'a> {
self.new
}
fn site_counter(&self, site: Site) -> SiteCounter {
self.old.site_counters.get(&site).copied().unwrap_or_else(|| {
// We should have already created all necessary site counters.
// But if we somehow didn't, avoid crashing in release builds,
// and just use an extra physical counter instead.
debug_assert!(false, "{site:?} should have a counter");
SiteCounter::Phys { site }
})
}
fn ensure_phys_counter(&mut self, site: Site) -> BcbCounter {
*self.phys_counter_for_site.entry(site).or_insert_with(|| self.new.make_phys_counter(site))
}
/// Resolves the given counter into flat lists of nodes/edges, whose counters
/// will then be added and subtracted to form a counter expression.
fn push_resolved_sites(&self, counter: SiteCounter, pos: &mut Vec<Site>, neg: &mut Vec<Site>) {
match counter {
SiteCounter::Phys { site } => pos.push(site),
SiteCounter::NodeSumExpr { bcb } => {
for &from_bcb in &self.old.graph.predecessors[bcb] {
let edge_counter = self.site_counter(Site::Edge { from_bcb, to_bcb: bcb });
self.push_resolved_sites(edge_counter, pos, neg);
}
}
SiteCounter::EdgeDiffExpr { from_bcb, to_bcb } => {
// First, add the count for `from_bcb`.
let node_counter = self.site_counter(Site::Node { bcb: from_bcb });
self.push_resolved_sites(node_counter, pos, neg);
// Then subtract the counts for the other out-edges.
for target in sibling_out_edge_targets(self.old.graph, from_bcb, to_bcb) {
let edge_counter = self.site_counter(Site::Edge { from_bcb, to_bcb: target });
// Swap `neg` and `pos` so that the counter is subtracted.
self.push_resolved_sites(edge_counter, neg, pos);
}
}
}
}
}
/// Given two lists:
/// - Sorts each list.
/// - Converts each list to `Vec<Option<T>>`.
/// - Scans for values that appear in both lists, and cancels them out by
/// replacing matching pairs of values with `None`.
fn sort_and_cancel<T: Ord>(mut pos: Vec<T>, mut neg: Vec<T>) -> (Vec<Option<T>>, Vec<Option<T>>) {
pos.sort();
neg.sort();
// Convert to `Vec<Option<T>>`. If `T` has a niche, this should be zero-cost.
let mut pos = pos.into_iter().map(Some).collect::<Vec<_>>();
let mut neg = neg.into_iter().map(Some).collect::<Vec<_>>();
// Scan through the lists using two cursors. When either cursor reaches the
// end of its list, there can be no more equal pairs, so stop.
let mut p = 0;
let mut n = 0;
while p < pos.len() && n < neg.len() {
// If the values are equal, remove them and advance both cursors.
// Otherwise, advance whichever cursor points to the lesser value.
// (Choosing which cursor to advance relies on both lists being sorted.)
match pos[p].cmp(&neg[n]) {
Ordering::Less => p += 1,
Ordering::Equal => {
pos[p] = None;
neg[n] = None;
p += 1;
n += 1;
}
Ordering::Greater => n += 1,
}
}
(pos, neg)
}

133
compiler/rustc_mir_transform/src/coverage/counters/balanced_flow.rs Normal file
View file

@ -0,0 +1,133 @@
//! A control-flow graph can be said to have “balanced flow” if the flow
//! (execution count) of each node is equal to the sum of its in-edge flows,
//! and also equal to the sum of its out-edge flows.
//!
//! Control-flow graphs typically have one or more nodes that don't satisfy the
//! balanced-flow property, e.g.:
//! - The start node has out-edges, but no in-edges.
//! - Return nodes have in-edges, but no out-edges.
//! - `Yield` nodes can have an out-flow that is less than their in-flow.
//! - Inescapable loops cause the in-flow/out-flow relationship to break down.
//!
//! Balanced-flow graphs are nevertheless useful for analysis, so this module
//! provides a wrapper type ([`BalancedFlowGraph`]) that imposes balanced flow
//! on an underlying graph. This is done by non-destructively adding synthetic
//! nodes and edges as necessary.
use rustc_data_structures::graph;
use rustc_data_structures::graph::iterate::DepthFirstSearch;
use rustc_data_structures::graph::reversed::ReversedGraph;
use rustc_index::Idx;
use rustc_index::bit_set::DenseBitSet;
use crate::coverage::counters::iter_nodes::IterNodes;
/// A view of an underlying graph that has been augmented to have “balanced flow”.
/// This means that the flow (execution count) of each node is equal to the
/// sum of its in-edge flows, and also equal to the sum of its out-edge flows.
///
/// To achieve this, a synthetic "sink" node is non-destructively added to the
/// graph, with synthetic in-edges from these nodes:
/// - Any node that has no out-edges.
/// - Any node that explicitly requires a sink edge, as indicated by a
/// caller-supplied `force_sink_edge` function.
/// - Any node that would otherwise be unable to reach the sink, because it is
/// part of an inescapable loop.
///
/// To make the graph fully balanced, there is also a synthetic edge from the
/// sink node back to the start node.
///
/// ---
/// The benefit of having a balanced-flow graph is that it can be subsequently
/// transformed in ways that are guaranteed to preserve balanced flow
/// (e.g. merging nodes together), which is useful for discovering relationships
/// between the node flows of different nodes in the graph.
pub(crate) struct BalancedFlowGraph<G: graph::DirectedGraph> {
graph: G,
sink_edge_nodes: DenseBitSet<G::Node>,
pub(crate) sink: G::Node,
}
impl<G: graph::DirectedGraph> BalancedFlowGraph<G> {
/// Creates a balanced view of an underlying graph, by adding a synthetic
/// sink node that has in-edges from nodes that need or request such an edge,
/// and a single out-edge to the start node.
///
/// Assumes that all nodes in the underlying graph are reachable from the
/// start node.
pub(crate) fn for_graph(graph: G, force_sink_edge: impl Fn(G::Node) -> bool) -> Self
where
G: graph::ControlFlowGraph,
{
let mut sink_edge_nodes = DenseBitSet::new_empty(graph.num_nodes());
let mut dfs = DepthFirstSearch::new(ReversedGraph::new(&graph));
// First, determine the set of nodes that explicitly request or require
// an out-edge to the sink.
for node in graph.iter_nodes() {
if force_sink_edge(node) || graph.successors(node).next().is_none() {
sink_edge_nodes.insert(node);
dfs.push_start_node(node);
}
}
// Next, find all nodes that are currently not reverse-reachable from
// `sink_edge_nodes`, and add them to the set as well.
dfs.complete_search();
sink_edge_nodes.union_not(dfs.visited_set());
// The sink node is 1 higher than the highest real node.
let sink = G::Node::new(graph.num_nodes());
BalancedFlowGraph { graph, sink_edge_nodes, sink }
}
}
impl<G> graph::DirectedGraph for BalancedFlowGraph<G>
where
G: graph::DirectedGraph,
{
type Node = G::Node;
/// Returns the number of nodes in this balanced-flow graph, which is 1
/// more than the number of nodes in the underlying graph, to account for
/// the synthetic sink node.
fn num_nodes(&self) -> usize {
// The sink node's index is already the size of the underlying graph,
// so just add 1 to that instead.
self.sink.index() + 1
}
}
impl<G> graph::StartNode for BalancedFlowGraph<G>
where
G: graph::StartNode,
{
fn start_node(&self) -> Self::Node {
self.graph.start_node()
}
}
impl<G> graph::Successors for BalancedFlowGraph<G>
where
G: graph::StartNode + graph::Successors,
{
fn successors(&self, node: Self::Node) -> impl Iterator<Item = Self::Node> {
let real_edges;
let sink_edge;
if node == self.sink {
// The sink node has no real out-edges, and one synthetic out-edge
// to the start node.
real_edges = None;
sink_edge = Some(self.graph.start_node());
} else {
// Real nodes have their real out-edges, and possibly one synthetic
// out-edge to the sink node.
real_edges = Some(self.graph.successors(node));
sink_edge = self.sink_edge_nodes.contains(node).then_some(self.sink);
}
real_edges.into_iter().flatten().chain(sink_edge)
}
}

16
compiler/rustc_mir_transform/src/coverage/counters/iter_nodes.rs Normal file
View file

@ -0,0 +1,16 @@
use rustc_data_structures::graph;
use rustc_index::Idx;
pub(crate) trait IterNodes: graph::DirectedGraph {
/// Iterates over all nodes of a graph in ascending numeric order.
/// Assumes that nodes are densely numbered, i.e. every index in
/// `0..num_nodes` is a valid node.
///
/// FIXME: Can this just be part of [`graph::DirectedGraph`]?
fn iter_nodes(
&self,
) -> impl Iterator<Item = Self::Node> + DoubleEndedIterator + ExactSizeIterator {
(0..self.num_nodes()).map(<Self::Node as Idx>::new)
}
}
impl<G: graph::DirectedGraph> IterNodes for G {}

290
compiler/rustc_mir_transform/src/coverage/counters/node_flow.rs Normal file
View file

@ -0,0 +1,290 @@
//! For each node in a control-flow graph, determines whether that node should
//! have a physical counter, or a counter expression that is derived from the
//! physical counters of other nodes.
//!
//! Based on the algorithm given in
//! "Optimal measurement points for program frequency counts"
//! (Knuth & Stevenson, 1973).
use rustc_data_structures::graph;
use rustc_index::bit_set::DenseBitSet;
use rustc_index::{Idx, IndexVec};
use rustc_middle::mir::coverage::Op;
use smallvec::SmallVec;
use crate::coverage::counters::iter_nodes::IterNodes;
use crate::coverage::counters::union_find::{FrozenUnionFind, UnionFind};
#[cfg(test)]
mod tests;
/// View of some underlying graph, in which each node's successors have been
/// merged into a single "supernode".
///
/// The resulting supernodes have no obvious meaning on their own.
/// However, merging successor nodes means that a node's out-edges can all
/// be combined into a single out-edge, whose flow is the same as the flow
/// (execution count) of its corresponding node in the original graph.
///
/// With all node flows now in the original graph now represented as edge flows
/// in the merged graph, it becomes possible to analyze the original node flows
/// using techniques for analyzing edge flows.
#[derive(Debug)]
pub(crate) struct MergedNodeFlowGraph<Node: Idx> {
/// Maps each node to the supernode that contains it, indicated by some
/// arbitrary "root" node that is part of that supernode.
supernodes: FrozenUnionFind<Node>,
/// For each node, stores the single supernode that all of its successors
/// have been merged into.
///
/// (Note that each node in a supernode can potentially have a _different_
/// successor supernode from its peers.)
succ_supernodes: IndexVec<Node, Node>,
}
impl<Node: Idx> MergedNodeFlowGraph<Node> {
/// Creates a "merged" view of an underlying graph.
///
/// The given graph is assumed to have [“balanced flow”](balanced-flow),
/// though it does not necessarily have to be a `BalancedFlowGraph`.
///
/// [balanced-flow]: `crate::coverage::counters::balanced_flow::BalancedFlowGraph`.
pub(crate) fn for_balanced_graph<G>(graph: G) -> Self
where
G: graph::DirectedGraph<Node = Node> + graph::Successors,
{
let mut supernodes = UnionFind::<G::Node>::new(graph.num_nodes());
// For each node, merge its successors into a single supernode, and
// arbitrarily choose one of those successors to represent all of them.
let successors = graph
.iter_nodes()
.map(|node| {
graph
.successors(node)
.reduce(|a, b| supernodes.unify(a, b))
.expect("each node in a balanced graph must have at least one out-edge")
})
.collect::<IndexVec<G::Node, G::Node>>();
// Now that unification is complete, freeze the supernode forest,
// and resolve each arbitrarily-chosen successor to its canonical root.
// (This avoids having to explicitly resolve them later.)
let supernodes = supernodes.freeze();
let succ_supernodes = successors.into_iter().map(|succ| supernodes.find(succ)).collect();
Self { supernodes, succ_supernodes }
}
fn num_nodes(&self) -> usize {
self.succ_supernodes.len()
}
fn is_supernode(&self, node: Node) -> bool {
self.supernodes.find(node) == node
}
/// Using the information in this merged graph, together with a given
/// permutation of all nodes in the graph, to create physical counters and
/// counter expressions for each node in the underlying graph.
///
/// The given list must contain exactly one copy of each node in the
/// underlying balanced-flow graph. The order of nodes is used as a hint to
/// influence counter allocation:
/// - Earlier nodes are more likely to receive counter expressions.
/// - Later nodes are more likely to receive physical counters.
pub(crate) fn make_node_counters(&self, all_nodes_permutation: &[Node]) -> NodeCounters<Node> {
let mut builder = SpantreeBuilder::new(self);
for &node in all_nodes_permutation {
builder.visit_node(node);
}
NodeCounters { counter_exprs: builder.finish() }
}
}
/// End result of allocating physical counters and counter expressions for the
/// nodes of a graph.
#[derive(Debug)]
pub(crate) struct NodeCounters<Node: Idx> {
counter_exprs: IndexVec<Node, CounterExprVec<Node>>,
}
impl<Node: Idx> NodeCounters<Node> {
/// For the given node, returns the finished list of terms that represent
/// its physical counter or counter expression. Always non-empty.
///
/// If a node was given a physical counter, its "expression" will contain
/// that counter as its sole element.
pub(crate) fn counter_expr(&self, this: Node) -> &[CounterTerm<Node>] {
self.counter_exprs[this].as_slice()
}
}
#[derive(Debug)]
struct SpantreeEdge<Node> {
/// If true, this edge in the spantree has been reversed an odd number of
/// times, so all physical counters added to its node's counter expression
/// need to be negated.
is_reversed: bool,
/// Each spantree edge is "claimed" by the (regular) node that caused it to
/// be created. When a node with a physical counter traverses this edge,
/// that counter is added to the claiming node's counter expression.
claiming_node: Node,
/// Supernode at the other end of this spantree edge. Transitively points
/// to the "root" of this supernode's spantree component.
span_parent: Node,
}
/// Part of a node's counter expression, which is a sum of counter terms.
#[derive(Debug)]
pub(crate) struct CounterTerm<Node> {
/// Whether to add or subtract the value of the node's physical counter.
pub(crate) op: Op,
/// The node whose physical counter is represented by this term.
pub(crate) node: Node,
}
/// Stores the list of counter terms that make up a node's counter expression.
type CounterExprVec<Node> = SmallVec<[CounterTerm<Node>; 2]>;
#[derive(Debug)]
struct SpantreeBuilder<'a, Node: Idx> {
graph: &'a MergedNodeFlowGraph<Node>,
is_unvisited: DenseBitSet<Node>,
/// Links supernodes to each other, gradually forming a spanning tree of
/// the merged-flow graph.
///
/// A supernode without a span edge is the root of its component of the
/// spantree. Nodes that aren't supernodes cannot have a spantree edge.
span_edges: IndexVec<Node, Option<SpantreeEdge<Node>>>,
/// An in-progress counter expression for each node. Each expression is
/// initially empty, and will be filled in as relevant nodes are visited.
counter_exprs: IndexVec<Node, CounterExprVec<Node>>,
}
impl<'a, Node: Idx> SpantreeBuilder<'a, Node> {
fn new(graph: &'a MergedNodeFlowGraph<Node>) -> Self {
let num_nodes = graph.num_nodes();
Self {
graph,
is_unvisited: DenseBitSet::new_filled(num_nodes),
span_edges: IndexVec::from_fn_n(|_| None, num_nodes),
counter_exprs: IndexVec::from_fn_n(|_| SmallVec::new(), num_nodes),
}
}
/// Given a supernode, finds the supernode that is the "root" of its
/// spantree component. Two nodes that have the same spantree root are
/// connected in the spantree.
fn spantree_root(&self, this: Node) -> Node {
debug_assert!(self.graph.is_supernode(this));
match self.span_edges[this] {
None => this,
Some(SpantreeEdge { span_parent, .. }) => self.spantree_root(span_parent),
}
}
/// Rotates edges in the spantree so that `this` is the root of its
/// spantree component.
fn yank_to_spantree_root(&mut self, this: Node) {
debug_assert!(self.graph.is_supernode(this));
// Temporarily remove this supernode (any any spantree-children) from its
// spantree component, by disconnecting the edge to its spantree-parent.
let Some(SpantreeEdge { is_reversed, claiming_node, span_parent }) =
self.span_edges[this].take()
else {
// This supernode has no spantree-parent edge, so it is already the
// root of its spantree component.
return;
};
// Recursively make our immediate spantree-parent the root of what's
// left of its component, so that only one more edge rotation is needed.
self.yank_to_spantree_root(span_parent);
// Recreate the removed edge, but in the opposite direction.
// Now `this` is the root of its spantree component.
self.span_edges[span_parent] =
Some(SpantreeEdge { is_reversed: !is_reversed, claiming_node, span_parent: this });
}
/// Must be called exactly once for each node in the balanced-flow graph.
fn visit_node(&mut self, this: Node) {
// Assert that this node was unvisited, and mark it visited.
assert!(self.is_unvisited.remove(this), "node has already been visited: {this:?}");
// Get the supernode containing `this`, and make it the root of its
// component of the spantree.
let this_supernode = self.graph.supernodes.find(this);
self.yank_to_spantree_root(this_supernode);
// Get the supernode containing all of this's successors.
let succ_supernode = self.graph.succ_supernodes[this];
debug_assert!(self.graph.is_supernode(succ_supernode));
// If two supernodes are already connected in the spantree, they will
// have the same spantree root. (Each supernode is connected to itself.)
if this_supernode != self.spantree_root(succ_supernode) {
// Adding this node's flow edge to the spantree would cause two
// previously-disconnected supernodes to become connected, so add
// it. That spantree-edge is now "claimed" by this node.
//
// Claiming a spantree-edge means that this node will get a counter
// expression instead of a physical counter. That expression is
// currently empty, but will be built incrementally as the other
// nodes are visited.
self.span_edges[this_supernode] = Some(SpantreeEdge {
is_reversed: false,
claiming_node: this,
span_parent: succ_supernode,
});
} else {
// This node's flow edge would join two supernodes that are already
// connected in the spantree (or are the same supernode). That would
// create a cycle in the spantree, so don't add an edge.
//
// Instead, create a physical counter for this node, and add that
// counter to all expressions on the path from `succ_supernode` to
// `this_supernode`.
// Instead of setting `this.measure = true` as in the original paper,
// we just add the node's ID to its own "expression".
self.counter_exprs[this].push(CounterTerm { node: this, op: Op::Add });
// Walk the spantree from `this.successor` back to `this`. For each
// spantree edge along the way, add this node's physical counter to
// the counter expression of the node that claimed the spantree edge.
let mut curr = succ_supernode;
while curr != this_supernode {
let &SpantreeEdge { is_reversed, claiming_node, span_parent } =
self.span_edges[curr].as_ref().unwrap();
let op = if is_reversed { Op::Subtract } else { Op::Add };
self.counter_exprs[claiming_node].push(CounterTerm { node: this, op });
curr = span_parent;
}
}
}
/// Asserts that all nodes have been visited, and returns the computed
/// counter expressions (made up of physical counters) for each node.
fn finish(self) -> IndexVec<Node, CounterExprVec<Node>> {
let Self { graph, is_unvisited, span_edges, counter_exprs } = self;
assert!(is_unvisited.is_empty(), "some nodes were never visited: {is_unvisited:?}");
debug_assert!(
span_edges
.iter_enumerated()
.all(|(node, span_edge)| { span_edge.is_some() <= graph.is_supernode(node) }),
"only supernodes can have a span edge",
);
debug_assert!(
counter_exprs.iter().all(|expr| !expr.is_empty()),
"after visiting all nodes, every node should have a non-empty expression",
);
counter_exprs
}
}

64
compiler/rustc_mir_transform/src/coverage/counters/node_flow/tests.rs Normal file
View file

@ -0,0 +1,64 @@
use itertools::Itertools;
use rustc_data_structures::graph;
use rustc_data_structures::graph::vec_graph::VecGraph;
use rustc_index::Idx;
use rustc_middle::mir::coverage::Op;
use super::{CounterTerm, MergedNodeFlowGraph, NodeCounters};
fn merged_node_flow_graph<G: graph::Successors>(graph: G) -> MergedNodeFlowGraph<G::Node> {
MergedNodeFlowGraph::for_balanced_graph(graph)
}
fn make_graph<Node: Idx + Ord>(num_nodes: usize, edge_pairs: Vec<(Node, Node)>) -> VecGraph<Node> {
VecGraph::new(num_nodes, edge_pairs)
}
/// Example used in "Optimal Measurement Points for Program Frequency Counts"
/// (Knuth & Stevenson, 1973), but with 0-based node IDs.
#[test]
fn example_driver() {
let graph = make_graph::<u32>(5, vec![
(0, 1),
(0, 3),
(1, 0),
(1, 2),
(2, 1),
(2, 4),
(3, 3),
(3, 4),
(4, 0),
]);
let merged = merged_node_flow_graph(&graph);
let counters = merged.make_node_counters(&[3, 1, 2, 0, 4]);
assert_eq!(format_counter_expressions(&counters), &[
// (comment to force vertical formatting for clarity)
"[0]: +c0",
"[1]: +c0 +c2 -c4",
"[2]: +c2",
"[3]: +c3",
"[4]: +c4",
]);
}
fn format_counter_expressions<Node: Idx>(counters: &NodeCounters<Node>) -> Vec<String> {
let format_item = |&CounterTerm { node, op }| {
let op = match op {
Op::Subtract => '-',
Op::Add => '+',
};
format!("{op}c{node:?}")
};
counters
.counter_exprs
.indices()
.map(|node| {
let mut expr = counters.counter_expr(node).iter().collect::<Vec<_>>();
expr.sort_by_key(|item| item.node.index());
format!("[{node:?}]: {}", expr.into_iter().map(format_item).join(" "))
})
.collect()
}

41
compiler/rustc_mir_transform/src/coverage/counters/tests.rs
View file

@ -1,41 +0,0 @@
use std::fmt::Debug;
use super::sort_and_cancel;
fn flatten<T>(input: Vec<Option<T>>) -> Vec<T> {
input.into_iter().flatten().collect()
}
fn sort_and_cancel_and_flatten<T: Clone + Ord>(pos: Vec<T>, neg: Vec<T>) -> (Vec<T>, Vec<T>) {
let (pos_actual, neg_actual) = sort_and_cancel(pos, neg);
(flatten(pos_actual), flatten(neg_actual))
}
#[track_caller]
fn check_test_case<T: Clone + Debug + Ord>(
pos: Vec<T>,
neg: Vec<T>,
pos_expected: Vec<T>,
neg_expected: Vec<T>,
) {
eprintln!("pos = {pos:?}; neg = {neg:?}");
let output = sort_and_cancel_and_flatten(pos, neg);
assert_eq!(output, (pos_expected, neg_expected));
}
#[test]
fn cancellation() {
let cases: &[(Vec<u32>, Vec<u32>, Vec<u32>, Vec<u32>)] = &[
(vec![], vec![], vec![], vec![]),
(vec![4, 2, 1, 5, 3], vec![], vec![1, 2, 3, 4, 5], vec![]),
(vec![5, 5, 5, 5, 5], vec![5], vec![5, 5, 5, 5], vec![]),
(vec![1, 1, 2, 2, 3, 3], vec![1, 2, 3], vec![1, 2, 3], vec![]),
(vec![1, 1, 2, 2, 3, 3], vec![2, 4, 2], vec![1, 1, 3, 3], vec![4]),
];
for (pos, neg, pos_expected, neg_expected) in cases {
check_test_case(pos.to_vec(), neg.to_vec(), pos_expected.to_vec(), neg_expected.to_vec());
// Same test case, but with its inputs flipped and its outputs flipped.
check_test_case(neg.to_vec(), pos.to_vec(), neg_expected.to_vec(), pos_expected.to_vec());
}
}

116
compiler/rustc_mir_transform/src/coverage/counters/union_find.rs Normal file
View file

@ -0,0 +1,116 @@
use std::cmp::Ordering;
use std::mem;
use rustc_index::{Idx, IndexVec};
#[cfg(test)]
mod tests;
/// Simple implementation of a union-find data structure, i.e. a disjoint-set
/// forest.
#[derive(Debug)]
pub(crate) struct UnionFind<Key: Idx> {
table: IndexVec<Key, UnionFindEntry<Key>>,
}
#[derive(Debug)]
struct UnionFindEntry<Key> {
/// Transitively points towards the "root" of the set containing this key.
///
/// Invariant: A root key is its own parent.
parent: Key,
/// When merging two "root" keys, their ranks determine which key becomes
/// the new root, to prevent the parent tree from becoming unnecessarily
/// tall. See [`UnionFind::unify`] for details.
rank: u32,
}
impl<Key: Idx> UnionFind<Key> {
/// Creates a new disjoint-set forest containing the keys `0..num_keys`.
/// Initially, every key is part of its own one-element set.
pub(crate) fn new(num_keys: usize) -> Self {
// Initially, every key is the root of its own set, so its parent is itself.
Self { table: IndexVec::from_fn_n(|key| UnionFindEntry { parent: key, rank: 0 }, num_keys) }
}
/// Returns the "root" key of the disjoint-set containing the given key.
/// If two keys have the same root, they belong to the same set.
///
/// Also updates internal data structures to make subsequent `find`
/// operations faster.
pub(crate) fn find(&mut self, key: Key) -> Key {
// Loop until we find a key that is its own parent.
let mut curr = key;
while let parent = self.table[curr].parent
&& curr != parent
{
// Perform "path compression" by peeking one layer ahead, and
// setting the current key's parent to that value.
// (This works even when `parent` is the root of its set, because
// of the invariant that a root is its own parent.)
let parent_parent = self.table[parent].parent;
self.table[curr].parent = parent_parent;
// Advance by one step and continue.
curr = parent;
}
curr
}
/// Merges the set containing `a` and the set containing `b` into one set.
///
/// Returns the common root of both keys, after the merge.
pub(crate) fn unify(&mut self, a: Key, b: Key) -> Key {
let mut a = self.find(a);
let mut b = self.find(b);
// If both keys have the same root, they're already in the same set,
// so there's nothing more to do.
if a == b {
return a;
};
// Ensure that `a` has strictly greater rank, swapping if necessary.
// If both keys have the same rank, increment the rank of `a` so that
// future unifications will also prefer `a`, leading to flatter trees.
match Ord::cmp(&self.table[a].rank, &self.table[b].rank) {
Ordering::Less => mem::swap(&mut a, &mut b),
Ordering::Equal => self.table[a].rank += 1,
Ordering::Greater => {}
}
debug_assert!(self.table[a].rank > self.table[b].rank);
debug_assert_eq!(self.table[b].parent, b);
// Make `a` the parent of `b`.
self.table[b].parent = a;
a
}
/// Creates a snapshot of this disjoint-set forest that can no longer be
/// mutated, but can be queried without mutation.
pub(crate) fn freeze(&mut self) -> FrozenUnionFind<Key> {
// Just resolve each key to its actual root.
let roots = self.table.indices().map(|key| self.find(key)).collect();
FrozenUnionFind { roots }
}
}
/// Snapshot of a disjoint-set forest that can no longer be mutated, but can be
/// queried in O(1) time without mutation.
///
/// This is really just a wrapper around a direct mapping from keys to roots,
/// but with a [`Self::find`] method that resembles [`UnionFind::find`].
#[derive(Debug)]
pub(crate) struct FrozenUnionFind<Key: Idx> {
roots: IndexVec<Key, Key>,
}
impl<Key: Idx> FrozenUnionFind<Key> {
/// Returns the "root" key of the disjoint-set containing the given key.
/// If two keys have the same root, they belong to the same set.
pub(crate) fn find(&self, key: Key) -> Key {
self.roots[key]
}
}

32
compiler/rustc_mir_transform/src/coverage/counters/union_find/tests.rs Normal file
View file

@ -0,0 +1,32 @@
use super::UnionFind;
#[test]
fn empty() {
let mut sets = UnionFind::<u32>::new(10);
for i in 1..10 {
assert_eq!(sets.find(i), i);
}
}
#[test]
fn transitive() {
let mut sets = UnionFind::<u32>::new(10);
sets.unify(3, 7);
sets.unify(4, 2);
assert_eq!(sets.find(7), sets.find(3));
assert_eq!(sets.find(2), sets.find(4));
assert_ne!(sets.find(3), sets.find(4));
sets.unify(7, 4);
assert_eq!(sets.find(7), sets.find(3));
assert_eq!(sets.find(2), sets.find(4));
assert_eq!(sets.find(3), sets.find(4));
for i in [0, 1, 5, 6, 8, 9] {
assert_eq!(sets.find(i), i);
}
}

174
compiler/rustc_mir_transform/src/coverage/graph.rs
View file

@ -1,7 +1,6 @@
use std::cmp::Ordering;
use std::collections::VecDeque;
use std::ops::{Index, IndexMut};
use std::{iter, mem, slice};
use std::{mem, slice};
use rustc_data_structures::captures::Captures;
use rustc_data_structures::fx::FxHashSet;
@ -211,54 +210,6 @@ impl CoverageGraph {
self.dominator_order_rank[a].cmp(&self.dominator_order_rank[b])
}
/// Returns the source of this node's sole in-edge, if it has exactly one.
/// That edge can be assumed to have the same execution count as the node
/// itself (in the absence of panics).
pub(crate) fn sole_predecessor(
&self,
to_bcb: BasicCoverageBlock,
) -> Option<BasicCoverageBlock> {
// Unlike `simple_successor`, there is no need for extra checks here.
if let &[from_bcb] = self.predecessors[to_bcb].as_slice() { Some(from_bcb) } else { None }
}
/// Returns the target of this node's sole out-edge, if it has exactly
/// one, but only if that edge can be assumed to have the same execution
/// count as the node itself (in the absence of panics).
pub(crate) fn simple_successor(
&self,
from_bcb: BasicCoverageBlock,
) -> Option<BasicCoverageBlock> {
// If a node's count is the sum of its out-edges, and it has exactly
// one out-edge, then that edge has the same count as the node.
if self.bcbs[from_bcb].is_out_summable
&& let &[to_bcb] = self.successors[from_bcb].as_slice()
{
Some(to_bcb)
} else {
None
}
}
/// For each loop that contains the given node, yields the "loop header"
/// node representing that loop, from innermost to outermost. If the given
/// node is itself a loop header, it is yielded first.
pub(crate) fn loop_headers_containing(
&self,
bcb: BasicCoverageBlock,
) -> impl Iterator<Item = BasicCoverageBlock> + Captures<'_> {
let self_if_loop_header = self.is_loop_header.contains(bcb).then_some(bcb).into_iter();
let mut curr = Some(bcb);
let strictly_enclosing = iter::from_fn(move || {
let enclosing = self.enclosing_loop_header[curr?];
curr = enclosing;
enclosing
});
self_if_loop_header.chain(strictly_enclosing)
}
/// For the given node, yields the subset of its predecessor nodes that
/// it dominates. If that subset is non-empty, the node is a "loop header",
/// and each of those predecessors represents an in-edge that jumps back to
@ -489,126 +440,3 @@ impl<'a, 'tcx> graph::Successors for CoverageRelevantSubgraph<'a, 'tcx> {
self.coverage_successors(bb).into_iter()
}
}
/// State of a node in the coverage graph during ready-first traversal.
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
enum ReadyState {
/// This node has not yet been added to the fallback queue or ready queue.
Unqueued,
/// This node is currently in the fallback queue.
InFallbackQueue,
/// This node's predecessors have all been visited, so it is in the ready queue.
/// (It might also have a stale entry in the fallback queue.)
InReadyQueue,
/// This node has been visited.
/// (It might also have a stale entry in the fallback queue.)
Visited,
}
/// Iterator that visits nodes in the coverage graph, in an order that always
/// prefers "ready" nodes whose predecessors have already been visited.
pub(crate) struct ReadyFirstTraversal<'a> {
graph: &'a CoverageGraph,
/// For each node, the number of its predecessor nodes that haven't been visited yet.
n_unvisited_preds: IndexVec<BasicCoverageBlock, u32>,
/// Indicates whether a node has been visited, or which queue it is in.
state: IndexVec<BasicCoverageBlock, ReadyState>,
/// Holds unvisited nodes whose predecessors have all been visited.
ready_queue: VecDeque<BasicCoverageBlock>,
/// Holds unvisited nodes with some unvisited predecessors.
/// Also contains stale entries for nodes that were upgraded to ready.
fallback_queue: VecDeque<BasicCoverageBlock>,
}
impl<'a> ReadyFirstTraversal<'a> {
pub(crate) fn new(graph: &'a CoverageGraph) -> Self {
let num_nodes = graph.num_nodes();
let n_unvisited_preds =
IndexVec::from_fn_n(|node| graph.predecessors[node].len() as u32, num_nodes);
let mut state = IndexVec::from_elem_n(ReadyState::Unqueued, num_nodes);
// We know from coverage graph construction that the start node is the
// only node with no predecessors.
debug_assert!(
n_unvisited_preds.iter_enumerated().all(|(node, &n)| (node == START_BCB) == (n == 0))
);
let ready_queue = VecDeque::from(vec![START_BCB]);
state[START_BCB] = ReadyState::InReadyQueue;
Self { graph, state, n_unvisited_preds, ready_queue, fallback_queue: VecDeque::new() }
}
/// Returns the next node from the ready queue, or else the next unvisited
/// node from the fallback queue.
fn next_inner(&mut self) -> Option<BasicCoverageBlock> {
// Always prefer to yield a ready node if possible.
if let Some(node) = self.ready_queue.pop_front() {
assert_eq!(self.state[node], ReadyState::InReadyQueue);
return Some(node);
}
while let Some(node) = self.fallback_queue.pop_front() {
match self.state[node] {
// This entry in the fallback queue is not stale, so yield it.
ReadyState::InFallbackQueue => return Some(node),
// This node was added to the fallback queue, but later became
// ready and was visited via the ready queue. Ignore it here.
ReadyState::Visited => {}
// Unqueued nodes can't be in the fallback queue, by definition.
// We know that the ready queue is empty at this point.
ReadyState::Unqueued | ReadyState::InReadyQueue => unreachable!(
"unexpected state for {node:?} in the fallback queue: {:?}",
self.state[node]
),
}
}
None
}
fn mark_visited_and_enqueue_successors(&mut self, node: BasicCoverageBlock) {
assert!(self.state[node] < ReadyState::Visited);
self.state[node] = ReadyState::Visited;
// For each of this node's successors, decrease the successor's
// "unvisited predecessors" count, and enqueue it if appropriate.
for &succ in &self.graph.successors[node] {
let is_unqueued = match self.state[succ] {
ReadyState::Unqueued => true,
ReadyState::InFallbackQueue => false,
ReadyState::InReadyQueue => {
unreachable!("nodes in the ready queue have no unvisited predecessors")
}
// The successor was already visited via one of its other predecessors.
ReadyState::Visited => continue,
};
self.n_unvisited_preds[succ] -= 1;
if self.n_unvisited_preds[succ] == 0 {
// This node's predecessors have all been visited, so add it to
// the ready queue. If it's already in the fallback queue, that
// fallback entry will be ignored later.
self.state[succ] = ReadyState::InReadyQueue;
self.ready_queue.push_back(succ);
} else if is_unqueued {
// This node has unvisited predecessors, so add it to the
// fallback queue in case we run out of ready nodes later.
self.state[succ] = ReadyState::InFallbackQueue;
self.fallback_queue.push_back(succ);
}
}
}
}
impl<'a> Iterator for ReadyFirstTraversal<'a> {
type Item = BasicCoverageBlock;
fn next(&mut self) -> Option<Self::Item> {
let node = self.next_inner()?;
self.mark_visited_and_enqueue_successors(node);
Some(node)
}
}

43
compiler/rustc_mir_transform/src/coverage/mod.rs
View file

@ -15,10 +15,7 @@ use rustc_middle::hir::nested_filter;
use rustc_middle::mir::coverage::{
CoverageKind, DecisionInfo, FunctionCoverageInfo, Mapping, MappingKind,
};
use rustc_middle::mir::{
self, BasicBlock, BasicBlockData, SourceInfo, Statement, StatementKind, Terminator,
TerminatorKind,
};
use rustc_middle::mir::{self, BasicBlock, Statement, StatementKind, TerminatorKind};
use rustc_middle::ty::TyCtxt;
use rustc_span::Span;
use rustc_span::def_id::LocalDefId;
@ -248,19 +245,6 @@ fn inject_coverage_statements<'tcx>(
// to create a new block between the two BCBs, and inject into that.
let target_bb = match site {
Site::Node { bcb } => graph[bcb].leader_bb(),
Site::Edge { from_bcb, to_bcb } => {
// Create a new block between the last block of `from_bcb` and
// the first block of `to_bcb`.
let from_bb = graph[from_bcb].last_bb();
let to_bb = graph[to_bcb].leader_bb();
let new_bb = inject_edge_counter_basic_block(mir_body, from_bb, to_bb);
debug!(
"Edge {from_bcb:?} (last {from_bb:?}) -> {to_bcb:?} (leader {to_bb:?}) \
requires a new MIR BasicBlock {new_bb:?} for counter increment {id:?}",
);
new_bb
}
};
inject_statement(mir_body, CoverageKind::CounterIncrement { id }, target_bb);
@ -335,31 +319,6 @@ fn inject_mcdc_statements<'tcx>(
}
}
/// Given two basic blocks that have a control-flow edge between them, creates
/// and returns a new block that sits between those blocks.
fn inject_edge_counter_basic_block(
mir_body: &mut mir::Body<'_>,
from_bb: BasicBlock,
to_bb: BasicBlock,
) -> BasicBlock {
let span = mir_body[from_bb].terminator().source_info.span.shrink_to_hi();
let new_bb = mir_body.basic_blocks_mut().push(BasicBlockData {
statements: vec![], // counter will be injected here
terminator: Some(Terminator {
source_info: SourceInfo::outermost(span),
kind: TerminatorKind::Goto { target: to_bb },
}),
is_cleanup: false,
});
let edge_ref = mir_body[from_bb]
.terminator_mut()
.successors_mut()
.find(|successor| **successor == to_bb)
.expect("from_bb should have a successor for to_bb");
*edge_ref = new_bb;
new_bb
}
fn inject_statement(mir_body: &mut mir::Body<'_>, counter_kind: CoverageKind, bb: BasicBlock) {
debug!(" injecting statement {counter_kind:?} for {bb:?}");
let data = &mut mir_body[bb];