coverage: Split `FunctionCoverage` into distinct collector/finished phases

This gives us a clearly-defined place to run code after the instance's MIR has
been traversed by codegen, but before we emit its `__llvm_covfun` record.

3 files changed, 40 insertions, 24 deletions

diff --git a/compiler/rustc_codegen_llvm/src/coverageinfo/map_data.rs b/compiler/rustc_codegen_llvm/src/coverageinfo/map_data.rs
index 84319b4ba2d..302e00b06ed 100644
--- a/compiler/rustc_codegen_llvm/src/coverageinfo/map_data.rs
+++ b/compiler/rustc_codegen_llvm/src/coverageinfo/map_data.rs
@@ -10,7 +10,7 @@ use rustc_middle::ty::Instance;
 /// Holds all of the coverage mapping data associated with a function instance,
 /// collected during traversal of `Coverage` statements in the function's MIR.
 #[derive(Debug)]
-pub struct FunctionCoverage<'tcx> {
+pub struct FunctionCoverageCollector<'tcx> {
     /// Coverage info that was attached to this function by the instrumentor.
     function_coverage_info: &'tcx FunctionCoverageInfo,
     is_used: bool,
@@ -26,7 +26,7 @@ pub struct FunctionCoverage<'tcx> {
     expressions_seen: BitSet<ExpressionId>,
 }
 
-impl<'tcx> FunctionCoverage<'tcx> {
+impl<'tcx> FunctionCoverageCollector<'tcx> {
     /// Creates a new set of coverage data for a used (called) function.
     pub fn new(
         instance: Instance<'tcx>,
@@ -76,11 +76,6 @@ impl<'tcx> FunctionCoverage<'tcx> {
         }
     }
 
-    /// Returns true for a used (called) function, and false for an unused function.
-    pub fn is_used(&self) -> bool {
-        self.is_used
-    }
-
     /// Marks a counter ID as having been seen in a counter-increment statement.
     #[instrument(level = "debug", skip(self))]
     pub(crate) fn mark_counter_id_seen(&mut self, id: CounterId) {
@@ -165,6 +160,28 @@ impl<'tcx> FunctionCoverage<'tcx> {
         ZeroExpressions(zero_expressions)
     }
 
+    pub(crate) fn into_finished(self) -> FunctionCoverage<'tcx> {
+        let zero_expressions = self.identify_zero_expressions();
+        let FunctionCoverageCollector { function_coverage_info, is_used, counters_seen, .. } = self;
+
+        FunctionCoverage { function_coverage_info, is_used, counters_seen, zero_expressions }
+    }
+}
+
+pub(crate) struct FunctionCoverage<'tcx> {
+    function_coverage_info: &'tcx FunctionCoverageInfo,
+    is_used: bool,
+
+    counters_seen: BitSet<CounterId>,
+    zero_expressions: ZeroExpressions,
+}
+
+impl<'tcx> FunctionCoverage<'tcx> {
+    /// Returns true for a used (called) function, and false for an unused function.
+    pub(crate) fn is_used(&self) -> bool {
+        self.is_used
+    }
+
     /// Return the source hash, generated from the HIR node structure, and used to indicate whether
     /// or not the source code structure changed between different compilations.
     pub fn source_hash(&self) -> u64 {
@@ -177,29 +194,27 @@ impl<'tcx> FunctionCoverage<'tcx> {
     pub fn get_expressions_and_counter_regions(
         &self,
     ) -> (Vec<CounterExpression>, impl Iterator<Item = (Counter, &CodeRegion)>) {
-        let zero_expressions = self.identify_zero_expressions();
-
-        let counter_expressions = self.counter_expressions(&zero_expressions);
+        let counter_expressions = self.counter_expressions();
         // Expression IDs are indices into `self.expressions`, and on the LLVM
         // side they will be treated as indices into `counter_expressions`, so
         // the two vectors should correspond 1:1.
         assert_eq!(self.function_coverage_info.expressions.len(), counter_expressions.len());
 
-        let counter_regions = self.counter_regions(zero_expressions);
+        let counter_regions = self.counter_regions();
 
         (counter_expressions, counter_regions)
     }
 
     /// Convert this function's coverage expression data into a form that can be
     /// passed through FFI to LLVM.
-    fn counter_expressions(&self, zero_expressions: &ZeroExpressions) -> Vec<CounterExpression> {
+    fn counter_expressions(&self) -> Vec<CounterExpression> {
         // We know that LLVM will optimize out any unused expressions before
         // producing the final coverage map, so there's no need to do the same
         // thing on the Rust side unless we're confident we can do much better.
         // (See `CounterExpressionsMinimizer` in `CoverageMappingWriter.cpp`.)
 
         let counter_from_operand = |operand: CovTerm| match operand {
-            CovTerm::Expression(id) if zero_expressions.contains(id) => Counter::ZERO,
+            CovTerm::Expression(id) if self.zero_expressions.contains(id) => Counter::ZERO,
             _ => Counter::from_term(operand),
         };
 
@@ -219,10 +234,7 @@ impl<'tcx> FunctionCoverage<'tcx> {
 
     /// Converts this function's coverage mappings into an intermediate form
     /// that will be used by `mapgen` when preparing for FFI.
-    fn counter_regions(
-        &self,
-        zero_expressions: ZeroExpressions,
-    ) -> impl Iterator<Item = (Counter, &CodeRegion)> {
+    fn counter_regions(&self) -> impl Iterator<Item = (Counter, &CodeRegion)> {
         // Historically, mappings were stored directly in counter/expression
         // statements in MIR, and MIR optimizations would sometimes remove them.
         // That's mostly no longer true, so now we detect cases where that would
@@ -230,7 +242,7 @@ impl<'tcx> FunctionCoverage<'tcx> {
         let counter_for_term = move |term: CovTerm| {
             let force_to_zero = match term {
                 CovTerm::Counter(id) => !self.counters_seen.contains(id),
-                CovTerm::Expression(id) => zero_expressions.contains(id),
+                CovTerm::Expression(id) => self.zero_expressions.contains(id),
                 CovTerm::Zero => false,
             };
             if force_to_zero { Counter::ZERO } else { Counter::from_term(term) }
diff --git a/compiler/rustc_codegen_llvm/src/coverageinfo/mapgen.rs b/compiler/rustc_codegen_llvm/src/coverageinfo/mapgen.rs
index 2f825b801ac..5b2dae2c0a4 100644
--- a/compiler/rustc_codegen_llvm/src/coverageinfo/mapgen.rs
+++ b/compiler/rustc_codegen_llvm/src/coverageinfo/mapgen.rs
@@ -1,7 +1,7 @@
 use crate::common::CodegenCx;
 use crate::coverageinfo;
 use crate::coverageinfo::ffi::CounterMappingRegion;
-use crate::coverageinfo::map_data::FunctionCoverage;
+use crate::coverageinfo::map_data::{FunctionCoverage, FunctionCoverageCollector};
 use crate::llvm;
 
 use rustc_codegen_ssa::traits::{BaseTypeMethods, ConstMethods};
@@ -62,6 +62,7 @@ pub fn finalize(cx: &CodegenCx<'_, '_>) {
     // Encode coverage mappings and generate function records
     let mut function_data = Vec::new();
     for (instance, function_coverage) in function_coverage_map {
+        let function_coverage = function_coverage.into_finished();
         debug!("Generate function coverage for {}, {:?}", cx.codegen_unit.name(), instance);
 
         let mangled_function_name = tcx.symbol_name(instance).name;
@@ -419,7 +420,7 @@ fn add_unused_function_coverage<'tcx>(
 ) {
     // An unused function's mappings will automatically be rewritten to map to
     // zero, because none of its counters/expressions are marked as seen.
-    let function_coverage = FunctionCoverage::unused(instance, function_coverage_info);
+    let function_coverage = FunctionCoverageCollector::unused(instance, function_coverage_info);
 
     if let Some(coverage_context) = cx.coverage_context() {
         coverage_context.function_coverage_map.borrow_mut().insert(instance, function_coverage);
diff --git a/compiler/rustc_codegen_llvm/src/coverageinfo/mod.rs b/compiler/rustc_codegen_llvm/src/coverageinfo/mod.rs
index cf7c7e6be60..7d69756181a 100644
--- a/compiler/rustc_codegen_llvm/src/coverageinfo/mod.rs
+++ b/compiler/rustc_codegen_llvm/src/coverageinfo/mod.rs
@@ -3,7 +3,7 @@ use crate::llvm;
 use crate::builder::Builder;
 use crate::common::CodegenCx;
 use crate::coverageinfo::ffi::{CounterExpression, CounterMappingRegion};
-use crate::coverageinfo::map_data::FunctionCoverage;
+use crate::coverageinfo::map_data::FunctionCoverageCollector;
 
 use libc::c_uint;
 use rustc_codegen_ssa::traits::{
@@ -29,7 +29,8 @@ const VAR_ALIGN_BYTES: usize = 8;
 /// A context object for maintaining all state needed by the coverageinfo module.
 pub struct CrateCoverageContext<'ll, 'tcx> {
     /// Coverage data for each instrumented function identified by DefId.
-    pub(crate) function_coverage_map: RefCell<FxHashMap<Instance<'tcx>, FunctionCoverage<'tcx>>>,
+    pub(crate) function_coverage_map:
+        RefCell<FxHashMap<Instance<'tcx>, FunctionCoverageCollector<'tcx>>>,
     pub(crate) pgo_func_name_var_map: RefCell<FxHashMap<Instance<'tcx>, &'ll llvm::Value>>,
 }
 
@@ -41,7 +42,9 @@ impl<'ll, 'tcx> CrateCoverageContext<'ll, 'tcx> {
         }
     }
 
-    pub fn take_function_coverage_map(&self) -> FxHashMap<Instance<'tcx>, FunctionCoverage<'tcx>> {
+    pub fn take_function_coverage_map(
+        &self,
+    ) -> FxHashMap<Instance<'tcx>, FunctionCoverageCollector<'tcx>> {
         self.function_coverage_map.replace(FxHashMap::default())
     }
 }
@@ -93,7 +96,7 @@ impl<'tcx> CoverageInfoBuilderMethods<'tcx> for Builder<'_, '_, 'tcx> {
         let mut coverage_map = coverage_context.function_coverage_map.borrow_mut();
         let func_coverage = coverage_map
             .entry(instance)
-            .or_insert_with(|| FunctionCoverage::new(instance, function_coverage_info));
+            .or_insert_with(|| FunctionCoverageCollector::new(instance, function_coverage_info));
 
         let Coverage { kind } = coverage;
         match *kind {