From 986db13c178126167b4c7e0c4eb76b1dad3908b7 Mon Sep 17 00:00:00 2001 From: Zalathar Date: Wed, 3 Dec 2025 17:28:47 +1100 Subject: [PATCH] coverage: Use the sorted expansion tree to determine min/max BCBs --- .../src/coverage/expansion.rs | 77 +++++++++++-------- .../rustc_mir_transform/src/coverage/spans.rs | 12 +-- compiler/rustc_mir_transform/src/lib.rs | 1 - 3 files changed, 49 insertions(+), 41 deletions(-) diff --git a/compiler/rustc_mir_transform/src/coverage/expansion.rs b/compiler/rustc_mir_transform/src/coverage/expansion.rs index 16c37d6e6cf0..77aec902faad 100644 --- a/compiler/rustc_mir_transform/src/coverage/expansion.rs +++ b/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 { - 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, /// Expansions whose call-site is in this expansion. pub(crate) child_expn_ids: FxIndexSet, + /// 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, /// Branch spans (recorded during MIR building) belonging to this expansion. pub(crate) branch_spans: Vec, @@ -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) -> 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, + node: &ExpnNode, +) -> Option { + 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 }) +} diff --git a/compiler/rustc_mir_transform/src/coverage/spans.rs b/compiler/rustc_mir_transform/src/coverage/spans.rs index bc26a8ccc47a..b1ce0069b43a 100644 --- a/compiler/rustc_mir_transform/src/coverage/spans.rs +++ b/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 { 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 diff --git a/compiler/rustc_mir_transform/src/lib.rs b/compiler/rustc_mir_transform/src/lib.rs index f12561492de4..afdfa35f97c0 100644 --- a/compiler/rustc_mir_transform/src/lib.rs +++ b/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)]