Auto merge of #124255 - RenjiSann:renji/mcdc-nested-expressions, r=Za…

…lathar

MCDC coverage: support nested decision coverage

#123409 provided the initial MCDC coverage implementation.

As referenced in #124144, it does not currently support "nested" decisions, like the following example :

```rust
fn nested_if_in_condition(a: bool, b: bool, c: bool) {
    if a && if b || c { true } else { false } {
        say("yes");
    } else {
        say("no");
    }
}
```

Note that there is an if-expression (`if b || c ...`) embedded inside a boolean expression in the decision of an outer if-expression.

This PR proposes a workaround for this cases, by introducing a Decision context stack, and by handing several `temporary condition bitmaps` instead of just one.
When instrumenting boolean expressions, if the current node is a leaf condition (i.e. not a `||`/`&&` logical operator nor a `!` not operator), we insert a new decision context, such that if there are more boolean expressions inside the condition, they are handled as separate expressions.

On the codegen LLVM side, we allocate as many `temp_cond_bitmap`s as necessary to handle the maximum encountered decision depth.

compiler/rustc_codegen_llvm/src/builder.rs

@@ -27,6 +27,7 @@ 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;
@@ -1709,7 +1710,8 @@ impl<'a, 'll, 'tcx> Builder<'a, 'll, 'tcx> {
         fn_name: &'ll Value,
         hash: &'ll Value,
         bitmap_bytes: &'ll Value,
-    ) -> &'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");
@@ -1722,6 +1724,8 @@ 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,
@@ -1733,17 +1737,22 @@ impl<'a, 'll, 'tcx> Builder<'a, 'll, 'tcx> {
                 0 as c_uint,
             );
             // Create condition bitmap named `mcdc.addr`.
-            let mut bx = Builder::with_cx(self.cx);
-            bx.position_at_start(llvm::LLVMGetFirstBasicBlock(self.llfn()));
-            let cond_bitmap = {
-                let alloca =
-                    llvm::LLVMBuildAlloca(bx.llbuilder, bx.cx.type_i32(), c"mcdc.addr".as_ptr());
-                llvm::LLVMSetAlignment(alloca, 4);
-                alloca
-            };
-            bx.store(self.const_i32(0), cond_bitmap, self.tcx().data_layout.i32_align.abi);
-            cond_bitmap
+            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(

compiler/rustc_codegen_llvm/src/coverageinfo/mod.rs

@@ -30,7 +30,7 @@ pub struct CrateCoverageContext<'ll, 'tcx> {
     pub(crate) function_coverage_map:
         RefCell<FxIndexMap<Instance<'tcx>, FunctionCoverageCollector<'tcx>>>,
     pub(crate) pgo_func_name_var_map: RefCell<FxHashMap<Instance<'tcx>, &'ll llvm::Value>>,
-    pub(crate) mcdc_condition_bitmap_map: RefCell<FxHashMap<Instance<'tcx>, &'ll llvm::Value>>,
+    pub(crate) mcdc_condition_bitmap_map: RefCell<FxHashMap<Instance<'tcx>, Vec<&'ll llvm::Value>>>,
 }
 
 impl<'ll, 'tcx> CrateCoverageContext<'ll, 'tcx> {
@@ -49,9 +49,20 @@ impl<'ll, 'tcx> CrateCoverageContext<'ll, 'tcx> {
     }
 
     /// LLVM use a temp value to record evaluated mcdc test vector of each decision, which is called condition bitmap.
-    /// This value is named `mcdc.addr` (same as clang) and is a 32-bit integer.
-    fn try_get_mcdc_condition_bitmap(&self, instance: &Instance<'tcx>) -> Option<&'ll llvm::Value> {
-        self.mcdc_condition_bitmap_map.borrow().get(instance).copied()
+    /// In order to handle nested decisions, several condition bitmaps can be
+    /// allocated for a function body.
+    /// These values are named `mcdc.addr.{i}` and are a 32-bit integers.
+    /// They respectively hold the condition bitmaps for decisions with a depth of `i`.
+    fn try_get_mcdc_condition_bitmap(
+        &self,
+        instance: &Instance<'tcx>,
+        decision_depth: u16,
+    ) -> Option<&'ll llvm::Value> {
+        self.mcdc_condition_bitmap_map
+            .borrow()
+            .get(instance)
+            .and_then(|bitmap_map| bitmap_map.get(decision_depth as usize))
+            .copied() // Dereference Option<&&Value> to Option<&Value>
     }
 }
 
@@ -143,26 +154,26 @@ impl<'tcx> CoverageInfoBuilderMethods<'tcx> for Builder<'_, '_, 'tcx> {
             CoverageKind::ExpressionUsed { id } => {
                 func_coverage.mark_expression_id_seen(id);
             }
-            CoverageKind::CondBitmapUpdate { id, value, .. } => {
+            CoverageKind::CondBitmapUpdate { id, value, decision_depth } => {
                 drop(coverage_map);
                 assert_ne!(
                     id.as_u32(),
                     0,
                     "ConditionId of evaluated conditions should never be zero"
                 );
                 let cond_bitmap = coverage_context
-                    .try_get_mcdc_condition_bitmap(&instance)
+                    .try_get_mcdc_condition_bitmap(&instance, decision_depth)
                     .expect("mcdc cond bitmap should have been allocated for updating");
                 let cond_loc = bx.const_i32(id.as_u32() as i32 - 1);
                 let bool_value = bx.const_bool(value);
                 let fn_name = bx.get_pgo_func_name_var(instance);
                 let hash = bx.const_u64(function_coverage_info.function_source_hash);
                 bx.mcdc_condbitmap_update(fn_name, hash, cond_loc, cond_bitmap, bool_value);
             }
-            CoverageKind::TestVectorBitmapUpdate { bitmap_idx } => {
+            CoverageKind::TestVectorBitmapUpdate { bitmap_idx, decision_depth } => {
                 drop(coverage_map);
                 let cond_bitmap = coverage_context
-                                    .try_get_mcdc_condition_bitmap(&instance)
+                                    .try_get_mcdc_condition_bitmap(&instance, decision_depth)
                                     .expect("mcdc cond bitmap should have been allocated for merging into the global bitmap");
                 let bitmap_bytes = bx.tcx().coverage_ids_info(instance.def).mcdc_bitmap_bytes;
                 assert!(bitmap_idx < bitmap_bytes, "bitmap index of the decision out of range");
@@ -195,7 +206,8 @@ fn ensure_mcdc_parameters<'ll, 'tcx>(
     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 cond_bitmap = bx.mcdc_parameters(fn_name, hash, 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

compiler/rustc_middle/src/mir/coverage.rs

@@ -132,15 +132,15 @@ pub enum CoverageKind {
     ///
     /// If this statement does not survive MIR optimizations, the condition would never be
     /// taken as evaluated.
-    CondBitmapUpdate { id: ConditionId, value: bool },
+    CondBitmapUpdate { id: ConditionId, value: bool, decision_depth: u16 },
 
     /// Marks the point in MIR control flow represented by a evaluated decision.
     ///
     /// This is eventually lowered to `llvm.instrprof.mcdc.tvbitmap.update` in LLVM IR.
     ///
     /// If this statement does not survive MIR optimizations, the decision would never be
     /// taken as evaluated.
-    TestVectorBitmapUpdate { bitmap_idx: u32 },
+    TestVectorBitmapUpdate { bitmap_idx: u32, decision_depth: u16 },
 }
 
 impl Debug for CoverageKind {
@@ -151,11 +151,17 @@ impl Debug for CoverageKind {
             BlockMarker { id } => write!(fmt, "BlockMarker({:?})", id.index()),
             CounterIncrement { id } => write!(fmt, "CounterIncrement({:?})", id.index()),
             ExpressionUsed { id } => write!(fmt, "ExpressionUsed({:?})", id.index()),
-            CondBitmapUpdate { id, value } => {
-                write!(fmt, "CondBitmapUpdate({:?}, {:?})", id.index(), value)
+            CondBitmapUpdate { id, value, decision_depth } => {
+                write!(
+                    fmt,
+                    "CondBitmapUpdate({:?}, {:?}, depth={:?})",
+                    id.index(),
+                    value,
+                    decision_depth
+                )
             }
-            TestVectorBitmapUpdate { bitmap_idx } => {
-                write!(fmt, "TestVectorUpdate({:?})", bitmap_idx)
+            TestVectorBitmapUpdate { bitmap_idx, decision_depth } => {
+                write!(fmt, "TestVectorUpdate({:?}, depth={:?})", bitmap_idx, decision_depth)
             }
         }
     }
@@ -269,6 +275,9 @@ pub struct FunctionCoverageInfo {
     pub mcdc_bitmap_bytes: u32,
     pub expressions: IndexVec<ExpressionId, Expression>,
     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,
 }
 
 /// Branch information recorded during THIR-to-MIR lowering, and stored in MIR.
@@ -319,6 +328,7 @@ pub struct MCDCBranchSpan {
     pub condition_info: Option<ConditionInfo>,
     pub true_marker: BlockMarkerId,
     pub false_marker: BlockMarkerId,
+    pub decision_depth: u16,
 }
 
 #[derive(Copy, Clone, Debug)]
@@ -334,4 +344,5 @@ pub struct MCDCDecisionSpan {
     pub span: Span,
     pub conditions_num: usize,
     pub end_markers: Vec<BlockMarkerId>,
+    pub decision_depth: u16,
 }

compiler/rustc_middle/src/mir/pretty.rs

@@ -496,20 +496,27 @@ fn write_coverage_branch_info(
         )?;
     }
 
-    for coverage::MCDCBranchSpan { span, condition_info, true_marker, false_marker } in
-        mcdc_branch_spans
+    for coverage::MCDCBranchSpan {
+        span,
+        condition_info,
+        true_marker,
+        false_marker,
+        decision_depth,
+    } in mcdc_branch_spans
     {
         writeln!(
             w,
-            "{INDENT}coverage mcdc branch {{ condition_id: {:?}, true: {true_marker:?}, false: {false_marker:?} }} => {span:?}",
+            "{INDENT}coverage mcdc branch {{ condition_id: {:?}, true: {true_marker:?}, false: {false_marker:?}, depth: {decision_depth:?} }} => {span:?}",
             condition_info.map(|info| info.condition_id)
         )?;
     }
 
-    for coverage::MCDCDecisionSpan { span, conditions_num, end_markers } in mcdc_decision_spans {
+    for coverage::MCDCDecisionSpan { span, conditions_num, end_markers, decision_depth } in
+        mcdc_decision_spans
+    {
         writeln!(
             w,
-            "{INDENT}coverage mcdc decision {{ conditions_num: {conditions_num:?}, end: {end_markers:?} }} => {span:?}"
+            "{INDENT}coverage mcdc decision {{ conditions_num: {conditions_num:?}, end: {end_markers:?}, depth: {decision_depth:?} }} => {span:?}"
         )?;
     }
 

compiler/rustc_mir_build/src/build/coverageinfo.rs

@@ -101,10 +101,14 @@ impl BranchInfoBuilder {
         tcx: TyCtxt<'_>,
         true_marker: BlockMarkerId,
         false_marker: BlockMarkerId,
-    ) -> Option<ConditionInfo> {
+    ) -> Option<(u16, ConditionInfo)> {
         let mcdc_state = self.mcdc_state.as_mut()?;
+        let decision_depth = mcdc_state.decision_depth();
         let (mut condition_info, decision_result) =
             mcdc_state.take_condition(true_marker, false_marker);
+        // take_condition() returns Some for decision_result when the decision stack
+        // is empty, i.e. when all the conditions of the decision were instrumented,
+        // and the decision is "complete".
         if let Some(decision) = decision_result {
             match decision.conditions_num {
                 0 => {
@@ -131,7 +135,7 @@ impl BranchInfoBuilder {
                 }
             }
         }
-        condition_info
+        condition_info.map(|cond_info| (decision_depth, cond_info))
     }
 
     fn add_two_way_branch<'tcx>(
@@ -199,17 +203,32 @@ impl BranchInfoBuilder {
 /// This limit may be relaxed if the [upstream change](https://github.com/llvm/llvm-project/pull/82448) is merged.
 const MAX_CONDITIONS_NUM_IN_DECISION: usize = 6;
 
-struct MCDCState {
+#[derive(Default)]
+struct MCDCDecisionCtx {
     /// To construct condition evaluation tree.
     decision_stack: VecDeque<ConditionInfo>,
     processing_decision: Option<MCDCDecisionSpan>,
 }
 
+struct MCDCState {
+    decision_ctx_stack: Vec<MCDCDecisionCtx>,
+}
+
 impl MCDCState {
     fn new_if_enabled(tcx: TyCtxt<'_>) -> Option<Self> {
         tcx.sess
             .instrument_coverage_mcdc()
-            .then(|| Self { decision_stack: VecDeque::new(), processing_decision: None })
+            .then(|| Self { decision_ctx_stack: vec![MCDCDecisionCtx::default()] })
+    }
+
+    /// Decision depth is given as a u16 to reduce the size of the `CoverageKind`,
+    /// as it is very unlikely that the depth ever reaches 2^16.
+    #[inline]
+    fn decision_depth(&self) -> u16 {
+        u16::try_from(
+            self.decision_ctx_stack.len().checked_sub(1).expect("Unexpected empty decision stack"),
+        )
+        .expect("decision depth did not fit in u16, this is likely to be an instrumentation error")
     }
 
     // At first we assign ConditionIds for each sub expression.
@@ -253,19 +272,23 @@ impl MCDCState {
     // - If the op is AND, the "false_next" of LHS and RHS should be the parent's "false_next". While "true_next" of the LHS is the RHS, the "true next" of RHS is the parent's "true_next".
     // - If the op is OR, the "true_next" of LHS and RHS should be the parent's "true_next". While "false_next" of the LHS is the RHS, the "false next" of RHS is the parent's "false_next".
     fn record_conditions(&mut self, op: LogicalOp, span: Span) {
-        let decision = match self.processing_decision.as_mut() {
+        let decision_depth = self.decision_depth();
+        let decision_ctx =
+            self.decision_ctx_stack.last_mut().expect("Unexpected empty decision_ctx_stack");
+        let decision = match decision_ctx.processing_decision.as_mut() {
             Some(decision) => {
                 decision.span = decision.span.to(span);
                 decision
             }
-            None => self.processing_decision.insert(MCDCDecisionSpan {
+            None => decision_ctx.processing_decision.insert(MCDCDecisionSpan {
                 span,
                 conditions_num: 0,
                 end_markers: vec![],
+                decision_depth,
             }),
         };
 
-        let parent_condition = self.decision_stack.pop_back().unwrap_or_default();
+        let parent_condition = decision_ctx.decision_stack.pop_back().unwrap_or_default();
         let lhs_id = if parent_condition.condition_id == ConditionId::NONE {
             decision.conditions_num += 1;
             ConditionId::from(decision.conditions_num)
@@ -305,19 +328,21 @@ impl MCDCState {
             }
         };
         // We visit expressions tree in pre-order, so place the left-hand side on the top.
-        self.decision_stack.push_back(rhs);
-        self.decision_stack.push_back(lhs);
+        decision_ctx.decision_stack.push_back(rhs);
+        decision_ctx.decision_stack.push_back(lhs);
     }
 
     fn take_condition(
         &mut self,
         true_marker: BlockMarkerId,
         false_marker: BlockMarkerId,
     ) -> (Option<ConditionInfo>, Option<MCDCDecisionSpan>) {
-        let Some(condition_info) = self.decision_stack.pop_back() else {
+        let decision_ctx =
+            self.decision_ctx_stack.last_mut().expect("Unexpected empty decision_ctx_stack");
+        let Some(condition_info) = decision_ctx.decision_stack.pop_back() else {
             return (None, None);
         };
-        let Some(decision) = self.processing_decision.as_mut() else {
+        let Some(decision) = decision_ctx.processing_decision.as_mut() else {
             bug!("Processing decision should have been created before any conditions are taken");
         };
         if condition_info.true_next_id == ConditionId::NONE {
@@ -327,8 +352,8 @@ impl MCDCState {
             decision.end_markers.push(false_marker);
         }
 
-        if self.decision_stack.is_empty() {
-            (Some(condition_info), self.processing_decision.take())
+        if decision_ctx.decision_stack.is_empty() {
+            (Some(condition_info), decision_ctx.processing_decision.take())
         } else {
             (Some(condition_info), None)
         }
@@ -364,13 +389,17 @@ impl Builder<'_, '_> {
                 |block| branch_info.inject_block_marker(&mut self.cfg, source_info, block);
             let true_marker = inject_block_marker(then_block);
             let false_marker = inject_block_marker(else_block);
-            let condition_info =
-                branch_info.fetch_mcdc_condition_info(self.tcx, true_marker, false_marker);
+            let (decision_depth, condition_info) = branch_info
+                .fetch_mcdc_condition_info(self.tcx, true_marker, false_marker)
+                .map_or((0, None), |(decision_depth, condition_info)| {
+                    (decision_depth, Some(condition_info))
+                });
             branch_info.mcdc_branch_spans.push(MCDCBranchSpan {
                 span: source_info.span,
                 condition_info,
                 true_marker,
                 false_marker,
+                decision_depth,
             });
             return;
         }
@@ -385,4 +414,20 @@ impl Builder<'_, '_> {
             mcdc_state.record_conditions(logical_op, span);
         }
     }
+
+    pub(crate) fn mcdc_increment_depth_if_enabled(&mut self) {
+        if let Some(branch_info) = self.coverage_branch_info.as_mut()
+            && let Some(mcdc_state) = branch_info.mcdc_state.as_mut()
+        {
+            mcdc_state.decision_ctx_stack.push(MCDCDecisionCtx::default());
+        };
+    }
+
+    pub(crate) fn mcdc_decrement_depth_if_enabled(&mut self) {
+        if let Some(branch_info) = self.coverage_branch_info.as_mut()
+            && let Some(mcdc_state) = branch_info.mcdc_state.as_mut()
+        {
+            mcdc_state.decision_ctx_stack.pop().expect("Unexpected empty decision stack");
+        };
+    }
 }