Rollup merge of #124555 - Zalathar:init-coverage, r=nnethercote

coverage: Clean up creation of MC/DC condition bitmaps

This PR improves the code for creating and initializing [MC/DC](https://en.wikipedia.org/wiki/Modified_condition/decision_coverage) condition bitmap variables, as introduced by #123409 and modified by #124255.

- The condition bitmap variables are now created eagerly at the start of per-function codegen, via a new `init_coverage` method in `CoverageInfoBuilderMethods`. This avoids having to retroactively create the bitmaps while doing codegen for an individual coverage statement.
- As a result, we can now create and initialize those bitmaps using existing safe APIs, instead of having to perform our own unsafe call to `llvm::LLVMBuildAlloca`.
- This PR also tweaks the way we count the number of condition bitmaps needed, by tracking the total number of bitmaps needed (max depth + 1), instead of only tracking the maximum depth. This reduces the potential for subtle off-by-one confusion.

8 files changed, 72 insertions, 56 deletions

diff --git a/compiler/rustc_abi/src/lib.rs b/compiler/rustc_abi/src/lib.rs
index 6a2943da4a3..5e3f64540e4 100644
--- a/compiler/rustc_abi/src/lib.rs
+++ b/compiler/rustc_abi/src/lib.rs
@@ -21,6 +21,8 @@ use rustc_macros::{Decodable_Generic, Encodable_Generic};
 use std::iter::Step;
 
 mod layout;
+#[cfg(test)]
+mod tests;
 
 pub use layout::LayoutCalculator;
 
diff --git a/compiler/rustc_abi/src/tests.rs b/compiler/rustc_abi/src/tests.rs
new file mode 100644
index 00000000000..d993012378c
--- /dev/null
+++ b/compiler/rustc_abi/src/tests.rs
@@ -0,0 +1,7 @@
+use super::*;
+
+#[test]
+fn align_constants() {
+    assert_eq!(Align::ONE, Align::from_bytes(1).unwrap());
+    assert_eq!(Align::EIGHT, Align::from_bytes(8).unwrap());
+}
diff --git a/compiler/rustc_codegen_llvm/src/builder.rs b/compiler/rustc_codegen_llvm/src/builder.rs
index 5c8f358d03a..1a1b4ae3831 100644
--- a/compiler/rustc_codegen_llvm/src/builder.rs
+++ b/compiler/rustc_codegen_llvm/src/builder.rs
@@ -17,7 +17,7 @@ use rustc_data_structures::small_c_str::SmallCStr;
 use rustc_hir::def_id::DefId;
 use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrs;
 use rustc_middle::ty::layout::{
-    FnAbiError, FnAbiOfHelpers, FnAbiRequest, HasTyCtxt, LayoutError, LayoutOfHelpers, TyAndLayout,
+    FnAbiError, FnAbiOfHelpers, FnAbiRequest, LayoutError, LayoutOfHelpers, TyAndLayout,
 };
 use rustc_middle::ty::{self, Instance, Ty, TyCtxt};
 use rustc_sanitizers::{cfi, kcfi};
@@ -27,7 +27,6 @@ use rustc_target::abi::{self, call::FnAbi, Align, Size, WrappingRange};
 use rustc_target::spec::{HasTargetSpec, SanitizerSet, Target};
 use smallvec::SmallVec;
 use std::borrow::Cow;
-use std::ffi::CString;
 use std::iter;
 use std::ops::Deref;
 use std::ptr;
@@ -1705,13 +1704,21 @@ impl<'a, 'll, 'tcx> Builder<'a, 'll, 'tcx> {
         kcfi_bundle
     }
 
+    /// Emits a call to `llvm.instrprof.mcdc.parameters`.
+    ///
+    /// This doesn't produce any code directly, but is used as input by
+    /// the LLVM pass that handles coverage instrumentation.
+    ///
+    /// (See clang's [`CodeGenPGO::emitMCDCParameters`] for comparison.)
+    ///
+    /// [`CodeGenPGO::emitMCDCParameters`]:
+    ///     https://github.com/rust-lang/llvm-project/blob/5399a24/clang/lib/CodeGen/CodeGenPGO.cpp#L1124
     pub(crate) fn mcdc_parameters(
         &mut self,
         fn_name: &'ll Value,
         hash: &'ll Value,
         bitmap_bytes: &'ll Value,
-        max_decision_depth: u32,
-    ) -> Vec<&'ll Value> {
+    ) {
         debug!("mcdc_parameters() with args ({:?}, {:?}, {:?})", fn_name, hash, bitmap_bytes);
 
         assert!(llvm_util::get_version() >= (18, 0, 0), "MCDC intrinsics require LLVM 18 or later");
@@ -1724,8 +1731,6 @@ impl<'a, 'll, 'tcx> Builder<'a, 'll, 'tcx> {
         let args = &[fn_name, hash, bitmap_bytes];
         let args = self.check_call("call", llty, llfn, args);
 
-        let mut cond_bitmaps = vec![];
-
         unsafe {
             let _ = llvm::LLVMRustBuildCall(
                 self.llbuilder,
@@ -1736,23 +1741,7 @@ impl<'a, 'll, 'tcx> Builder<'a, 'll, 'tcx> {
                 [].as_ptr(),
                 0 as c_uint,
             );
-            // Create condition bitmap named `mcdc.addr`.
-            for i in 0..=max_decision_depth {
-                let mut bx = Builder::with_cx(self.cx);
-                bx.position_at_start(llvm::LLVMGetFirstBasicBlock(self.llfn()));
-
-                let name = CString::new(format!("mcdc.addr.{i}")).unwrap();
-                let cond_bitmap = {
-                    let alloca =
-                        llvm::LLVMBuildAlloca(bx.llbuilder, bx.cx.type_i32(), name.as_ptr());
-                    llvm::LLVMSetAlignment(alloca, 4);
-                    alloca
-                };
-                bx.store(self.const_i32(0), cond_bitmap, self.tcx().data_layout.i32_align.abi);
-                cond_bitmaps.push(cond_bitmap);
-            }
         }
-        cond_bitmaps
     }
 
     pub(crate) fn mcdc_tvbitmap_update(
@@ -1794,8 +1783,7 @@ impl<'a, 'll, 'tcx> Builder<'a, 'll, 'tcx> {
                 0 as c_uint,
             );
         }
-        let i32_align = self.tcx().data_layout.i32_align.abi;
-        self.store(self.const_i32(0), mcdc_temp, i32_align);
+        self.store(self.const_i32(0), mcdc_temp, self.tcx.data_layout.i32_align.abi);
     }
 
     pub(crate) fn mcdc_condbitmap_update(
diff --git a/compiler/rustc_codegen_llvm/src/coverageinfo/mod.rs b/compiler/rustc_codegen_llvm/src/coverageinfo/mod.rs
index 679c6e1a2ff..c51a7744a30 100644
--- a/compiler/rustc_codegen_llvm/src/coverageinfo/mod.rs
+++ b/compiler/rustc_codegen_llvm/src/coverageinfo/mod.rs
@@ -13,10 +13,10 @@ 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;
+use rustc_target::abi::{Align, Size};
 
 use std::cell::RefCell;
 
@@ -91,6 +91,42 @@ 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);
+        self.mcdc_parameters(fn_name, hash, bitmap_bytes);
+
+        // Create pointers named `mcdc.addr.{i}` to stack-allocated condition bitmaps.
+        let mut cond_bitmaps = vec![];
+        for i in 0..function_coverage_info.mcdc_num_condition_bitmaps {
+            // MC/DC intrinsics will perform loads/stores that use the ABI default
+            // alignment for i32, so our variable declaration should match.
+            let align = self.tcx.data_layout.i32_align.abi;
+            let cond_bitmap = self.alloca(Size::from_bytes(4), align);
+            llvm::set_value_name(cond_bitmap, format!("mcdc.addr.{i}").as_bytes());
+            self.store(self.const_i32(0), cond_bitmap, align);
+            cond_bitmaps.push(cond_bitmap);
+        }
+
+        self.coverage_context()
+            .expect("always present when coverage is enabled")
+            .mcdc_condition_bitmap_map
+            .borrow_mut()
+            .insert(instance, cond_bitmaps);
+    }
+
     #[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 +145,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 +225,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
diff --git a/compiler/rustc_codegen_ssa/src/mir/mod.rs b/compiler/rustc_codegen_ssa/src/mir/mod.rs
index 0064c16f5d9..cf6e2e8d14c 100644
--- a/compiler/rustc_codegen_ssa/src/mir/mod.rs
+++ b/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);
diff --git a/compiler/rustc_codegen_ssa/src/traits/coverageinfo.rs b/compiler/rustc_codegen_ssa/src/traits/coverageinfo.rs
index d1d813bd389..906d8b87d3b 100644
--- a/compiler/rustc_codegen_ssa/src/traits/coverageinfo.rs
+++ b/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
diff --git a/compiler/rustc_middle/src/mir/coverage.rs b/compiler/rustc_middle/src/mir/coverage.rs
index 9d9ca22247a..477303e2434 100644
--- a/compiler/rustc_middle/src/mir/coverage.rs
+++ b/compiler/rustc_middle/src/mir/coverage.rs
@@ -277,7 +277,7 @@ pub struct FunctionCoverageInfo {
     pub mappings: Vec<Mapping>,
     /// The depth of the deepest decision is used to know how many
     /// temp condbitmaps should be allocated for the function.
-    pub mcdc_max_decision_depth: u16,
+    pub mcdc_num_condition_bitmaps: usize,
 }
 
 /// Branch information recorded during THIR-to-MIR lowering, and stored in MIR.
diff --git a/compiler/rustc_mir_transform/src/coverage/mod.rs b/compiler/rustc_mir_transform/src/coverage/mod.rs
index 9edde666246..ffe61e761c5 100644
--- a/compiler/rustc_mir_transform/src/coverage/mod.rs
+++ b/compiler/rustc_mir_transform/src/coverage/mod.rs
@@ -102,7 +102,7 @@ fn instrument_function_for_coverage<'tcx>(tcx: TyCtxt<'tcx>, mir_body: &mut mir:
 
     inject_mcdc_statements(mir_body, &basic_coverage_blocks, &coverage_spans);
 
-    let mcdc_max_decision_depth = coverage_spans
+    let mcdc_num_condition_bitmaps = coverage_spans
         .mappings
         .iter()
         .filter_map(|bcb_mapping| match bcb_mapping.kind {
@@ -110,7 +110,7 @@ fn instrument_function_for_coverage<'tcx>(tcx: TyCtxt<'tcx>, mir_body: &mut mir:
             _ => None,
         })
         .max()
-        .unwrap_or(0);
+        .map_or(0, |max| usize::from(max) + 1);
 
     mir_body.function_coverage_info = Some(Box::new(FunctionCoverageInfo {
         function_source_hash: hir_info.function_source_hash,
@@ -118,7 +118,7 @@ fn instrument_function_for_coverage<'tcx>(tcx: TyCtxt<'tcx>, mir_body: &mut mir:
         mcdc_bitmap_bytes: coverage_spans.test_vector_bitmap_bytes(),
         expressions: coverage_counters.into_expressions(),
         mappings,
-        mcdc_max_decision_depth,
+        mcdc_num_condition_bitmaps,
     }));
 }