coverage: Store BCB node IDs in mappings, and resolve them in codegen

Even though the coverage graph itself is no longer available during codegen,
its nodes can still be used as opaque IDs.

Author: Zalathar

compiler/rustc_codegen_llvm/src/coverageinfo/mapgen/covfun.rs

@@ -11,7 +11,8 @@ use rustc_codegen_ssa::traits::{
     BaseTypeCodegenMethods, ConstCodegenMethods, StaticCodegenMethods,
 };
 use rustc_middle::mir::coverage::{
-    CovTerm, CoverageIdsInfo, Expression, FunctionCoverageInfo, Mapping, MappingKind, Op,
+    BasicCoverageBlock, CovTerm, CoverageIdsInfo, Expression, FunctionCoverageInfo, Mapping,
+    MappingKind, Op,
 };
 use rustc_middle::ty::{Instance, TyCtxt};
 use rustc_span::Span;
@@ -140,7 +141,12 @@ fn fill_region_tables<'tcx>(
     for &Mapping { ref kind, span } in &fn_cov_info.mappings {
         // If the mapping refers to counters/expressions that were removed by
         // MIR opts, replace those occurrences with zero.
-        let kind = kind.map_terms(|term| if is_zero_term(term) { CovTerm::Zero } else { term });
+        let counter_for_bcb = |bcb: BasicCoverageBlock| -> ffi::Counter {
+            let term =
+                fn_cov_info.term_for_bcb[bcb].expect("every BCB in a mapping was given a term");
+            let term = if is_zero_term(term) { CovTerm::Zero } else { term };
+            ffi::Counter::from_term(term)
+        };
 
         // Convert the `Span` into coordinates that we can pass to LLVM, or
         // discard the span if conversion fails. In rare, cases _all_ of a
@@ -154,23 +160,22 @@ fn fill_region_tables<'tcx>(
             continue;
         }
 
-        match kind {
-            MappingKind::Code(term) => {
-                code_regions
-                    .push(ffi::CodeRegion { cov_span, counter: ffi::Counter::from_term(term) });
+        match *kind {
+            MappingKind::Code { bcb } => {
+                code_regions.push(ffi::CodeRegion { cov_span, counter: counter_for_bcb(bcb) });
             }
-            MappingKind::Branch { true_term, false_term } => {
+            MappingKind::Branch { true_bcb, false_bcb } => {
                 branch_regions.push(ffi::BranchRegion {
                     cov_span,
-                    true_counter: ffi::Counter::from_term(true_term),
-                    false_counter: ffi::Counter::from_term(false_term),
+                    true_counter: counter_for_bcb(true_bcb),
+                    false_counter: counter_for_bcb(false_bcb),
                 });
             }
-            MappingKind::MCDCBranch { true_term, false_term, mcdc_params } => {
+            MappingKind::MCDCBranch { true_bcb, false_bcb, mcdc_params } => {
                 mcdc_branch_regions.push(ffi::MCDCBranchRegion {
                     cov_span,
-                    true_counter: ffi::Counter::from_term(true_term),
-                    false_counter: ffi::Counter::from_term(false_term),
+                    true_counter: counter_for_bcb(true_bcb),
+                    false_counter: counter_for_bcb(false_bcb),
                     mcdc_branch_params: ffi::mcdc::BranchParameters::from(mcdc_params),
                 });
             }

compiler/rustc_middle/src/mir/coverage.rs

@@ -179,34 +179,19 @@ pub struct Expression {
 #[derive(TyEncodable, TyDecodable, Hash, HashStable)]
 pub enum MappingKind {
     /// Associates a normal region of code with a counter/expression/zero.
-    Code(CovTerm),
+    Code { bcb: BasicCoverageBlock },
     /// Associates a branch region with separate counters for true and false.
-    Branch { true_term: CovTerm, false_term: CovTerm },
+    Branch { true_bcb: BasicCoverageBlock, false_bcb: BasicCoverageBlock },
     /// Associates a branch region with separate counters for true and false.
-    MCDCBranch { true_term: CovTerm, false_term: CovTerm, mcdc_params: ConditionInfo },
+    MCDCBranch {
+        true_bcb: BasicCoverageBlock,
+        false_bcb: BasicCoverageBlock,
+        mcdc_params: ConditionInfo,
+    },
     /// Associates a decision region with a bitmap and number of conditions.
     MCDCDecision(DecisionInfo),
 }
 
-impl MappingKind {
-    /// Returns a copy of this mapping kind, in which all coverage terms have
-    /// been replaced with ones returned by the given function.
-    pub fn map_terms(&self, map_fn: impl Fn(CovTerm) -> CovTerm) -> Self {
-        match *self {
-            Self::Code(term) => Self::Code(map_fn(term)),
-            Self::Branch { true_term, false_term } => {
-                Self::Branch { true_term: map_fn(true_term), false_term: map_fn(false_term) }
-            }
-            Self::MCDCBranch { true_term, false_term, mcdc_params } => Self::MCDCBranch {
-                true_term: map_fn(true_term),
-                false_term: map_fn(false_term),
-                mcdc_params,
-            },
-            Self::MCDCDecision(param) => Self::MCDCDecision(param),
-        }
-    }
-}
-
 #[derive(Clone, Debug)]
 #[derive(TyEncodable, TyDecodable, Hash, HashStable)]
 pub struct Mapping {
@@ -222,10 +207,14 @@ pub struct Mapping {
 pub struct FunctionCoverageInfo {
     pub function_source_hash: u64,
     pub body_span: Span,
+
     pub num_counters: usize,
-    pub mcdc_bitmap_bits: usize,
     pub expressions: IndexVec<ExpressionId, Expression>,
+
     pub mappings: Vec<Mapping>,
+    pub term_for_bcb: IndexVec<BasicCoverageBlock, Option<CovTerm>>,
+
+    pub mcdc_bitmap_bits: usize,
     /// The depth of the deepest decision is used to know how many
     /// temp condbitmaps should be allocated for the function.
     pub mcdc_num_condition_bitmaps: usize,
@@ -329,3 +318,19 @@ impl CoverageIdsInfo {
         }
     }
 }
+
+rustc_index::newtype_index! {
+    /// During the `InstrumentCoverage` MIR pass, a BCB is a node in the
+    /// "coverage graph", which is a refinement of the MIR control-flow graph
+    /// that merges or omits some blocks that aren't relevant to coverage.
+    ///
+    /// After that pass is complete, the coverage graph no longer exists, so a
+    /// BCB is effectively an opaque ID.
+    #[derive(HashStable)]
+    #[encodable]
+    #[orderable]
+    #[debug_format = "bcb{}"]
+    pub struct BasicCoverageBlock {
+        const START_BCB = 0;
+    }
+}

compiler/rustc_middle/src/ty/structural_impls.rs

@@ -246,6 +246,7 @@ TrivialTypeTraversalImpls! {
     ::rustc_hir::HirId,
     ::rustc_hir::MatchSource,
     ::rustc_target::asm::InlineAsmRegOrRegClass,
+    crate::mir::coverage::BasicCoverageBlock,
     crate::mir::coverage::BlockMarkerId,
     crate::mir::coverage::CounterId,
     crate::mir::coverage::ExpressionId,

compiler/rustc_mir_transform/src/coverage/counters.rs

@@ -133,7 +133,7 @@ pub(super) struct CoverageCounters {
     next_counter_id: CounterId,
 
     /// Coverage counters/expressions that are associated with individual BCBs.
-    node_counters: IndexVec<BasicCoverageBlock, Option<CovTerm>>,
+    pub(crate) node_counters: IndexVec<BasicCoverageBlock, Option<CovTerm>>,
 
     /// Table of expression data, associating each expression ID with its
     /// corresponding operator (+ or -) and its LHS/RHS operands.
@@ -203,10 +203,6 @@ impl CoverageCounters {
         counter
     }
 
-    pub(super) fn term_for_bcb(&self, bcb: BasicCoverageBlock) -> Option<CovTerm> {
-        self.node_counters[bcb]
-    }
-
     /// Returns an iterator over all the nodes in the coverage graph that
     /// should have a counter-increment statement injected into MIR, along with
     /// each site's corresponding counter ID.

compiler/rustc_mir_transform/src/coverage/graph.rs

@@ -8,6 +8,7 @@ use rustc_data_structures::graph::dominators::Dominators;
 use rustc_data_structures::graph::{self, DirectedGraph, StartNode};
 use rustc_index::IndexVec;
 use rustc_index::bit_set::DenseBitSet;
+pub(crate) use rustc_middle::mir::coverage::{BasicCoverageBlock, START_BCB};
 use rustc_middle::mir::{self, BasicBlock, Terminator, TerminatorKind};
 use tracing::debug;
 
@@ -269,15 +270,6 @@ impl graph::Predecessors for CoverageGraph {
     }
 }
 
-rustc_index::newtype_index! {
-    /// A node in the control-flow graph of CoverageGraph.
-    #[orderable]
-    #[debug_format = "bcb{}"]
-    pub(crate) struct BasicCoverageBlock {
-        const START_BCB = 0;
-    }
-}
-
 /// `BasicCoverageBlockData` holds the data indexed by a `BasicCoverageBlock`.
 ///
 /// A `BasicCoverageBlock` (BCB) represents the maximal-length sequence of MIR `BasicBlock`s without

compiler/rustc_mir_transform/src/coverage/mod.rs

@@ -95,12 +95,13 @@ fn instrument_function_for_coverage<'tcx>(tcx: TyCtxt<'tcx>, mir_body: &mut mir:
 
     let coverage_counters = counters::make_bcb_counters(&graph, &bcbs_with_counter_mappings);
 
-    let mappings = create_mappings(&extracted_mappings, &coverage_counters);
+    let mappings = create_mappings(&extracted_mappings);
     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");
         return;
     }
+    let term_for_bcb = coverage_counters.node_counters.clone();
 
     inject_coverage_statements(mir_body, &graph, &extracted_mappings, &coverage_counters);
 
@@ -116,26 +117,24 @@ fn instrument_function_for_coverage<'tcx>(tcx: TyCtxt<'tcx>, mir_body: &mut mir:
     mir_body.function_coverage_info = Some(Box::new(FunctionCoverageInfo {
         function_source_hash: hir_info.function_source_hash,
         body_span: hir_info.body_span,
+
         num_counters: coverage_counters.num_counters(),
-        mcdc_bitmap_bits: extracted_mappings.mcdc_bitmap_bits,
         expressions: coverage_counters.into_expressions(),
+
         mappings,
+        term_for_bcb,
+
+        mcdc_bitmap_bits: extracted_mappings.mcdc_bitmap_bits,
         mcdc_num_condition_bitmaps,
     }));
 }
 
-/// For each coverage span extracted from MIR, create a corresponding
-/// mapping.
+/// For each coverage span extracted from MIR, create a corresponding mapping.
 ///
-/// Precondition: All BCBs corresponding to those spans have been given
-/// coverage counters.
-fn create_mappings(
-    extracted_mappings: &ExtractedMappings,
-    coverage_counters: &CoverageCounters,
-) -> Vec<Mapping> {
-    let term_for_bcb =
-        |bcb| coverage_counters.term_for_bcb(bcb).expect("all BCBs with spans were given counters");
-
+/// FIXME(Zalathar): This used to be where BCBs in the extracted mappings were
+/// resolved to a `CovTerm`. But that is now handled elsewhere, so this
+/// function can potentially be simplified even further.
+fn create_mappings(extracted_mappings: &ExtractedMappings) -> Vec<Mapping> {
     // Fully destructure the mappings struct to make sure we don't miss any kinds.
     let ExtractedMappings {
         num_bcbs: _,
@@ -150,23 +149,18 @@ fn create_mappings(
     mappings.extend(code_mappings.iter().map(
         // Ordinary code mappings are the simplest kind.
         |&mappings::CodeMapping { span, bcb }| {
-            let kind = MappingKind::Code(term_for_bcb(bcb));
+            let kind = MappingKind::Code { bcb };
             Mapping { kind, span }
         },
     ));
 
     mappings.extend(branch_pairs.iter().map(
         |&mappings::BranchPair { span, true_bcb, false_bcb }| {
-            let true_term = term_for_bcb(true_bcb);
-            let false_term = term_for_bcb(false_bcb);
-            let kind = MappingKind::Branch { true_term, false_term };
+            let kind = MappingKind::Branch { true_bcb, false_bcb };
             Mapping { kind, span }
         },
     ));
 
-    let term_for_bcb =
-        |bcb| coverage_counters.term_for_bcb(bcb).expect("all BCBs with spans were given counters");
-
     // MCDC branch mappings are appended with their decisions in case decisions were ignored.
     mappings.extend(mcdc_degraded_branches.iter().map(
         |&mappings::MCDCBranch {
@@ -176,11 +170,7 @@ fn create_mappings(
              condition_info: _,
              true_index: _,
              false_index: _,
-         }| {
-            let true_term = term_for_bcb(true_bcb);
-            let false_term = term_for_bcb(false_bcb);
-            Mapping { kind: MappingKind::Branch { true_term, false_term }, span }
-        },
+         }| { Mapping { kind: MappingKind::Branch { true_bcb, false_bcb }, span } },
     ));
 
     for (decision, branches) in mcdc_mappings {
@@ -201,12 +191,10 @@ fn create_mappings(
                      true_index: _,
                      false_index: _,
                  }| {
-                    let true_term = term_for_bcb(true_bcb);
-                    let false_term = term_for_bcb(false_bcb);
                     Mapping {
                         kind: MappingKind::MCDCBranch {
-                            true_term,
-                            false_term,
+                            true_bcb,
+                            false_bcb,
                             mcdc_params: condition_info,
                         },
                         span,

compiler/rustc_mir_transform/src/coverage/query.rs

@@ -103,7 +103,9 @@ fn coverage_ids_info<'tcx>(
         // to any particular point in the control-flow graph.
         // (Keep this in sync with the injection of `ExpressionUsed`
         // statements in the `InstrumentCoverage` MIR pass.)
-        if let MappingKind::Code(CovTerm::Expression(id)) = mapping.kind {
+        if let MappingKind::Code { bcb } = mapping.kind
+            && let Some(CovTerm::Expression(id)) = fn_cov_info.term_for_bcb[bcb]
+        {
             expressions_seen.remove(id);
         }
     }

tests/mir-opt/coverage/branch_match_arms.main.InstrumentCoverage.diff

@@ -30,12 +30,12 @@
 +     coverage ExpressionId(0) => Expression { lhs: Counter(1), op: Add, rhs: Counter(2) };
 +     coverage ExpressionId(1) => Expression { lhs: Expression(0), op: Add, rhs: Counter(3) };
 +     coverage ExpressionId(2) => Expression { lhs: Counter(0), op: Subtract, rhs: Expression(1) };
-+     coverage Code(Counter(0)) => $DIR/branch_match_arms.rs:14:1: 15:21 (#0);
-+     coverage Code(Counter(1)) => $DIR/branch_match_arms.rs:16:17: 16:33 (#0);
-+     coverage Code(Counter(2)) => $DIR/branch_match_arms.rs:17:17: 17:33 (#0);
-+     coverage Code(Counter(3)) => $DIR/branch_match_arms.rs:18:17: 18:33 (#0);
-+     coverage Code(Expression(2)) => $DIR/branch_match_arms.rs:19:17: 19:33 (#0);
-+     coverage Code(Counter(0)) => $DIR/branch_match_arms.rs:21:2: 21:2 (#0);
++     coverage Code { bcb: bcb0 } => $DIR/branch_match_arms.rs:14:1: 15:21 (#0);
++     coverage Code { bcb: bcb1 } => $DIR/branch_match_arms.rs:16:17: 16:33 (#0);
++     coverage Code { bcb: bcb3 } => $DIR/branch_match_arms.rs:17:17: 17:33 (#0);
++     coverage Code { bcb: bcb4 } => $DIR/branch_match_arms.rs:18:17: 18:33 (#0);
++     coverage Code { bcb: bcb5 } => $DIR/branch_match_arms.rs:19:17: 19:33 (#0);
++     coverage Code { bcb: bcb2 } => $DIR/branch_match_arms.rs:21:2: 21:2 (#0);
 + 
       bb0: {
 +         Coverage::CounterIncrement(0);

tests/mir-opt/coverage/instrument_coverage.bar.InstrumentCoverage.diff

@@ -5,7 +5,7 @@
       let mut _0: bool;
 
 +     coverage body span: $DIR/instrument_coverage.rs:29:18: 31:2 (#0)
-+     coverage Code(Counter(0)) => $DIR/instrument_coverage.rs:29:1: 31:2 (#0);
++     coverage Code { bcb: bcb0 } => $DIR/instrument_coverage.rs:29:1: 31:2 (#0);
 + 
       bb0: {
 +         Coverage::CounterIncrement(0);

tests/mir-opt/coverage/instrument_coverage.main.InstrumentCoverage.diff

@@ -9,11 +9,11 @@
 
 +     coverage body span: $DIR/instrument_coverage.rs:14:11: 20:2 (#0)
 +     coverage ExpressionId(0) => Expression { lhs: Counter(1), op: Subtract, rhs: Counter(0) };
-+     coverage Code(Counter(0)) => $DIR/instrument_coverage.rs:14:1: 14:11 (#0);
-+     coverage Code(Counter(1)) => $DIR/instrument_coverage.rs:16:12: 16:17 (#0);
-+     coverage Code(Counter(0)) => $DIR/instrument_coverage.rs:17:13: 17:18 (#0);
-+     coverage Code(Expression(0)) => $DIR/instrument_coverage.rs:18:10: 18:10 (#0);
-+     coverage Code(Counter(0)) => $DIR/instrument_coverage.rs:20:2: 20:2 (#0);
++     coverage Code { bcb: bcb0 } => $DIR/instrument_coverage.rs:14:1: 14:11 (#0);
++     coverage Code { bcb: bcb1 } => $DIR/instrument_coverage.rs:16:12: 16:17 (#0);
++     coverage Code { bcb: bcb2 } => $DIR/instrument_coverage.rs:17:13: 17:18 (#0);
++     coverage Code { bcb: bcb3 } => $DIR/instrument_coverage.rs:18:10: 18:10 (#0);
++     coverage Code { bcb: bcb2 } => $DIR/instrument_coverage.rs:20:2: 20:2 (#0);
 + 
       bb0: {
 +         Coverage::CounterIncrement(0);