Add CGU size heuristic for partitioning

This addresses the concern of #47316 by estimating CGU size based on
the size of its MIR. Looking at the size estimate differences for a
small selection of crates, this heuristic produces different orderings,
which should more accurately reflect optimisation time.

Fixes #47316.

src/librustc_mir/monomorphize/partitioning.rs

@@ -110,11 +110,12 @@ use rustc::mir::mono::{Linkage, Visibility};
 use rustc::ty::{self, TyCtxt, InstanceDef};
 use rustc::ty::item_path::characteristic_def_id_of_type;
 use rustc::util::nodemap::{FxHashMap, FxHashSet};
-use std::collections::hash_map::Entry;
+use std::collections::hash_map::{HashMap, Entry};
 use syntax::ast::NodeId;
 use syntax::symbol::{Symbol, InternedString};
 use rustc::mir::mono::MonoItem;
 use monomorphize::item::{MonoItemExt, InstantiationMode};
+use core::usize;
 
 pub use rustc::mir::mono::CodegenUnit;
 
@@ -229,7 +230,7 @@ pub fn partition<'a, 'tcx, I>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
     // If the partitioning should produce a fixed count of codegen units, merge
     // until that count is reached.
     if let PartitioningStrategy::FixedUnitCount(count) = strategy {
-        merge_codegen_units(&mut initial_partitioning, count, &tcx.crate_name.as_str());
+        merge_codegen_units(tcx, &mut initial_partitioning, count, &tcx.crate_name.as_str());
 
         debug_dump(tcx, "POST MERGING:", initial_partitioning.codegen_units.iter());
     }
@@ -404,7 +405,8 @@ fn place_root_translation_items<'a, 'tcx, I>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
     }
 }
 
-fn merge_codegen_units<'tcx>(initial_partitioning: &mut PreInliningPartitioning<'tcx>,
+fn merge_codegen_units<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
+                             initial_partitioning: &mut PreInliningPartitioning<'tcx>,
                              target_cgu_count: usize,
                              crate_name: &str) {
     assert!(target_cgu_count >= 1);
@@ -421,12 +423,48 @@ fn merge_codegen_units<'tcx>(initial_partitioning: &mut PreInliningPartitioning<
     // the stable sort below will keep everything nice and deterministic.
     codegen_units.sort_by_key(|cgu| cgu.name().clone());
 
+    // Estimate the size of a codegen unit as (approximately) the number of MIR
+    // statements it corresponds to.
+    fn codegen_unit_size_estimate<'a, 'tcx>(cgu: &CodegenUnit<'tcx>,
+                                            mono_item_sizes: &HashMap<MonoItem, usize>)
+                                            -> usize {
+        cgu.items().keys().map(|mi| mono_item_sizes.get(mi).unwrap()).sum()
+    }
+
+    fn mono_item_size_estimate<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
+                                         item: &MonoItem<'tcx>)
+                                         -> usize {
+        match item {
+            MonoItem::Fn(instance) => {
+                // Estimate the size of a function based on how many statements
+                // it contains.
+                let mir = tcx.instance_mir(instance.def);
+                mir.basic_blocks().iter().map(|bb| bb.statements.len()).sum()
+            },
+            // Conservatively estimate the size of a static declaration
+            // or assembly to be 1.
+            MonoItem::Static(_) | MonoItem::GlobalAsm(_) => 1,
+        }
+    }
+
+    // Since `sort_by_key` currently recomputes the keys for each comparison,
+    // we can save unnecessary recomputations by storing size estimates for
+    // each `MonoItem`. Storing estimates for `CodegenUnit` might be preferable,
+    // but its structure makes it awkward to use as a key and additionally their
+    // sizes change as the merging occurs, requiring the map to be updated.
+    let mut sizes: HashMap<MonoItem, usize> = HashMap::new();
+    for mis in codegen_units.iter().map(|cgu| cgu.items().keys()) {
+        mis.for_each(|mi| {
+            sizes.entry(*mi).or_insert_with(|| mono_item_size_estimate(tcx, mi));
+        });
+    }
+
     // Merge the two smallest codegen units until the target size is reached.
     // Note that "size" is estimated here rather inaccurately as the number of
     // translation items in a given unit. This could be improved on.
     while codegen_units.len() > target_cgu_count {
         // Sort small cgus to the back
-        codegen_units.sort_by_key(|cgu| -(cgu.items().len() as i64));
+        codegen_units.sort_by_key(|cgu| usize::MAX - codegen_unit_size_estimate(cgu, &sizes));
         let mut smallest = codegen_units.pop().unwrap();
         let second_smallest = codegen_units.last_mut().unwrap();