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,14 @@ 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 mapping in function_coverage_info.mappings.iter() {
+        for mapping 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 { .. }))
+        {
             // Currently we only worry about ordinary code mappings.
             // For branch and MC/DC mappings, expressions might not correspond
             // to any particular point in the control-flow graph.
@@ -196,7 +203,11 @@ impl<'tcx> FunctionCoverage<'tcx> {
     /// Return the source hash, generated from the HIR node structure, and used to indicate whether
     /// or not the source code structure changed between different compilations.
     pub(crate) fn source_hash(&self) -> u64 {
-        if self.is_used { self.function_coverage_info.function_source_hash } else { 0 }
+        if self.is_used {
+            self.function_coverage_info.function_source_hash
+        } else {
+            0
+        }
     }
 
     /// Returns an iterator over all filenames used by this function's mappings.
@@ -240,7 +251,11 @@ impl<'tcx> FunctionCoverage<'tcx> {
     }
 
     fn counter_for_term(&self, term: CovTerm) -> Counter {
-        if self.is_zero_term(term) { Counter::ZERO } else { Counter::from_term(term) }
+        if self.is_zero_term(term) {
+            Counter::ZERO
+        } else {
+            Counter::from_term(term)
+        }
     }
 
     fn is_zero_term(&self, term: CovTerm) -> bool {

compiler/rustc_mir_transform/src/coverage/counters.rs

@@ -122,7 +122,12 @@ impl CoverageCounters {
         self.set_bcb_edge_counter(from_bcb, to_bcb, counter)
     }
 
-    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

@@ -22,14 +22,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.
@@ -61,7 +58,7 @@ pub(super) struct ExtractedMappings {
     /// only public so that other code can still use exhaustive destructuring.
     pub(super) num_bcbs: usize,
     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>,
@@ -76,7 +73,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![];
@@ -98,7 +95,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,
@@ -112,7 +109,7 @@ pub(super) fn extract_all_mapping_info_from_mir<'tcx>(
     ExtractedMappings {
         num_bcbs: basic_coverage_blocks.num_nodes(),
         code_mappings,
-        branch_pairs,
+        branch_arm_lists,
         mcdc_bitmap_bytes,
         mcdc_branches,
         mcdc_decisions,
@@ -125,7 +122,7 @@ impl ExtractedMappings {
         let Self {
             num_bcbs,
             code_mappings,
-            branch_pairs,
+            branch_arm_lists,
             mcdc_bitmap_bytes: _,
             mcdc_branches,
             mcdc_decisions,
@@ -140,9 +137,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);
@@ -192,16 +191,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(coverage_info_hi) = mir_body.coverage_info_hi.as_deref() else { return vec![] };
 
     let block_markers = resolve_block_markers(coverage_info_hi, mir_body);
@@ -223,7 +222,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

@@ -13,7 +13,7 @@ use rustc_hir::intravisit::{Visitor, walk_expr};
 use rustc_middle::hir::map::Map;
 use rustc_middle::hir::nested_filter;
 use rustc_middle::mir::coverage::{
-    CoverageKind, DecisionInfo, FunctionCoverageInfo, Mapping, MappingKind, SourceRegion,
+    CoverageKind, DecisionInfo, FunctionCoverageInfo, Mapping, MappingKind, Op, SourceRegion,
 };
 use rustc_middle::mir::{
     self, BasicBlock, BasicBlockData, SourceInfo, Statement, StatementKind, Terminator,
@@ -27,7 +27,7 @@ use tracing::{debug, debug_span, instrument, trace};
 
 use crate::coverage::counters::{CounterIncrementSite, CoverageCounters};
 use crate::coverage::graph::CoverageGraph;
-use crate::coverage::mappings::ExtractedMappings;
+use crate::coverage::mappings::{BranchArm, ExtractedMappings};
 
 /// Inserts `StatementKind::Coverage` statements that either instrument the binary with injected
 /// counters, via intrinsic `llvm.instrprof.increment`, and/or inject metadata used during codegen
@@ -95,10 +95,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");
@@ -140,7 +140,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;
@@ -153,25 +153,67 @@ 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_source_region(source_map, file_name, span, body_span);
 
     // Fully destructure the mappings struct to make sure we don't miss any kinds.
     let ExtractedMappings {
         num_bcbs: _,
         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(source_region) = region_for_span(span) {
+                mappings.push(Mapping { kind, source_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 }| {
@@ -181,16 +223,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 source_region = region_for_span(span)?;
-            Some(Mapping { kind, source_region })
-        },
-    ));
-
     mappings.extend(mcdc_branches.iter().filter_map(
         |&mappings::MCDCBranch { span, true_bcb, false_bcb, condition_info, decision_depth: _ }| {
             let source_region = region_for_span(span)?;