coverage: Store the number of counters/expressions in function coverage info

Coverage codegen can now allocate arrays based on the number of
counters/expressions originally used by the instrumentor.

The existing query that inspects coverage statements is still used for
determining the number of counters passed to `llvm.instrprof.increment`. If
some high-numbered counters were removed by MIR optimizations, the instrumented
binary can potentially use less memory and disk space at runtime.

compiler/rustc_codegen_llvm/src/coverageinfo/map_data.rs

@@ -6,7 +6,6 @@ use rustc_middle::mir::coverage::{
     CodeRegion, CounterId, ExpressionId, FunctionCoverageInfo, Op, Operand,
 };
 use rustc_middle::ty::Instance;
-use rustc_middle::ty::TyCtxt;
 
 #[derive(Clone, Debug, PartialEq)]
 pub struct Expression {
@@ -40,38 +39,36 @@ pub struct FunctionCoverage<'tcx> {
 impl<'tcx> FunctionCoverage<'tcx> {
     /// Creates a new set of coverage data for a used (called) function.
     pub fn new(
-        tcx: TyCtxt<'tcx>,
         instance: Instance<'tcx>,
         function_coverage_info: &'tcx FunctionCoverageInfo,
     ) -> Self {
-        Self::create(tcx, instance, function_coverage_info, true)
+        Self::create(instance, function_coverage_info, true)
     }
 
     /// Creates a new set of coverage data for an unused (never called) function.
     pub fn unused(
-        tcx: TyCtxt<'tcx>,
         instance: Instance<'tcx>,
         function_coverage_info: &'tcx FunctionCoverageInfo,
     ) -> Self {
-        Self::create(tcx, instance, function_coverage_info, false)
+        Self::create(instance, function_coverage_info, false)
     }
 
     fn create(
-        tcx: TyCtxt<'tcx>,
         instance: Instance<'tcx>,
         function_coverage_info: &'tcx FunctionCoverageInfo,
         is_used: bool,
     ) -> Self {
-        let coverageinfo = tcx.coverageinfo(instance.def);
+        let num_counters = function_coverage_info.num_counters;
+        let num_expressions = function_coverage_info.num_expressions;
         debug!(
-            "FunctionCoverage::create(instance={:?}) has coverageinfo={:?}. is_used={}",
-            instance, coverageinfo, is_used
+            "FunctionCoverage::create(instance={instance:?}) has \
+            num_counters={num_counters}, num_expressions={num_expressions}, is_used={is_used}"
         );
         Self {
             function_coverage_info,
             is_used,
-            counters: IndexVec::from_elem_n(None, coverageinfo.num_counters as usize),
-            expressions: IndexVec::from_elem_n(None, coverageinfo.num_expressions as usize),
+            counters: IndexVec::from_elem_n(None, num_counters),
+            expressions: IndexVec::from_elem_n(None, num_expressions),
             unreachable_regions: Vec::new(),
         }
     }

compiler/rustc_codegen_llvm/src/coverageinfo/mod.rs

@@ -114,7 +114,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(bx.tcx(), instance, function_coverage_info));
+            .or_insert_with(|| FunctionCoverage::new(instance, function_coverage_info));
 
         let Coverage { kind, code_regions } = coverage;
         match *kind {
@@ -124,11 +124,21 @@ impl<'tcx> CoverageInfoBuilderMethods<'tcx> for Builder<'_, '_, 'tcx> {
                 // as that needs an exclusive borrow.
                 drop(coverage_map);
 
-                let coverageinfo = bx.tcx().coverageinfo(instance.def);
+                // The number of counters passed to `llvm.instrprof.increment` might
+                // be smaller than the number originally inserted by the instrumentor,
+                // if some high-numbered counters were removed by MIR optimizations.
+                // If so, LLVM's profiler runtime will use fewer physical counters.
+                let num_counters =
+                    bx.tcx().coverage_ids_info(instance.def).max_counter_id.as_u32() + 1;
+                assert!(
+                    num_counters as usize <= function_coverage_info.num_counters,
+                    "num_counters disagreement: query says {num_counters} but function info only has {}",
+                    function_coverage_info.num_counters
+                );
 
                 let fn_name = bx.get_pgo_func_name_var(instance);
                 let hash = bx.const_u64(function_coverage_info.function_source_hash);
-                let num_counters = bx.const_u32(coverageinfo.num_counters);
+                let num_counters = bx.const_u32(num_counters);
                 let index = bx.const_u32(id.as_u32());
                 debug!(
                     "codegen intrinsic instrprof.increment(fn_name={:?}, hash={:?}, num_counters={:?}, index={:?})",
@@ -208,7 +218,7 @@ fn add_unused_function_coverage<'tcx>(
 ) {
     let tcx = cx.tcx;
 
-    let mut function_coverage = FunctionCoverage::unused(tcx, instance, function_coverage_info);
+    let mut function_coverage = FunctionCoverage::unused(instance, function_coverage_info);
     for &code_region in tcx.covered_code_regions(def_id) {
         let code_region = std::slice::from_ref(code_region);
         function_coverage.add_unreachable_regions(code_region);