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

@@ -65,7 +65,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

@@ -88,7 +88,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 expression = Expression { lhs: lhs.as_term(), op, rhs: rhs.as_term() };
         let id = self.expressions.push(expression);
         BcbCounter::Expression { id }

compiler/rustc_mir_transform/src/coverage/mod.rs

@@ -9,7 +9,7 @@ mod tests;
 
 use self::counters::{CounterIncrementSite, CoverageCounters};
 use self::graph::{BasicCoverageBlock, CoverageGraph};
-use self::spans::{BcbBranchPair, BcbMapping, BcbMappingKind, CoverageSpans};
+use self::spans::{BcbBranchArm, BcbMapping, BcbMappingKind, CoverageSpans};
 
 use crate::MirPass;
 
@@ -83,10 +83,10 @@ fn instrument_function_for_coverage<'tcx>(tcx: TyCtxt<'tcx>, mir_body: &mut mir:
     // and all `Expression` dependencies (operands) are also generated, for any other
     // `BasicCoverageBlock`s not already associated with a coverage span.
     let bcb_has_coverage_spans = |bcb| coverage_spans.bcb_has_coverage_spans(bcb);
-    let coverage_counters =
+    let mut coverage_counters =
         CoverageCounters::make_bcb_counters(&basic_coverage_blocks, bcb_has_coverage_spans);
 
-    let mappings = create_mappings(tcx, &hir_info, &coverage_spans, &coverage_counters);
+    let mappings = create_mappings(tcx, &hir_info, &coverage_spans, &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");
@@ -120,7 +120,7 @@ fn create_mappings<'tcx>(
     tcx: TyCtxt<'tcx>,
     hir_info: &ExtractedHirInfo,
     coverage_spans: &CoverageSpans,
-    coverage_counters: &CoverageCounters,
+    coverage_counters: &mut CoverageCounters,
 ) -> Vec<Mapping> {
     let source_map = tcx.sess.source_map();
     let body_span = hir_info.body_span;
@@ -169,15 +169,32 @@ fn create_mappings<'tcx>(
         },
     ));
 
-    mappings.extend(coverage_spans.branch_pairs.iter().filter_map(
-        |&BcbBranchPair { 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 = make_code_region(source_map, file_name, span, body_span)?;
-            Some(Mapping { kind, code_region })
-        },
-    ));
+    for arm_list in &coverage_spans.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(&BcbBranchArm { span: _, bcb }) = arms_rev.next() else { continue };
+            coverage_counters.bcb_counter(bcb).expect("all relevant BCBs have counters")
+        };
+
+        for &BcbBranchArm { span, bcb } in arms_rev {
+            let true_counter =
+                coverage_counters.bcb_counter(bcb).expect("all relevant BCBs have counters");
+            let kind = MappingKind::Branch {
+                true_term: true_counter.as_term(),
+                false_term: rest_counter.as_term(),
+            };
+
+            if let Some(code_region) = make_code_region(source_map, file_name, span, body_span) {
+                mappings.push(Mapping { kind, code_region });
+            }
+
+            // FIXME: Avoid creating an unused expression on the last iteration.
+            rest_counter = coverage_counters.make_expression(true_counter, Op::Add, rest_counter);
+        }
+    }
 
     mappings
 }

compiler/rustc_mir_transform/src/coverage/spans.rs

@@ -37,20 +37,16 @@ pub(super) struct BcbMapping {
     pub(super) span: Span,
 }
 
-/// This is separate from [`BcbMappingKind`] 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 BcbBranchPair {
+pub(super) struct BcbBranchArm {
     pub(super) span: Span,
-    pub(super) true_bcb: BasicCoverageBlock,
-    pub(super) false_bcb: BasicCoverageBlock,
+    pub(super) bcb: BasicCoverageBlock,
 }
 
 pub(super) struct CoverageSpans {
     bcb_has_mappings: BitSet<BasicCoverageBlock>,
     pub(super) mappings: Vec<BcbMapping>,
-    pub(super) branch_pairs: Vec<BcbBranchPair>,
+    pub(super) branch_arm_lists: Vec<Vec<BcbBranchArm>>,
     test_vector_bitmap_bytes: u32,
 }
 
@@ -74,7 +70,7 @@ pub(super) fn generate_coverage_spans(
     basic_coverage_blocks: &CoverageGraph,
 ) -> Option<CoverageSpans> {
     let mut mappings = vec![];
-    let mut branch_pairs = vec![];
+    let mut branch_arm_lists = vec![];
 
     if hir_info.is_async_fn {
         // An async function desugars into a function that returns a future,
@@ -96,7 +92,7 @@ pub(super) fn generate_coverage_spans(
             BcbMapping { kind: BcbMappingKind::Code(bcb), span }
         }));
 
-        branch_pairs.extend(from_mir::extract_branch_pairs(
+        branch_arm_lists.extend(from_mir::extract_branch_arm_lists(
             mir_body,
             hir_info,
             basic_coverage_blocks,
@@ -109,7 +105,7 @@ pub(super) fn generate_coverage_spans(
         ));
     }
 
-    if mappings.is_empty() && branch_pairs.is_empty() {
+    if mappings.is_empty() && branch_arm_lists.is_empty() {
         return None;
     }
 
@@ -135,12 +131,11 @@ pub(super) fn generate_coverage_spans(
             }
         }
     }
-    for &BcbBranchPair { true_bcb, false_bcb, .. } in &branch_pairs {
-        insert(true_bcb);
-        insert(false_bcb);
+    for &BcbBranchArm { bcb, .. } in branch_arm_lists.iter().flatten() {
+        insert(bcb);
     }
 
-    Some(CoverageSpans { bcb_has_mappings, mappings, branch_pairs, test_vector_bitmap_bytes })
+    Some(CoverageSpans { bcb_has_mappings, mappings, branch_arm_lists, test_vector_bitmap_bytes })
 }
 
 #[derive(Debug)]

compiler/rustc_mir_transform/src/coverage/spans/from_mir.rs

@@ -13,7 +13,7 @@ use rustc_span::{ExpnKind, MacroKind, Span, Symbol};
 use crate::coverage::graph::{
     BasicCoverageBlock, BasicCoverageBlockData, CoverageGraph, START_BCB,
 };
-use crate::coverage::spans::{BcbBranchPair, BcbMapping, BcbMappingKind};
+use crate::coverage::spans::{BcbBranchArm, BcbMapping, BcbMappingKind};
 use crate::coverage::ExtractedHirInfo;
 
 /// Traverses the MIR body to produce an initial collection of coverage-relevant
@@ -388,16 +388,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<BcbBranchPair> {
+) -> Vec<Vec<BcbBranchArm>> {
     let Some(branch_info) = mir_body.coverage_branch_info.as_deref() else { return vec![] };
 
     let block_markers = resolve_block_markers(branch_info, mir_body);
@@ -420,7 +420,7 @@ pub(super) fn extract_branch_pairs(
             let true_bcb = bcb_from_marker(true_marker)?;
             let false_bcb = bcb_from_marker(false_marker)?;
 
-            Some(BcbBranchPair { span, true_bcb, false_bcb })
+            Some(vec![BcbBranchArm { span, bcb: true_bcb }, BcbBranchArm { span, bcb: false_bcb }])
         })
         .collect::<Vec<_>>()
 }