mcdc-coverage: Refactor the MCDC info building

Author: RenjiSann

compiler/rustc_mir_build/src/build/coverageinfo.rs

@@ -4,7 +4,7 @@ use std::collections::hash_map::Entry;
 use rustc_data_structures::fx::FxHashMap;
 use rustc_index::IndexVec;
 use rustc_middle::mir::coverage::{
-    BlockMarkerId, BranchSpan, CoverageKind, DecisionId, DecisionSpan,
+    BlockMarkerId, BranchSpan, ConditionId, CoverageKind, DecisionId, DecisionSpan,
 };
 use rustc_middle::mir::{self, BasicBlock, UnOp};
 use rustc_middle::thir::{ExprId, ExprKind, Thir};
@@ -15,14 +15,7 @@ use rustc_span::Span;
 use crate::build::Builder;
 use crate::errors::MCDCNestedDecision;
 
-pub(crate) struct BranchInfoBuilder {
-    /// Maps condition expressions to their enclosing `!`, for better instrumentation.
-    nots: FxHashMap<ExprId, NotInfo>,
-
-    num_block_markers: usize,
-    branch_spans: Vec<BranchSpan>,
-
-    // MCDC decision stuff
+struct MCDCInfoBuilder {
     /// ID of the current decision.
     /// Do not use directly. Use the function instead, as it will hide
     /// the decision in the scope of nested decisions.
@@ -31,7 +24,54 @@ pub(crate) struct BranchInfoBuilder {
     /// nested decisions.
     nested_decision_level: u32,
     /// Vector for storing all the decisions with their span
-    decisions: IndexVec<DecisionId, Span>,
+    decision_spans: IndexVec<DecisionId, Span>,
+
+    next_condition_id: u32,
+}
+
+impl MCDCInfoBuilder {
+    /// Increase the nested decision level and return true if the
+    /// decision can be instrumented (not in a nested condition).
+    pub fn enter_decision(&mut self, span: Span) -> bool {
+        self.nested_decision_level += 1;
+        let can_mcdc = !self.in_nested_condition();
+
+        if can_mcdc {
+            self.current_decision_id = Some(self.decision_spans.push(span));
+        }
+
+        can_mcdc
+    }
+
+    pub fn exit_decision(&mut self) {
+        self.nested_decision_level -= 1;
+    }
+
+    /// Return true if the current decision is located inside another decision.
+    pub fn in_nested_condition(&self) -> bool {
+        self.nested_decision_level > 1
+    }
+
+    pub fn current_decision_id(&self) -> Option<DecisionId> {
+        if self.in_nested_condition() { None } else { self.current_decision_id }
+    }
+
+    pub fn next_condition_id(&mut self) -> u32 {
+        let res = self.next_condition_id;
+        self.next_condition_id += 1;
+        res
+    }
+}
+
+pub(crate) struct BranchInfoBuilder {
+    /// Maps condition expressions to their enclosing `!`, for better instrumentation.
+    nots: FxHashMap<ExprId, NotInfo>,
+
+    num_block_markers: usize,
+    branch_spans: Vec<BranchSpan>,
+
+    // MCDC decision stuff
+    mcdc_info: Option<MCDCInfoBuilder>,
 }
 
 #[derive(Clone, Copy)]
@@ -51,13 +91,21 @@ impl BranchInfoBuilder {
         if (tcx.sess.instrument_coverage_branch() || tcx.sess.instrument_coverage_mcdc())
             && tcx.is_eligible_for_coverage(def_id)
         {
+            let mcdc_info = if tcx.sess.instrument_coverage_mcdc() {
+                Some(MCDCInfoBuilder {
+                    current_decision_id: None,
+                    nested_decision_level: 0,
+                    decision_spans: IndexVec::new(),
+                    next_condition_id: 0,
+                })
+            } else {
+                None
+            };
             Some(Self {
                 nots: FxHashMap::default(),
                 num_block_markers: 0,
                 branch_spans: vec![],
-                current_decision_id: None,
-                nested_decision_level: 0,
-                decisions: IndexVec::new(),
+                mcdc_info,
             })
         } else {
             None
@@ -111,15 +159,18 @@ impl BranchInfoBuilder {
     }
 
     pub(crate) fn into_done(self) -> Option<Box<mir::coverage::BranchInfo>> {
-        let Self { nots: _, num_block_markers, branch_spans, decisions, .. } = self;
+        let Self { nots: _, num_block_markers, branch_spans, mcdc_info, .. } = self;
 
         if num_block_markers == 0 {
             assert!(branch_spans.is_empty());
             return None;
         }
 
         let mut decision_spans = IndexVec::from_iter(
-            decisions.into_iter().map(|span| DecisionSpan { span, num_conditions: 0 }),
+            mcdc_info
+                .map_or(IndexVec::new(), |mcdc_info| mcdc_info.decision_spans)
+                .into_iter()
+                .map(|span| DecisionSpan { span, num_conditions: 0 }),
         );
 
         // Count the number of conditions linked to each decision.
@@ -135,32 +186,6 @@ impl BranchInfoBuilder {
             decision_spans,
         }))
     }
-
-    /// Increase the nested decision level and return true if the
-    /// decision can be instrumented (not in a nested condition).
-    pub fn enter_decision(&mut self, span: Span) -> bool {
-        self.nested_decision_level += 1;
-        let can_mcdc = !self.in_nested_condition();
-
-        if can_mcdc {
-            self.current_decision_id = Some(self.decisions.push(span));
-        }
-
-        can_mcdc
-    }
-
-    pub fn exit_decision(&mut self) {
-        self.nested_decision_level -= 1;
-    }
-
-    /// Return true if the current decision is located inside another decision.
-    pub fn in_nested_condition(&self) -> bool {
-        self.nested_decision_level > 1
-    }
-
-    pub fn current_decision_id(&self) -> Option<DecisionId> {
-        if self.in_nested_condition() { None } else { self.current_decision_id }
-    }
 }
 
 impl Builder<'_, '_> {
@@ -206,7 +231,11 @@ impl Builder<'_, '_> {
         branch_info.branch_spans.push(BranchSpan {
             span: source_info.span,
             // FIXME(dprn): Handle case when MCDC is disabled better than just putting 0.
-            decision_id: branch_info.current_decision_id.unwrap_or(DecisionId::from_u32(0)),
+            decision_id: branch_info
+                .mcdc_info
+                .as_ref()
+                .and_then(|mcdc_info| mcdc_info.current_decision_id())
+                .unwrap_or(DecisionId::from_u32(0)),
             true_marker,
             false_marker,
         });
@@ -215,25 +244,24 @@ impl Builder<'_, '_> {
     /// If MCDC coverage is enabled, inject a decision entry marker in the given decision.
     /// return true
     pub(crate) fn begin_mcdc_decision_coverage(&mut self, expr_id: ExprId, block: BasicBlock) {
-        // Early return if MCDC coverage is not enabled.
-        if !self.tcx.sess.instrument_coverage_mcdc() {
-            return;
-        }
-        let Some(branch_info) = self.coverage_branch_info.as_mut() else {
+        // Get the MCDCInfoBuilder object, which existence implies that MCDC is enabled.
+        let Some(BranchInfoBuilder { mcdc_info: Some(mcdc_info), .. }) =
+            self.coverage_branch_info.as_mut()
+        else {
             return;
         };
 
         let span = self.thir[expr_id].span;
 
         // enter_decision returns false if it detects nested decisions.
-        if !branch_info.enter_decision(span) {
-            // FIXME(dprn): do WARNING for nested decision.
+        if !mcdc_info.enter_decision(span) {
+            // FIXME(dprn): WARNING for nested decision does not seem to work properly
             debug!("MCDC: Unsupported nested decision");
             self.tcx.dcx().emit_warn(MCDCNestedDecision { span });
             return;
         }
 
-        let decision_id = branch_info.current_decision_id().expect("Should have returned.");
+        let decision_id = mcdc_info.current_decision_id().expect("Should have returned.");
 
         // Inject a decision marker
         let source_info = self.source_info(span);
@@ -248,16 +276,15 @@ impl Builder<'_, '_> {
 
     /// If MCDC is enabled, and function is instrumented,
     pub(crate) fn end_mcdc_decision_coverage(&mut self) {
-        // Early return if MCDC coverage is not enabled.
-        if !self.tcx.sess.instrument_coverage_mcdc() {
-            return;
-        }
-        let Some(branch_info) = self.coverage_branch_info.as_mut() else {
+        // Get the MCDCInfoBuilder object, which existence implies that MCDC is enabled.
+        let Some(BranchInfoBuilder { mcdc_info: Some(mcdc_info), .. }) =
+            self.coverage_branch_info.as_mut()
+        else {
             return;
         };
 
         // Exit decision now so we can drop &mut to branch info
-        branch_info.exit_decision();
+        mcdc_info.exit_decision();
     }
 
     /// If MCDC is enabled and the current decision is being instrumented,
@@ -268,16 +295,19 @@ impl Builder<'_, '_> {
         bb: BasicBlock,
         outcome: bool,
     ) -> BasicBlock {
-        let Some(branch_info) = self.coverage_branch_info.as_mut() else {
-            // Coverage instrumentation is not enabled.
+        // Get the MCDCInfoBuilder object, which existence implies that MCDC is enabled.
+        let Some(BranchInfoBuilder { mcdc_info: Some(mcdc_info), .. }) =
+            self.coverage_branch_info.as_mut()
+        else {
             return bb;
         };
-        let Some(decision_id) = branch_info.current_decision_id() else {
+
+        let Some(decision_id) = mcdc_info.current_decision_id() else {
             // Decision is not instrumented
             return bb;
         };
 
-        let span = branch_info.decisions[decision_id];
+        let span = mcdc_info.decision_spans[decision_id];
         let source_info = self.source_info(span);
         let marker_statement = mir::Statement {
             source_info,
@@ -298,4 +328,49 @@ impl Builder<'_, '_> {
 
         new_bb
     }
+
+    /// Add markers on the condition's basic blocks to ease the later MCDC instrumentation.
+    pub(crate) fn visit_mcdc_condition(
+        &mut self,
+        condition_expr: ExprId,
+        condition_block: BasicBlock,
+        then_block: BasicBlock,
+        else_block: BasicBlock,
+    ) {
+        // Get the MCDCInfoBuilder object, which existence implies that MCDC is enabled.
+        let Some(BranchInfoBuilder { mcdc_info: Some(mcdc_info), .. }) =
+            self.coverage_branch_info.as_mut()
+        else {
+            return;
+        };
+
+        let Some(decision_id) = mcdc_info.current_decision_id() else {
+            // If current_decision_id() is None, the decision is not instrumented.
+            return;
+        };
+
+
+        let id = ConditionId::from_u32(mcdc_info.next_condition_id());
+        let span = self.thir[condition_expr].span;
+        let source_info = self.source_info(span);
+
+        let mut inject_statement = |bb, cov_kind| {
+            let statement =
+                mir::Statement { source_info, kind: mir::StatementKind::Coverage(cov_kind) };
+            self.cfg.basic_blocks[bb].statements.insert(0, statement);
+        };
+
+        inject_statement(
+            condition_block,
+            CoverageKind::MCDCConditionEntryMarker { decision_id, id },
+        );
+        inject_statement(
+            then_block,
+            CoverageKind::MCDCConditionOutputMarker { decision_id, id, outcome: true },
+        );
+        inject_statement(
+            else_block,
+            CoverageKind::MCDCConditionOutputMarker { decision_id, id, outcome: false },
+        );
+    }
 }

compiler/rustc_mir_build/src/build/matches/mod.rs

@@ -169,6 +169,14 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                     then_block,
                     else_block,
                 );
+                // Record MCDC coverage info
+                // (Does nothing if mcdc coverage is not enabled.)
+                this.visit_mcdc_condition(
+                    expr_id,
+                    block,
+                    then_block,
+                    else_block,
+                );
 
                 let source_info = this.source_info(expr_span);
                 this.cfg.terminate(block, source_info, term);