Fixed four bugs found while compiling a larger set of real world crates

Compiling with coverage, with the expression optimization, now works on
the json5format crate and its dependencies.

Author: richkadel

compiler/rustc_codegen_llvm/src/coverageinfo/mod.rs

@@ -28,15 +28,15 @@ const COVMAP_VAR_ALIGN_BYTES: usize = 8;
 /// A context object for maintaining all state needed by the coverageinfo module.
 pub struct CrateCoverageContext<'tcx> {
     // Coverage data for each instrumented function identified by DefId.
-    pub(crate) function_coverage_map: RefCell<FxHashMap<Instance<'tcx>, FunctionCoverage>>,
+    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,6 +58,18 @@ impl CoverageInfoBuilderMethods<'tcx> for Builder<'a, 'll, 'tcx> {
         unsafe { llvm::LLVMRustCoverageCreatePGOFuncNameVar(llfn, mangled_fn_name.as_ptr()) }
     }
 
+    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>,

compiler/rustc_codegen_ssa/src/coverageinfo/map.rs

@@ -28,28 +28,43 @@ pub struct Expression {
 /// 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<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.
@@ -111,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();

compiler/rustc_codegen_ssa/src/mir/coverageinfo.rs

@@ -24,6 +24,8 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
 
                 if let Some(code_region) = code_region {
                     bx.add_coverage_counter(self.instance, function_source_hash, id, code_region);
+                } else {
+                    bx.set_function_source_hash(self.instance, function_source_hash);
                 }
                 // Note: Some counters do not have code regions, but may still be referenced from
                 // expressions.

compiler/rustc_codegen_ssa/src/traits/coverageinfo.rs

@@ -9,6 +9,8 @@ pub trait CoverageInfoMethods: BackendTypes {
 pub trait CoverageInfoBuilderMethods<'tcx>: BackendTypes {
     fn create_pgo_func_name_var(&self, instance: Instance<'tcx>) -> Self::Value;
 
+    fn set_function_source_hash(&mut self, instance: Instance<'tcx>, function_source_hash: u64);
+
     fn add_coverage_counter(
         &mut self,
         instance: Instance<'tcx>,

compiler/rustc_mir/src/transform/coverage/debug.rs

@@ -7,6 +7,8 @@ use rustc_middle::mir::{BasicBlock, TerminatorKind};
 
 use std::lazy::SyncOnceCell;
 
+pub const NESTED_INDENT: &str = "    ";
+
 const RUSTC_COVERAGE_DEBUG_OPTIONS: &str = "RUSTC_COVERAGE_DEBUG_OPTIONS";
 
 pub(crate) fn debug_options<'a>() -> &'a DebugOptions {

compiler/rustc_mir/src/transform/coverage/graph.rs

@@ -1,3 +1,5 @@
+use super::Error;
+
 use rustc_data_structures::fx::FxHashMap;
 use rustc_data_structures::graph::dominators::{self, Dominators};
 use rustc_data_structures::graph::{self, GraphSuccessors, WithNumNodes, WithStartNode};
@@ -313,16 +315,28 @@ impl BasicCoverageBlockData {
     }
 
     #[inline(always)]
-    pub fn set_counter(&mut self, counter_kind: CoverageKind) -> ExpressionOperandId {
+    pub fn set_counter(
+        &mut self,
+        counter_kind: CoverageKind,
+    ) -> Result<ExpressionOperandId, Error> {
         debug_assert!(
+            // If the BCB has an edge counter (to be injected into a new `BasicBlock`), it can also
+            // have an expression (to be injected into an existing `BasicBlock` represented by this
+            // `BasicCoverageBlock`).
             self.edge_from_bcbs.is_none() || counter_kind.is_expression(),
             "attempt to add a `Counter` to a BCB target with existing incoming edge counters"
         );
         let operand = counter_kind.as_operand_id();
-        self.counter_kind
-            .replace(counter_kind)
-            .expect_none("attempt to set a BasicCoverageBlock coverage counter more than once");
-        operand
+        let expect_none = self.counter_kind.replace(counter_kind);
+        if expect_none.is_some() {
+            return Error::from_string(format!(
+                "attempt to set a BasicCoverageBlock coverage counter more than once; \
+                {:?} already had counter {:?}",
+                self,
+                expect_none.unwrap(),
+            ));
+        }
+        Ok(operand)
     }
 
     #[inline(always)]
@@ -340,19 +354,33 @@ impl BasicCoverageBlockData {
         &mut self,
         from_bcb: BasicCoverageBlock,
         counter_kind: CoverageKind,
-    ) -> ExpressionOperandId {
-        debug_assert!(
-            self.counter_kind.as_ref().map_or(true, |c| c.is_expression()),
-            "attempt to add an incoming edge counter from {:?} when the target BCB already has a \
-            `Counter`",
-            from_bcb
-        );
+    ) -> Result<ExpressionOperandId, Error> {
+        if level_enabled!(tracing::Level::DEBUG) {
+            // If the BCB has an edge counter (to be injected into a new `BasicBlock`), it can also
+            // have an expression (to be injected into an existing `BasicBlock` represented by this
+            // `BasicCoverageBlock`).
+            if !self.counter_kind.as_ref().map_or(true, |c| c.is_expression()) {
+                return Error::from_string(format!(
+                    "attempt to add an incoming edge counter from {:?} when the target BCB already \
+                    has a `Counter`",
+                    from_bcb
+                ));
+            }
+        }
         let operand = counter_kind.as_operand_id();
-        self.edge_from_bcbs
+        let expect_none = self
+            .edge_from_bcbs
             .get_or_insert_with(|| FxHashMap::default())
-            .insert(from_bcb, counter_kind)
-            .expect_none("attempt to set an edge counter more than once");
-        operand
+            .insert(from_bcb, counter_kind);
+        if expect_none.is_some() {
+            return Error::from_string(format!(
+                "attempt to set an edge counter more than once; from_bcb: \
+                {:?} already had counter {:?}",
+                from_bcb,
+                expect_none.unwrap(),
+            ));
+        }
+        Ok(operand)
     }
 
     #[inline(always)]
@@ -383,6 +411,56 @@ impl BasicCoverageBlockData {
     }
 }
 
+/// Represents a successor from a branching BasicCoverageBlock (such as the arms of a `SwitchInt`)
+/// as either the successor BCB itself, if it has only one incoming edge, or the successor _plus_
+/// the specific branching BCB, representing the edge between the two. The latter case
+/// distinguishes this incoming edge from other incoming edges to the same `target_bcb`.
+#[derive(Clone, Copy, PartialEq, Eq)]
+pub(crate) struct BcbBranch {
+    pub edge_from_bcb: Option<BasicCoverageBlock>,
+    pub target_bcb: BasicCoverageBlock,
+}
+
+impl BcbBranch {
+    pub fn from_to(
+        from_bcb: BasicCoverageBlock,
+        to_bcb: BasicCoverageBlock,
+        basic_coverage_blocks: &CoverageGraph,
+    ) -> Self {
+        let edge_from_bcb = if basic_coverage_blocks.predecessors[to_bcb].len() > 1 {
+            Some(from_bcb)
+        } else {
+            None
+        };
+        Self { edge_from_bcb, target_bcb: to_bcb }
+    }
+
+    pub fn counter<'a>(
+        &self,
+        basic_coverage_blocks: &'a CoverageGraph,
+    ) -> Option<&'a CoverageKind> {
+        if let Some(from_bcb) = self.edge_from_bcb {
+            basic_coverage_blocks[self.target_bcb].edge_counter_from(from_bcb)
+        } else {
+            basic_coverage_blocks[self.target_bcb].counter()
+        }
+    }
+
+    pub fn is_only_path_to_target(&self) -> bool {
+        self.edge_from_bcb.is_none()
+    }
+}
+
+impl std::fmt::Debug for BcbBranch {
+    fn fmt(&self, fmt: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        if let Some(from_bcb) = self.edge_from_bcb {
+            write!(fmt, "{:?}->{:?}", from_bcb, self.target_bcb)
+        } else {
+            write!(fmt, "{:?}", self.target_bcb)
+        }
+    }
+}
+
 fn bcb_filtered_successors<'a, 'tcx>(
     body: &'tcx &'a mir::Body<'tcx>,
     term_kind: &'tcx TerminatorKind<'tcx>,
@@ -437,14 +515,18 @@ impl TraverseCoverageGraphWithLoops {
     }
 
     pub fn next(&mut self) -> Option<BasicCoverageBlock> {
-        // Strip contexts with empty worklists from the top of the stack
-        while self.context_stack.last().map_or(false, |context| context.worklist.is_empty()) {
-            self.context_stack.pop();
-        }
-        // Pop the next bcb off of the current context_stack. If none, all BCBs were visited.
-        while let Some(next_bcb) =
+        debug!(
+            "TraverseCoverageGraphWithLoops::next - context_stack: {:?}",
+            self.context_stack.iter().rev().collect::<Vec<_>>()
+        );
+        while let Some(next_bcb) = {
+            // Strip contexts with empty worklists from the top of the stack
+            while self.context_stack.last().map_or(false, |context| context.worklist.is_empty()) {
+                self.context_stack.pop();
+            }
+            // Pop the next bcb off of the current context_stack. If none, all BCBs were visited.
             self.context_stack.last_mut().map_or(None, |context| context.worklist.pop())
-        {
+        } {
             if !self.visited.insert(next_bcb) {
                 debug!("Already visited: {:?}", next_bcb);
                 continue;
@@ -459,9 +541,19 @@ impl TraverseCoverageGraphWithLoops {
             }
             return Some(next_bcb);
         }
-        debug_assert_eq!(self.visited.count(), self.visited.domain_size());
         None
     }
+
+    pub fn is_complete(&self) -> bool {
+        self.visited.count() == self.visited.domain_size()
+    }
+
+    pub fn unvisited(&self) -> Vec<BasicCoverageBlock> {
+        let mut unvisited_set: BitSet<BasicCoverageBlock> =
+            BitSet::new_filled(self.visited.domain_size());
+        unvisited_set.subtract(&self.visited);
+        unvisited_set.iter().collect::<Vec<_>>()
+    }
 }
 
 fn find_loop_backedges(