Rollup merge of #139397 - Zalathar:virtual, r=jieyouxu

coverage: Build the CGU's global file table as late as possible

Embedding coverage metadata in the output binary is a delicate dance, because per-function records need to embed references to the per-CGU filename table, but we only want to include files in that table if they are successfully used by at least one function.

The way that we build the file tables has changed a few times over the last few years. This particular change is motivated by experimental work on properly supporting macro-expansion regions, which adds some additional constraints that our previous implementation wasn't equipped to deal with.

LLVM is very strict about not allowing unused entries in local file tables. Currently that's not much of an issue, because we assume one source file per function, but to support expansion regions we need the flexibility to avoid committing to the use of a file until we're completely sure that we are able and willing to produce at least one coverage mapping region for it. In particular, when preparing a function's covfun record, we need the flexibility to decide at a late stage that a particular file isn't needed/usable after all.

(It's OK for the *global* file table to contain unused entries, but we would still prefer to avoid that if possible, and this implementation also achieves that.)

Author: Zalathar

compiler/rustc_codegen_llvm/src/coverageinfo/mapgen.rs

@@ -53,31 +53,20 @@ pub(crate) fn finalize(cx: &CodegenCx<'_, '_>) {
         None => return,
     };
 
-    // The order of entries in this global file table needs to be deterministic,
-    // and ideally should also be independent of the details of stable-hashing,
-    // because coverage tests snapshots (`.cov-map`) can observe the order and
-    // would need to be re-blessed if it changes. As long as those requirements
-    // are satisfied, the order can be arbitrary.
-    let mut global_file_table = GlobalFileTable::new();
-
     let mut covfun_records = instances_used
         .iter()
         .copied()
         // Sort by symbol name, so that the global file table is built in an
         // order that doesn't depend on the stable-hash-based order in which
         // instances were visited during codegen.
         .sorted_by_cached_key(|&instance| tcx.symbol_name(instance).name)
-        .filter_map(|instance| prepare_covfun_record(tcx, &mut global_file_table, instance, true))
+        .filter_map(|instance| prepare_covfun_record(tcx, instance, true))
         .collect::<Vec<_>>();
 
     // In a single designated CGU, also prepare covfun records for functions
     // in this crate that were instrumented for coverage, but are unused.
     if cx.codegen_unit.is_code_coverage_dead_code_cgu() {
-        unused::prepare_covfun_records_for_unused_functions(
-            cx,
-            &mut global_file_table,
-            &mut covfun_records,
-        );
+        unused::prepare_covfun_records_for_unused_functions(cx, &mut covfun_records);
     }
 
     // If there are no covfun records for this CGU, don't generate a covmap record.
@@ -89,68 +78,88 @@ pub(crate) fn finalize(cx: &CodegenCx<'_, '_>) {
         return;
     }
 
-    // Encode all filenames referenced by coverage mappings in this CGU.
-    let filenames_buffer = global_file_table.make_filenames_buffer(tcx);
-    // The `llvm-cov` tool uses this hash to associate each covfun record with
-    // its corresponding filenames table, since the final binary will typically
-    // contain multiple covmap records from different compilation units.
-    let filenames_hash = llvm_cov::hash_bytes(&filenames_buffer);
+    // Prepare the global file table for this CGU, containing all paths needed
+    // by one or more covfun records.
+    let global_file_table =
+        GlobalFileTable::build(tcx, covfun_records.iter().flat_map(|c| c.all_source_files()));
 
     for covfun in &covfun_records {
-        covfun::generate_covfun_record(cx, filenames_hash, covfun)
+        covfun::generate_covfun_record(cx, &global_file_table, covfun)
     }
 
     // Generate the coverage map header, which contains the filenames used by
     // this CGU's coverage mappings, and store it in a well-known global.
     // (This is skipped if we returned early due to having no covfun records.)
-    generate_covmap_record(cx, covmap_version, &filenames_buffer);
+    generate_covmap_record(cx, covmap_version, &global_file_table.filenames_buffer);
 }
 
-/// Maps "global" (per-CGU) file ID numbers to their underlying source files.
+/// Maps "global" (per-CGU) file ID numbers to their underlying source file paths.
+#[derive(Debug)]
 struct GlobalFileTable {
     /// This "raw" table doesn't include the working dir, so a file's
     /// global ID is its index in this set **plus one**.
-    raw_file_table: FxIndexMap<StableSourceFileId, Arc<SourceFile>>,
+    raw_file_table: FxIndexMap<StableSourceFileId, String>,
+
+    /// The file table in encoded form (possibly compressed), which can be
+    /// included directly in this CGU's `__llvm_covmap` record.
+    filenames_buffer: Vec<u8>,
+
+    /// Truncated hash of the bytes in `filenames_buffer`.
+    ///
+    /// The `llvm-cov` tool uses this hash to associate each covfun record with
+    /// its corresponding filenames table, since the final binary will typically
+    /// contain multiple covmap records from different compilation units.
+    filenames_hash: u64,
 }
 
 impl GlobalFileTable {
-    fn new() -> Self {
-        Self { raw_file_table: FxIndexMap::default() }
-    }
+    /// Builds a "global file table" for this CGU, mapping numeric IDs to
+    /// path strings.
+    fn build<'a>(tcx: TyCtxt<'_>, all_files: impl Iterator<Item = &'a SourceFile>) -> Self {
+        let mut raw_file_table = FxIndexMap::default();
+
+        for file in all_files {
+            raw_file_table.entry(file.stable_id).or_insert_with(|| {
+                file.name
+                    .for_scope(tcx.sess, RemapPathScopeComponents::MACRO)
+                    .to_string_lossy()
+                    .into_owned()
+            });
+        }
+
+        // FIXME(Zalathar): Consider sorting the file table here, but maybe
+        // only after adding filename support to coverage-dump, so that the
+        // table order isn't directly visible in `.coverage-map` snapshots.
+
+        let mut table = Vec::with_capacity(raw_file_table.len() + 1);
+
+        // Since version 6 of the LLVM coverage mapping format, the first entry
+        // in the global file table is treated as a base directory, used to
+        // resolve any other entries that are stored as relative paths.
+        let base_dir = tcx
+            .sess
+            .opts
+            .working_dir
+            .for_scope(tcx.sess, RemapPathScopeComponents::MACRO)
+            .to_string_lossy();
+        table.push(base_dir.as_ref());
 
-    fn global_file_id_for_file(&mut self, file: &Arc<SourceFile>) -> GlobalFileId {
-        // Ensure the given file has a table entry, and get its index.
-        let entry = self.raw_file_table.entry(file.stable_id);
-        let raw_id = entry.index();
-        entry.or_insert_with(|| Arc::clone(file));
-
-        // The raw file table doesn't include an entry for the working dir
-        // (which has ID 0), so add 1 to get the correct ID.
-        GlobalFileId::from_usize(raw_id + 1)
-    }
+        // Add the regular entries after the base directory.
+        table.extend(raw_file_table.values().map(|name| name.as_str()));
 
-    fn make_filenames_buffer(&self, tcx: TyCtxt<'_>) -> Vec<u8> {
-        let mut table = Vec::with_capacity(self.raw_file_table.len() + 1);
-
-        // LLVM Coverage Mapping Format version 6 (zero-based encoded as 5)
-        // requires setting the first filename to the compilation directory.
-        // Since rustc generates coverage maps with relative paths, the
-        // compilation directory can be combined with the relative paths
-        // to get absolute paths, if needed.
-        table.push(
-            tcx.sess
-                .opts
-                .working_dir
-                .for_scope(tcx.sess, RemapPathScopeComponents::MACRO)
-                .to_string_lossy(),
-        );
+        // Encode the file table into a buffer, and get the hash of its encoded
+        // bytes, so that we can embed that hash in `__llvm_covfun` records.
+        let filenames_buffer = llvm_cov::write_filenames_to_buffer(&table);
+        let filenames_hash = llvm_cov::hash_bytes(&filenames_buffer);
 
-        // Add the regular entries after the base directory.
-        table.extend(self.raw_file_table.values().map(|file| {
-            file.name.for_scope(tcx.sess, RemapPathScopeComponents::MACRO).to_string_lossy()
-        }));
+        Self { raw_file_table, filenames_buffer, filenames_hash }
+    }
 
-        llvm_cov::write_filenames_to_buffer(&table)
+    fn get_existing_id(&self, file: &SourceFile) -> Option<GlobalFileId> {
+        let raw_id = self.raw_file_table.get_index_of(&file.stable_id)?;
+        // The raw file table doesn't include an entry for the base dir
+        // (which has ID 0), so add 1 to get the correct ID.
+        Some(GlobalFileId::from_usize(raw_id + 1))
     }
 }
 
@@ -166,26 +175,31 @@ rustc_index::newtype_index! {
     struct LocalFileId {}
 }
 
-/// Holds a mapping from "local" (per-function) file IDs to "global" (per-CGU)
-/// file IDs.
+/// Holds a mapping from "local" (per-function) file IDs to their corresponding
+/// source files.
 #[derive(Debug, Default)]
 struct VirtualFileMapping {
-    local_to_global: IndexVec<LocalFileId, GlobalFileId>,
-    global_to_local: FxIndexMap<GlobalFileId, LocalFileId>,
+    local_file_table: IndexVec<LocalFileId, Arc<SourceFile>>,
 }
 
 impl VirtualFileMapping {
-    fn local_id_for_global(&mut self, global_file_id: GlobalFileId) -> LocalFileId {
-        *self
-            .global_to_local
-            .entry(global_file_id)
-            .or_insert_with(|| self.local_to_global.push(global_file_id))
+    fn push_file(&mut self, source_file: &Arc<SourceFile>) -> LocalFileId {
+        self.local_file_table.push(Arc::clone(source_file))
     }
 
-    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()
+    /// Resolves all of the filenames in this local file mapping to a list of
+    /// global file IDs in its CGU, for inclusion in this function's
+    /// `__llvm_covfun` record.
+    ///
+    /// The global file IDs are returned as `u32` to make FFI easier.
+    fn resolve_all(&self, global_file_table: &GlobalFileTable) -> Option<Vec<u32>> {
+        self.local_file_table
+            .iter()
+            .map(|file| try {
+                let id = global_file_table.get_existing_id(file)?;
+                GlobalFileId::as_u32(id)
+            })
+            .collect::<Option<Vec<_>>>()
     }
 }
 

compiler/rustc_codegen_llvm/src/coverageinfo/mapgen/covfun.rs

@@ -5,6 +5,7 @@
 //! [^win]: On Windows the section name is `.lcovfun`.
 
 use std::ffi::CString;
+use std::sync::Arc;
 
 use rustc_abi::Align;
 use rustc_codegen_ssa::traits::{
@@ -15,7 +16,7 @@ use rustc_middle::mir::coverage::{
     MappingKind, Op,
 };
 use rustc_middle::ty::{Instance, TyCtxt};
-use rustc_span::Span;
+use rustc_span::{SourceFile, Span};
 use rustc_target::spec::HasTargetSpec;
 use tracing::debug;
 
@@ -37,9 +38,16 @@ pub(crate) struct CovfunRecord<'tcx> {
     regions: ffi::Regions,
 }
 
+impl<'tcx> CovfunRecord<'tcx> {
+    /// Iterator that yields all source files referred to by this function's
+    /// coverage mappings. Used to build the global file table for the CGU.
+    pub(crate) fn all_source_files(&self) -> impl Iterator<Item = &SourceFile> {
+        self.virtual_file_mapping.local_file_table.iter().map(Arc::as_ref)
+    }
+}
+
 pub(crate) fn prepare_covfun_record<'tcx>(
     tcx: TyCtxt<'tcx>,
-    global_file_table: &mut GlobalFileTable,
     instance: Instance<'tcx>,
     is_used: bool,
 ) -> Option<CovfunRecord<'tcx>> {
@@ -57,7 +65,7 @@ pub(crate) fn prepare_covfun_record<'tcx>(
         regions: ffi::Regions::default(),
     };
 
-    fill_region_tables(tcx, global_file_table, fn_cov_info, ids_info, &mut covfun);
+    fill_region_tables(tcx, fn_cov_info, ids_info, &mut covfun);
 
     if covfun.regions.has_no_regions() {
         debug!(?covfun, "function has no mappings to embed; skipping");
@@ -92,7 +100,6 @@ fn prepare_expressions(ids_info: &CoverageIdsInfo) -> Vec<ffi::CounterExpression
 /// Populates the mapping region tables in the current function's covfun record.
 fn fill_region_tables<'tcx>(
     tcx: TyCtxt<'tcx>,
-    global_file_table: &mut GlobalFileTable,
     fn_cov_info: &'tcx FunctionCoverageInfo,
     ids_info: &'tcx CoverageIdsInfo,
     covfun: &mut CovfunRecord<'tcx>,
@@ -106,11 +113,7 @@ fn fill_region_tables<'tcx>(
     };
     let source_file = source_map.lookup_source_file(first_span.lo());
 
-    // Look up the global file ID for that file.
-    let global_file_id = global_file_table.global_file_id_for_file(&source_file);
-
-    // Associate that global file ID with a local file ID for this function.
-    let local_file_id = covfun.virtual_file_mapping.local_id_for_global(global_file_id);
+    let local_file_id = covfun.virtual_file_mapping.push_file(&source_file);
 
     // In rare cases, _all_ of a function's spans are discarded, and coverage
     // codegen needs to handle that gracefully to avoid #133606.
@@ -179,7 +182,7 @@ fn fill_region_tables<'tcx>(
 /// as a global variable in the `__llvm_covfun` section.
 pub(crate) fn generate_covfun_record<'tcx>(
     cx: &CodegenCx<'_, 'tcx>,
-    filenames_hash: u64,
+    global_file_table: &GlobalFileTable,
     covfun: &CovfunRecord<'tcx>,
 ) {
     let &CovfunRecord {
@@ -191,12 +194,19 @@ pub(crate) fn generate_covfun_record<'tcx>(
         ref regions,
     } = covfun;
 
+    let Some(local_file_table) = virtual_file_mapping.resolve_all(global_file_table) else {
+        debug_assert!(
+            false,
+            "all local files should be present in the global file table: \
+                global_file_table = {global_file_table:?}, \
+                virtual_file_mapping = {virtual_file_mapping:?}"
+        );
+        return;
+    };
+
     // 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,
-    );
+    let coverage_mapping_buffer =
+        llvm_cov::write_function_mappings_to_buffer(&local_file_table, expressions, regions);
 
     // A covfun record consists of four target-endian integers, followed by the
     // encoded mapping data in bytes. Note that the length field is 32 bits.
@@ -209,7 +219,7 @@ pub(crate) fn generate_covfun_record<'tcx>(
             cx.const_u64(func_name_hash),
             cx.const_u32(coverage_mapping_buffer.len() as u32),
             cx.const_u64(source_hash),
-            cx.const_u64(filenames_hash),
+            cx.const_u64(global_file_table.filenames_hash),
             cx.const_bytes(&coverage_mapping_buffer),
         ],
         // This struct needs to be packed, so that the 32-bit length field

compiler/rustc_codegen_llvm/src/coverageinfo/mapgen/unused.rs

@@ -7,7 +7,6 @@ use rustc_middle::ty::{self, TyCtxt};
 use rustc_span::def_id::DefIdSet;
 
 use crate::common::CodegenCx;
-use crate::coverageinfo::mapgen::GlobalFileTable;
 use crate::coverageinfo::mapgen::covfun::{CovfunRecord, prepare_covfun_record};
 use crate::llvm;
 
@@ -21,7 +20,6 @@ use crate::llvm;
 /// its embedded coverage data.
 pub(crate) fn prepare_covfun_records_for_unused_functions<'tcx>(
     cx: &CodegenCx<'_, 'tcx>,
-    global_file_table: &mut GlobalFileTable,
     covfun_records: &mut Vec<CovfunRecord<'tcx>>,
 ) {
     assert!(cx.codegen_unit.is_code_coverage_dead_code_cgu());
@@ -33,7 +31,7 @@ pub(crate) fn prepare_covfun_records_for_unused_functions<'tcx>(
     // Try to create a covfun record for each unused function.
     let mut name_globals = Vec::with_capacity(unused_instances.len());
     covfun_records.extend(unused_instances.into_iter().filter_map(|unused| try {
-        let record = prepare_covfun_record(cx.tcx, global_file_table, unused.instance, false)?;
+        let record = prepare_covfun_record(cx.tcx, unused.instance, false)?;
         // If successful, also store its symbol name in a global constant.
         name_globals.push(cx.const_str(unused.symbol_name.name).0);
         record