Querify MonoItem size estimation

Author: saethlin

compiler/rustc_middle/src/mir/mono.rs

@@ -19,6 +19,7 @@ use rustc_target::spec::SymbolVisibility;
 use tracing::debug;
 
 use crate::dep_graph::{DepNode, WorkProduct, WorkProductId};
+use crate::query::Providers;
 use crate::ty::{GenericArgs, Instance, InstanceKind, SymbolName, TyCtxt};
 
 /// Describes how a monomorphization will be instantiated in object files.
@@ -62,30 +63,6 @@ impl<'tcx> MonoItem<'tcx> {
         }
     }
 
-    // Note: if you change how item size estimates work, you might need to
-    // change NON_INCR_MIN_CGU_SIZE as well.
-    pub fn size_estimate(&self, tcx: TyCtxt<'tcx>) -> usize {
-        match *self {
-            MonoItem::Fn(instance) => {
-                match instance.def {
-                    // "Normal" functions size estimate: the number of
-                    // statements, plus one for the terminator.
-                    InstanceKind::Item(..)
-                    | InstanceKind::DropGlue(..)
-                    | InstanceKind::AsyncDropGlueCtorShim(..) => {
-                        let mir = tcx.instance_mir(instance.def);
-                        mir.basic_blocks.iter().map(|bb| bb.statements.len() + 1).sum()
-                    }
-                    // Other compiler-generated shims size estimate: 1
-                    _ => 1,
-                }
-            }
-            // Conservatively estimate the size of a static declaration or
-            // assembly item to be 1.
-            MonoItem::Static(_) | MonoItem::GlobalAsm(_) => 1,
-        }
-    }
-
     pub fn is_generic_fn(&self) -> bool {
         match self {
             MonoItem::Fn(instance) => instance.args.non_erasable_generics().next().is_some(),
@@ -574,3 +551,31 @@ pub enum CollectionMode {
     /// that it is optimization-independent.
     MentionedItems,
 }
+
+// Note: if you change how item size estimates work, you might need to
+// change NON_INCR_MIN_CGU_SIZE as well.
+fn size_estimate<'tcx>(tcx: TyCtxt<'tcx>, item: MonoItem<'tcx>) -> usize {
+    match item {
+        MonoItem::Fn(instance) => {
+            match instance.def {
+                // "Normal" functions size estimate: the number of
+                // statements, plus one for the terminator.
+                InstanceKind::Item(..)
+                | InstanceKind::DropGlue(..)
+                | InstanceKind::AsyncDropGlueCtorShim(..) => {
+                    let mir = tcx.instance_mir(instance.def);
+                    mir.basic_blocks.iter().map(|bb| bb.statements.len() + 1).sum()
+                }
+                // Other compiler-generated shims size estimate: 1
+                _ => 1,
+            }
+        }
+        // Conservatively estimate the size of a static declaration or
+        // assembly item to be 1.
+        MonoItem::Static(_) | MonoItem::GlobalAsm(_) => 1,
+    }
+}
+
+pub fn provide(providers: &mut Providers) {
+    providers.size_estimate = size_estimate;
+}

compiler/rustc_middle/src/query/keys.rs

@@ -9,6 +9,7 @@ use rustc_target::abi;
 
 use crate::infer::canonical::CanonicalQueryInput;
 use crate::mir::mono::CollectionMode;
+use crate::query::MonoItem;
 use crate::ty::fast_reject::SimplifiedType;
 use crate::ty::layout::{TyAndLayout, ValidityRequirement};
 use crate::ty::{self, GenericArg, GenericArgsRef, Ty, TyCtxt};
@@ -87,6 +88,14 @@ impl<'tcx> Key for ty::Instance<'tcx> {
     }
 }
 
+impl<'tcx> Key for MonoItem<'tcx> {
+    type Cache<V> = DefaultCache<Self, V>;
+
+    fn default_span(&self, tcx: TyCtxt<'_>) -> Span {
+        tcx.def_span(self.def_id())
+    }
+}
+
 impl<'tcx> Key for mir::interpret::GlobalId<'tcx> {
     type Cache<V> = DefaultCache<Self, V>;
 

compiler/rustc_middle/src/query/mod.rs

@@ -2320,6 +2320,11 @@ rustc_queries! {
     query items_of_instance(key: (ty::Instance<'tcx>, CollectionMode)) -> (&'tcx [Spanned<MonoItem<'tcx>>], &'tcx [Spanned<MonoItem<'tcx>>]) {
         desc { "collecting used items" }
     }
+
+    query size_estimate(key: MonoItem<'tcx>) -> usize {
+        desc { "estimating codegen size" }
+        cache_on_disk_if { true }
+    }
 }
 
 rustc_query_append! { define_callbacks! }

compiler/rustc_middle/src/ty/mod.rs

@@ -2155,6 +2155,7 @@ pub fn provide(providers: &mut Providers) {
     print::provide(providers);
     super::util::bug::provide(providers);
     super::middle::provide(providers);
+    super::mir::mono::provide(providers);
     *providers = Providers {
         trait_impls_of: trait_def::trait_impls_of_provider,
         incoherent_impls: trait_def::incoherent_impls_provider,

compiler/rustc_monomorphize/src/partitioning.rs

@@ -254,7 +254,7 @@ where
         if visibility == Visibility::Hidden && can_be_internalized {
             internalization_candidates.insert(mono_item);
         }
-        let size_estimate = mono_item.size_estimate(cx.tcx);
+        let size_estimate = cx.tcx.size_estimate(mono_item);
 
         cgu.items_mut().insert(mono_item, MonoItemData {
             inlined: false,
@@ -279,7 +279,7 @@ where
                 inlined: true,
                 linkage: Linkage::Internal,
                 visibility: Visibility::Default,
-                size_estimate: inlined_item.size_estimate(cx.tcx),
+                size_estimate: cx.tcx.size_estimate(inlined_item),
             });
         }
     }
@@ -1274,7 +1274,7 @@ fn dump_mono_items_stats<'tcx>(
         .map(|(def_id, items)| {
             let name = with_no_trimmed_paths!(tcx.def_path_str(def_id));
             let instantiation_count = items.len();
-            let size_estimate = items[0].size_estimate(tcx);
+            let size_estimate = tcx.size_estimate(*items[0]);
             let total_estimate = instantiation_count * size_estimate;
             MonoItem { name, instantiation_count, size_estimate, total_estimate }
         })