coverage: Hoist expansion tree creation out of span refinement

This is an incremental step towards making the expansion tree central to
coverage mapping creation, which will be needed for proper expansion region
support.

Author: Zalathar

compiler/rustc_mir_transform/src/coverage/expansion.rs

@@ -1,7 +1,11 @@
 use rustc_data_structures::fx::{FxIndexMap, FxIndexSet, IndexEntry};
+use rustc_middle::mir;
 use rustc_middle::mir::coverage::BasicCoverageBlock;
 use rustc_span::{ExpnId, ExpnKind, Span};
 
+use crate::coverage::from_mir;
+use crate::coverage::graph::CoverageGraph;
+
 #[derive(Clone, Copy, Debug)]
 pub(crate) struct SpanWithBcb {
     pub(crate) span: Span,
@@ -92,13 +96,17 @@ impl ExpnNode {
     }
 }
 
-/// Given a collection of span/BCB pairs from potentially-different syntax contexts,
+/// Extracts raw span/BCB pairs from potentially-different syntax contexts, and
 /// arranges them into an "expansion tree" based on their expansion call-sites.
-pub(crate) fn build_expn_tree(spans: impl IntoIterator<Item = SpanWithBcb>) -> ExpnTree {
+pub(crate) fn build_expn_tree(mir_body: &mir::Body<'_>, graph: &CoverageGraph) -> ExpnTree {
+    let raw_spans = from_mir::extract_raw_spans_from_mir(mir_body, graph);
+
     let mut nodes = FxIndexMap::default();
     let new_node = |&expn_id: &ExpnId| ExpnNode::new(expn_id);
 
-    for span_with_bcb in spans {
+    for from_mir::RawSpanFromMir { raw_span, bcb } in raw_spans {
+        let span_with_bcb = SpanWithBcb { span: raw_span, bcb };
+
         // Create a node for this span's enclosing expansion, and add the span to it.
         let expn_id = span_with_bcb.span.ctxt().outer_expn();
         let node = nodes.entry(expn_id).or_insert_with_key(new_node);

compiler/rustc_mir_transform/src/coverage/spans/from_mir.rs renamed to compiler/rustc_mir_transform/src/coverage/from_mir.rs

@@ -5,6 +5,7 @@ use rustc_middle::mir::coverage::{
 use rustc_middle::mir::{self, BasicBlock, StatementKind};
 use rustc_middle::ty::TyCtxt;
 
+use crate::coverage::expansion;
 use crate::coverage::graph::CoverageGraph;
 use crate::coverage::hir_info::ExtractedHirInfo;
 use crate::coverage::spans::extract_refined_covspans;
@@ -23,10 +24,12 @@ pub(crate) fn extract_mappings_from_mir<'tcx>(
     hir_info: &ExtractedHirInfo,
     graph: &CoverageGraph,
 ) -> ExtractedMappings {
+    let expn_tree = expansion::build_expn_tree(mir_body, graph);
+
     let mut mappings = vec![];
 
     // Extract ordinary code mappings from MIR statement/terminator spans.
-    extract_refined_covspans(tcx, mir_body, hir_info, graph, &mut mappings);
+    extract_refined_covspans(tcx, hir_info, graph, &expn_tree, &mut mappings);
 
     extract_branch_mappings(mir_body, hir_info, graph, &mut mappings);
 

compiler/rustc_mir_transform/src/coverage/mappings.rs

@@ -9,6 +9,7 @@ use crate::coverage::mappings::ExtractedMappings;
 
 mod counters;
 mod expansion;
+mod from_mir;
 mod graph;
 mod hir_info;
 mod mappings;

compiler/rustc_mir_transform/src/coverage/mod.rs

@@ -1,22 +1,19 @@
-use rustc_middle::mir;
 use rustc_middle::mir::coverage::{Mapping, MappingKind, START_BCB};
 use rustc_middle::ty::TyCtxt;
 use rustc_span::source_map::SourceMap;
 use rustc_span::{BytePos, DesugaringKind, ExpnId, ExpnKind, MacroKind, Span};
 use tracing::instrument;
 
-use crate::coverage::expansion::{self, ExpnTree, SpanWithBcb};
+use crate::coverage::expansion::{ExpnTree, SpanWithBcb};
+use crate::coverage::from_mir::Hole;
 use crate::coverage::graph::{BasicCoverageBlock, CoverageGraph};
 use crate::coverage::hir_info::ExtractedHirInfo;
-use crate::coverage::spans::from_mir::{Hole, RawSpanFromMir};
-
-mod from_mir;
 
 pub(super) fn extract_refined_covspans<'tcx>(
     tcx: TyCtxt<'tcx>,
-    mir_body: &mir::Body<'tcx>,
     hir_info: &ExtractedHirInfo,
     graph: &CoverageGraph,
+    expn_tree: &ExpnTree,
     mappings: &mut Vec<Mapping>,
 ) {
     if hir_info.is_async_fn {
@@ -32,14 +29,6 @@ pub(super) fn extract_refined_covspans<'tcx>(
 
     let &ExtractedHirInfo { body_span, .. } = hir_info;
 
-    let raw_spans = from_mir::extract_raw_spans_from_mir(mir_body, graph);
-    // Use the raw spans to build a tree of expansions for this function.
-    let expn_tree = expansion::build_expn_tree(
-        raw_spans
-            .into_iter()
-            .map(|RawSpanFromMir { raw_span, bcb }| SpanWithBcb { span: raw_span, bcb }),
-    );
-
     let mut covspans = vec![];
     let mut push_covspan = |covspan: Covspan| {
         let covspan_span = covspan.span;