coverage: Store intermediate region tables in `CovfunRecord`

This defers the call to `llvm_cov::write_function_mappings_to_buffer` until
just before its enclosing global variable is created.

4 files changed, 77 insertions, 55 deletions

diff --git a/compiler/rustc_codegen_llvm/src/coverageinfo/ffi.rs b/compiler/rustc_codegen_llvm/src/coverageinfo/ffi.rs
index a6e07ea2a60..1f133141c18 100644
--- a/compiler/rustc_codegen_llvm/src/coverageinfo/ffi.rs
+++ b/compiler/rustc_codegen_llvm/src/coverageinfo/ffi.rs
@@ -152,6 +152,34 @@ impl CoverageSpan {
     }
 }
 
+/// Holds tables of the various region types in one struct.
+///
+/// Don't pass this struct across FFI; pass the individual region tables as
+/// pointer/length pairs instead.
+///
+/// Each field name has a `_regions` suffix for improved readability after
+/// exhaustive destructing, which ensures that all region types are handled.
+#[derive(Clone, Debug, Default)]
+pub(crate) struct Regions {
+    pub(crate) code_regions: Vec<CodeRegion>,
+    pub(crate) branch_regions: Vec<BranchRegion>,
+    pub(crate) mcdc_branch_regions: Vec<MCDCBranchRegion>,
+    pub(crate) mcdc_decision_regions: Vec<MCDCDecisionRegion>,
+}
+
+impl Regions {
+    /// Returns true if none of this structure's tables contain any regions.
+    pub(crate) fn has_no_regions(&self) -> bool {
+        let Self { code_regions, branch_regions, mcdc_branch_regions, mcdc_decision_regions } =
+            self;
+
+        code_regions.is_empty()
+            && branch_regions.is_empty()
+            && mcdc_branch_regions.is_empty()
+            && mcdc_decision_regions.is_empty()
+    }
+}
+
 /// Must match the layout of `LLVMRustCoverageCodeRegion`.
 #[derive(Clone, Debug)]
 #[repr(C)]
diff --git a/compiler/rustc_codegen_llvm/src/coverageinfo/llvm_cov.rs b/compiler/rustc_codegen_llvm/src/coverageinfo/llvm_cov.rs
index 99c2d12b261..086cf1f44a0 100644
--- a/compiler/rustc_codegen_llvm/src/coverageinfo/llvm_cov.rs
+++ b/compiler/rustc_codegen_llvm/src/coverageinfo/llvm_cov.rs
@@ -62,11 +62,10 @@ pub(crate) fn write_filenames_to_buffer<'a>(
 pub(crate) fn write_function_mappings_to_buffer(
     virtual_file_mapping: &[u32],
     expressions: &[ffi::CounterExpression],
-    code_regions: &[ffi::CodeRegion],
-    branch_regions: &[ffi::BranchRegion],
-    mcdc_branch_regions: &[ffi::MCDCBranchRegion],
-    mcdc_decision_regions: &[ffi::MCDCDecisionRegion],
+    regions: &ffi::Regions,
 ) -> Vec<u8> {
+    let ffi::Regions { code_regions, branch_regions, mcdc_branch_regions, mcdc_decision_regions } =
+        regions;
     llvm::build_byte_buffer(|buffer| unsafe {
         llvm::LLVMRustCoverageWriteFunctionMappingsToBuffer(
             virtual_file_mapping.as_ptr(),
diff --git a/compiler/rustc_codegen_llvm/src/coverageinfo/mapgen.rs b/compiler/rustc_codegen_llvm/src/coverageinfo/mapgen.rs
index f7c3193a449..a573a37beb3 100644
--- a/compiler/rustc_codegen_llvm/src/coverageinfo/mapgen.rs
+++ b/compiler/rustc_codegen_llvm/src/coverageinfo/mapgen.rs
@@ -194,7 +194,7 @@ rustc_index::newtype_index! {
 
 /// Holds a mapping from "local" (per-function) file IDs to "global" (per-CGU)
 /// file IDs.
-#[derive(Default)]
+#[derive(Debug, Default)]
 struct VirtualFileMapping {
     local_to_global: IndexVec<LocalFileId, GlobalFileId>,
     global_to_local: FxIndexMap<GlobalFileId, LocalFileId>,
@@ -208,10 +208,10 @@ impl VirtualFileMapping {
             .or_insert_with(|| self.local_to_global.push(global_file_id))
     }
 
-    fn into_vec(self) -> Vec<u32> {
-        // This conversion should be optimized away to ~zero overhead.
-        // In any case, it's probably not hot enough to worry about.
-        self.local_to_global.into_iter().map(|global| global.as_u32()).collect()
+    fn to_vec(&self) -> Vec<u32> {
+        // This clone could be avoided by transmuting `&[GlobalFileId]` to `&[u32]`,
+        // but it isn't hot or expensive enough to justify the extra unsafety.
+        self.local_to_global.iter().map(|&global| GlobalFileId::as_u32(global)).collect()
     }
 }
 
diff --git a/compiler/rustc_codegen_llvm/src/coverageinfo/mapgen/covfun.rs b/compiler/rustc_codegen_llvm/src/coverageinfo/mapgen/covfun.rs
index a68b026a4c0..530e6827f55 100644
--- a/compiler/rustc_codegen_llvm/src/coverageinfo/mapgen/covfun.rs
+++ b/compiler/rustc_codegen_llvm/src/coverageinfo/mapgen/covfun.rs
@@ -29,7 +29,10 @@ pub(crate) struct CovfunRecord<'tcx> {
     mangled_function_name: &'tcx str,
     source_hash: u64,
     is_used: bool,
-    coverage_mapping_buffer: Vec<u8>,
+
+    virtual_file_mapping: VirtualFileMapping,
+    expressions: Vec<ffi::CounterExpression>,
+    regions: ffi::Regions,
 }
 
 impl<'tcx> CovfunRecord<'tcx> {
@@ -46,51 +49,41 @@ pub(crate) fn prepare_covfun_record<'tcx>(
     instance: Instance<'tcx>,
     function_coverage: &FunctionCoverage<'tcx>,
 ) -> Option<CovfunRecord<'tcx>> {
-    let mangled_function_name = tcx.symbol_name(instance).name;
-    let source_hash = function_coverage.source_hash();
-    let is_used = function_coverage.is_used();
-
-    let coverage_mapping_buffer =
-        encode_mappings_for_function(tcx, global_file_table, function_coverage);
-
-    if coverage_mapping_buffer.is_empty() {
-        if function_coverage.is_used() {
-            bug!(
-                "A used function should have had coverage mapping data but did not: {}",
-                mangled_function_name
-            );
+    let mut covfun = CovfunRecord {
+        mangled_function_name: tcx.symbol_name(instance).name,
+        source_hash: function_coverage.source_hash(),
+        is_used: function_coverage.is_used(),
+        virtual_file_mapping: VirtualFileMapping::default(),
+        expressions: function_coverage.counter_expressions().collect::<Vec<_>>(),
+        regions: ffi::Regions::default(),
+    };
+
+    fill_region_tables(tcx, global_file_table, function_coverage, &mut covfun);
+
+    if covfun.regions.has_no_regions() {
+        if covfun.is_used {
+            bug!("a used function should have had coverage mapping data but did not: {covfun:?}");
         } else {
-            debug!("unused function had no coverage mapping data: {}", mangled_function_name);
+            debug!(?covfun, "unused function had no coverage mapping data");
             return None;
         }
     }
 
-    Some(CovfunRecord { mangled_function_name, source_hash, is_used, coverage_mapping_buffer })
+    Some(covfun)
 }
 
-/// Using the expressions and counter regions collected for a single function,
-/// generate the variable-sized payload of its corresponding `__llvm_covfun`
-/// entry. The payload is returned as a vector of bytes.
-///
-/// Newly-encountered filenames will be added to the global file table.
-fn encode_mappings_for_function(
-    tcx: TyCtxt<'_>,
+/// Populates the mapping region tables in the current function's covfun record.
+fn fill_region_tables<'tcx>(
+    tcx: TyCtxt<'tcx>,
     global_file_table: &GlobalFileTable,
-    function_coverage: &FunctionCoverage<'_>,
-) -> Vec<u8> {
+    function_coverage: &FunctionCoverage<'tcx>,
+    covfun: &mut CovfunRecord<'tcx>,
+) {
     let counter_regions = function_coverage.counter_regions();
     if counter_regions.is_empty() {
-        return Vec::new();
+        return;
     }
 
-    let expressions = function_coverage.counter_expressions().collect::<Vec<_>>();
-
-    let mut virtual_file_mapping = VirtualFileMapping::default();
-    let mut code_regions = vec![];
-    let mut branch_regions = vec![];
-    let mut mcdc_branch_regions = vec![];
-    let mut mcdc_decision_regions = vec![];
-
     // Currently a function's mappings must all be in the same file as its body span.
     let file_name = span_file_name(tcx, function_coverage.function_coverage_info.body_span);
 
@@ -98,9 +91,12 @@ fn encode_mappings_for_function(
     let global_file_id = global_file_table.global_file_id_for_file_name(file_name);
 
     // Associate that global file ID with a local file ID for this function.
-    let local_file_id = virtual_file_mapping.local_id_for_global(global_file_id);
+    let local_file_id = covfun.virtual_file_mapping.local_id_for_global(global_file_id);
     debug!("  file id: {local_file_id:?} => {global_file_id:?} = '{file_name:?}'");
 
+    let ffi::Regions { code_regions, branch_regions, mcdc_branch_regions, mcdc_decision_regions } =
+        &mut covfun.regions;
+
     // For each counter/region pair in this function+file, convert it to a
     // form suitable for FFI.
     for (mapping_kind, region) in counter_regions {
@@ -133,16 +129,6 @@ fn encode_mappings_for_function(
             }
         }
     }
-
-    // Encode the function's coverage mappings into a buffer.
-    llvm_cov::write_function_mappings_to_buffer(
-        &virtual_file_mapping.into_vec(),
-        &expressions,
-        &code_regions,
-        &branch_regions,
-        &mcdc_branch_regions,
-        &mcdc_decision_regions,
-    )
 }
 
 /// Generates the contents of the covfun record for this function, which
@@ -157,9 +143,18 @@ pub(crate) fn generate_covfun_record<'tcx>(
         mangled_function_name,
         source_hash,
         is_used,
-        ref coverage_mapping_buffer, // Previously-encoded coverage mappings
+        ref virtual_file_mapping,
+        ref expressions,
+        ref regions,
     } = covfun;
 
+    // Encode the function's coverage mappings into a buffer.
+    let coverage_mapping_buffer = llvm_cov::write_function_mappings_to_buffer(
+        &virtual_file_mapping.to_vec(),
+        expressions,
+        regions,
+    );
+
     // Concatenate the encoded coverage mappings
     let coverage_mapping_size = coverage_mapping_buffer.len();
     let coverage_mapping_val = cx.const_bytes(&coverage_mapping_buffer);