coverage: Pass around basic blocks instead of references

For code that already has access to the MIR body, it's simpler and easier to
pass around basic block indices, and look up details in the body only where
they are needed.

Author: Zalathar

compiler/rustc_mir_transform/src/coverage/graph.rs

@@ -2,7 +2,7 @@ use rustc_data_structures::graph::dominators::{self, Dominators};
 use rustc_data_structures::graph::{self, GraphSuccessors, WithNumNodes, WithStartNode};
 use rustc_index::bit_set::BitSet;
 use rustc_index::{IndexSlice, IndexVec};
-use rustc_middle::mir::{self, BasicBlock, BasicBlockData, Terminator, TerminatorKind};
+use rustc_middle::mir::{self, BasicBlock, TerminatorKind};
 
 use std::cmp::Ordering;
 use std::ops::{Index, IndexMut};
@@ -36,9 +36,8 @@ impl CoverageGraph {
                 }
                 let bcb_data = &bcbs[bcb];
                 let mut bcb_successors = Vec::new();
-                for successor in
-                    bcb_filtered_successors(&mir_body, &bcb_data.terminator(mir_body).kind)
-                        .filter_map(|successor_bb| bb_to_bcb[successor_bb])
+                for successor in bcb_filtered_successors(&mir_body, bcb_data.last_bb())
+                    .filter_map(|successor_bb| bb_to_bcb[successor_bb])
                 {
                     if !seen[successor] {
                         seen[successor] = true;
@@ -83,7 +82,7 @@ impl CoverageGraph {
         let mir_cfg_without_unwind = ShortCircuitPreorder::new(&mir_body, bcb_filtered_successors);
 
         let mut basic_blocks = Vec::new();
-        for (bb, data) in mir_cfg_without_unwind {
+        for bb in mir_cfg_without_unwind {
             if let Some(last) = basic_blocks.last() {
                 let predecessors = &mir_body.basic_blocks.predecessors()[bb];
                 if predecessors.len() > 1 || !predecessors.contains(last) {
@@ -109,7 +108,7 @@ impl CoverageGraph {
             }
             basic_blocks.push(bb);
 
-            let term = data.terminator();
+            let term = mir_body[bb].terminator();
 
             match term.kind {
                 TerminatorKind::Return { .. }
@@ -316,11 +315,6 @@ impl BasicCoverageBlockData {
     pub fn last_bb(&self) -> BasicBlock {
         *self.basic_blocks.last().unwrap()
     }
-
-    #[inline(always)]
-    pub fn terminator<'a, 'tcx>(&self, mir_body: &'a mir::Body<'tcx>) -> &'a Terminator<'tcx> {
-        &mir_body[self.last_bb()].terminator()
-    }
 }
 
 /// Represents a successor from a branching BasicCoverageBlock (such as the arms of a `SwitchInt`)
@@ -362,13 +356,15 @@ impl std::fmt::Debug for BcbBranch {
     }
 }
 
-// Returns the `Terminator`s non-unwind successors.
+// Returns the subset of a block's successors that are relevant for the coverage
+// graph, i.e. those that do not represent unwinds or unreachable branches.
 // FIXME(#78544): MIR InstrumentCoverage: Improve coverage of `#[should_panic]` tests and
 // `catch_unwind()` handlers.
 fn bcb_filtered_successors<'a, 'tcx>(
     body: &'a mir::Body<'tcx>,
-    term_kind: &'a TerminatorKind<'tcx>,
+    bb: BasicBlock,
 ) -> Box<dyn Iterator<Item = BasicBlock> + 'a> {
+    let term_kind = &body[bb].terminator().kind;
     Box::new(
         match &term_kind {
             // SwitchInt successors are never unwind, and all of them should be traversed.
@@ -562,10 +558,7 @@ struct ShortCircuitPreorder<'a, 'tcx, F> {
 
 impl<'a, 'tcx, F> ShortCircuitPreorder<'a, 'tcx, F>
 where
-    F: Fn(
-        &'a mir::Body<'tcx>,
-        &'a TerminatorKind<'tcx>,
-    ) -> Box<dyn Iterator<Item = BasicBlock> + 'a>,
+    F: Fn(&'a mir::Body<'tcx>, BasicBlock) -> Box<dyn Iterator<Item = BasicBlock> + 'a>,
 {
     fn new(body: &'a mir::Body<'tcx>, filtered_successors: F) -> Self {
         Self {
@@ -579,26 +572,19 @@ where
 
 impl<'a, 'tcx, F> Iterator for ShortCircuitPreorder<'a, 'tcx, F>
 where
-    F: Fn(
-        &'a mir::Body<'tcx>,
-        &'a TerminatorKind<'tcx>,
-    ) -> Box<dyn Iterator<Item = BasicBlock> + 'a>,
+    F: Fn(&'a mir::Body<'tcx>, BasicBlock) -> Box<dyn Iterator<Item = BasicBlock> + 'a>,
 {
-    type Item = (BasicBlock, &'a BasicBlockData<'tcx>);
+    type Item = BasicBlock;
 
-    fn next(&mut self) -> Option<(BasicBlock, &'a BasicBlockData<'tcx>)> {
-        while let Some(idx) = self.worklist.pop() {
-            if !self.visited.insert(idx) {
+    fn next(&mut self) -> Option<BasicBlock> {
+        while let Some(bb) = self.worklist.pop() {
+            if !self.visited.insert(bb) {
                 continue;
             }
 
-            let data = &self.body[idx];
-
-            if let Some(ref term) = data.terminator {
-                self.worklist.extend((self.filtered_successors)(&self.body, &term.kind));
-            }
+            self.worklist.extend((self.filtered_successors)(self.body, bb));
 
-            return Some((idx, data));
+            return Some(bb);
         }
 
         None

compiler/rustc_mir_transform/src/coverage/tests.rs

@@ -241,7 +241,7 @@ fn print_coverage_graphviz(
                         "    {:?} [label=\"{:?}: {}\"];\n{}",
                         bcb,
                         bcb,
-                        bcb_data.terminator(mir_body).kind.name(),
+                        mir_body[bcb_data.last_bb()].terminator().kind.name(),
                         basic_coverage_blocks
                             .successors(bcb)
                             .map(|successor| { format!("    {:?} -> {:?};", bcb, successor) })