coverage: Represent branches as a list of arms

Within the `InstrumentCoverage` pass, we now represent branches as a list of
arms, instead of a true/false pair, until we prepare the final table of
mappings to be attached to the MIR body.

(We then flatten the list into two-way branches by treating each arm as a
branch between its success block, and the total of all later arms.)

Currently all of the branches produced by MIR building are still two-way, but
this is a step towards allowing many-way branches.

Author: Zalathar

compiler/rustc_codegen_llvm/src/coverageinfo/map_data.rs

@@ -66,7 +66,15 @@ impl<'tcx> FunctionCoverageCollector<'tcx> {
         // For each expression ID that is directly used by one or more mappings,
         // mark it as not-yet-seen. This indicates that we expect to see a
         // corresponding `ExpressionUsed` statement during MIR traversal.
-        for term in function_coverage_info.mappings.iter().flat_map(|m| m.kind.terms()) {
+        for term in function_coverage_info
+            .mappings
+            .iter()
+            // For many-armed branches, some branch mappings will have expressions
+            // that don't correspond to any node in the control-flow graph, so don't
+            // expect to see `ExpressionUsed` statements for them.
+            .filter(|m| !matches!(m.kind, MappingKind::Branch { .. }))
+            .flat_map(|m| m.kind.terms())
+        {
             if let CovTerm::Expression(id) = term {
                 expressions_seen.remove(id);
             }

compiler/rustc_mir_transform/src/coverage/counters.rs

@@ -101,7 +101,12 @@ impl CoverageCounters {
         BcbCounter::Counter { id }
     }
 
-    fn make_expression(&mut self, lhs: BcbCounter, op: Op, rhs: BcbCounter) -> BcbCounter {
+    pub(super) fn make_expression(
+        &mut self,
+        lhs: BcbCounter,
+        op: Op,
+        rhs: BcbCounter,
+    ) -> BcbCounter {
         let new_expr = BcbExpression { lhs, op, rhs };
         *self
             .expressions_memo

compiler/rustc_mir_transform/src/coverage/mappings.rs

@@ -21,14 +21,11 @@ pub(super) struct CodeMapping {
     pub(super) bcb: BasicCoverageBlock,
 }
 
-/// This is separate from [`MCDCBranch`] to help prepare for larger changes
-/// that will be needed for improved branch coverage in the future.
-/// (See <https://github.com/rust-lang/rust/pull/124217>.)
 #[derive(Debug)]
-pub(super) struct BranchPair {
+pub(super) struct BranchArm {
     pub(super) span: Span,
-    pub(super) true_bcb: BasicCoverageBlock,
-    pub(super) false_bcb: BasicCoverageBlock,
+    pub(super) pre_guard_bcb: BasicCoverageBlock,
+    pub(super) arm_taken_bcb: BasicCoverageBlock,
 }
 
 /// Associates an MC/DC branch span with condition info besides fields for normal branch.
@@ -56,7 +53,7 @@ pub(super) struct MCDCDecision {
 #[derive(Default)]
 pub(super) struct ExtractedMappings {
     pub(super) code_mappings: Vec<CodeMapping>,
-    pub(super) branch_pairs: Vec<BranchPair>,
+    pub(super) branch_arm_lists: Vec<Vec<BranchArm>>,
     pub(super) mcdc_bitmap_bytes: u32,
     pub(super) mcdc_branches: Vec<MCDCBranch>,
     pub(super) mcdc_decisions: Vec<MCDCDecision>,
@@ -71,7 +68,7 @@ pub(super) fn extract_all_mapping_info_from_mir<'tcx>(
     basic_coverage_blocks: &CoverageGraph,
 ) -> ExtractedMappings {
     let mut code_mappings = vec![];
-    let mut branch_pairs = vec![];
+    let mut branch_arm_lists = vec![];
     let mut mcdc_bitmap_bytes = 0;
     let mut mcdc_branches = vec![];
     let mut mcdc_decisions = vec![];
@@ -93,7 +90,7 @@ pub(super) fn extract_all_mapping_info_from_mir<'tcx>(
         extract_refined_covspans(mir_body, hir_info, basic_coverage_blocks, &mut code_mappings);
     }
 
-    branch_pairs.extend(extract_branch_pairs(mir_body, hir_info, basic_coverage_blocks));
+    branch_arm_lists.extend(extract_branch_arm_lists(mir_body, hir_info, basic_coverage_blocks));
 
     extract_mcdc_mappings(
         mir_body,
@@ -106,7 +103,7 @@ pub(super) fn extract_all_mapping_info_from_mir<'tcx>(
 
     ExtractedMappings {
         code_mappings,
-        branch_pairs,
+        branch_arm_lists,
         mcdc_bitmap_bytes,
         mcdc_branches,
         mcdc_decisions,
@@ -121,7 +118,7 @@ impl ExtractedMappings {
         // Fully destructure self to make sure we don't miss any fields that have mappings.
         let Self {
             code_mappings,
-            branch_pairs,
+            branch_arm_lists,
             mcdc_bitmap_bytes: _,
             mcdc_branches,
             mcdc_decisions,
@@ -136,9 +133,11 @@ impl ExtractedMappings {
         for &CodeMapping { span: _, bcb } in code_mappings {
             insert(bcb);
         }
-        for &BranchPair { true_bcb, false_bcb, .. } in branch_pairs {
-            insert(true_bcb);
-            insert(false_bcb);
+        for &BranchArm { span: _, pre_guard_bcb, arm_taken_bcb } in
+            branch_arm_lists.iter().flatten()
+        {
+            insert(pre_guard_bcb);
+            insert(arm_taken_bcb);
         }
         for &MCDCBranch { true_bcb, false_bcb, .. } in mcdc_branches {
             insert(true_bcb);
@@ -179,16 +178,16 @@ fn resolve_block_markers(
 }
 
 // FIXME: There is currently a lot of redundancy between
-// `extract_branch_pairs` and `extract_mcdc_mappings`. This is needed so
+// `extract_branch_arm_lists` and `extract_mcdc_mappings`. This is needed so
 // that they can each be modified without interfering with the other, but in
 // the long term we should try to bring them together again when branch coverage
 // and MC/DC coverage support are more mature.
 
-pub(super) fn extract_branch_pairs(
+pub(super) fn extract_branch_arm_lists(
     mir_body: &mir::Body<'_>,
     hir_info: &ExtractedHirInfo,
     basic_coverage_blocks: &CoverageGraph,
-) -> Vec<BranchPair> {
+) -> Vec<Vec<BranchArm>> {
     let Some(branch_info) = mir_body.coverage_branch_info.as_deref() else { return vec![] };
 
     let block_markers = resolve_block_markers(branch_info, mir_body);
@@ -211,7 +210,8 @@ pub(super) fn extract_branch_pairs(
             let true_bcb = bcb_from_marker(true_marker)?;
             let false_bcb = bcb_from_marker(false_marker)?;
 
-            Some(BranchPair { span, true_bcb, false_bcb })
+            let arm = |bcb| BranchArm { span, pre_guard_bcb: bcb, arm_taken_bcb: bcb };
+            Some(vec![arm(true_bcb), arm(false_bcb)])
         })
         .collect::<Vec<_>>()
 }

compiler/rustc_mir_transform/src/coverage/mod.rs

@@ -8,7 +8,7 @@ mod spans;
 mod tests;
 
 use rustc_middle::mir::coverage::{
-    CodeRegion, CoverageKind, DecisionInfo, FunctionCoverageInfo, Mapping, MappingKind,
+    CodeRegion, CoverageKind, DecisionInfo, FunctionCoverageInfo, Mapping, MappingKind, Op,
 };
 use rustc_middle::mir::{
     self, BasicBlock, BasicBlockData, SourceInfo, Statement, StatementKind, Terminator,
@@ -21,7 +21,7 @@ use rustc_span::{BytePos, Pos, RelativeBytePos, Span, Symbol};
 
 use crate::coverage::counters::{CounterIncrementSite, CoverageCounters};
 use crate::coverage::graph::{BasicCoverageBlock, CoverageGraph};
-use crate::coverage::mappings::ExtractedMappings;
+use crate::coverage::mappings::{BranchArm, ExtractedMappings};
 use crate::MirPass;
 
 /// Inserts `StatementKind::Coverage` statements that either instrument the binary with injected
@@ -91,10 +91,10 @@ fn instrument_function_for_coverage<'tcx>(tcx: TyCtxt<'tcx>, mir_body: &mut mir:
     }
 
     let bcb_has_counter_mappings = |bcb| bcbs_with_counter_mappings.contains(bcb);
-    let coverage_counters =
+    let mut coverage_counters =
         CoverageCounters::make_bcb_counters(&basic_coverage_blocks, bcb_has_counter_mappings);
 
-    let mappings = create_mappings(tcx, &hir_info, &extracted_mappings, &coverage_counters);
+    let mappings = create_mappings(tcx, &hir_info, &extracted_mappings, &mut coverage_counters);
     if mappings.is_empty() {
         // No spans could be converted into valid mappings, so skip this function.
         debug!("no spans could be converted into valid mappings; skipping");
@@ -136,7 +136,7 @@ fn create_mappings<'tcx>(
     tcx: TyCtxt<'tcx>,
     hir_info: &ExtractedHirInfo,
     extracted_mappings: &ExtractedMappings,
-    coverage_counters: &CoverageCounters,
+    coverage_counters: &mut CoverageCounters,
 ) -> Vec<Mapping> {
     let source_map = tcx.sess.source_map();
     let body_span = hir_info.body_span;
@@ -148,24 +148,66 @@ fn create_mappings<'tcx>(
         &source_file.name.for_scope(tcx.sess, RemapPathScopeComponents::MACRO).to_string_lossy(),
     );
 
-    let term_for_bcb = |bcb| {
-        coverage_counters
-            .bcb_counter(bcb)
-            .expect("all BCBs with spans were given counters")
-            .as_term()
-    };
     let region_for_span = |span: Span| make_code_region(source_map, file_name, span, body_span);
 
     // Fully destructure the mappings struct to make sure we don't miss any kinds.
     let ExtractedMappings {
         code_mappings,
-        branch_pairs,
+        branch_arm_lists,
         mcdc_bitmap_bytes: _,
         mcdc_branches,
         mcdc_decisions,
     } = extracted_mappings;
     let mut mappings = Vec::new();
 
+    // Process branch arms first, because they might need to mutate `coverage_counters`
+    // to create new expressions.
+    for arm_list in branch_arm_lists {
+        let mut arms_rev = arm_list.iter().rev();
+
+        let mut rest_counter = {
+            // The last arm's span is ignored, because its BCB is only used as the
+            // false branch of the second-last arm; it's not a branch of its own.
+            let Some(&BranchArm { span: _, pre_guard_bcb, arm_taken_bcb }) = arms_rev.next() else {
+                continue;
+            };
+            debug_assert_eq!(pre_guard_bcb, arm_taken_bcb, "last arm should not have a guard");
+            coverage_counters.bcb_counter(pre_guard_bcb).expect("all relevant BCBs have counters")
+        };
+
+        // All relevant BCBs should have counters, so we can `.unwrap()` them.
+        for &BranchArm { span, pre_guard_bcb, arm_taken_bcb } in arms_rev {
+            // Number of times the pattern matched.
+            let matched_counter = coverage_counters.bcb_counter(pre_guard_bcb).unwrap();
+            // Number of times the pattern matched and the guard succeeded.
+            let arm_taken_counter = coverage_counters.bcb_counter(arm_taken_bcb).unwrap();
+            // Total number of times execution logically reached this pattern.
+            let reached_counter =
+                coverage_counters.make_expression(rest_counter, Op::Add, arm_taken_counter);
+            // Number of times execution reached this pattern, but didn't match it.
+            let unmatched_counter =
+                coverage_counters.make_expression(reached_counter, Op::Subtract, matched_counter);
+
+            let kind = MappingKind::Branch {
+                true_term: matched_counter.as_term(),
+                false_term: unmatched_counter.as_term(),
+            };
+
+            if let Some(code_region) = region_for_span(span) {
+                mappings.push(Mapping { kind, code_region });
+            }
+
+            rest_counter = reached_counter;
+        }
+    }
+
+    let term_for_bcb = |bcb| {
+        coverage_counters
+            .bcb_counter(bcb)
+            .expect("all BCBs with spans were given counters")
+            .as_term()
+    };
+
     mappings.extend(code_mappings.iter().filter_map(
         // Ordinary code mappings are the simplest kind.
         |&mappings::CodeMapping { span, bcb }| {
@@ -175,16 +217,6 @@ fn create_mappings<'tcx>(
         },
     ));
 
-    mappings.extend(branch_pairs.iter().filter_map(
-        |&mappings::BranchPair { span, true_bcb, false_bcb }| {
-            let true_term = term_for_bcb(true_bcb);
-            let false_term = term_for_bcb(false_bcb);
-            let kind = MappingKind::Branch { true_term, false_term };
-            let code_region = region_for_span(span)?;
-            Some(Mapping { kind, code_region })
-        },
-    ));
-
     mappings.extend(mcdc_branches.iter().filter_map(
         |&mappings::MCDCBranch { span, true_bcb, false_bcb, condition_info, decision_depth: _ }| {
             let code_region = region_for_span(span)?;