diff --git a/compiler/rustc_middle/src/mir/coverage.rs b/compiler/rustc_middle/src/mir/coverage.rs index cd6b2d65bf184..65f51ae9d39c6 100644 --- a/compiler/rustc_middle/src/mir/coverage.rs +++ b/compiler/rustc_middle/src/mir/coverage.rs @@ -71,11 +71,7 @@ impl ConditionId { /// Enum that can hold a constant zero value, the ID of an physical coverage /// counter, or the ID of a coverage-counter expression. -/// -/// This was originally only used for expression operands (and named `Operand`), -/// but the zero/counter/expression distinction is also useful for representing -/// the value of code/gap mappings, and the true/false arms of branch mappings. -#[derive(Copy, Clone, PartialEq, Eq)] +#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord)] #[derive(TyEncodable, TyDecodable, Hash, HashStable, TypeFoldable, TypeVisitable)] pub enum CovTerm { Zero, @@ -171,7 +167,7 @@ impl Op { } } -#[derive(Clone, Debug)] +#[derive(Clone, Debug, PartialEq, Eq)] #[derive(TyEncodable, TyDecodable, Hash, HashStable, TypeFoldable, TypeVisitable)] pub struct Expression { pub lhs: CovTerm, diff --git a/compiler/rustc_mir_transform/src/coverage/counters.rs b/compiler/rustc_mir_transform/src/coverage/counters.rs index 1a9323329f61f..8d397f63cc7e0 100644 --- a/compiler/rustc_mir_transform/src/coverage/counters.rs +++ b/compiler/rustc_mir_transform/src/coverage/counters.rs @@ -1,10 +1,9 @@ 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::fx::{FxHashMap, FxIndexMap}; use rustc_data_structures::graph::DirectedGraph; use rustc_index::IndexVec; use rustc_index::bit_set::DenseBitSet; @@ -20,134 +19,163 @@ mod iter_nodes; mod node_flow; mod union_find; -/// The coverage counter or counter expression associated with a particular -/// BCB node or BCB edge. -#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] -enum BcbCounter { - Counter { id: CounterId }, - Expression { id: ExpressionId }, +/// Ensures that each BCB node needing a counter has one, by creating physical +/// counters or counter expressions for nodes as required. +pub(super) fn make_bcb_counters( + graph: &CoverageGraph, + bcb_needs_counter: &DenseBitSet, +) -> CoverageCounters { + // Create the derived graphs that are necessary for subsequent steps. + let balanced_graph = BalancedFlowGraph::for_graph(graph, |n| !graph[n].is_out_summable); + let merged_graph = MergedNodeFlowGraph::for_balanced_graph(&balanced_graph); + + // Use those graphs to determine which nodes get physical counters, and how + // to compute the execution counts of other nodes from those counters. + let nodes = make_node_counter_priority_list(graph, balanced_graph); + let node_counters = merged_graph.make_node_counters(&nodes); + + // Convert the counters into a form suitable for embedding into MIR. + transcribe_counters(&node_counters, bcb_needs_counter) } -impl BcbCounter { - fn as_term(&self) -> CovTerm { - match *self { - BcbCounter::Counter { id, .. } => CovTerm::Counter(id), - BcbCounter::Expression { id, .. } => CovTerm::Expression(id), - } - } +/// Arranges the nodes in `balanced_graph` into a list, such that earlier nodes +/// take priority in being given a counter expression instead of a physical counter. +fn make_node_counter_priority_list( + graph: &CoverageGraph, + balanced_graph: BalancedFlowGraph<&CoverageGraph>, +) -> Vec { + // 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::>(); + // 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()) + }); + nodes } -impl Debug for BcbCounter { - fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - Self::Counter { id, .. } => write!(fmt, "Counter({:?})", id.index()), - Self::Expression { id } => write!(fmt, "Expression({:?})", id.index()), +// Converts node counters into a form suitable for embedding into MIR. +fn transcribe_counters( + old: &NodeCounters, + bcb_needs_counter: &DenseBitSet, +) -> CoverageCounters { + let mut new = CoverageCounters::with_num_bcbs(bcb_needs_counter.domain_size()); + + for bcb in bcb_needs_counter.iter() { + // Our counter-creation algorithm doesn't guarantee that a counter + // expression starts or ends with a positive term, so partition the + // counters into "positive" and "negative" lists for easier handling. + let (mut pos, mut neg): (Vec<_>, Vec<_>) = + 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, 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, "{bcb:?} has no positive counter terms"); + pos = vec![bcb]; + neg = vec![]; } - } -} -#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] -struct BcbExpression { - lhs: BcbCounter, - op: Op, - rhs: BcbCounter, -} + // These intermediate sorts are not strictly necessary, but were helpful + // in reducing churn when switching to the current counter-creation scheme. + // They also help to slightly decrease the overall size of the expression + // table, due to more subexpressions being shared. + pos.sort(); + neg.sort(); + + let mut new_counters_for_sites = |sites: Vec| { + sites.into_iter().map(|node| new.ensure_phys_counter(node)).collect::>() + }; + let mut pos = new_counters_for_sites(pos); + let mut neg = new_counters_for_sites(neg); + + // These sorts are also not strictly necessary; see above. + pos.sort(); + neg.sort(); + + let pos_counter = new.make_sum(&pos).expect("`pos` should not be empty"); + let new_counter = new.make_subtracted_sum(pos_counter, &neg); + new.set_node_counter(bcb, new_counter); + } -/// 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 }, + new } /// Generates and stores coverage counter and coverage expression information -/// associated with nodes/edges in the BCB graph. +/// associated with nodes in the coverage graph. pub(super) struct CoverageCounters { /// List of places where a counter-increment statement should be injected /// into MIR, each with its corresponding counter ID. - counter_increment_sites: IndexVec, + phys_counter_for_node: FxIndexMap, + next_counter_id: CounterId, /// Coverage counters/expressions that are associated with individual BCBs. - node_counters: IndexVec>, + node_counters: IndexVec>, /// Table of expression data, associating each expression ID with its /// corresponding operator (+ or -) and its LHS/RHS operands. - expressions: IndexVec, + expressions: IndexVec, /// Remember expressions that have already been created (or simplified), /// so that we don't create unnecessary duplicates. - expressions_memo: FxHashMap, + expressions_memo: FxHashMap, } impl CoverageCounters { - /// Ensures that each BCB node needing a counter has one, by creating physical - /// counters or counter expressions for nodes and edges as required. - pub(super) fn make_bcb_counters( - graph: &CoverageGraph, - bcb_needs_counter: &DenseBitSet, - ) -> Self { - let balanced_graph = BalancedFlowGraph::for_graph(graph, |n| !graph[n].is_out_summable); - let merged_graph = MergedNodeFlowGraph::for_balanced_graph(&balanced_graph); - - // 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::>(); - // 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 { Self { - counter_increment_sites: IndexVec::new(), + phys_counter_for_node: FxIndexMap::default(), + next_counter_id: CounterId::ZERO, node_counters: IndexVec::from_elem_n(None, num_bcbs), expressions: IndexVec::new(), expressions_memo: FxHashMap::default(), } } - /// Creates a new physical counter for a BCB node or edge. - fn make_phys_counter(&mut self, site: Site) -> BcbCounter { - let id = self.counter_increment_sites.push(site); - BcbCounter::Counter { id } + /// Returns the physical counter for the given node, creating it if necessary. + fn ensure_phys_counter(&mut self, bcb: BasicCoverageBlock) -> CovTerm { + let id = *self.phys_counter_for_node.entry(bcb).or_insert_with(|| { + let id = self.next_counter_id; + self.next_counter_id = id + 1; + id + }); + CovTerm::Counter(id) } - fn make_expression(&mut self, lhs: BcbCounter, op: Op, rhs: BcbCounter) -> BcbCounter { - let new_expr = BcbExpression { lhs, op, rhs }; - *self.expressions_memo.entry(new_expr).or_insert_with(|| { + fn make_expression(&mut self, lhs: CovTerm, op: Op, rhs: CovTerm) -> CovTerm { + let new_expr = Expression { lhs, op, rhs }; + *self.expressions_memo.entry(new_expr.clone()).or_insert_with(|| { let id = self.expressions.push(new_expr); - BcbCounter::Expression { id } + CovTerm::Expression(id) }) } /// Creates a counter that is the sum of the given counters. /// /// Returns `None` if the given list of counters was empty. - fn make_sum(&mut self, counters: &[BcbCounter]) -> Option { + fn make_sum(&mut self, counters: &[CovTerm]) -> Option { counters .iter() .copied() @@ -155,16 +183,18 @@ impl CoverageCounters { } /// Creates a counter whose value is `lhs - SUM(rhs)`. - fn make_subtracted_sum(&mut self, lhs: BcbCounter, rhs: &[BcbCounter]) -> BcbCounter { + fn make_subtracted_sum(&mut self, lhs: CovTerm, rhs: &[CovTerm]) -> CovTerm { let Some(rhs_sum) = self.make_sum(rhs) else { return lhs }; self.make_expression(lhs, Op::Subtract, rhs_sum) } pub(super) fn num_counters(&self) -> usize { - self.counter_increment_sites.len() + let num_counters = self.phys_counter_for_node.len(); + assert_eq!(num_counters, self.next_counter_id.as_usize()); + num_counters } - fn set_node_counter(&mut self, bcb: BasicCoverageBlock, counter: BcbCounter) -> BcbCounter { + fn set_node_counter(&mut self, bcb: BasicCoverageBlock, counter: CovTerm) -> CovTerm { let existing = self.node_counters[bcb].replace(counter); assert!( existing.is_none(), @@ -174,16 +204,16 @@ impl CoverageCounters { } pub(super) fn term_for_bcb(&self, bcb: BasicCoverageBlock) -> Option { - self.node_counters[bcb].map(|counter| counter.as_term()) + self.node_counters[bcb] } - /// Returns an iterator over all the nodes/edges in the coverage graph that + /// Returns an iterator over all the nodes in the coverage graph that /// should have a counter-increment statement injected into MIR, along with /// each site's corresponding counter ID. pub(super) fn counter_increment_sites( &self, - ) -> impl Iterator + Captures<'_> { - self.counter_increment_sites.iter_enumerated().map(|(id, &site)| (id, site)) + ) -> impl Iterator + Captures<'_> { + self.phys_counter_for_node.iter().map(|(&site, &id)| (id, site)) } /// Returns an iterator over the subset of BCB nodes that have been associated @@ -193,93 +223,13 @@ impl CoverageCounters { ) -> impl Iterator + Captures<'_> { self.node_counters.iter_enumerated().filter_map(|(bcb, &counter)| match counter { // Yield the BCB along with its associated expression ID. - Some(BcbCounter::Expression { id }) => Some((bcb, id)), + Some(CovTerm::Expression(id)) => Some((bcb, id)), // This BCB is associated with a counter or nothing, so skip it. - Some(BcbCounter::Counter { .. }) | None => None, + Some(CovTerm::Counter { .. } | CovTerm::Zero) | None => None, }) } pub(super) fn into_expressions(self) -> IndexVec { - let old_len = self.expressions.len(); - let expressions = self - .expressions - .into_iter() - .map(|BcbExpression { lhs, op, rhs }| Expression { - lhs: lhs.as_term(), - op, - rhs: rhs.as_term(), - }) - .collect::>(); - - // Expression IDs are indexes into this vector, so make sure we didn't - // accidentally invalidate them by changing its length. - assert_eq!(old_len, expressions.len()); - expressions - } -} - -struct Transcriber { - old: NodeCounters, - new: CoverageCounters, - phys_counter_for_site: FxHashMap, -} - -impl Transcriber { - fn new(num_nodes: usize, old: NodeCounters) -> Self { - Self { - old, - new: CoverageCounters::with_num_bcbs(num_nodes), - phys_counter_for_site: FxHashMap::default(), - } - } - - fn transcribe_counters( - mut self, - bcb_needs_counter: &DenseBitSet, - ) -> CoverageCounters { - for bcb in bcb_needs_counter.iter() { - let site = Site::Node { bcb }; - 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, 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![bcb]; - neg = vec![]; - } - - pos.sort(); - neg.sort(); - - let mut new_counters_for_sites = |sites: Vec| { - sites - .into_iter() - .map(|node| self.ensure_phys_counter(Site::Node { bcb: node })) - .collect::>() - }; - let mut pos = new_counters_for_sites(pos); - let mut neg = new_counters_for_sites(neg); - - pos.sort(); - neg.sort(); - - let pos_counter = self.new.make_sum(&pos).expect("`pos` should not be empty"); - let new_counter = self.new.make_subtracted_sum(pos_counter, &neg); - self.new.set_node_counter(bcb, new_counter); - } - - self.new - } - - 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)) + self.expressions } } diff --git a/compiler/rustc_mir_transform/src/coverage/mod.rs b/compiler/rustc_mir_transform/src/coverage/mod.rs index b1b609595b74a..19568735df76e 100644 --- a/compiler/rustc_mir_transform/src/coverage/mod.rs +++ b/compiler/rustc_mir_transform/src/coverage/mod.rs @@ -21,7 +21,7 @@ use rustc_span::Span; use rustc_span::def_id::LocalDefId; use tracing::{debug, debug_span, trace}; -use crate::coverage::counters::{CoverageCounters, Site}; +use crate::coverage::counters::CoverageCounters; use crate::coverage::graph::CoverageGraph; use crate::coverage::mappings::ExtractedMappings; @@ -89,8 +89,7 @@ fn instrument_function_for_coverage<'tcx>(tcx: TyCtxt<'tcx>, mir_body: &mut mir: return; } - let coverage_counters = - CoverageCounters::make_bcb_counters(&graph, &bcbs_with_counter_mappings); + let coverage_counters = counters::make_bcb_counters(&graph, &bcbs_with_counter_mappings); let mappings = create_mappings(&extracted_mappings, &coverage_counters); if mappings.is_empty() { @@ -239,14 +238,8 @@ fn inject_coverage_statements<'tcx>( coverage_counters: &CoverageCounters, ) { // Inject counter-increment statements into MIR. - for (id, site) in coverage_counters.counter_increment_sites() { - // Determine the block to inject a counter-increment statement into. - // For BCB nodes this is just their first block, but for edges we need - // 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(), - }; - + for (id, bcb) in coverage_counters.counter_increment_sites() { + let target_bb = graph[bcb].leader_bb(); inject_statement(mir_body, CoverageKind::CounterIncrement { id }, target_bb); }