coverage: Sort the expansion tree in depth-first order

This makes it possible for subsequent operations to iterate over all nodes,
while assuming that every node occurs before all of its descendants.

Author: Zalathar

compiler/rustc_mir_transform/src/coverage/expansion.rs

@@ -6,6 +6,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 {
@@ -62,6 +63,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,
@@ -100,6 +103,7 @@ impl ExpnNode {
 
         Self {
             expn_id,
+            dfs_rank: usize::MAX,
 
             expn_kind: expn_data.kind,
             call_site,
@@ -124,7 +128,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 +161,9 @@ pub(crate) fn build_expn_tree(
         }
     }
 
+    // Sort the tree nodes into depth-first order.
+    sort_nodes_depth_first(&mut nodes)?;
+
     // 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 +196,32 @@ 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(())
 }

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