Avoid intermediate Vec allocations in coverage mapping creation

Author: Swatinem

compiler/rustc_codegen_llvm/src/coverageinfo/map_data.rs

@@ -180,12 +180,12 @@ impl<'tcx> FunctionCoverage<'tcx> {
         self.source_hash
     }
 
-    /// Generate an array of CounterExpressions, and an iterator over all `Counter`s and their
-    /// associated `Regions` (from which the LLVM-specific `CoverageMapGenerator` will create
-    /// `CounterMappingRegion`s.
+    /// Generate an array of CounterExpressions, and an array of all `Counter`s and their
+    /// associated `Regions` sorted by `Region`, from which the LLVM-specific
+    /// `CoverageMapGenerator` will create `CounterMappingRegion`s.
     pub fn get_expressions_and_counter_regions(
         &self,
-    ) -> (Vec<CounterExpression>, impl Iterator<Item = (Counter, &CodeRegion)>) {
+    ) -> (Vec<CounterExpression>, Vec<(Counter, &CodeRegion)>) {
         assert!(
             self.source_hash != 0 || !self.is_used,
             "No counters provided the source_hash for used function: {:?}",
@@ -198,21 +198,31 @@ impl<'tcx> FunctionCoverage<'tcx> {
         // the two vectors should correspond 1:1.
         assert_eq!(self.expressions.len(), counter_expressions.len());
 
-        let counter_regions = self.counter_regions();
-        let expression_regions = self.expression_regions();
-        let unreachable_regions = self.unreachable_regions();
-
-        let counter_regions =
-            counter_regions.chain(expression_regions.into_iter().chain(unreachable_regions));
-        (counter_expressions, counter_regions)
-    }
-
-    fn counter_regions(&self) -> impl Iterator<Item = (Counter, &CodeRegion)> {
-        self.counters.iter_enumerated().filter_map(|(index, entry)| {
+        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))
-        })
+        });
+
+        // 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 unreachable_regions =
+            self.unreachable_regions.iter().map(|region| (Counter::ZERO, region));
+
+        let mut all_regions: Vec<_> = expression_regions.collect();
+        all_regions.extend(counter_regions);
+        all_regions.extend(unreachable_regions);
+
+        // make sure all the regions are sorted
+        all_regions.sort_unstable_by_key(|(_counter, region)| *region);
+
+        (counter_expressions, all_regions)
     }
 
     /// Convert this function's coverage expression data into a form that can be
@@ -244,21 +254,4 @@ impl<'tcx> FunctionCoverage<'tcx> {
             })
             .collect::<Vec<_>>()
     }
-
-    fn expression_regions(&self) -> Vec<(Counter, &CodeRegion)> {
-        // 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.
-        self.expressions
-            .iter_enumerated()
-            .filter_map(|(id, expression)| {
-                let code_region = expression.as_ref()?.region.as_ref()?;
-                Some((Counter::expression(id), code_region))
-            })
-            .collect::<Vec<_>>()
-    }
-
-    fn unreachable_regions(&self) -> impl Iterator<Item = (Counter, &CodeRegion)> {
-        self.unreachable_regions.iter().map(|region| (Counter::ZERO, region))
-    }
 }

compiler/rustc_codegen_llvm/src/coverageinfo/mapgen.rs

@@ -142,14 +142,13 @@ impl CoverageMapGenerator {
     /// Using the `expressions` and `counter_regions` collected for the current function, generate
     /// the `mapping_regions` and `virtual_file_mapping`, and capture any new filenames. Then use
     /// LLVM APIs to encode the `virtual_file_mapping`, `expressions`, and `mapping_regions` into
-    /// the given `coverage_mapping` byte buffer, compliant with the LLVM Coverage Mapping format.
-    fn write_coverage_mapping<'a>(
+    /// the given `coverage_mapping_buffer` byte buffer, compliant with the LLVM Coverage Mapping format.
+    fn write_coverage_mapping(
         &mut self,
         expressions: Vec<CounterExpression>,
-        counter_regions: impl Iterator<Item = (Counter, &'a CodeRegion)>,
+        counter_regions: Vec<(Counter, &CodeRegion)>,
         coverage_mapping_buffer: &RustString,
     ) {
-        let mut counter_regions = counter_regions.collect::<Vec<_>>();
         if counter_regions.is_empty() {
             return;
         }
@@ -164,7 +163,6 @@ impl CoverageMapGenerator {
         // `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.
-        counter_regions.sort_unstable_by_key(|(_counter, region)| *region);
         for (counter, region) in counter_regions {
             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);

compiler/rustc_mir_transform/src/coverage/debug.rs

@@ -44,15 +44,15 @@
 //! points, which can be enabled via environment variable:
 //!
 //! ```shell
-//! RUSTC_LOG=rustc_mir_transform::transform::coverage=debug
+//! RUSTC_LOG=rustc_mir_transform::coverage=debug
 //! ```
 //!
 //! Other module paths with coverage-related debug logs may also be of interest, particularly for
 //! debugging the coverage map data, injected as global variables in the LLVM IR (during rustc's
 //! code generation pass). For example:
 //!
 //! ```shell
-//! RUSTC_LOG=rustc_mir_transform::transform::coverage,rustc_codegen_ssa::coverageinfo,rustc_codegen_llvm::coverageinfo=debug
+//! RUSTC_LOG=rustc_mir_transform::coverage,rustc_codegen_ssa::coverageinfo,rustc_codegen_llvm::coverageinfo=debug
 //! ```
 //!
 //! Coverage Debug Options

compiler/rustc_mir_transform/src/coverage/mod.rs

@@ -216,7 +216,7 @@ impl<'a, 'tcx> Instrumentor<'a, 'tcx> {
             }
 
             ////////////////////////////////////////////////////
-            // Remove the counter or edge counter from of each `CoverageSpan`s associated
+            // Remove the counter or edge counter from each `CoverageSpan`s associated
             // `BasicCoverageBlock`, and inject a `Coverage` statement into the MIR.
             //
             // `Coverage` statements injected from `CoverageSpan`s will include the code regions

compiler/rustc_mir_transform/src/coverage/tests.rs

@@ -2,7 +2,7 @@
 //! pass.
 //!
 //! ```shell
-//! ./x.py test --keep-stage 1 compiler/rustc_mir --test-args '--show-output coverage'
+//! ./x.py test --keep-stage 1 compiler/rustc_mir_transform --test-args '--show-output coverage'
 //! ```
 //!
 //! The tests construct a few "mock" objects, as needed, to support the `InstrumentCoverage`