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

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
+