locally caching the sizes of each mono-item

Author: andjo403

Diff for: src/librustc/mir/mono.rs

@@ -47,20 +47,6 @@ pub enum MonoItem<'tcx> {
 }
 
 impl<'tcx> MonoItem<'tcx> {
-    pub fn size_estimate(&self, tcx: TyCtxt<'tcx>) -> usize {
-        match *self {
-            MonoItem::Fn(instance) => {
-                // Estimate the size of a function based on how many statements
-                // it contains.
-                tcx.instance_def_size_estimate(instance.def)
-            },
-            // Conservatively estimate the size of a static declaration
-            // or assembly to be 1.
-            MonoItem::Static(_) |
-            MonoItem::GlobalAsm(_) => 1,
-        }
-    }
-
     pub fn is_generic_fn(&self) -> bool {
         match *self {
             MonoItem::Fn(ref instance) => {
@@ -248,7 +234,7 @@ pub struct CodegenUnit<'tcx> {
     /// as well as the crate name and disambiguator.
     name: InternedString,
     items: FxHashMap<MonoItem<'tcx>, (Linkage, Visibility)>,
-    size_estimate: Option<usize>,
+    size_estimate: usize,
 }
 
 #[derive(Copy, Clone, PartialEq, Eq, Hash, Debug, RustcEncodable, RustcDecodable)]
@@ -298,7 +284,7 @@ impl<'tcx> CodegenUnit<'tcx> {
         CodegenUnit {
             name: name,
             items: Default::default(),
-            size_estimate: None,
+            size_estimate: 0,
         }
     }
 
@@ -330,21 +316,17 @@ impl<'tcx> CodegenUnit<'tcx> {
         base_n::encode(hash, base_n::CASE_INSENSITIVE)
     }
 
-    pub fn estimate_size(&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());
-    }
-
     pub fn size_estimate(&self) -> usize {
-        // Should only be called if `estimate_size` has previously been called.
-        self.size_estimate.expect("estimate_size must be called before getting a size_estimate")
+        self.size_estimate
     }
 
-    pub fn modify_size_estimate(&mut self, delta: usize) {
-        assert!(self.size_estimate.is_some());
-        if let Some(size_estimate) = self.size_estimate {
-            self.size_estimate = Some(size_estimate + delta);
+    pub fn add_item(
+            &mut self,
+            item: MonoItem<'tcx>,
+            linkage: (Linkage, Visibility),
+            item_size: usize ) {
+        if self.items.insert(item, linkage).is_none() {
+            self.size_estimate += item_size;
         }
     }
 

Diff for: src/librustc_mir/monomorphize/partitioning.rs

@@ -136,15 +136,21 @@ where
 {
     let _prof_timer = tcx.prof.generic_activity("cgu_partitioning");
 
+    let mut estimator = SizeEstimator::new();
+
     // In the first step, we place all regular monomorphizations into their
     // respective 'home' codegen unit. Regular monomorphizations are all
     // functions and statics defined in the local crate.
     let initial_partitioning = {
         let _prof_timer = tcx.prof.generic_activity("cgu_partitioning_place_roots");
-        place_root_mono_items(tcx, mono_items)
+        place_root_mono_items(tcx, mono_items, &mut estimator)
     };
 
-    debug_dump(tcx, "INITIAL PARTITIONING:", initial_partitioning.codegen_units.iter());
+    debug_dump(
+        tcx,
+        "INITIAL PARTITIONING:",
+        initial_partitioning.codegen_units.iter(),
+        &mut estimator);
 
     // In the next step, we use the inlining map to determine which additional
     // monomorphizations have to go into each codegen unit. These additional
@@ -153,19 +159,17 @@ where
     let mut post_inlining = {
         let _prof_timer =
             tcx.prof.generic_activity("cgu_partitioning_place_inline_items");
-        place_inlined_mono_items(initial_partitioning, inlining_map)
+        place_inlined_mono_items(tcx, initial_partitioning, inlining_map, &mut estimator)
     };
 
-    post_inlining.codegen_units.iter_mut().for_each(|cgu| cgu.estimate_size(tcx));
-
-    debug_dump(tcx, "POST INLINING:", post_inlining.codegen_units.iter());
+    debug_dump(tcx, "POST INLINING:", post_inlining.codegen_units.iter(), &mut estimator);
 
     // If the partitioning should produce a fixed count of codegen units, merge
     // until that count is reached.
     if let PartitioningStrategy::FixedUnitCount(count) = strategy {
         let _prof_timer = tcx.prof.generic_activity("cgu_partitioning_merge_cgus");
-        merge_codegen_units(tcx, &mut post_inlining, count);
-        debug_dump(tcx, "POST MERGING:", post_inlining.codegen_units.iter());
+        merge_codegen_units(tcx, &mut post_inlining, count, &mut estimator);
+        debug_dump(tcx, "POST MERGING:", post_inlining.codegen_units.iter(), &mut estimator);
     }
 
     // Next we try to make as many symbols "internal" as possible, so LLVM has
@@ -189,6 +193,27 @@ where
     result
 }
 
+struct SizeEstimator<'tcx> {
+    cache: FxHashMap<MonoItem<'tcx>, usize>,
+}
+impl<'tcx> SizeEstimator<'tcx> {
+    pub fn new() -> SizeEstimator<'tcx> {
+        SizeEstimator { cache: Default::default() }
+    }
+    pub fn size_estimate(&mut self, tcx: TyCtxt<'tcx>, mono_item: MonoItem<'tcx>) -> usize {
+        *self.cache.entry(mono_item).or_insert_with(|| match mono_item {
+            MonoItem::Fn(instance) => {
+                // Estimate the size of a function based on how many statements
+                // it contains.
+                tcx.instance_def_size_estimate(instance.def)
+            }
+            // Conservatively estimate the size of a static declaration
+            // or assembly to be 1.
+            MonoItem::Static(_) | MonoItem::GlobalAsm(_) => 1,
+        })
+    }
+}
+
 struct PreInliningPartitioning<'tcx> {
     codegen_units: Vec<CodegenUnit<'tcx>>,
     roots: FxHashSet<MonoItem<'tcx>>,
@@ -209,7 +234,11 @@ struct PostInliningPartitioning<'tcx> {
     internalization_candidates: FxHashSet<MonoItem<'tcx>>,
 }
 
-fn place_root_mono_items<'tcx, I>(tcx: TyCtxt<'tcx>, mono_items: I) -> PreInliningPartitioning<'tcx>
+fn place_root_mono_items<'tcx, I>(
+    tcx: TyCtxt<'tcx>,
+    mono_items: I,
+    estimator: &mut SizeEstimator<'tcx>,
+) -> PreInliningPartitioning<'tcx>
 where
     I: Iterator<Item = MonoItem<'tcx>>,
 {
@@ -260,8 +289,8 @@ where
         if visibility == Visibility::Hidden && can_be_internalized {
             internalization_candidates.insert(mono_item);
         }
-
-        codegen_unit.items_mut().insert(mono_item, (linkage, visibility));
+        let item_size = estimator.size_estimate(tcx, mono_item);
+        codegen_unit.add_item(mono_item, (linkage, visibility), item_size);
         roots.insert(mono_item);
     }
 
@@ -475,6 +504,7 @@ fn merge_codegen_units<'tcx>(
     tcx: TyCtxt<'tcx>,
     partitioning: &mut PostInliningPartitioning<'tcx>,
     target_cgu_count: usize,
+    estimator: &mut SizeEstimator<'tcx>,
 ) {
     assert!(target_cgu_count >= 1);
     let codegen_units = &mut partitioning.codegen_units;
@@ -497,10 +527,10 @@ fn merge_codegen_units<'tcx>(
         let mut smallest = codegen_units.pop().unwrap();
         let second_smallest = codegen_units.last_mut().unwrap();
 
-        for (k, v) in smallest.items_mut().drain() {
-            second_smallest.items_mut().insert(k, v);
+        for (mono_item, v) in smallest.items_mut().drain() {
+            let item_size = estimator.size_estimate(tcx, mono_item);
+            second_smallest.add_item(mono_item, v, item_size);
         }
-        second_smallest.estimate_size(tcx);
         debug!("CodegenUnit {} merged in to CodegenUnit {}",
                smallest.name(),
                second_smallest.name());
@@ -512,9 +542,11 @@ fn merge_codegen_units<'tcx>(
     }
 }
 
-fn place_inlined_mono_items<'tcx>(initial_partitioning: PreInliningPartitioning<'tcx>,
-                                  inlining_map: &InliningMap<'tcx>)
-                                  -> PostInliningPartitioning<'tcx> {
+fn place_inlined_mono_items<'tcx>(tcx: TyCtxt<'tcx>,
+                                  initial_partitioning: PreInliningPartitioning<'tcx>,
+                                  inlining_map: &InliningMap<'tcx>,
+                                  estimator: &mut SizeEstimator<'tcx>,
+) -> PostInliningPartitioning<'tcx> {
     let mut new_partitioning = Vec::new();
 
     let PreInliningPartitioning {
@@ -536,18 +568,20 @@ fn place_inlined_mono_items<'tcx>(initial_partitioning: PreInliningPartitioning<
         for mono_item in reachable {
             if let Some(linkage) = old_codegen_unit.items().get(&mono_item) {
                 // This is a root, just copy it over.
-                new_codegen_unit.items_mut().insert(mono_item, *linkage);
+                let item_size = estimator.size_estimate(tcx, mono_item);
+                new_codegen_unit.add_item(mono_item, *linkage, item_size);
             } else {
                 if roots.contains(&mono_item) {
                     bug!("GloballyShared mono-item inlined into other CGU: \
                           {:?}", mono_item);
                 }
 
                 // This is a CGU-private copy.
-                new_codegen_unit.items_mut().insert(
+                let item_size = estimator.size_estimate(tcx, mono_item);
+                new_codegen_unit.add_item(
                     mono_item,
                     (Linkage::Internal, Visibility::Default),
-                );
+                    item_size);
             }
         }
 
@@ -768,7 +802,10 @@ fn numbered_codegen_unit_name(
     name_builder.build_cgu_name_no_mangle(LOCAL_CRATE, &["cgu"], Some(index))
 }
 
-fn debug_dump<'a, 'tcx, I>(tcx: TyCtxt<'tcx>, label: &str, cgus: I)
+fn debug_dump<'a, 'tcx, I>(tcx: TyCtxt<'tcx>,
+    label: &str,
+    cgus: I,
+    estimator: &mut SizeEstimator<'tcx>)
 where
     I: Iterator<Item = &'a CodegenUnit<'tcx>>,
     'tcx: 'a,
@@ -788,7 +825,7 @@ where
                        mono_item.to_string(tcx, true),
                        linkage,
                        symbol_hash,
-                       mono_item.size_estimate(tcx));
+                       estimator.size_estimate(tcx, *mono_item));
             }
 
             debug!("");