Rollup merge of #149587 - tree-sort, r=davidtwco

coverage: Sort the expansion tree to help choose a single BCB for child expansions

This PR is another incremental step on the road towards proper coverage instrumentation of expansion regions.

When creating coverage mappings for each relevant span in a function body, the current implementation also needs to do a separate tree traversal for each child expansion (e.g. macro calls like `println!("foo")`), in order to associate a single control-flow point (BCB) with that macro call's span.

This PR changes things so that we now keep track of the “minimum” and ”maximum” BCB associated with each node in the expansion tree, which makes it much easier for the parent node to get a single BCB (min or max) for each of its child expansions.

In order to make this (relatively) easy, we first need to sort the tree nodes into depth-first order. Once that's taken care of, we can iterate over all the nodes in reverse order, knowing that each node's children will have been visited before visiting the node itself.

Author: Zalathar

compiler/rustc_mir_transform/src/coverage/expansion.rs

@@ -1,3 +1,4 @@
+use itertools::Itertools;
 use rustc_data_structures::fx::{FxIndexMap, FxIndexSet, IndexEntry};
 use rustc_middle::mir;
 use rustc_middle::mir::coverage::{BasicCoverageBlock, BranchSpan};
@@ -6,6 +7,7 @@ use rustc_span::{ExpnId, ExpnKind, Span};
 use crate::coverage::from_mir;
 use crate::coverage::graph::CoverageGraph;
 use crate::coverage::hir_info::ExtractedHirInfo;
+use crate::coverage::mappings::MappingsError;
 
 #[derive(Clone, Copy, Debug)]
 pub(crate) struct SpanWithBcb {
@@ -22,38 +24,6 @@ impl ExpnTree {
     pub(crate) fn get(&self, expn_id: ExpnId) -> Option<&ExpnNode> {
         self.nodes.get(&expn_id)
     }
-
-    /// Yields the tree node for the given expansion ID (if present), followed
-    /// by the nodes of all of its descendants in depth-first order.
-    pub(crate) fn iter_node_and_descendants(
-        &self,
-        root_expn_id: ExpnId,
-    ) -> impl Iterator<Item = &ExpnNode> {
-        gen move {
-            let Some(root_node) = self.get(root_expn_id) else { return };
-            yield root_node;
-
-            // Stack of child-node-ID iterators that drives the depth-first traversal.
-            let mut iter_stack = vec![root_node.child_expn_ids.iter()];
-
-            while let Some(curr_iter) = iter_stack.last_mut() {
-                // Pull the next ID from the top of the stack.
-                let Some(&curr_id) = curr_iter.next() else {
-                    iter_stack.pop();
-                    continue;
-                };
-
-                // Yield this node.
-                let Some(node) = self.get(curr_id) else { continue };
-                yield node;
-
-                // Push the node's children, to be traversed next.
-                if !node.child_expn_ids.is_empty() {
-                    iter_stack.push(node.child_expn_ids.iter());
-                }
-            }
-        }
-    }
 }
 
 #[derive(Debug)]
@@ -62,6 +32,8 @@ pub(crate) struct ExpnNode {
     /// but is helpful for debugging and might be useful later.
     #[expect(dead_code)]
     pub(crate) expn_id: ExpnId,
+    /// Index of this node in a depth-first traversal from the root.
+    pub(crate) dfs_rank: usize,
 
     // Useful info extracted from `ExpnData`.
     pub(crate) expn_kind: ExpnKind,
@@ -82,6 +54,10 @@ pub(crate) struct ExpnNode {
     pub(crate) spans: Vec<SpanWithBcb>,
     /// Expansions whose call-site is in this expansion.
     pub(crate) child_expn_ids: FxIndexSet<ExpnId>,
+    /// The "minimum" and "maximum" BCBs (in dominator order) of ordinary spans
+    /// belonging to this tree node and all of its descendants. Used when
+    /// creating a single code mapping representing an entire child expansion.
+    pub(crate) minmax_bcbs: Option<MinMaxBcbs>,
 
     /// Branch spans (recorded during MIR building) belonging to this expansion.
     pub(crate) branch_spans: Vec<BranchSpan>,
@@ -100,6 +76,7 @@ impl ExpnNode {
 
         Self {
             expn_id,
+            dfs_rank: usize::MAX,
 
             expn_kind: expn_data.kind,
             call_site,
@@ -110,6 +87,7 @@ impl ExpnNode {
 
             spans: vec![],
             child_expn_ids: FxIndexSet::default(),
+            minmax_bcbs: None,
 
             branch_spans: vec![],
 
@@ -124,7 +102,7 @@ pub(crate) fn build_expn_tree(
     mir_body: &mir::Body<'_>,
     hir_info: &ExtractedHirInfo,
     graph: &CoverageGraph,
-) -> ExpnTree {
+) -> Result<ExpnTree, MappingsError> {
     let raw_spans = from_mir::extract_raw_spans_from_mir(mir_body, graph);
 
     let mut nodes = FxIndexMap::default();
@@ -157,6 +135,20 @@ pub(crate) fn build_expn_tree(
         }
     }
 
+    // Sort the tree nodes into depth-first order.
+    sort_nodes_depth_first(&mut nodes)?;
+
+    // For each node, determine its "minimum" and "maximum" BCBs, based on its
+    // own spans and its immediate children. This relies on the nodes having
+    // been sorted, so that each node's children are processed before the node
+    // itself.
+    for i in (0..nodes.len()).rev() {
+        // Computing a node's min/max BCBs requires a shared ref to other nodes.
+        let minmax_bcbs = minmax_bcbs_for_expn_tree_node(graph, &nodes, &nodes[i]);
+        // Now we can mutate the current node to set its min/max BCBs.
+        nodes[i].minmax_bcbs = minmax_bcbs;
+    }
+
     // If we have a span for the function signature, associate it with the
     // corresponding expansion tree node.
     if let Some(fn_sig_span) = hir_info.fn_sig_span
@@ -189,5 +181,60 @@ pub(crate) fn build_expn_tree(
         }
     }
 
-    ExpnTree { nodes }
+    Ok(ExpnTree { nodes })
+}
+
+/// Sorts the tree nodes in the map into depth-first order.
+///
+/// This allows subsequent operations to iterate over all nodes, while assuming
+/// that every node occurs before all of its descendants.
+fn sort_nodes_depth_first(nodes: &mut FxIndexMap<ExpnId, ExpnNode>) -> Result<(), MappingsError> {
+    let mut dfs_stack = vec![ExpnId::root()];
+    let mut next_dfs_rank = 0usize;
+    while let Some(expn_id) = dfs_stack.pop() {
+        if let Some(node) = nodes.get_mut(&expn_id) {
+            node.dfs_rank = next_dfs_rank;
+            next_dfs_rank += 1;
+            dfs_stack.extend(node.child_expn_ids.iter().rev().copied());
+        }
+    }
+    nodes.sort_by_key(|_expn_id, node| node.dfs_rank);
+
+    // Verify that the depth-first search visited each node exactly once.
+    for (i, &ExpnNode { dfs_rank, .. }) in nodes.values().enumerate() {
+        if dfs_rank != i {
+            tracing::debug!(dfs_rank, i, "expansion tree node's rank does not match its index");
+            return Err(MappingsError::TreeSortFailure);
+        }
+    }
+
+    Ok(())
+}
+
+#[derive(Clone, Copy, Debug)]
+pub(crate) struct MinMaxBcbs {
+    pub(crate) min: BasicCoverageBlock,
+    pub(crate) max: BasicCoverageBlock,
+}
+
+/// For a single node in the expansion tree, compute its "minimum" and "maximum"
+/// BCBs (in dominator order), from among the BCBs of its immediate spans,
+/// and the min/max of its immediate children.
+fn minmax_bcbs_for_expn_tree_node(
+    graph: &CoverageGraph,
+    nodes: &FxIndexMap<ExpnId, ExpnNode>,
+    node: &ExpnNode,
+) -> Option<MinMaxBcbs> {
+    let immediate_span_bcbs = node.spans.iter().map(|sp: &SpanWithBcb| sp.bcb);
+    let child_minmax_bcbs = node
+        .child_expn_ids
+        .iter()
+        .flat_map(|id| nodes.get(id))
+        .flat_map(|child| child.minmax_bcbs)
+        .flat_map(|MinMaxBcbs { min, max }| [min, max]);
+
+    let (min, max) = Iterator::chain(immediate_span_bcbs, child_minmax_bcbs)
+        .minmax_by(|&a, &b| graph.cmp_in_dominator_order(a, b))
+        .into_option()?;
+    Some(MinMaxBcbs { min, max })
 }

compiler/rustc_mir_transform/src/coverage/mappings.rs

@@ -11,6 +11,13 @@ use crate::coverage::graph::CoverageGraph;
 use crate::coverage::hir_info::ExtractedHirInfo;
 use crate::coverage::spans::extract_refined_covspans;
 
+/// Indicates why mapping extraction failed, for debug-logging purposes.
+#[derive(Debug)]
+pub(crate) enum MappingsError {
+    NoMappings,
+    TreeSortFailure,
+}
+
 #[derive(Default)]
 pub(crate) struct ExtractedMappings {
     pub(crate) mappings: Vec<Mapping>,
@@ -23,8 +30,8 @@ pub(crate) fn extract_mappings_from_mir<'tcx>(
     mir_body: &mir::Body<'tcx>,
     hir_info: &ExtractedHirInfo,
     graph: &CoverageGraph,
-) -> ExtractedMappings {
-    let expn_tree = expansion::build_expn_tree(mir_body, hir_info, graph);
+) -> Result<ExtractedMappings, MappingsError> {
+    let expn_tree = expansion::build_expn_tree(mir_body, hir_info, graph)?;
 
     let mut mappings = vec![];
 
@@ -33,7 +40,11 @@ pub(crate) fn extract_mappings_from_mir<'tcx>(
 
     extract_branch_mappings(mir_body, hir_info, graph, &expn_tree, &mut mappings);
 
-    ExtractedMappings { mappings }
+    if mappings.is_empty() {
+        tracing::debug!("no mappings were extracted");
+        return Err(MappingsError::NoMappings);
+    }
+    Ok(ExtractedMappings { mappings })
 }
 
 fn resolve_block_markers(

compiler/rustc_mir_transform/src/coverage/mod.rs

@@ -73,12 +73,13 @@ fn instrument_function_for_coverage<'tcx>(tcx: TyCtxt<'tcx>, mir_body: &mut mir:
     ////////////////////////////////////////////////////
     // Extract coverage spans and other mapping info from MIR.
     let ExtractedMappings { mappings } =
-        mappings::extract_mappings_from_mir(tcx, mir_body, &hir_info, &graph);
-    if mappings.is_empty() {
-        // No spans could be converted into valid mappings, so skip this function.
-        debug!("no spans could be converted into valid mappings; skipping");
-        return;
-    }
+        match mappings::extract_mappings_from_mir(tcx, mir_body, &hir_info, &graph) {
+            Ok(m) => m,
+            Err(error) => {
+                tracing::debug!(?error, "mapping extraction failed; skipping this function");
+                return;
+            }
+        };
 
     // Use the coverage graph to prepare intermediate data that will eventually
     // be used to assign physical counters and counter expressions to points in

compiler/rustc_mir_transform/src/coverage/spans.rs

@@ -39,8 +39,7 @@ pub(super) fn extract_refined_covspans<'tcx>(
     // For each expansion with its call-site in the body span, try to
     // distill a corresponding covspan.
     for &child_expn_id in &node.child_expn_ids {
-        if let Some(covspan) = single_covspan_for_child_expn(tcx, graph, &expn_tree, child_expn_id)
-        {
+        if let Some(covspan) = single_covspan_for_child_expn(tcx, &expn_tree, child_expn_id) {
             covspans.push(covspan);
         }
     }
@@ -127,24 +126,21 @@ pub(super) fn extract_refined_covspans<'tcx>(
 /// For a single child expansion, try to distill it into a single span+BCB mapping.
 fn single_covspan_for_child_expn(
     tcx: TyCtxt<'_>,
-    graph: &CoverageGraph,
     expn_tree: &ExpnTree,
     expn_id: ExpnId,
 ) -> Option<Covspan> {
     let node = expn_tree.get(expn_id)?;
-
-    let bcbs =
-        expn_tree.iter_node_and_descendants(expn_id).flat_map(|n| n.spans.iter().map(|s| s.bcb));
+    let minmax_bcbs = node.minmax_bcbs?;
 
     let bcb = match node.expn_kind {
         // For bang-macros (e.g. `assert!`, `trace!`) and for `await`, taking
         // the "first" BCB in dominator order seems to give good results.
         ExpnKind::Macro(MacroKind::Bang, _) | ExpnKind::Desugaring(DesugaringKind::Await) => {
-            bcbs.min_by(|&a, &b| graph.cmp_in_dominator_order(a, b))?
+            minmax_bcbs.min
         }
         // For other kinds of expansion, taking the "last" (most-dominated) BCB
         // seems to give good results.
-        _ => bcbs.max_by(|&a, &b| graph.cmp_in_dominator_order(a, b))?,
+        _ => minmax_bcbs.max,
     };
 
     // For bang-macro expansions, limit the call-site span to just the macro

compiler/rustc_mir_transform/src/lib.rs

@@ -5,7 +5,6 @@
 #![feature(const_type_name)]
 #![feature(cow_is_borrowed)]
 #![feature(file_buffered)]
-#![feature(gen_blocks)]
 #![feature(if_let_guard)]
 #![feature(impl_trait_in_assoc_type)]
 #![feature(try_blocks)]

tests/coverage/macros/pass-through.cov-map

@@ -0,0 +1,57 @@
+Function name: pass_through::uses_inner_macro
+Raw bytes (51): 0x[01, 01, 01, 01, 05, 09, 01, 24, 01, 00, 16, 01, 01, 08, 00, 1d, 05, 01, 09, 00, 0c, 05, 00, 0d, 00, 21, 05, 01, 09, 00, 15, 05, 01, 09, 00, 0c, 05, 00, 0d, 00, 20, 02, 01, 05, 00, 06, 01, 01, 01, 00, 02]
+Number of files: 1
+- file 0 => $DIR/pass-through.rs
+Number of expressions: 1
+- expression 0 operands: lhs = Counter(0), rhs = Counter(1)
+Number of file 0 mappings: 9
+- Code(Counter(0)) at (prev + 36, 1) to (start + 0, 22)
+- Code(Counter(0)) at (prev + 1, 8) to (start + 0, 29)
+- Code(Counter(1)) at (prev + 1, 9) to (start + 0, 12)
+- Code(Counter(1)) at (prev + 0, 13) to (start + 0, 33)
+- Code(Counter(1)) at (prev + 1, 9) to (start + 0, 21)
+- Code(Counter(1)) at (prev + 1, 9) to (start + 0, 12)
+- Code(Counter(1)) at (prev + 0, 13) to (start + 0, 32)
+- Code(Expression(0, Sub)) at (prev + 1, 5) to (start + 0, 6)
+    = (c0 - c1)
+- Code(Counter(0)) at (prev + 1, 1) to (start + 0, 2)
+Highest counter ID seen: c1
+
+Function name: pass_through::uses_middle_macro
+Raw bytes (51): 0x[01, 01, 01, 01, 05, 09, 01, 2c, 01, 00, 17, 01, 01, 08, 00, 1d, 05, 01, 09, 00, 0c, 05, 00, 0d, 00, 22, 05, 01, 09, 00, 16, 05, 01, 09, 00, 0c, 05, 00, 0d, 00, 21, 02, 01, 05, 00, 06, 01, 01, 01, 00, 02]
+Number of files: 1
+- file 0 => $DIR/pass-through.rs
+Number of expressions: 1
+- expression 0 operands: lhs = Counter(0), rhs = Counter(1)
+Number of file 0 mappings: 9
+- Code(Counter(0)) at (prev + 44, 1) to (start + 0, 23)
+- Code(Counter(0)) at (prev + 1, 8) to (start + 0, 29)
+- Code(Counter(1)) at (prev + 1, 9) to (start + 0, 12)
+- Code(Counter(1)) at (prev + 0, 13) to (start + 0, 34)
+- Code(Counter(1)) at (prev + 1, 9) to (start + 0, 22)
+- Code(Counter(1)) at (prev + 1, 9) to (start + 0, 12)
+- Code(Counter(1)) at (prev + 0, 13) to (start + 0, 33)
+- Code(Expression(0, Sub)) at (prev + 1, 5) to (start + 0, 6)
+    = (c0 - c1)
+- Code(Counter(0)) at (prev + 1, 1) to (start + 0, 2)
+Highest counter ID seen: c1
+
+Function name: pass_through::uses_outer_macro
+Raw bytes (51): 0x[01, 01, 01, 01, 05, 09, 01, 34, 01, 00, 16, 01, 01, 08, 00, 1d, 05, 01, 09, 00, 0c, 05, 00, 0d, 00, 21, 05, 01, 09, 00, 15, 05, 01, 09, 00, 0c, 05, 00, 0d, 00, 20, 02, 01, 05, 00, 06, 01, 01, 01, 00, 02]
+Number of files: 1
+- file 0 => $DIR/pass-through.rs
+Number of expressions: 1
+- expression 0 operands: lhs = Counter(0), rhs = Counter(1)
+Number of file 0 mappings: 9
+- Code(Counter(0)) at (prev + 52, 1) to (start + 0, 22)
+- Code(Counter(0)) at (prev + 1, 8) to (start + 0, 29)
+- Code(Counter(1)) at (prev + 1, 9) to (start + 0, 12)
+- Code(Counter(1)) at (prev + 0, 13) to (start + 0, 33)
+- Code(Counter(1)) at (prev + 1, 9) to (start + 0, 21)
+- Code(Counter(1)) at (prev + 1, 9) to (start + 0, 12)
+- Code(Counter(1)) at (prev + 0, 13) to (start + 0, 32)
+- Code(Expression(0, Sub)) at (prev + 1, 5) to (start + 0, 6)
+    = (c0 - c1)
+- Code(Counter(0)) at (prev + 1, 1) to (start + 0, 2)
+Highest counter ID seen: c1
+