From ab154e69993f15c4f37842dd2c2dfc6750af6a8a Mon Sep 17 00:00:00 2001
From: Zalathar <Zalathar@users.noreply.github.com>
Date: Sun, 6 Oct 2024 22:44:14 +1100
Subject: [PATCH] coverage: Store `bcb_needs_counter` in a field as a bitset

This makes it possible for other parts of counter-assignment to check whether a
node is guaranteed to end up with some kind of counter.

Switching from `impl Fn` to a concrete `&BitSet` just avoids the hassle of
trying to store a closure in a struct field, and currently there's no
foreseeable need for this information to not be a bitset.
---
 .../src/coverage/counters.rs                  | 19 +++++++++++++------
 .../rustc_mir_transform/src/coverage/mod.rs   |  3 +--
 2 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/compiler/rustc_mir_transform/src/coverage/counters.rs b/compiler/rustc_mir_transform/src/coverage/counters.rs
index 95c4d9f88fbe9..9408815675692 100644
--- a/compiler/rustc_mir_transform/src/coverage/counters.rs
+++ b/compiler/rustc_mir_transform/src/coverage/counters.rs
@@ -4,6 +4,7 @@ use rustc_data_structures::captures::Captures;
 use rustc_data_structures::fx::FxHashMap;
 use rustc_data_structures::graph::DirectedGraph;
 use rustc_index::IndexVec;
+use rustc_index::bit_set::BitSet;
 use rustc_middle::bug;
 use rustc_middle::mir::coverage::{CounterId, CovTerm, Expression, ExpressionId, Op};
 use tracing::{debug, debug_span, instrument};
@@ -78,10 +79,10 @@ impl CoverageCounters {
     /// counters or counter expressions for nodes and edges as required.
     pub(super) fn make_bcb_counters(
         basic_coverage_blocks: &CoverageGraph,
-        bcb_needs_counter: impl Fn(BasicCoverageBlock) -> bool,
+        bcb_needs_counter: &BitSet<BasicCoverageBlock>,
     ) -> Self {
-        let mut counters = MakeBcbCounters::new(basic_coverage_blocks);
-        counters.make_bcb_counters(bcb_needs_counter);
+        let mut counters = MakeBcbCounters::new(basic_coverage_blocks, bcb_needs_counter);
+        counters.make_bcb_counters();
 
         counters.coverage_counters
     }
@@ -268,17 +269,23 @@ impl CoverageCounters {
 struct MakeBcbCounters<'a> {
     coverage_counters: CoverageCounters,
     basic_coverage_blocks: &'a CoverageGraph,
+    bcb_needs_counter: &'a BitSet<BasicCoverageBlock>,
 }
 
 impl<'a> MakeBcbCounters<'a> {
-    fn new(basic_coverage_blocks: &'a CoverageGraph) -> Self {
+    fn new(
+        basic_coverage_blocks: &'a CoverageGraph,
+        bcb_needs_counter: &'a BitSet<BasicCoverageBlock>,
+    ) -> Self {
+        assert_eq!(basic_coverage_blocks.num_nodes(), bcb_needs_counter.domain_size());
         Self {
             coverage_counters: CoverageCounters::with_num_bcbs(basic_coverage_blocks.num_nodes()),
             basic_coverage_blocks,
+            bcb_needs_counter,
         }
     }
 
-    fn make_bcb_counters(&mut self, bcb_needs_counter: impl Fn(BasicCoverageBlock) -> bool) {
+    fn make_bcb_counters(&mut self) {
         debug!("make_bcb_counters(): adding a counter or expression to each BasicCoverageBlock");
 
         // Traverse the coverage graph, ensuring that every node that needs a
@@ -291,7 +298,7 @@ impl<'a> MakeBcbCounters<'a> {
         let mut traversal = TraverseCoverageGraphWithLoops::new(self.basic_coverage_blocks);
         while let Some(bcb) = traversal.next() {
             let _span = debug_span!("traversal", ?bcb).entered();
-            if bcb_needs_counter(bcb) {
+            if self.bcb_needs_counter.contains(bcb) {
                 self.make_node_counter_and_out_edge_counters(&traversal, bcb);
             }
         }
diff --git a/compiler/rustc_mir_transform/src/coverage/mod.rs b/compiler/rustc_mir_transform/src/coverage/mod.rs
index 97684d796b726..79482ba39198e 100644
--- a/compiler/rustc_mir_transform/src/coverage/mod.rs
+++ b/compiler/rustc_mir_transform/src/coverage/mod.rs
@@ -94,9 +94,8 @@ fn instrument_function_for_coverage<'tcx>(tcx: TyCtxt<'tcx>, mir_body: &mut mir:
         return;
     }
 
-    let bcb_has_counter_mappings = |bcb| bcbs_with_counter_mappings.contains(bcb);
     let coverage_counters =
-        CoverageCounters::make_bcb_counters(&basic_coverage_blocks, bcb_has_counter_mappings);
+        CoverageCounters::make_bcb_counters(&basic_coverage_blocks, &bcbs_with_counter_mappings);
 
     let mappings = create_mappings(tcx, &hir_info, &extracted_mappings, &coverage_counters);
     if mappings.is_empty() {