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

@@ -181,9 +181,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

@@ -14,6 +14,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_: Operand,
+    false_: 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`,
@@ -32,6 +49,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>,
 }
 
@@ -58,6 +76,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(),
         }
     }
@@ -84,6 +103,11 @@ impl<'tcx> FunctionCoverage<'tcx> {
         }
     }
 
+    /// Adds a branch region using the two provided true/false operands
+    pub fn add_branch_counter(&mut self, true_: Operand, false_: Operand, region: CodeRegion) {
+        self.branches.push(CoverageBranch { true_, false_, 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.
@@ -185,7 +209,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: {:?}",
@@ -201,26 +225,43 @@ impl<'tcx> FunctionCoverage<'tcx> {
         let counter_regions = 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,
+            })
         });
 
         // Find all of the expression IDs that weren't optimized out AND have
         // an attached code region, and return the corresponding mapping as a
         // counter/region pair.
         let expression_regions =
             self.expressions.iter_enumerated().filter_map(|(id, expression)| {
-                let code_region = expression.as_ref()?.region.as_ref()?;
-                Some((Counter::expression(id), code_region))
+                let region = expression.as_ref()?.region.as_ref()?;
+                Some(CoverageCounterAndRegion {
+                    kind: CoverageCounterKind::Counter(Counter::expression(id)),
+                    region,
+                })
             });
-        let unreachable_regions =
-            self.unreachable_regions.iter().map(|region| (Counter::ZERO, region));
+
+        let unreachable_regions = self.unreachable_regions.iter().map(|region| {
+            CoverageCounterAndRegion { kind: CoverageCounterKind::Counter(Counter::ZERO), region }
+        });
+
+        let branch_regions = self.branches.iter().map(|branch| CoverageCounterAndRegion {
+            kind: CoverageCounterKind::Branch {
+                true_counter: Counter::from_operand(branch.true_),
+                false_counter: Counter::from_operand(branch.false_),
+            },
+            region: &branch.region,
+        });
 
         let mut all_regions: Vec<_> = expression_regions.collect();
         all_regions.extend(counter_regions);
         all_regions.extend(unreachable_regions);
+        all_regions.extend(branch_regions);
 
         // make sure all the regions are sorted
-        all_regions.sort_unstable_by_key(|(_counter, region)| *region);
+        all_regions.sort_unstable_by_key(|coverage_region| coverage_region.region);
 
         (counter_expressions, all_regions)
     }

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
@@ -146,7 +148,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() {
@@ -158,12 +160,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 {
@@ -175,15 +178,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, code_region);
             }
+            CoverageKind::Branch { true_, false_ } => {
+                let code_region = code_region.expect("branch regions always have code regions");
+                debug!(
+                    "adding branch to coverage_map: instance={:?}, {:?} / {:?}; region: {:?}",
+                    instance, true_, false_, code_region,
+                );
+                func_coverage.add_branch_counter(true_, false_, code_region);
+            }
             CoverageKind::Unreachable => {
                 let code_region =
                     code_region.expect("unreachable regions always have code regions");

compiler/rustc_middle/src/mir/coverage.rs

@@ -83,6 +83,10 @@ pub enum CoverageKind {
         op: Op,
         rhs: Operand,
     },
+    Branch {
+        true_: Operand,
+        false_: Operand,
+    },
     Unreachable,
 }
 
@@ -92,6 +96,7 @@ impl CoverageKind {
         match *self {
             Counter { id, .. } => Operand::Counter(id),
             Expression { id, .. } => Operand::Expression(id),
+            Branch { .. } => bug!("Branch coverage cannot be part of an expression"),
             Unreachable => bug!("Unreachable coverage cannot be part of an expression"),
         }
     }
@@ -117,6 +122,9 @@ impl Debug for CoverageKind {
                 },
                 rhs,
             ),
+            Branch { true_, false_ } => {
+                write!(fmt, "Branch: {true_:?} / {false_:?}")
+            }
             Unreachable => write!(fmt, "Unreachable"),
         }
     }

compiler/rustc_mir_transform/src/coverage/debug.rs

@@ -299,6 +299,9 @@ impl DebugCounters {
             CoverageKind::Expression { .. } => {
                 format!("Expression({})", self.format_counter_kind(counter_kind))
             }
+            CoverageKind::Branch { true_, false_ } => {
+                format!("Branch({} / {})", self.format_operand(true_), self.format_operand(false_))
+            }
             CoverageKind::Unreachable { .. } => "Unreachable".to_owned(),
         }
     }

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_, false_ } => stable_mir::mir::CoverageKind::Branch {
+                true_: opaque(true_),
+                false_: opaque(false_),
+            },
             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_: ExpressionOperandId,
+        false_: ExpressionOperandId,
+    },
     Unreachable,
 }