Add a new CoverageKind::Branch to MIR

This adds a couple intermediate types, and threads branch coverage from MIR through codegen.

Author: Swatinem

compiler/rustc_codegen_llvm/src/coverageinfo/ffi.rs

@@ -183,9 +183,6 @@ impl CounterMappingRegion {
         }
     }
 
-    // This function might be used in the future; the LLVM API is still evolving, as is coverage
-    // support.
-    #[allow(dead_code)]
     pub(crate) fn branch_region(
         counter: Counter,
         false_counter: Counter,

compiler/rustc_codegen_llvm/src/coverageinfo/map_data.rs

@@ -16,6 +16,23 @@ pub struct Expression {
     region: Option<CodeRegion>,
 }
 
+pub struct CoverageCounterAndRegion<'a> {
+    pub(crate) kind: CoverageCounterKind,
+    pub(crate) region: &'a CodeRegion,
+}
+
+pub enum CoverageCounterKind {
+    Counter(Counter),
+    Branch { true_counter: Counter, false_counter: Counter },
+}
+
+#[derive(Debug)]
+struct CoverageBranch {
+    true_op: Operand,
+    false_op: Operand,
+    region: CodeRegion,
+}
+
 /// Collects all of the coverage regions associated with (a) injected counters, (b) counter
 /// expressions (additions or subtraction), and (c) unreachable regions (always counted as zero),
 /// for a given Function. This struct also stores the `function_source_hash`,
@@ -34,6 +51,7 @@ pub struct FunctionCoverage<'tcx> {
     is_used: bool,
     counters: IndexVec<CounterId, Option<CodeRegion>>,
     expressions: IndexVec<ExpressionId, Option<Expression>>,
+    branches: Vec<CoverageBranch>,
     unreachable_regions: Vec<CodeRegion>,
 }
 
@@ -60,6 +78,7 @@ impl<'tcx> FunctionCoverage<'tcx> {
             is_used,
             counters: IndexVec::from_elem_n(None, coverageinfo.num_counters as usize),
             expressions: IndexVec::from_elem_n(None, coverageinfo.num_expressions as usize),
+            branches: Vec::new(),
             unreachable_regions: Vec::new(),
         }
     }
@@ -86,6 +105,11 @@ impl<'tcx> FunctionCoverage<'tcx> {
         }
     }
 
+    /// Adds a branch region using the two provided true/false operands
+    pub fn add_branch_counter(&mut self, true_op: Operand, false_op: Operand, region: CodeRegion) {
+        self.branches.push(CoverageBranch { true_op, false_op, region })
+    }
+
     /// Both counters and "counter expressions" (or simply, "expressions") can be operands in other
     /// expressions. These are tracked as separate variants of `Operand`, so there is no ambiguity
     /// between operands that are counter IDs and operands that are expression IDs.
@@ -139,7 +163,7 @@ impl<'tcx> FunctionCoverage<'tcx> {
     /// `CoverageMapGenerator` will create `CounterMappingRegion`s.
     pub fn get_expressions_and_counter_regions(
         &self,
-    ) -> (Vec<CounterExpression>, Vec<(Counter, &CodeRegion)>) {
+    ) -> (Vec<CounterExpression>, Vec<CoverageCounterAndRegion<'_>>) {
         assert!(
             self.source_hash != 0 || !self.is_used,
             "No counters provided the source_hash for used function: {:?}",
@@ -153,22 +177,25 @@ impl<'tcx> FunctionCoverage<'tcx> {
         let mut collected_counter_regions = expression_regions;
         collected_counter_regions.extend(counter_regions);
         collected_counter_regions.extend(unreachable_regions);
-        collected_counter_regions.sort_unstable_by_key(|(_counter, region)| *region);
+        collected_counter_regions.sort_unstable_by_key(|counter| counter.region);
 
         (counter_expressions, collected_counter_regions)
     }
 
-    fn counter_regions(&self) -> impl Iterator<Item = (Counter, &CodeRegion)> {
+    fn counter_regions(&self) -> impl Iterator<Item = CoverageCounterAndRegion<'_>> {
         self.counters.iter_enumerated().filter_map(|(index, entry)| {
             // Option::map() will return None to filter out missing counters. This may happen
             // if, for example, a MIR-instrumented counter is removed during an optimization.
-            entry.as_ref().map(|region| (Counter::counter_value_reference(index), region))
+            entry.as_ref().map(|region| CoverageCounterAndRegion {
+                kind: CoverageCounterKind::Counter(Counter::counter_value_reference(index)),
+                region,
+            })
         })
     }
 
     fn expressions_with_regions(
         &self,
-    ) -> (Vec<CounterExpression>, Vec<(Counter, &CodeRegion)>) {
+    ) -> (Vec<CounterExpression>, Vec<CoverageCounterAndRegion<'_>>) {
         let mut counter_expressions = Vec::with_capacity(self.expressions.len());
         let mut expression_regions = Vec::with_capacity(self.expressions.len());
         let mut new_indexes = IndexVec::from_elem_n(None, self.expressions.len());
@@ -291,7 +318,12 @@ impl<'tcx> FunctionCoverage<'tcx> {
                 counter_expressions.push(expression);
                 new_indexes[original_index] = Some(mapped_expression_index);
                 if let Some(region) = optional_region {
-                    expression_regions.push((Counter::expression(mapped_expression_index), region));
+                    expression_regions.push(CoverageCounterAndRegion {
+                        kind: CoverageCounterKind::Counter(Counter::expression(
+                            mapped_expression_index,
+                        )),
+                        region,
+                    });
                 }
             } else {
                 bug!(
@@ -305,10 +337,36 @@ impl<'tcx> FunctionCoverage<'tcx> {
                 );
             }
         }
+
+        for branch in &self.branches {
+            if let Some(Some((true_counter, false_counter))) =
+                id_to_counter(&new_indexes, branch.true_op).map(|true_counter| {
+                    id_to_counter(&new_indexes, branch.false_op)
+                        .map(|false_counter| (true_counter, false_counter))
+                })
+            {
+                expression_regions.push(CoverageCounterAndRegion {
+                    kind: CoverageCounterKind::Branch { true_counter, false_counter },
+                    region: &branch.region,
+                });
+            } else {
+                bug!(
+                    "branch has one or more missing operands \
+                      true={:?}, false={:?}, region={:?}",
+                    branch.true_op,
+                    branch.false_op,
+                    branch.region,
+                );
+            }
+        }
+
         (counter_expressions, expression_regions)
     }
 
-    fn unreachable_regions(&self) -> impl Iterator<Item = (Counter, &CodeRegion)> {
-        self.unreachable_regions.iter().map(|region| (Counter::zero(), region))
+    fn unreachable_regions(&self) -> impl Iterator<Item = CoverageCounterAndRegion<'_>> {
+        self.unreachable_regions.iter().map(|region| CoverageCounterAndRegion {
+            kind: CoverageCounterKind::Counter(Counter::zero()),
+            region,
+        })
     }
 }

compiler/rustc_codegen_llvm/src/coverageinfo/mapgen.rs

@@ -1,6 +1,6 @@
 use crate::common::CodegenCx;
 use crate::coverageinfo;
-use crate::coverageinfo::ffi::{Counter, CounterExpression, CounterMappingRegion};
+use crate::coverageinfo::ffi::{CounterExpression, CounterMappingRegion};
 use crate::llvm;
 
 use rustc_codegen_ssa::traits::ConstMethods;
@@ -14,6 +14,8 @@ use rustc_middle::mir::coverage::CodeRegion;
 use rustc_middle::ty::TyCtxt;
 use rustc_span::Symbol;
 
+use super::map_data::{CoverageCounterAndRegion, CoverageCounterKind};
+
 /// Generates and exports the Coverage Map.
 ///
 /// Rust Coverage Map generation supports LLVM Coverage Mapping Format version
@@ -144,7 +146,7 @@ impl CoverageMapGenerator {
     fn write_coverage_mapping(
         &mut self,
         expressions: Vec<CounterExpression>,
-        counter_regions: Vec<(Counter, &CodeRegion)>,
+        counter_regions: Vec<CoverageCounterAndRegion<'_>>,
         coverage_mapping_buffer: &RustString,
     ) {
         if counter_regions.is_empty() {
@@ -156,12 +158,13 @@ impl CoverageMapGenerator {
         let mut current_file_name = None;
         let mut current_file_id = 0;
 
-        // Convert the list of (Counter, CodeRegion) pairs to an array of `CounterMappingRegion`, sorted
+        // Convert the list of `CoverageCounterAndRegion` to an array of `CounterMappingRegion`, sorted
         // by filename and position. Capture any new files to compute the `CounterMappingRegion`s
         // `file_id` (indexing files referenced by the current function), and construct the
         // function-specific `virtual_file_mapping` from `file_id` to its index in the module's
         // `filenames` array.
-        for (counter, region) in counter_regions {
+        for counter_region in counter_regions {
+            let region = counter_region.region;
             let CodeRegion { file_name, start_line, start_col, end_line, end_col } = *region;
             let same_file = current_file_name.is_some_and(|p| p == file_name);
             if !same_file {
@@ -173,15 +176,34 @@ impl CoverageMapGenerator {
                 let (filenames_index, _) = self.filenames.insert_full(file_name);
                 virtual_file_mapping.push(filenames_index as u32);
             }
-            debug!("Adding counter {:?} to map for {:?}", counter, region);
-            mapping_regions.push(CounterMappingRegion::code_region(
-                counter,
-                current_file_id,
-                start_line,
-                start_col,
-                end_line,
-                end_col,
-            ));
+            match counter_region.kind {
+                CoverageCounterKind::Counter(counter) => {
+                    debug!("Adding counter {:?} to map for {:?}", counter, region);
+                    mapping_regions.push(CounterMappingRegion::code_region(
+                        counter,
+                        current_file_id,
+                        start_line,
+                        start_col,
+                        end_line,
+                        end_col,
+                    ));
+                }
+                CoverageCounterKind::Branch { true_counter, false_counter } => {
+                    debug!(
+                        "Adding branch ({:?} / {:?}) to map for {:?}",
+                        true_counter, false_counter, region
+                    );
+                    mapping_regions.push(CounterMappingRegion::branch_region(
+                        true_counter,
+                        false_counter,
+                        current_file_id,
+                        start_line,
+                        start_col,
+                        end_line,
+                        end_col,
+                    ));
+                }
+            }
         }
 
         // Encode and append the current function's coverage mapping data

compiler/rustc_codegen_llvm/src/coverageinfo/mod.rs

@@ -152,6 +152,14 @@ impl<'tcx> CoverageInfoBuilderMethods<'tcx> for Builder<'_, '_, 'tcx> {
                 );
                 func_coverage.add_counter_expression(id, lhs, op, rhs, region);
             }
+            CoverageKind::Branch { true_op, false_op } => {
+                let region = region.expect("branch regions always have code regions");
+                debug!(
+                    "adding branch to coverage_map: instance={:?}, {:?} / {:?}; region: {:?}",
+                    instance, true_op, false_op, region,
+                );
+                func_coverage.add_branch_counter(true_op, false_op, region);
+            }
             CoverageKind::Unreachable => {
                 let region = region.expect("unreachable regions always have code regions");
                 debug!(

compiler/rustc_middle/src/mir/coverage.rs

@@ -93,6 +93,10 @@ pub enum CoverageKind {
         op: Op,
         rhs: Operand,
     },
+    Branch {
+        true_op: Operand,
+        false_op: Operand,
+    },
     Unreachable,
 }
 
@@ -112,6 +116,7 @@ impl Debug for CoverageKind {
                 },
                 rhs,
             ),
+            Branch { true_op, false_op } => write!(fmt, "Branch: {true_op:?} / {false_op:?}"),
             Unreachable => write!(fmt, "Unreachable"),
         }
     }

compiler/rustc_smir/src/rustc_smir/mod.rs

@@ -491,6 +491,10 @@ impl<'tcx> Stable<'tcx> for mir::coverage::CoverageKind {
                     rhs: opaque(rhs),
                 }
             }
+            CoverageKind::Branch { true_op, false_op } => stable_mir::mir::CoverageKind::Branch {
+                true_op: opaque(true_op),
+                false_op: opaque(false_op),
+            },
             CoverageKind::Unreachable => stable_mir::mir::CoverageKind::Unreachable,
         }
     }

compiler/rustc_smir/src/stable_mir/mir/body.rs

@@ -184,6 +184,10 @@ pub enum CoverageKind {
         op: Op,
         rhs: ExpressionOperandId,
     },
+    Branch {
+        true_op: ExpressionOperandId,
+        false_op: ExpressionOperandId,
+    },
     Unreachable,
 }