From abde9ba527704449ae84848673b52a8ab0f5cbd1 Mon Sep 17 00:00:00 2001
From: Nicholas Nethercote <n.nethercote@gmail.com>
Date: Wed, 21 Jun 2023 16:38:47 +1000
Subject: [PATCH] Tweak CGU size estimate code.

- Rename `create_size_estimate` as `compute_size_estimate`, because that
  makes more sense for the second and subsequent calls for each CGU.
- Change `CodegenUnit::size_estimate` from `Option<usize>` to `usize`.
  We can still assert that `compute_size_estimate` is called first.
- Move the size estimation for `place_mono_items` inside the function,
  for consistency with `merge_codegen_units`.
---
 compiler/rustc_middle/src/mir/mono.rs           | 16 +++++++++-------
 compiler/rustc_monomorphize/src/partitioning.rs | 15 ++++++++-------
 2 files changed, 17 insertions(+), 14 deletions(-)

diff --git a/compiler/rustc_middle/src/mir/mono.rs b/compiler/rustc_middle/src/mir/mono.rs
index 1cbf1a36283..ca735d52314 100644
--- a/compiler/rustc_middle/src/mir/mono.rs
+++ b/compiler/rustc_middle/src/mir/mono.rs
@@ -231,7 +231,7 @@ pub struct CodegenUnit<'tcx> {
     /// as well as the crate name and disambiguator.
     name: Symbol,
     items: FxHashMap<MonoItem<'tcx>, (Linkage, Visibility)>,
-    size_estimate: Option<usize>,
+    size_estimate: usize,
     primary: bool,
     /// True if this is CGU is used to hold code coverage information for dead code,
     /// false otherwise.
@@ -269,7 +269,7 @@ pub fn new(name: Symbol) -> CodegenUnit<'tcx> {
         CodegenUnit {
             name,
             items: Default::default(),
-            size_estimate: None,
+            size_estimate: 0,
             primary: false,
             is_code_coverage_dead_code_cgu: false,
         }
@@ -320,19 +320,21 @@ pub fn mangle_name(human_readable_name: &str) -> String {
         base_n::encode(hash, base_n::CASE_INSENSITIVE)
     }
 
-    pub fn create_size_estimate(&mut self, tcx: TyCtxt<'tcx>) {
+    pub fn compute_size_estimate(&mut self, tcx: TyCtxt<'tcx>) {
         // Estimate the size of a codegen unit as (approximately) the number of MIR
         // statements it corresponds to.
-        self.size_estimate = Some(self.items.keys().map(|mi| mi.size_estimate(tcx)).sum());
+        self.size_estimate = self.items.keys().map(|mi| mi.size_estimate(tcx)).sum();
     }
 
     #[inline]
-    /// Should only be called if [`create_size_estimate`] has previously been called.
+    /// Should only be called if [`compute_size_estimate`] has previously been called.
     ///
-    /// [`create_size_estimate`]: Self::create_size_estimate
+    /// [`compute_size_estimate`]: Self::compute_size_estimate
     pub fn size_estimate(&self) -> usize {
+        // Items are never zero-sized, so if we have items the estimate must be
+        // non-zero, unless we forgot to call `compute_size_estimate` first.
+        assert!(self.items.is_empty() || self.size_estimate != 0);
         self.size_estimate
-            .expect("create_size_estimate must be called before getting a size_estimate")
     }
 
     pub fn contains_item(&self, item: &MonoItem<'tcx>) -> bool {
diff --git a/compiler/rustc_monomorphize/src/partitioning.rs b/compiler/rustc_monomorphize/src/partitioning.rs
index 4e9a4f91ddf..97e3748b8b8 100644
--- a/compiler/rustc_monomorphize/src/partitioning.rs
+++ b/compiler/rustc_monomorphize/src/partitioning.rs
@@ -150,14 +150,11 @@ fn partition<'tcx, I>(
 
     let cx = &PartitioningCx { tcx, usage_map };
 
-    // Place all mono items into a codegen unit.
+    // Place all mono items into a codegen unit. `place_mono_items` is
+    // responsible for initializing the CGU size estimates.
     let PlacedMonoItems { mut codegen_units, internalization_candidates, unique_inlined_stats } = {
         let _prof_timer = tcx.prof.generic_activity("cgu_partitioning_place_items");
-        let mut placed = place_mono_items(cx, mono_items);
-
-        for cgu in &mut placed.codegen_units {
-            cgu.create_size_estimate(tcx);
-        }
+        let placed = place_mono_items(cx, mono_items);
 
         debug_dump(tcx, "PLACE", &placed.codegen_units, placed.unique_inlined_stats);
 
@@ -282,6 +279,10 @@ fn place_mono_items<'tcx, I>(cx: &PartitioningCx<'_, 'tcx>, mono_items: I) -> Pl
     let mut codegen_units: Vec<_> = codegen_units.into_values().collect();
     codegen_units.sort_by(|a, b| a.name().as_str().cmp(b.name().as_str()));
 
+    for cgu in codegen_units.iter_mut() {
+        cgu.compute_size_estimate(cx.tcx);
+    }
+
     return PlacedMonoItems {
         codegen_units,
         internalization_candidates,
@@ -351,7 +352,7 @@ fn merge_codegen_units<'tcx>(
         // may be duplicate inlined items, in which case the destination CGU is
         // unaffected. Recalculate size estimates afterwards.
         second_smallest.items_mut().extend(smallest.items_mut().drain());
-        second_smallest.create_size_estimate(cx.tcx);
+        second_smallest.compute_size_estimate(cx.tcx);
 
         // Record that `second_smallest` now contains all the stuff that was
         // in `smallest` before.