coverage: Use the sorted expansion tree to determine min/max BCBs

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};
@@ -23,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)]
@@ -85,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>,
@@ -114,6 +87,7 @@ impl ExpnNode {
 
             spans: vec![],
             child_expn_ids: FxIndexSet::default(),
+            minmax_bcbs: None,
 
             branch_spans: vec![],
 
@@ -164,6 +138,17 @@ 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
@@ -225,3 +210,31 @@ fn sort_nodes_depth_first(nodes: &mut FxIndexMap<ExpnId, ExpnNode>) -> Result<()
 
     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/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)]