mcdc-coverage: Get decision_depth from THIR lowering

Use decision context stack to handle nested decisions:
- Introduce MCDCDecisionCtx
- Use a stack of MCDCDecisionCtx to handle nested decisions

2 files changed, 66 insertions, 17 deletions

diff --git a/compiler/rustc_mir_build/src/build/coverageinfo.rs b/compiler/rustc_mir_build/src/build/coverageinfo.rs
index 1818be2c1bd..a1c20fcfd83 100644
--- a/compiler/rustc_mir_build/src/build/coverageinfo.rs
+++ b/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,20 +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: 0,
+                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)
@@ -306,8 +328,8 @@ 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(
@@ -315,10 +337,12 @@ impl MCDCState {
         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 {
@@ -328,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)
         }
@@ -365,14 +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: 0,
+                decision_depth,
             });
             return;
         }
@@ -387,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");
+        };
+    }
 }
diff --git a/compiler/rustc_mir_build/src/build/matches/mod.rs b/compiler/rustc_mir_build/src/build/matches/mod.rs
index f46dceeeedf..b36e97194c9 100644
--- a/compiler/rustc_mir_build/src/build/matches/mod.rs
+++ b/compiler/rustc_mir_build/src/build/matches/mod.rs
@@ -151,8 +151,14 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                 let mut block = block;
                 let temp_scope = args.temp_scope_override.unwrap_or_else(|| this.local_scope());
                 let mutability = Mutability::Mut;
+
+                // Increment the decision depth, in case we encounter boolean expressions
+                // further down.
+                this.mcdc_increment_depth_if_enabled();
                 let place =
                     unpack!(block = this.as_temp(block, Some(temp_scope), expr_id, mutability));
+                this.mcdc_decrement_depth_if_enabled();
+
                 let operand = Operand::Move(Place::from(place));
 
                 let then_block = this.cfg.start_new_block();