rustc_mir_transform/coverage/
counters.rs

1use std::cmp::Ordering;
2
3use either::Either;
4use itertools::Itertools;
5use rustc_data_structures::fx::{FxHashMap, FxIndexMap};
6use rustc_data_structures::graph::DirectedGraph;
7use rustc_index::IndexVec;
8use rustc_index::bit_set::DenseBitSet;
9use rustc_middle::mir::coverage::{CounterId, CovTerm, Expression, ExpressionId, Op};
10
11use crate::coverage::counters::balanced_flow::BalancedFlowGraph;
12use crate::coverage::counters::node_flow::{
13    CounterTerm, NodeCounters, NodeFlowData, node_flow_data_for_balanced_graph,
14};
15use crate::coverage::graph::{BasicCoverageBlock, CoverageGraph};
16
17mod balanced_flow;
18pub(crate) mod node_flow;
19
20/// Struct containing the results of [`prepare_bcb_counters_data`].
21pub(crate) struct BcbCountersData {
22    pub(crate) node_flow_data: NodeFlowData<BasicCoverageBlock>,
23    pub(crate) priority_list: Vec<BasicCoverageBlock>,
24}
25
26/// Analyzes the coverage graph to create intermediate data structures that
27/// will later be used (during codegen) to create physical counters or counter
28/// expressions for each BCB node that needs one.
29pub(crate) fn prepare_bcb_counters_data(graph: &CoverageGraph) -> BcbCountersData {
30    // Create the derived graphs that are necessary for subsequent steps.
31    let balanced_graph = BalancedFlowGraph::for_graph(graph, |n| !graph[n].is_out_summable);
32    let node_flow_data = node_flow_data_for_balanced_graph(&balanced_graph);
33
34    // Also create a "priority list" of coverage graph nodes, to help determine
35    // which ones get physical counters or counter expressions. This needs to
36    // be done now, because the later parts of the counter-creation process
37    // won't have access to the original coverage graph.
38    let priority_list = make_node_flow_priority_list(graph, balanced_graph);
39
40    BcbCountersData { node_flow_data, priority_list }
41}
42
43/// Arranges the nodes in `balanced_graph` into a list, such that earlier nodes
44/// take priority in being given a counter expression instead of a physical counter.
45fn make_node_flow_priority_list(
46    graph: &CoverageGraph,
47    balanced_graph: BalancedFlowGraph<&CoverageGraph>,
48) -> Vec<BasicCoverageBlock> {
49    // A "reloop" node has exactly one out-edge, which jumps back to the top
50    // of an enclosing loop. Reloop nodes are typically visited more times
51    // than loop-exit nodes, so try to avoid giving them physical counters.
52    let is_reloop_node = IndexVec::<BasicCoverageBlock, _>::from_fn_n(
53        |node| match graph.successors[node].as_slice() {
54            &[succ] => graph.dominates(succ, node),
55            _ => false,
56        },
57        graph.num_nodes(),
58    );
59
60    let mut nodes = balanced_graph.iter_nodes().rev().collect::<Vec<_>>();
61    // The first node is the sink, which must not get a physical counter.
62    assert_eq!(nodes[0], balanced_graph.sink);
63    // Sort the real nodes, such that earlier (lesser) nodes take priority
64    // in being given a counter expression instead of a physical counter.
65    nodes[1..].sort_by(|&a, &b| {
66        // Start with a dummy `Equal` to make the actual tests line up nicely.
67        Ordering::Equal
68            // Prefer a physical counter for return/yield nodes.
69            .then_with(|| Ord::cmp(&graph[a].is_out_summable, &graph[b].is_out_summable))
70            // Prefer an expression for reloop nodes (see definition above).
71            .then_with(|| Ord::cmp(&is_reloop_node[a], &is_reloop_node[b]).reverse())
72            // Otherwise, prefer a physical counter for dominating nodes.
73            .then_with(|| graph.cmp_in_dominator_order(a, b).reverse())
74    });
75    nodes
76}
77
78// Converts node counters into a form suitable for embedding into MIR.
79pub(crate) fn transcribe_counters(
80    old: &NodeCounters<BasicCoverageBlock>,
81    bcb_needs_counter: &DenseBitSet<BasicCoverageBlock>,
82    bcbs_seen: &DenseBitSet<BasicCoverageBlock>,
83) -> CoverageCounters {
84    let mut new = CoverageCounters::with_num_bcbs(bcb_needs_counter.domain_size());
85
86    for bcb in bcb_needs_counter.iter() {
87        if !bcbs_seen.contains(bcb) {
88            // This BCB's code was removed by MIR opts, so its counter is always zero.
89            new.set_node_counter(bcb, CovTerm::Zero);
90            continue;
91        }
92
93        // Our counter-creation algorithm doesn't guarantee that a node's list
94        // of terms starts or ends with a positive term, so partition the
95        // counters into "positive" and "negative" lists for easier handling.
96        let (mut pos, mut neg): (Vec<_>, Vec<_>) = old.counter_terms[bcb]
97            .iter()
98            // Filter out any BCBs that were removed by MIR opts;
99            // this treats them as having an execution count of 0.
100            .filter(|term| bcbs_seen.contains(term.node))
101            .partition_map(|&CounterTerm { node, op }| match op {
102                Op::Add => Either::Left(node),
103                Op::Subtract => Either::Right(node),
104            });
105
106        // These intermediate sorts are not strictly necessary, but were helpful
107        // in reducing churn when switching to the current counter-creation scheme.
108        // They also help to slightly decrease the overall size of the expression
109        // table, due to more subexpressions being shared.
110        pos.sort();
111        neg.sort();
112
113        let mut new_counters_for_sites = |sites: Vec<BasicCoverageBlock>| {
114            sites.into_iter().map(|node| new.ensure_phys_counter(node)).collect::<Vec<_>>()
115        };
116        let pos = new_counters_for_sites(pos);
117        let neg = new_counters_for_sites(neg);
118
119        let pos_counter = new.make_sum(&pos).unwrap_or(CovTerm::Zero);
120        let new_counter = new.make_subtracted_sum(pos_counter, &neg);
121        new.set_node_counter(bcb, new_counter);
122    }
123
124    new
125}
126
127/// Generates and stores coverage counter and coverage expression information
128/// associated with nodes in the coverage graph.
129pub(super) struct CoverageCounters {
130    /// List of places where a counter-increment statement should be injected
131    /// into MIR, each with its corresponding counter ID.
132    pub(crate) phys_counter_for_node: FxIndexMap<BasicCoverageBlock, CounterId>,
133    pub(crate) next_counter_id: CounterId,
134
135    /// Coverage counters/expressions that are associated with individual BCBs.
136    pub(crate) node_counters: IndexVec<BasicCoverageBlock, Option<CovTerm>>,
137
138    /// Table of expression data, associating each expression ID with its
139    /// corresponding operator (+ or -) and its LHS/RHS operands.
140    pub(crate) expressions: IndexVec<ExpressionId, Expression>,
141    /// Remember expressions that have already been created (or simplified),
142    /// so that we don't create unnecessary duplicates.
143    expressions_memo: FxHashMap<Expression, CovTerm>,
144}
145
146impl CoverageCounters {
147    fn with_num_bcbs(num_bcbs: usize) -> Self {
148        Self {
149            phys_counter_for_node: FxIndexMap::default(),
150            next_counter_id: CounterId::ZERO,
151            node_counters: IndexVec::from_elem_n(None, num_bcbs),
152            expressions: IndexVec::new(),
153            expressions_memo: FxHashMap::default(),
154        }
155    }
156
157    /// Returns the physical counter for the given node, creating it if necessary.
158    fn ensure_phys_counter(&mut self, bcb: BasicCoverageBlock) -> CovTerm {
159        let id = *self.phys_counter_for_node.entry(bcb).or_insert_with(|| {
160            let id = self.next_counter_id;
161            self.next_counter_id = id + 1;
162            id
163        });
164        CovTerm::Counter(id)
165    }
166
167    fn make_expression(&mut self, lhs: CovTerm, op: Op, rhs: CovTerm) -> CovTerm {
168        let new_expr = Expression { lhs, op, rhs };
169        *self.expressions_memo.entry(new_expr.clone()).or_insert_with(|| {
170            let id = self.expressions.push(new_expr);
171            CovTerm::Expression(id)
172        })
173    }
174
175    /// Creates a counter that is the sum of the given counters.
176    ///
177    /// Returns `None` if the given list of counters was empty.
178    fn make_sum(&mut self, counters: &[CovTerm]) -> Option<CovTerm> {
179        counters
180            .iter()
181            .copied()
182            .reduce(|accum, counter| self.make_expression(accum, Op::Add, counter))
183    }
184
185    /// Creates a counter whose value is `lhs - SUM(rhs)`.
186    fn make_subtracted_sum(&mut self, lhs: CovTerm, rhs: &[CovTerm]) -> CovTerm {
187        let Some(rhs_sum) = self.make_sum(rhs) else { return lhs };
188        self.make_expression(lhs, Op::Subtract, rhs_sum)
189    }
190
191    fn set_node_counter(&mut self, bcb: BasicCoverageBlock, counter: CovTerm) -> CovTerm {
192        let existing = self.node_counters[bcb].replace(counter);
193        assert!(
194            existing.is_none(),
195            "node {bcb:?} already has a counter: {existing:?} => {counter:?}"
196        );
197        counter
198    }
199}
rustc_mir_transform/coverage/counters.rs

rustc_mir_transform/coverage/
counters.rs
