coverage: Eagerly do start-of-function codegen for coverage

Author: Zalathar

compiler/rustc_codegen_llvm/src/coverageinfo/mod.rs

@@ -13,7 +13,7 @@ use rustc_codegen_ssa::traits::{
 use rustc_data_structures::fx::{FxHashMap, FxIndexMap};
 use rustc_llvm::RustString;
 use rustc_middle::bug;
-use rustc_middle::mir::coverage::{CoverageKind, FunctionCoverageInfo};
+use rustc_middle::mir::coverage::CoverageKind;
 use rustc_middle::ty::layout::HasTyCtxt;
 use rustc_middle::ty::Instance;
 use rustc_target::abi::Align;
@@ -91,6 +91,32 @@ impl<'ll, 'tcx> CodegenCx<'ll, 'tcx> {
 }
 
 impl<'tcx> CoverageInfoBuilderMethods<'tcx> for Builder<'_, '_, 'tcx> {
+    fn init_coverage(&mut self, instance: Instance<'tcx>) {
+        let Some(function_coverage_info) =
+            self.tcx.instance_mir(instance.def).function_coverage_info.as_deref()
+        else {
+            return;
+        };
+
+        // If there are no MC/DC bitmaps to set up, return immediately.
+        if function_coverage_info.mcdc_bitmap_bytes == 0 {
+            return;
+        }
+
+        let fn_name = self.get_pgo_func_name_var(instance);
+        let hash = self.const_u64(function_coverage_info.function_source_hash);
+        let bitmap_bytes = self.const_u32(function_coverage_info.mcdc_bitmap_bytes);
+        let max_decision_depth = function_coverage_info.mcdc_max_decision_depth;
+        let cond_bitmap =
+            self.mcdc_parameters(fn_name, hash, bitmap_bytes, max_decision_depth as u32);
+
+        self.coverage_context()
+            .expect("always present when coverage is enabled")
+            .mcdc_condition_bitmap_map
+            .borrow_mut()
+            .insert(instance, cond_bitmap);
+    }
+
     #[instrument(level = "debug", skip(self))]
     fn add_coverage(&mut self, instance: Instance<'tcx>, kind: &CoverageKind) {
         // Our caller should have already taken care of inlining subtleties,
@@ -109,10 +135,6 @@ impl<'tcx> CoverageInfoBuilderMethods<'tcx> for Builder<'_, '_, 'tcx> {
             return;
         };
 
-        if function_coverage_info.mcdc_bitmap_bytes > 0 {
-            ensure_mcdc_parameters(bx, instance, function_coverage_info);
-        }
-
         let Some(coverage_context) = bx.coverage_context() else { return };
         let mut coverage_map = coverage_context.function_coverage_map.borrow_mut();
         let func_coverage = coverage_map
@@ -193,28 +215,6 @@ impl<'tcx> CoverageInfoBuilderMethods<'tcx> for Builder<'_, '_, 'tcx> {
     }
 }
 
-fn ensure_mcdc_parameters<'ll, 'tcx>(
-    bx: &mut Builder<'_, 'll, 'tcx>,
-    instance: Instance<'tcx>,
-    function_coverage_info: &FunctionCoverageInfo,
-) {
-    let Some(cx) = bx.coverage_context() else { return };
-    if cx.mcdc_condition_bitmap_map.borrow().contains_key(&instance) {
-        return;
-    }
-
-    let fn_name = bx.get_pgo_func_name_var(instance);
-    let hash = bx.const_u64(function_coverage_info.function_source_hash);
-    let bitmap_bytes = bx.const_u32(function_coverage_info.mcdc_bitmap_bytes);
-    let max_decision_depth = function_coverage_info.mcdc_max_decision_depth;
-    let cond_bitmap = bx.mcdc_parameters(fn_name, hash, bitmap_bytes, max_decision_depth as u32);
-    bx.coverage_context()
-        .expect("already checked above")
-        .mcdc_condition_bitmap_map
-        .borrow_mut()
-        .insert(instance, cond_bitmap);
-}
-
 /// Calls llvm::createPGOFuncNameVar() with the given function instance's
 /// mangled function name. The LLVM API returns an llvm::GlobalVariable
 /// containing the function name, with the specific variable name and linkage

compiler/rustc_codegen_ssa/src/mir/mod.rs

@@ -259,6 +259,10 @@ pub fn codegen_mir<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
     // Apply debuginfo to the newly allocated locals.
     fx.debug_introduce_locals(&mut start_bx);
 
+    // If the backend supports coverage, and coverage is enabled for this function,
+    // do any necessary start-of-function codegen (e.g. locals for MC/DC bitmaps).
+    start_bx.init_coverage(instance);
+
     // The builders will be created separately for each basic block at `codegen_block`.
     // So drop the builder of `start_llbb` to avoid having two at the same time.
     drop(start_bx);

compiler/rustc_codegen_ssa/src/traits/coverageinfo.rs

@@ -3,6 +3,11 @@ use rustc_middle::mir::coverage::CoverageKind;
 use rustc_middle::ty::Instance;
 
 pub trait CoverageInfoBuilderMethods<'tcx>: BackendTypes {
+    /// Performs any start-of-function codegen needed for coverage instrumentation.
+    ///
+    /// Can be a no-op in backends that don't support coverage instrumentation.
+    fn init_coverage(&mut self, _instance: Instance<'tcx>) {}
+
     /// Handle the MIR coverage info in a backend-specific way.
     ///
     /// This can potentially be a no-op in backends that don't support