coverage: Use a standard depth-first search on a custom subgraph

compiler/rustc_mir_transform/src/coverage/graph.rs

@@ -145,18 +145,17 @@ impl CoverageGraph {
             bcbs.push(bcb_data);
         };
 
-        // Walk the MIR CFG using a Preorder traversal, which starts from `START_BLOCK` and follows
+        // Traverse the MIR control-flow graph, accumulating chains of blocks
-        // each block terminator's `successors()`. Coverage spans must map to actual source code,
+        // that can be combined into a single node in the coverage graph.
-        // so compiler generated blocks and paths can be ignored. To that end, the CFG traversal
+        // A depth-first search ensures that if two nodes can be chained
-        // intentionally omits unwind paths.
+        // together, they will be adjacent in the traversal order.
-        // FIXME(#78544): MIR InstrumentCoverage: Improve coverage of `#[should_panic]` tests and
-        // `catch_unwind()` handlers.
 
         // Accumulates a chain of blocks that will be combined into one BCB.
         let mut basic_blocks = Vec::new();
 
         let filtered_successors = |bb| bcb_filtered_successors(mir_body[bb].terminator());
-        for bb in short_circuit_preorder(mir_body, filtered_successors)
+        let subgraph = CoverageRelevantSubgraph::new(&mir_body.basic_blocks);
+        for bb in graph::depth_first_search(subgraph, mir::START_BLOCK)
             .filter(|&bb| mir_body[bb].terminator().kind != TerminatorKind::Unreachable)
         {
             // If the previous block can't be chained into `bb`, flush the accumulated
@@ -599,28 +598,31 @@ impl<'a> TraverseCoverageGraphWithLoops<'a> {
     }
 }
 
-fn short_circuit_preorder<'a, 'tcx, F, Iter>(
+/// Wrapper around a [`mir::BasicBlocks`] graph that restricts each node's
-    body: &'a mir::Body<'tcx>,
+/// successors to only the ones considered "relevant" when building a coverage
-    filtered_successors: F,
+/// graph.
-) -> impl Iterator<Item = BasicBlock> + Captures<'a> + Captures<'tcx>
+#[derive(Clone, Copy)]
-where
+struct CoverageRelevantSubgraph<'a, 'tcx> {
-    F: Fn(BasicBlock) -> Iter,
+    basic_blocks: &'a mir::BasicBlocks<'tcx>,
-    Iter: IntoIterator<Item = BasicBlock>,
+}
-{
+impl<'a, 'tcx> CoverageRelevantSubgraph<'a, 'tcx> {
-    let mut visited = BitSet::new_empty(body.basic_blocks.len());
+    fn new(basic_blocks: &'a mir::BasicBlocks<'tcx>) -> Self {
-    let mut worklist = vec![mir::START_BLOCK];
+        Self { basic_blocks }
-
+    }
-    std::iter::from_fn(move || {
+
-        while let Some(bb) = worklist.pop() {
+    fn coverage_successors(&self, bb: BasicBlock) -> CoverageSuccessors<'_> {
-            if !visited.insert(bb) {
+        bcb_filtered_successors(self.basic_blocks[bb].terminator())
-                continue;
+    }
-            }
+}
-
+impl<'a, 'tcx> graph::DirectedGraph for CoverageRelevantSubgraph<'a, 'tcx> {
-            worklist.extend(filtered_successors(bb));
+    type Node = BasicBlock;
-
+
-            return Some(bb);
+    fn num_nodes(&self) -> usize {
-        }
+        self.basic_blocks.num_nodes()
-
+    }
-        None
+}
-    })
+impl<'a, 'tcx> graph::Successors for CoverageRelevantSubgraph<'a, 'tcx> {
+    fn successors(&self, bb: Self::Node) -> impl Iterator<Item = Self::Node> {
+        self.coverage_successors(bb).into_iter()
+    }
 }