[WIP] Working expression optimization, and some improvements to branch-level source coverage

compiler/rustc_codegen_llvm/src/coverageinfo/mapgen.rs

@@ -126,7 +126,7 @@ impl CoverageMapGenerator {
                 let (filenames_index, _) = self.filenames.insert_full(c_filename);
                 virtual_file_mapping.push(filenames_index as u32);
             }
-            debug!("Adding counter {:?} to map for {:?}", counter, region,);
+            debug!("Adding counter {:?} to map for {:?}", counter, region);
             mapping_regions.push(CounterMappingRegion::code_region(
                 counter,
                 current_file_id,

compiler/rustc_codegen_llvm/src/coverageinfo/mod.rs

@@ -12,7 +12,7 @@ use rustc_codegen_ssa::traits::{
 use rustc_data_structures::fx::FxHashMap;
 use rustc_llvm::RustString;
 use rustc_middle::mir::coverage::{
-    CodeRegion, CounterValueReference, ExpressionOperandId, InjectedExpressionIndex, Op,
+    CodeRegion, CounterValueReference, ExpressionOperandId, InjectedExpressionId, Op,
 };
 use rustc_middle::ty::Instance;
 
@@ -27,16 +27,16 @@ const COVMAP_VAR_ALIGN_BYTES: usize = 8;
 
 /// A context object for maintaining all state needed by the coverageinfo module.
 pub struct CrateCoverageContext<'tcx> {
-    // Coverage region data for each instrumented function identified by DefId.
-    pub(crate) function_coverage_map: RefCell<FxHashMap<Instance<'tcx>, FunctionCoverage>>,
+    // Coverage data for each instrumented function identified by DefId.
+    pub(crate) function_coverage_map: RefCell<FxHashMap<Instance<'tcx>, FunctionCoverage<'tcx>>>,
 }
 
 impl<'tcx> CrateCoverageContext<'tcx> {
     pub fn new() -> Self {
         Self { function_coverage_map: Default::default() }
     }
 
-    pub fn take_function_coverage_map(&self) -> FxHashMap<Instance<'tcx>, FunctionCoverage> {
+    pub fn take_function_coverage_map(&self) -> FxHashMap<Instance<'tcx>, FunctionCoverage<'tcx>> {
         self.function_coverage_map.replace(FxHashMap::default())
     }
 }
@@ -58,53 +58,62 @@ impl CoverageInfoBuilderMethods<'tcx> for Builder<'a, 'll, 'tcx> {
         unsafe { llvm::LLVMRustCoverageCreatePGOFuncNameVar(llfn, mangled_fn_name.as_ptr()) }
     }
 
-    fn add_counter_region(
+    fn set_function_source_hash(&mut self, instance: Instance<'tcx>, function_source_hash: u64) {
+        debug!(
+            "ensuring function source hash is set for instance={:?}; function_source_hash={}",
+            instance, function_source_hash,
+        );
+        let mut coverage_map = self.coverage_context().function_coverage_map.borrow_mut();
+        coverage_map
+            .entry(instance)
+            .or_insert_with(|| FunctionCoverage::new(self.tcx, instance))
+            .set_function_source_hash(function_source_hash);
+    }
+
+    fn add_coverage_counter(
         &mut self,
         instance: Instance<'tcx>,
         function_source_hash: u64,
         id: CounterValueReference,
         region: CodeRegion,
     ) {
         debug!(
-            "adding counter to coverage_regions: instance={:?}, function_source_hash={}, id={:?}, \
+            "adding counter to coverage_map: instance={:?}, function_source_hash={}, id={:?}, \
              at {:?}",
             instance, function_source_hash, id, region,
         );
-        let mut coverage_regions = self.coverage_context().function_coverage_map.borrow_mut();
-        coverage_regions
+        let mut coverage_map = self.coverage_context().function_coverage_map.borrow_mut();
+        coverage_map
             .entry(instance)
             .or_insert_with(|| FunctionCoverage::new(self.tcx, instance))
             .add_counter(function_source_hash, id, region);
     }
 
-    fn add_counter_expression_region(
+    fn add_coverage_counter_expression(
         &mut self,
         instance: Instance<'tcx>,
-        id: InjectedExpressionIndex,
+        id: InjectedExpressionId,
         lhs: ExpressionOperandId,
         op: Op,
         rhs: ExpressionOperandId,
-        region: CodeRegion,
+        region: Option<CodeRegion>,
     ) {
         debug!(
-            "adding counter expression to coverage_regions: instance={:?}, id={:?}, {:?} {:?} {:?}, \
-             at {:?}",
+            "adding counter expression to coverage_map: instance={:?}, id={:?}, {:?} {:?} {:?}; \
+             region: {:?}",
             instance, id, lhs, op, rhs, region,
         );
-        let mut coverage_regions = self.coverage_context().function_coverage_map.borrow_mut();
-        coverage_regions
+        let mut coverage_map = self.coverage_context().function_coverage_map.borrow_mut();
+        coverage_map
             .entry(instance)
             .or_insert_with(|| FunctionCoverage::new(self.tcx, instance))
             .add_counter_expression(id, lhs, op, rhs, region);
     }
 
-    fn add_unreachable_region(&mut self, instance: Instance<'tcx>, region: CodeRegion) {
-        debug!(
-            "adding unreachable code to coverage_regions: instance={:?}, at {:?}",
-            instance, region,
-        );
-        let mut coverage_regions = self.coverage_context().function_coverage_map.borrow_mut();
-        coverage_regions
+    fn add_coverage_unreachable(&mut self, instance: Instance<'tcx>, region: CodeRegion) {
+        debug!("adding unreachable code to coverage_map: instance={:?}, at {:?}", instance, region,);
+        let mut coverage_map = self.coverage_context().function_coverage_map.borrow_mut();
+        coverage_map
             .entry(instance)
             .or_insert_with(|| FunctionCoverage::new(self.tcx, instance))
             .add_unreachable_region(region);

compiler/rustc_codegen_ssa/src/coverageinfo/ffi.rs

@@ -3,7 +3,7 @@ use rustc_middle::mir::coverage::{CounterValueReference, MappedExpressionIndex};
 /// Aligns with [llvm::coverage::Counter::CounterKind](https://github.com/rust-lang/llvm-project/blob/rustc/10.0-2020-05-05/llvm/include/llvm/ProfileData/Coverage/CoverageMapping.h#L91)
 #[derive(Copy, Clone, Debug)]
 #[repr(C)]
-enum CounterKind {
+pub enum CounterKind {
     Zero = 0,
     CounterValueReference = 1,
     Expression = 2,
@@ -23,8 +23,8 @@ enum CounterKind {
 #[repr(C)]
 pub struct Counter {
     // Important: The layout (order and types of fields) must match its C++ counterpart.
-    kind: CounterKind,
-    id: u32,
+    pub kind: CounterKind,
+    pub id: u32,
 }
 
 impl Counter {
@@ -55,9 +55,9 @@ pub enum ExprKind {
 #[derive(Copy, Clone, Debug)]
 #[repr(C)]
 pub struct CounterExpression {
-    kind: ExprKind,
-    lhs: Counter,
-    rhs: Counter,
+    pub kind: ExprKind,
+    pub lhs: Counter,
+    pub rhs: Counter,
 }
 
 impl CounterExpression {

compiler/rustc_codegen_ssa/src/coverageinfo/map.rs

@@ -2,18 +2,18 @@ pub use super::ffi::*;
 
 use rustc_index::vec::IndexVec;
 use rustc_middle::mir::coverage::{
-    CodeRegion, CounterValueReference, ExpressionOperandId, InjectedExpressionIndex,
-    MappedExpressionIndex, Op,
+    CodeRegion, CounterValueReference, ExpressionOperandId, InjectedExpressionId,
+    InjectedExpressionIndex, MappedExpressionIndex, Op,
 };
 use rustc_middle::ty::Instance;
 use rustc_middle::ty::TyCtxt;
 
 #[derive(Clone, Debug)]
-pub struct ExpressionRegion {
+pub struct Expression {
     lhs: ExpressionOperandId,
     op: Op,
     rhs: ExpressionOperandId,
-    region: CodeRegion,
+    region: Option<CodeRegion>,
 }
 
 /// Collects all of the coverage regions associated with (a) injected counters, (b) counter
@@ -28,24 +28,43 @@ pub struct ExpressionRegion {
 /// only whitespace or comments). According to LLVM Code Coverage Mapping documentation, "A count
 /// for a gap area is only used as the line execution count if there are no other regions on a
 /// line."
-pub struct FunctionCoverage {
+pub struct FunctionCoverage<'tcx> {
+    instance: Instance<'tcx>,
     source_hash: u64,
     counters: IndexVec<CounterValueReference, Option<CodeRegion>>,
-    expressions: IndexVec<InjectedExpressionIndex, Option<ExpressionRegion>>,
+    expressions: IndexVec<InjectedExpressionIndex, Option<Expression>>,
     unreachable_regions: Vec<CodeRegion>,
 }
 
-impl FunctionCoverage {
-    pub fn new<'tcx>(tcx: TyCtxt<'tcx>, instance: Instance<'tcx>) -> Self {
+impl<'tcx> FunctionCoverage<'tcx> {
+    pub fn new(tcx: TyCtxt<'tcx>, instance: Instance<'tcx>) -> Self {
         let coverageinfo = tcx.coverageinfo(instance.def_id());
+        debug!(
+            "FunctionCoverage::new(instance={:?}) has coverageinfo={:?}",
+            instance, coverageinfo
+        );
         Self {
+            instance,
             source_hash: 0, // will be set with the first `add_counter()`
             counters: IndexVec::from_elem_n(None, coverageinfo.num_counters as usize),
             expressions: IndexVec::from_elem_n(None, coverageinfo.num_expressions as usize),
             unreachable_regions: Vec::new(),
         }
     }
 
+    /// Although every function should have at least one `Counter`, the `Counter` isn't required to
+    /// have a `CodeRegion`. (The `CodeRegion` may be associated only with `Expressions`.) This
+    /// method supports the ability to ensure the `function_source_hash` is set from `Counters` that
+    /// do not trigger the call to `add_counter()` because they don't have an associated
+    /// `CodeRegion` to add.
+    pub fn set_function_source_hash(&mut self, source_hash: u64) {
+        if self.source_hash == 0 {
+            self.source_hash = source_hash;
+        } else {
+            debug_assert_eq!(source_hash, self.source_hash);
+        }
+    }
+
     /// Adds a code region to be counted by an injected counter intrinsic.
     /// The source_hash (computed during coverage instrumentation) should also be provided, and
     /// should be the same for all counters in a given function.
@@ -74,15 +93,19 @@ impl FunctionCoverage {
     /// counters and expressions have been added.
     pub fn add_counter_expression(
         &mut self,
-        expression_id: InjectedExpressionIndex,
+        expression_id: InjectedExpressionId,
         lhs: ExpressionOperandId,
         op: Op,
         rhs: ExpressionOperandId,
-        region: CodeRegion,
+        region: Option<CodeRegion>,
     ) {
+        debug!(
+            "add_counter_expression({:?}, lhs={:?}, op={:?}, rhs={:?} at {:?}",
+            expression_id, lhs, op, rhs, region
+        );
         let expression_index = self.expression_index(u32::from(expression_id));
         self.expressions[expression_index]
-            .replace(ExpressionRegion { lhs, op, rhs, region })
+            .replace(Expression { lhs, op, rhs, region })
             .expect_none("add_counter_expression called with duplicate `id_descending_from_max`");
     }
 
@@ -103,7 +126,11 @@ impl FunctionCoverage {
     pub fn get_expressions_and_counter_regions<'a>(
         &'a self,
     ) -> (Vec<CounterExpression>, impl Iterator<Item = (Counter, &'a CodeRegion)>) {
-        assert!(self.source_hash != 0);
+        assert!(
+            self.source_hash != 0,
+            "No counters provided the source_hash for function: {:?}",
+            self.instance
+        );
 
         let counter_regions = self.counter_regions();
         let (counter_expressions, expression_regions) = self.expressions_with_regions();
@@ -129,54 +156,60 @@ impl FunctionCoverage {
     ) -> (Vec<CounterExpression>, impl Iterator<Item = (Counter, &'a CodeRegion)>) {
         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(MappedExpressionIndex::from(u32::MAX), self.expressions.len());
-        // Note, the initial value shouldn't matter since every index in use in `self.expressions`
-        // will be set, and after that, `new_indexes` will only be accessed using those same
-        // indexes.
-
-        // Note that an `ExpressionRegion`s at any given index can include other expressions as
+        let mut new_indexes = IndexVec::from_elem_n(None, self.expressions.len());
+        // Note that an `Expression`s at any given index can include other expressions as
         // operands, but expression operands can only come from the subset of expressions having
-        // `expression_index`s lower than the referencing `ExpressionRegion`. Therefore, it is
+        // `expression_index`s lower than the referencing `Expression`. Therefore, it is
         // reasonable to look up the new index of an expression operand while the `new_indexes`
         // vector is only complete up to the current `ExpressionIndex`.
         let id_to_counter =
-            |new_indexes: &IndexVec<InjectedExpressionIndex, MappedExpressionIndex>,
+            |new_indexes: &IndexVec<InjectedExpressionIndex, Option<MappedExpressionIndex>>,
              id: ExpressionOperandId| {
                 if id == ExpressionOperandId::ZERO {
                     Some(Counter::zero())
                 } else if id.index() < self.counters.len() {
+                    // Note: Some codegen-injected Counters may be only referenced by `Expression`s,
+                    // and may not have their own `CodeRegion`s,
                     let index = CounterValueReference::from(id.index());
-                    self.counters
-                        .get(index)
-                        .unwrap() // pre-validated
-                        .as_ref()
-                        .map(|_| Counter::counter_value_reference(index))
+                    Some(Counter::counter_value_reference(index))
                 } else {
                     let index = self.expression_index(u32::from(id));
                     self.expressions
                         .get(index)
                         .expect("expression id is out of range")
                         .as_ref()
-                        .map(|_| Counter::expression(new_indexes[index]))
+                        // If an expression was optimized out, assume it would have produced a count
+                        // of zero. This ensures that expressions dependent on optimized-out
+                        // expressions are still valid.
+                        .map_or(Some(Counter::zero()), |_| {
+                            new_indexes[index].map(|new_index| Counter::expression(new_index))
+                        })
                 }
             };
 
-        for (original_index, expression_region) in
+        for (original_index, expression) in
             self.expressions.iter_enumerated().filter_map(|(original_index, entry)| {
                 // Option::map() will return None to filter out missing expressions. This may happen
                 // if, for example, a MIR-instrumented expression is removed during an optimization.
-                entry.as_ref().map(|region| (original_index, region))
+                entry.as_ref().map(|expression| (original_index, expression))
             })
         {
-            let region = &expression_region.region;
-            let ExpressionRegion { lhs, op, rhs, .. } = *expression_region;
+            let optional_region = &expression.region;
+            let Expression { lhs, op, rhs, .. } = *expression;
 
             if let Some(Some((lhs_counter, rhs_counter))) =
                 id_to_counter(&new_indexes, lhs).map(|lhs_counter| {
                     id_to_counter(&new_indexes, rhs).map(|rhs_counter| (lhs_counter, rhs_counter))
                 })
             {
+                debug_assert!(
+                    (lhs_counter.id as usize)
+                        < usize::max(self.counters.len(), self.expressions.len())
+                );
+                debug_assert!(
+                    (rhs_counter.id as usize)
+                        < usize::max(self.counters.len(), self.expressions.len())
+                );
                 // Both operands exist. `Expression` operands exist in `self.expressions` and have
                 // been assigned a `new_index`.
                 let mapped_expression_index =
@@ -190,12 +223,20 @@ impl FunctionCoverage {
                     rhs_counter,
                 );
                 debug!(
-                    "Adding expression {:?} = {:?} at {:?}",
-                    mapped_expression_index, expression, region
+                    "Adding expression {:?} = {:?}, region: {:?}",
+                    mapped_expression_index, expression, optional_region
                 );
                 counter_expressions.push(expression);
-                new_indexes[original_index] = mapped_expression_index;
-                expression_regions.push((Counter::expression(mapped_expression_index), region));
+                new_indexes[original_index] = Some(mapped_expression_index);
+                if let Some(region) = optional_region {
+                    expression_regions.push((Counter::expression(mapped_expression_index), region));
+                }
+            } else {
+                debug!(
+                    "Ignoring expression with one or more missing operands: \
+                    original_index={:?}, lhs={:?}, op={:?}, rhs={:?}, region={:?}",
+                    original_index, lhs, op, rhs, optional_region,
+                )
             }
         }
         (counter_expressions, expression_regions.into_iter())