Store BCB counters externally, not directly in the BCB graph

Storing coverage counter information in `CoverageCounters` has a few advantages
over storing it directly inside BCB graph nodes:

- The graph doesn't need to be mutable when making the counters, making it
easier to see that the graph itself is not modified during this step.

- All of the counter data is clearly visible in one place.

- It becomes possible to use a representation that doesn't correspond 1:1 to
graph nodes, e.g. storing all the edge counters in a single hashmap instead of
several.

5 files changed, 157 insertions, 158 deletions

diff --git a/compiler/rustc_mir_transform/src/coverage/counters.rs b/compiler/rustc_mir_transform/src/coverage/counters.rs
index ed362a72e92..d1f2f0c76c8 100644
--- a/compiler/rustc_mir_transform/src/coverage/counters.rs
+++ b/compiler/rustc_mir_transform/src/coverage/counters.rs
@@ -8,17 +8,28 @@ use debug::{DebugCounters, NESTED_INDENT};
 use graph::{BasicCoverageBlock, BcbBranch, CoverageGraph, TraverseCoverageGraphWithLoops};
 use spans::CoverageSpan;
 
+use rustc_data_structures::fx::FxHashMap;
 use rustc_data_structures::graph::WithNumNodes;
 use rustc_index::bit_set::BitSet;
+use rustc_index::IndexVec;
 use rustc_middle::mir::coverage::*;
 
-/// Manages the counter and expression indexes/IDs to generate `CoverageKind` components for MIR
-/// `Coverage` statements.
+/// Generates and stores coverage counter and coverage expression information
+/// associated with nodes/edges in the BCB graph.
 pub(super) struct CoverageCounters {
     function_source_hash: u64,
     next_counter_id: CounterId,
     next_expression_id: ExpressionId,
 
+    /// Coverage counters/expressions that are associated with individual BCBs.
+    bcb_counters: IndexVec<BasicCoverageBlock, Option<CoverageKind>>,
+    /// Coverage counters/expressions that are associated with the control-flow
+    /// edge between two BCBs.
+    bcb_edge_counters: FxHashMap<(BasicCoverageBlock, BasicCoverageBlock), CoverageKind>,
+    /// Tracks which BCBs have a counter associated with some incoming edge.
+    /// Only used by debug assertions, to verify that BCBs with incoming edge
+    /// counters do not have their own physical counters (expressions are allowed).
+    bcb_has_incoming_edge_counters: BitSet<BasicCoverageBlock>,
     /// Expression nodes that are not directly associated with any particular
     /// BCB/edge, but are needed as operands to more complex expressions.
     /// These are always `CoverageKind::Expression`.
@@ -28,12 +39,17 @@ pub(super) struct CoverageCounters {
 }
 
 impl CoverageCounters {
-    pub fn new(function_source_hash: u64) -> Self {
+    pub(super) fn new(function_source_hash: u64, basic_coverage_blocks: &CoverageGraph) -> Self {
+        let num_bcbs = basic_coverage_blocks.num_nodes();
+
         Self {
             function_source_hash,
             next_counter_id: CounterId::START,
             next_expression_id: ExpressionId::START,
 
+            bcb_counters: IndexVec::from_elem_n(None, num_bcbs),
+            bcb_edge_counters: FxHashMap::default(),
+            bcb_has_incoming_edge_counters: BitSet::new_empty(num_bcbs),
             intermediate_expressions: Vec::new(),
 
             debug_counters: DebugCounters::new(),
@@ -51,7 +67,7 @@ impl CoverageCounters {
     /// representing intermediate values.
     pub fn make_bcb_counters(
         &mut self,
-        basic_coverage_blocks: &mut CoverageGraph,
+        basic_coverage_blocks: &CoverageGraph,
         coverage_spans: &[CoverageSpan],
     ) -> Result<(), Error> {
         MakeBcbCounters::new(self, basic_coverage_blocks).make_bcb_counters(coverage_spans)
@@ -114,6 +130,80 @@ impl CoverageCounters {
         self.next_expression_id = next.next_id();
         next
     }
+
+    fn set_bcb_counter(
+        &mut self,
+        bcb: BasicCoverageBlock,
+        counter_kind: CoverageKind,
+    ) -> Result<Operand, Error> {
+        debug_assert!(
+            // If the BCB has an edge counter (to be injected into a new `BasicBlock`), it can also
+            // have an expression (to be injected into an existing `BasicBlock` represented by this
+            // `BasicCoverageBlock`).
+            counter_kind.is_expression() || !self.bcb_has_incoming_edge_counters.contains(bcb),
+            "attempt to add a `Counter` to a BCB target with existing incoming edge counters"
+        );
+        let operand = counter_kind.as_operand();
+        if let Some(replaced) = self.bcb_counters[bcb].replace(counter_kind) {
+            Error::from_string(format!(
+                "attempt to set a BasicCoverageBlock coverage counter more than once; \
+                {bcb:?} already had counter {replaced:?}",
+            ))
+        } else {
+            Ok(operand)
+        }
+    }
+
+    fn set_bcb_edge_counter(
+        &mut self,
+        from_bcb: BasicCoverageBlock,
+        to_bcb: BasicCoverageBlock,
+        counter_kind: CoverageKind,
+    ) -> Result<Operand, Error> {
+        if level_enabled!(tracing::Level::DEBUG) {
+            // If the BCB has an edge counter (to be injected into a new `BasicBlock`), it can also
+            // have an expression (to be injected into an existing `BasicBlock` represented by this
+            // `BasicCoverageBlock`).
+            if self.bcb_counter(to_bcb).is_some_and(|c| !c.is_expression()) {
+                return Error::from_string(format!(
+                    "attempt to add an incoming edge counter from {from_bcb:?} when the target BCB already \
+                    has a `Counter`"
+                ));
+            }
+        }
+        self.bcb_has_incoming_edge_counters.insert(to_bcb);
+        let operand = counter_kind.as_operand();
+        if let Some(replaced) = self.bcb_edge_counters.insert((from_bcb, to_bcb), counter_kind) {
+            Error::from_string(format!(
+                "attempt to set an edge counter more than once; from_bcb: \
+                {from_bcb:?} already had counter {replaced:?}",
+            ))
+        } else {
+            Ok(operand)
+        }
+    }
+
+    pub(super) fn bcb_counter(&self, bcb: BasicCoverageBlock) -> Option<&CoverageKind> {
+        self.bcb_counters[bcb].as_ref()
+    }
+
+    pub(super) fn take_bcb_counter(&mut self, bcb: BasicCoverageBlock) -> Option<CoverageKind> {
+        self.bcb_counters[bcb].take()
+    }
+
+    pub(super) fn drain_bcb_counters(
+        &mut self,
+    ) -> impl Iterator<Item = (BasicCoverageBlock, CoverageKind)> + '_ {
+        self.bcb_counters
+            .iter_enumerated_mut()
+            .filter_map(|(bcb, counter)| Some((bcb, counter.take()?)))
+    }
+
+    pub(super) fn drain_bcb_edge_counters(
+        &mut self,
+    ) -> impl Iterator<Item = ((BasicCoverageBlock, BasicCoverageBlock), CoverageKind)> + '_ {
+        self.bcb_edge_counters.drain()
+    }
 }
 
 /// Traverse the `CoverageGraph` and add either a `Counter` or `Expression` to every BCB, to be
@@ -122,13 +212,13 @@ impl CoverageCounters {
 /// embedded counter, an `Expression` should be used.
 struct MakeBcbCounters<'a> {
     coverage_counters: &'a mut CoverageCounters,
-    basic_coverage_blocks: &'a mut CoverageGraph,
+    basic_coverage_blocks: &'a CoverageGraph,
 }
 
 impl<'a> MakeBcbCounters<'a> {
     fn new(
         coverage_counters: &'a mut CoverageCounters,
-        basic_coverage_blocks: &'a mut CoverageGraph,
+        basic_coverage_blocks: &'a CoverageGraph,
     ) -> Self {
         Self { coverage_counters, basic_coverage_blocks }
     }
@@ -202,9 +292,7 @@ impl<'a> MakeBcbCounters<'a> {
             branching_bcb,
             branches
                 .iter()
-                .map(|branch| {
-                    format!("{:?}: {:?}", branch, branch.counter(&self.basic_coverage_blocks))
-                })
+                .map(|branch| { format!("{:?}: {:?}", branch, self.branch_counter(branch)) })
                 .collect::<Vec<_>>()
                 .join("\n  "),
         );
@@ -274,9 +362,9 @@ impl<'a> MakeBcbCounters<'a> {
         debug!("{:?} gets an expression: {}", expression_branch, self.format_counter(&expression));
         let bcb = expression_branch.target_bcb;
         if expression_branch.is_only_path_to_target() {
-            self.basic_coverage_blocks[bcb].set_counter(expression)?;
+            self.coverage_counters.set_bcb_counter(bcb, expression)?;
         } else {
-            self.basic_coverage_blocks[bcb].set_edge_counter_from(branching_bcb, expression)?;
+            self.coverage_counters.set_bcb_edge_counter(branching_bcb, bcb, expression)?;
         }
         Ok(())
     }
@@ -291,7 +379,7 @@ impl<'a> MakeBcbCounters<'a> {
         debug_indent_level: usize,
     ) -> Result<Operand, Error> {
         // If the BCB already has a counter, return it.
-        if let Some(counter_kind) = self.basic_coverage_blocks[bcb].counter() {
+        if let Some(counter_kind) = &self.coverage_counters.bcb_counters[bcb] {
             debug!(
                 "{}{:?} already has a counter: {}",
                 NESTED_INDENT.repeat(debug_indent_level),
@@ -324,7 +412,7 @@ impl<'a> MakeBcbCounters<'a> {
                     self.format_counter(&counter_kind),
                 );
             }
-            return self.basic_coverage_blocks[bcb].set_counter(counter_kind);
+            return self.coverage_counters.set_bcb_counter(bcb, counter_kind);
         }
 
         // A BCB with multiple incoming edges can compute its count by `Expression`, summing up the
@@ -380,7 +468,7 @@ impl<'a> MakeBcbCounters<'a> {
             bcb,
             self.format_counter(&counter_kind)
         );
-        self.basic_coverage_blocks[bcb].set_counter(counter_kind)
+        self.coverage_counters.set_bcb_counter(bcb, counter_kind)
     }
 
     fn get_or_make_edge_counter_operand(
@@ -405,7 +493,9 @@ impl<'a> MakeBcbCounters<'a> {
         }
 
         // If the edge already has a counter, return it.
-        if let Some(counter_kind) = self.basic_coverage_blocks[to_bcb].edge_counter_from(from_bcb) {
+        if let Some(counter_kind) =
+            self.coverage_counters.bcb_edge_counters.get(&(from_bcb, to_bcb))
+        {
             debug!(
                 "{}Edge {:?}->{:?} already has a counter: {}",
                 NESTED_INDENT.repeat(debug_indent_level),
@@ -426,7 +516,7 @@ impl<'a> MakeBcbCounters<'a> {
             to_bcb,
             self.format_counter(&counter_kind)
         );
-        self.basic_coverage_blocks[to_bcb].set_edge_counter_from(from_bcb, counter_kind)
+        self.coverage_counters.set_bcb_edge_counter(from_bcb, to_bcb, counter_kind)
     }
 
     /// Select a branch for the expression, either the recommended `reloop_branch`, or if none was
@@ -436,8 +526,7 @@ impl<'a> MakeBcbCounters<'a> {
         traversal: &TraverseCoverageGraphWithLoops,
         branches: &[BcbBranch],
     ) -> BcbBranch {
-        let branch_needs_a_counter =
-            |branch: &BcbBranch| branch.counter(&self.basic_coverage_blocks).is_none();
+        let branch_needs_a_counter = |branch: &BcbBranch| self.branch_has_no_counter(branch);
 
         let some_reloop_branch = self.find_some_reloop_branch(traversal, &branches);
         if let Some(reloop_branch_without_counter) =
@@ -450,10 +539,8 @@ impl<'a> MakeBcbCounters<'a> {
             );
             reloop_branch_without_counter
         } else {
-            let &branch_without_counter = branches
-                .iter()
-                .find(|&&branch| branch.counter(&self.basic_coverage_blocks).is_none())
-                .expect(
+            let &branch_without_counter =
+                branches.iter().find(|&branch| self.branch_has_no_counter(branch)).expect(
                     "needs_branch_counters was `true` so there should be at least one \
                     branch",
                 );
@@ -480,8 +567,7 @@ impl<'a> MakeBcbCounters<'a> {
         traversal: &TraverseCoverageGraphWithLoops,
         branches: &[BcbBranch],
     ) -> Option<BcbBranch> {
-        let branch_needs_a_counter =
-            |branch: &BcbBranch| branch.counter(&self.basic_coverage_blocks).is_none();
+        let branch_needs_a_counter = |branch: &BcbBranch| self.branch_has_no_counter(branch);
 
         let mut some_reloop_branch: Option<BcbBranch> = None;
         for context in traversal.context_stack.iter().rev() {
@@ -492,7 +578,7 @@ impl<'a> MakeBcbCounters<'a> {
                         self.bcb_dominates(branch.target_bcb, backedge_from_bcb)
                     }) {
                         if let Some(reloop_branch) = some_reloop_branch {
-                            if reloop_branch.counter(&self.basic_coverage_blocks).is_none() {
+                            if self.branch_has_no_counter(&reloop_branch) {
                                 // we already found a candidate reloop_branch that still
                                 // needs a counter
                                 continue;
@@ -558,12 +644,24 @@ impl<'a> MakeBcbCounters<'a> {
     }
 
     fn bcb_needs_branch_counters(&self, bcb: BasicCoverageBlock) -> bool {
-        let branch_needs_a_counter =
-            |branch: &BcbBranch| branch.counter(&self.basic_coverage_blocks).is_none();
+        let branch_needs_a_counter = |branch: &BcbBranch| self.branch_has_no_counter(branch);
         let branches = self.bcb_branches(bcb);
         branches.len() > 1 && branches.iter().any(branch_needs_a_counter)
     }
 
+    fn branch_has_no_counter(&self, branch: &BcbBranch) -> bool {
+        self.branch_counter(branch).is_none()
+    }
+
+    fn branch_counter(&self, branch: &BcbBranch) -> Option<&CoverageKind> {
+        let to_bcb = branch.target_bcb;
+        if let Some(from_bcb) = branch.edge_from_bcb {
+            self.coverage_counters.bcb_edge_counters.get(&(from_bcb, to_bcb))
+        } else {
+            self.coverage_counters.bcb_counters[to_bcb].as_ref()
+        }
+    }
+
     /// Returns true if the BasicCoverageBlock has zero or one incoming edge. (If zero, it should be
     /// the entry point for the function.)
     #[inline]
diff --git a/compiler/rustc_mir_transform/src/coverage/debug.rs b/compiler/rustc_mir_transform/src/coverage/debug.rs
index 26f9cfd0b86..d2c0c4ba069 100644
--- a/compiler/rustc_mir_transform/src/coverage/debug.rs
+++ b/compiler/rustc_mir_transform/src/coverage/debug.rs
@@ -108,6 +108,7 @@
 //!         recursively, generating labels with nested operations, enclosed in parentheses
 //!         (for example: `bcb2 + (bcb0 - bcb1)`).
 
+use super::counters::CoverageCounters;
 use super::graph::{BasicCoverageBlock, BasicCoverageBlockData, CoverageGraph};
 use super::spans::CoverageSpan;
 
@@ -659,18 +660,21 @@ pub(super) fn dump_coverage_graphviz<'tcx>(
     mir_body: &mir::Body<'tcx>,
     pass_name: &str,
     basic_coverage_blocks: &CoverageGraph,
-    debug_counters: &DebugCounters,
+    coverage_counters: &CoverageCounters,
     graphviz_data: &GraphvizData,
     intermediate_expressions: &[CoverageKind],
     debug_used_expressions: &UsedExpressions,
 ) {
+    let debug_counters = &coverage_counters.debug_counters;
+
     let mir_source = mir_body.source;
     let def_id = mir_source.def_id();
     let node_content = |bcb| {
         bcb_to_string_sections(
             tcx,
             mir_body,
-            debug_counters,
+            coverage_counters,
+            bcb,
             &basic_coverage_blocks[bcb],
             graphviz_data.get_bcb_coverage_spans_with_counters(bcb),
             graphviz_data.get_bcb_dependency_counters(bcb),
@@ -736,12 +740,15 @@ pub(super) fn dump_coverage_graphviz<'tcx>(
 fn bcb_to_string_sections<'tcx>(
     tcx: TyCtxt<'tcx>,
     mir_body: &mir::Body<'tcx>,
-    debug_counters: &DebugCounters,
+    coverage_counters: &CoverageCounters,
+    bcb: BasicCoverageBlock,
     bcb_data: &BasicCoverageBlockData,
     some_coverage_spans_with_counters: Option<&[(CoverageSpan, CoverageKind)]>,
     some_dependency_counters: Option<&[CoverageKind]>,
     some_intermediate_expressions: Option<&[CoverageKind]>,
 ) -> Vec<String> {
+    let debug_counters = &coverage_counters.debug_counters;
+
     let len = bcb_data.basic_blocks.len();
     let mut sections = Vec::new();
     if let Some(collect_intermediate_expressions) = some_intermediate_expressions {
@@ -777,7 +784,7 @@ fn bcb_to_string_sections<'tcx>(
                 .join("  \n"),
         ));
     }
-    if let Some(counter_kind) = &bcb_data.counter_kind {
+    if let Some(counter_kind) = coverage_counters.bcb_counter(bcb) {
         sections.push(format!("{counter_kind:?}"));
     }
     let non_term_blocks = bcb_data.basic_blocks[0..len - 1]
diff --git a/compiler/rustc_mir_transform/src/coverage/graph.rs b/compiler/rustc_mir_transform/src/coverage/graph.rs
index f94dad4c8da..59b01ffec0f 100644
--- a/compiler/rustc_mir_transform/src/coverage/graph.rs
+++ b/compiler/rustc_mir_transform/src/coverage/graph.rs
@@ -1,12 +1,8 @@
-use super::Error;
-
 use itertools::Itertools;
-use rustc_data_structures::fx::FxHashMap;
 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::coverage::*;
 use rustc_middle::mir::{self, BasicBlock, BasicBlockData, Terminator, TerminatorKind};
 
 use std::cmp::Ordering;
@@ -15,10 +11,7 @@ use std::ops::{Index, IndexMut};
 const ID_SEPARATOR: &str = ",";
 
 /// A coverage-specific simplification of the MIR control flow graph (CFG). The `CoverageGraph`s
-/// nodes are `BasicCoverageBlock`s, which encompass one or more MIR `BasicBlock`s, plus a
-/// `CoverageKind` counter (to be added by `CoverageCounters::make_bcb_counters`), and an optional
-/// set of additional counters--if needed--to count incoming edges, if there are more than one.
-/// (These "edge counters" are eventually converted into new MIR `BasicBlock`s.)
+/// nodes are `BasicCoverageBlock`s, which encompass one or more MIR `BasicBlock`s.
 #[derive(Debug)]
 pub(super) struct CoverageGraph {
     bcbs: IndexVec<BasicCoverageBlock, BasicCoverageBlockData>,
@@ -196,13 +189,6 @@ impl CoverageGraph {
     }
 
     #[inline(always)]
-    pub fn iter_enumerated_mut(
-        &mut self,
-    ) -> impl Iterator<Item = (BasicCoverageBlock, &mut BasicCoverageBlockData)> {
-        self.bcbs.iter_enumerated_mut()
-    }
-
-    #[inline(always)]
     pub fn bcb_from_bb(&self, bb: BasicBlock) -> Option<BasicCoverageBlock> {
         if bb.index() < self.bb_to_bcb.len() { self.bb_to_bcb[bb] } else { None }
     }
@@ -320,14 +306,12 @@ rustc_index::newtype_index! {
 #[derive(Debug, Clone)]
 pub(super) struct BasicCoverageBlockData {
     pub basic_blocks: Vec<BasicBlock>,
-    pub counter_kind: Option<CoverageKind>,
-    edge_from_bcbs: Option<FxHashMap<BasicCoverageBlock, CoverageKind>>,
 }
 
 impl BasicCoverageBlockData {
     pub fn from(basic_blocks: Vec<BasicBlock>) -> Self {
         assert!(basic_blocks.len() > 0);
-        Self { basic_blocks, counter_kind: None, edge_from_bcbs: None }
+        Self { basic_blocks }
     }
 
     #[inline(always)]
@@ -345,80 +329,6 @@ impl BasicCoverageBlockData {
         &mir_body[self.last_bb()].terminator()
     }
 
-    pub fn set_counter(&mut self, counter_kind: CoverageKind) -> Result<Operand, Error> {
-        debug_assert!(
-            // If the BCB has an edge counter (to be injected into a new `BasicBlock`), it can also
-            // have an expression (to be injected into an existing `BasicBlock` represented by this
-            // `BasicCoverageBlock`).
-            self.edge_from_bcbs.is_none() || counter_kind.is_expression(),
-            "attempt to add a `Counter` to a BCB target with existing incoming edge counters"
-        );
-        let operand = counter_kind.as_operand();
-        if let Some(replaced) = self.counter_kind.replace(counter_kind) {
-            Error::from_string(format!(
-                "attempt to set a BasicCoverageBlock coverage counter more than once; \
-                {self:?} already had counter {replaced:?}",
-            ))
-        } else {
-            Ok(operand)
-        }
-    }
-
-    #[inline(always)]
-    pub fn counter(&self) -> Option<&CoverageKind> {
-        self.counter_kind.as_ref()
-    }
-
-    #[inline(always)]
-    pub fn take_counter(&mut self) -> Option<CoverageKind> {
-        self.counter_kind.take()
-    }
-
-    pub fn set_edge_counter_from(
-        &mut self,
-        from_bcb: BasicCoverageBlock,
-        counter_kind: CoverageKind,
-    ) -> Result<Operand, Error> {
-        if level_enabled!(tracing::Level::DEBUG) {
-            // If the BCB has an edge counter (to be injected into a new `BasicBlock`), it can also
-            // have an expression (to be injected into an existing `BasicBlock` represented by this
-            // `BasicCoverageBlock`).
-            if self.counter_kind.as_ref().is_some_and(|c| !c.is_expression()) {
-                return Error::from_string(format!(
-                    "attempt to add an incoming edge counter from {from_bcb:?} when the target BCB already \
-                    has a `Counter`"
-                ));
-            }
-        }
-        let operand = counter_kind.as_operand();
-        if let Some(replaced) =
-            self.edge_from_bcbs.get_or_insert_default().insert(from_bcb, counter_kind)
-        {
-            Error::from_string(format!(
-                "attempt to set an edge counter more than once; from_bcb: \
-                {from_bcb:?} already had counter {replaced:?}",
-            ))
-        } else {
-            Ok(operand)
-        }
-    }
-
-    #[inline]
-    pub fn edge_counter_from(&self, from_bcb: BasicCoverageBlock) -> Option<&CoverageKind> {
-        if let Some(edge_from_bcbs) = &self.edge_from_bcbs {
-            edge_from_bcbs.get(&from_bcb)
-        } else {
-            None
-        }
-    }
-
-    #[inline]
-    pub fn take_edge_counters(
-        &mut self,
-    ) -> Option<impl Iterator<Item = (BasicCoverageBlock, CoverageKind)>> {
-        self.edge_from_bcbs.take().map(|m| m.into_iter())
-    }
-
     pub fn id(&self) -> String {
         format!("@{}", self.basic_blocks.iter().map(|bb| bb.index().to_string()).join(ID_SEPARATOR))
     }
@@ -448,17 +358,6 @@ impl BcbBranch {
         Self { edge_from_bcb, target_bcb: to_bcb }
     }
 
-    pub fn counter<'a>(
-        &self,
-        basic_coverage_blocks: &'a CoverageGraph,
-    ) -> Option<&'a CoverageKind> {
-        if let Some(from_bcb) = self.edge_from_bcb {
-            basic_coverage_blocks[self.target_bcb].edge_counter_from(from_bcb)
-        } else {
-            basic_coverage_blocks[self.target_bcb].counter()
-        }
-    }
-
     pub fn is_only_path_to_target(&self) -> bool {
         self.edge_from_bcb.is_none()
     }
diff --git a/compiler/rustc_mir_transform/src/coverage/mod.rs b/compiler/rustc_mir_transform/src/coverage/mod.rs
index 750e2a2a215..e08b6d6f6e8 100644
--- a/compiler/rustc_mir_transform/src/coverage/mod.rs
+++ b/compiler/rustc_mir_transform/src/coverage/mod.rs
@@ -137,6 +137,8 @@ impl<'a, 'tcx> Instrumentor<'a, 'tcx> {
 
         let function_source_hash = hash_mir_source(tcx, hir_body);
         let basic_coverage_blocks = CoverageGraph::from_mir(mir_body);
+        let coverage_counters = CoverageCounters::new(function_source_hash, &basic_coverage_blocks);
+
         Self {
             pass_name,
             tcx,
@@ -145,7 +147,7 @@ impl<'a, 'tcx> Instrumentor<'a, 'tcx> {
             fn_sig_span,
             body_span,
             basic_coverage_blocks,
-            coverage_counters: CoverageCounters::new(function_source_hash),
+            coverage_counters,
         }
     }
 
@@ -250,7 +252,7 @@ impl<'a, 'tcx> Instrumentor<'a, 'tcx> {
                 self.mir_body,
                 self.pass_name,
                 &self.basic_coverage_blocks,
-                &self.coverage_counters.debug_counters,
+                &self.coverage_counters,
                 &graphviz_data,
                 &self.coverage_counters.intermediate_expressions,
                 &debug_used_expressions,
@@ -298,7 +300,7 @@ impl<'a, 'tcx> Instrumentor<'a, 'tcx> {
             let span = covspan.span;
             let counter_kind = if let Some(&counter_operand) = bcb_counters[bcb].as_ref() {
                 self.coverage_counters.make_identity_counter(counter_operand)
-            } else if let Some(counter_kind) = self.bcb_data_mut(bcb).take_counter() {
+            } else if let Some(counter_kind) = self.coverage_counters.take_bcb_counter(bcb) {
                 bcb_counters[bcb] = Some(counter_kind.as_operand());
                 debug_used_expressions.add_expression_operands(&counter_kind);
                 counter_kind
@@ -338,19 +340,17 @@ impl<'a, 'tcx> Instrumentor<'a, 'tcx> {
         debug_used_expressions: &mut debug::UsedExpressions,
     ) {
         let mut bcb_counters_without_direct_coverage_spans = Vec::new();
-        for (target_bcb, target_bcb_data) in self.basic_coverage_blocks.iter_enumerated_mut() {
-            if let Some(counter_kind) = target_bcb_data.take_counter() {
-                bcb_counters_without_direct_coverage_spans.push((None, target_bcb, counter_kind));
-            }
-            if let Some(edge_counters) = target_bcb_data.take_edge_counters() {
-                for (from_bcb, counter_kind) in edge_counters {
-                    bcb_counters_without_direct_coverage_spans.push((
-                        Some(from_bcb),
-                        target_bcb,
-                        counter_kind,
-                    ));
-                }
-            }
+        for (target_bcb, counter_kind) in self.coverage_counters.drain_bcb_counters() {
+            bcb_counters_without_direct_coverage_spans.push((None, target_bcb, counter_kind));
+        }
+        for ((from_bcb, target_bcb), counter_kind) in
+            self.coverage_counters.drain_bcb_edge_counters()
+        {
+            bcb_counters_without_direct_coverage_spans.push((
+                Some(from_bcb),
+                target_bcb,
+                counter_kind,
+            ));
         }
 
         // If debug is enabled, validate that every BCB or edge counter not directly associated
@@ -426,11 +426,6 @@ impl<'a, 'tcx> Instrumentor<'a, 'tcx> {
     }
 
     #[inline]
-    fn bcb_data_mut(&mut self, bcb: BasicCoverageBlock) -> &mut BasicCoverageBlockData {
-        &mut self.basic_coverage_blocks[bcb]
-    }
-
-    #[inline]
     fn format_counter(&self, counter_kind: &CoverageKind) -> String {
         self.coverage_counters.debug_counters.format_counter(counter_kind)
     }
diff --git a/compiler/rustc_mir_transform/src/coverage/tests.rs b/compiler/rustc_mir_transform/src/coverage/tests.rs
index 0699723fdc4..d797a6057a7 100644
--- a/compiler/rustc_mir_transform/src/coverage/tests.rs
+++ b/compiler/rustc_mir_transform/src/coverage/tests.rs
@@ -675,7 +675,7 @@ fn test_make_bcb_counters() {
                 ));
             }
         }
-        let mut coverage_counters = counters::CoverageCounters::new(0);
+        let mut coverage_counters = counters::CoverageCounters::new(0, &basic_coverage_blocks);
         let () = coverage_counters
             .make_bcb_counters(&mut basic_coverage_blocks, &coverage_spans)
             .expect("should be Ok");
@@ -684,7 +684,7 @@ fn test_make_bcb_counters() {
         let_bcb!(1);
         assert_eq!(
             0, // bcb1 has a `Counter` with id = 0
-            match basic_coverage_blocks[bcb1].counter().expect("should have a counter") {
+            match coverage_counters.bcb_counter(bcb1).expect("should have a counter") {
                 CoverageKind::Counter { id, .. } => id,
                 _ => panic!("expected a Counter"),
             }
@@ -694,7 +694,7 @@ fn test_make_bcb_counters() {
         let_bcb!(2);
         assert_eq!(
             1, // bcb2 has a `Counter` with id = 1
-            match basic_coverage_blocks[bcb2].counter().expect("should have a counter") {
+            match coverage_counters.bcb_counter(bcb2).expect("should have a counter") {
                 CoverageKind::Counter { id, .. } => id,
                 _ => panic!("expected a Counter"),
             }