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_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::{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");
@@ -117,7 +117,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;
@@ -136,20 +136,48 @@ fn create_mappings<'tcx>(
             .as_term()
     };
 
-    coverage_spans
-        .all_bcb_mappings()
-        .filter_map(|&BcbMapping { kind: bcb_mapping_kind, span }| {
+    let mut mappings = vec![];
+
+    mappings.extend(coverage_spans.all_bcb_mappings().filter_map(
+        |&BcbMapping { kind: bcb_mapping_kind, span }| {
             let kind = match bcb_mapping_kind {
                 BcbMappingKind::Code(bcb) => MappingKind::Code(term_for_bcb(bcb)),
-                BcbMappingKind::Branch { true_bcb, false_bcb } => MappingKind::Branch {
-                    true_term: term_for_bcb(true_bcb),
-                    false_term: term_for_bcb(false_bcb),
-                },
             };
             let code_region = make_code_region(source_map, file_name, span, body_span)?;
             Some(Mapping { kind, code_region })
-        })
-        .collect::<Vec<_>>()
+        },
+    ));
+
+    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.
+            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(),
+            };
+
+            mappings.extend::<Option<_>>(
+                try {
+                    let code_region = make_code_region(source_map, file_name, span, body_span)?;
+                    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
 }
 
 /// For each BCB node or BCB edge that has an associated coverage counter,

compiler/rustc_mir_transform/src/coverage/spans.rs

@@ -13,8 +13,8 @@ mod from_mir;
 pub(super) enum BcbMappingKind {
     /// Associates an ordinary executable code span with its corresponding BCB.
     Code(BasicCoverageBlock),
-    /// Associates a branch span with BCBs for its true and false arms.
-    Branch { true_bcb: BasicCoverageBlock, false_bcb: BasicCoverageBlock },
+    // Branches have additional structure that can't be represented here, so
+    // they are handled separately.
 }
 
 #[derive(Debug)]
@@ -23,9 +23,15 @@ pub(super) struct BcbMapping {
     pub(super) span: Span,
 }
 
+pub(super) struct BcbBranchArm {
+    pub(super) span: Span,
+    pub(super) bcb: BasicCoverageBlock,
+}
+
 pub(super) struct CoverageSpans {
     bcb_has_mappings: BitSet<BasicCoverageBlock>,
     mappings: Vec<BcbMapping>,
+    pub(super) branch_arm_lists: Vec<Vec<BcbBranchArm>>,
 }
 
 impl CoverageSpans {
@@ -48,6 +54,7 @@ pub(super) fn generate_coverage_spans(
     basic_coverage_blocks: &CoverageGraph,
 ) -> Option<CoverageSpans> {
     let mut mappings = vec![];
+    let mut branch_arm_lists = vec![];
 
     if hir_info.is_async_fn {
         // An async function desugars into a function that returns a future,
@@ -69,14 +76,11 @@ pub(super) fn generate_coverage_spans(
             BcbMapping { kind: BcbMappingKind::Code(bcb), span }
         }));
 
-        mappings.extend(from_mir::extract_branch_mappings(
-            mir_body,
-            hir_info.body_span,
-            basic_coverage_blocks,
-        ));
+        branch_arm_lists =
+            from_mir::extract_branch_arm_lists(mir_body, hir_info, basic_coverage_blocks);
     }
 
-    if mappings.is_empty() {
+    if mappings.is_empty() && branch_arm_lists.is_empty() {
         return None;
     }
 
@@ -88,14 +92,13 @@ pub(super) fn generate_coverage_spans(
     for &BcbMapping { kind, span: _ } in &mappings {
         match kind {
             BcbMappingKind::Code(bcb) => insert(bcb),
-            BcbMappingKind::Branch { true_bcb, false_bcb } => {
-                insert(true_bcb);
-                insert(false_bcb);
-            }
         }
     }
+    for &BcbBranchArm { bcb, .. } in branch_arm_lists.iter().flatten() {
+        insert(bcb);
+    }
 
-    Some(CoverageSpans { bcb_has_mappings, mappings })
+    Some(CoverageSpans { bcb_has_mappings, mappings, branch_arm_lists })
 }
 
 #[derive(Debug)]

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

@@ -11,7 +11,7 @@ use rustc_span::{ExpnKind, MacroKind, Span, Symbol};
 use crate::coverage::graph::{
     BasicCoverageBlock, BasicCoverageBlockData, CoverageGraph, START_BCB,
 };
-use crate::coverage::spans::{BcbMapping, BcbMappingKind};
+use crate::coverage::spans::BcbBranchArm;
 use crate::coverage::ExtractedHirInfo;
 
 /// Traverses the MIR body to produce an initial collection of coverage-relevant
@@ -361,20 +361,17 @@ impl SpanFromMir {
     }
 }
 
-pub(super) fn extract_branch_mappings(
+pub(super) fn extract_branch_arm_lists(
     mir_body: &mir::Body<'_>,
-    body_span: Span,
+    hir_info: &ExtractedHirInfo,
     basic_coverage_blocks: &CoverageGraph,
-) -> Vec<BcbMapping> {
-    let Some(branch_info) = mir_body.coverage_branch_info.as_deref() else {
-        return vec![];
-    };
+) -> Vec<Vec<BcbBranchArm>> {
+    let Some(branch_info) = mir_body.coverage_branch_info.as_deref() else { return vec![] };
 
     let mut block_markers = IndexVec::<BlockMarkerId, Option<BasicBlock>>::from_elem_n(
         None,
         branch_info.num_block_markers,
     );
-
     // Fill out the mapping from block marker IDs to their enclosing blocks.
     for (bb, data) in mir_body.basic_blocks.iter_enumerated() {
         for statement in &data.statements {
@@ -393,15 +390,16 @@ pub(super) fn extract_branch_mappings(
             if !raw_span.ctxt().outer_expn_data().is_root() {
                 return None;
             }
-            let (span, _) = unexpand_into_body_span_with_visible_macro(raw_span, body_span)?;
+            let (span, _) =
+                unexpand_into_body_span_with_visible_macro(raw_span, hir_info.body_span)?;
 
             let bcb_from_marker =
                 |marker: BlockMarkerId| basic_coverage_blocks.bcb_from_bb(block_markers[marker]?);
 
             let true_bcb = bcb_from_marker(true_marker)?;
             let false_bcb = bcb_from_marker(false_marker)?;
 
-            Some(BcbMapping { kind: BcbMappingKind::Branch { true_bcb, false_bcb }, span })
+            Some(vec![BcbBranchArm { span, bcb: true_bcb }, BcbBranchArm { span, bcb: false_bcb }])
         })
         .collect::<Vec<_>>()
 }