Auto merge of #132566 - saethlin:querify-mir-collection, r=<try>

Querify mir collection

Factored out of #131650, these changes are required for post-mono MIR opts but I want to benchmark them on their own so that I can tune the implementation.

r? ghost

Author: bors

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.
@@ -46,7 +47,7 @@ pub enum InstantiationMode {
     LocalCopy,
 }
 
-#[derive(PartialEq, Eq, Clone, Copy, Debug, Hash, HashStable)]
+#[derive(PartialEq, Eq, Clone, Copy, Debug, Hash, HashStable, TyEncodable, TyDecodable)]
 pub enum MonoItem<'tcx> {
     Fn(Instance<'tcx>),
     Static(DefId),
@@ -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(),
@@ -556,3 +533,49 @@ impl<'tcx> CodegenUnitNameBuilder<'tcx> {
         Symbol::intern(&cgu_name)
     }
 }
+
+/// See module-level docs on some contect for "mentioned" items.
+#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash, HashStable)]
+pub enum CollectionMode {
+    /// Collect items that are used, i.e., actually needed for codegen.
+    ///
+    /// Which items are used can depend on optimization levels, as MIR optimizations can remove
+    /// uses.
+    UsedItems,
+    /// Collect items that are mentioned. The goal of this mode is that it is independent of
+    /// optimizations: the set of "mentioned" items is computed before optimizations are run.
+    ///
+    /// The exact contents of this set are *not* a stable guarantee. (For instance, it is currently
+    /// computed after drop-elaboration. If we ever do some optimizations even in debug builds, we
+    /// might decide to run them before computing mentioned items.) The key property of this set is
+    /// 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/erase.rs

@@ -218,6 +218,10 @@ impl<T0, T1> EraseType for (&'_ T0, &'_ [T1]) {
     type Result = [u8; size_of::<(&'static (), &'static [()])>()];
 }
 
+impl<T0, T1> EraseType for (&'_ [T0], &'_ [T1]) {
+    type Result = [u8; size_of::<(&'static [()], &'static [()])>()];
+}
+
 impl<T0> EraseType for (&'_ T0, Result<(), ErrorGuaranteed>) {
     type Result = [u8; size_of::<(&'static (), Result<(), ErrorGuaranteed>)>()];
 }

compiler/rustc_middle/src/query/keys.rs

@@ -8,6 +8,8 @@ use rustc_span::{DUMMY_SP, Span};
 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};
@@ -86,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>;
 
@@ -591,3 +601,11 @@ impl<'tcx> Key for (ValidityRequirement, ty::ParamEnvAnd<'tcx, Ty<'tcx>>) {
         }
     }
 }
+
+impl<'tcx> Key for (ty::Instance<'tcx>, CollectionMode) {
+    type Cache<V> = DefaultCache<Self, V>;
+
+    fn default_span(&self, tcx: TyCtxt<'_>) -> Span {
+        self.0.default_span(tcx)
+    }
+}

compiler/rustc_middle/src/query/mod.rs

@@ -40,6 +40,7 @@ use rustc_session::cstore::{
 };
 use rustc_session::lint::LintExpectationId;
 use rustc_span::def_id::LOCAL_CRATE;
+use rustc_span::source_map::Spanned;
 use rustc_span::symbol::Symbol;
 use rustc_span::{DUMMY_SP, Span};
 use rustc_target::abi;
@@ -60,7 +61,7 @@ use crate::mir::interpret::{
     EvalStaticInitializerRawResult, EvalToAllocationRawResult, EvalToConstValueResult,
     EvalToValTreeResult, GlobalId, LitToConstError, LitToConstInput,
 };
-use crate::mir::mono::CodegenUnit;
+use crate::mir::mono::{CodegenUnit, CollectionMode, MonoItem};
 use crate::query::erase::{Erase, erase, restore};
 use crate::query::plumbing::{
     CyclePlaceholder, DynamicQuery, query_ensure, query_ensure_error_guaranteed, query_get_at,
@@ -2315,6 +2316,16 @@ rustc_queries! {
         desc { "whether the item should be made inlinable across crates" }
         separate_provide_extern
     }
+
+    query items_of_instance(key: (ty::Instance<'tcx>, CollectionMode)) -> (&'tcx [Spanned<MonoItem<'tcx>>], &'tcx [Spanned<MonoItem<'tcx>>]) {
+        desc { "collecting items used by `{}`", key.0 }
+        cache_on_disk_if { true }
+    }
+
+    query size_estimate(key: MonoItem<'tcx>) -> usize {
+        desc { "estimating codegen size of `{}`", key }
+        cache_on_disk_if { true }
+    }
 }
 
 rustc_query_append! { define_callbacks! }

compiler/rustc_middle/src/query/on_disk_cache.rs

@@ -12,6 +12,7 @@ use rustc_index::{Idx, IndexVec};
 use rustc_macros::{Decodable, Encodable};
 use rustc_middle::dep_graph::{DepNodeIndex, SerializedDepNodeIndex};
 use rustc_middle::mir::interpret::{AllocDecodingSession, AllocDecodingState};
+use rustc_middle::mir::mono::MonoItem;
 use rustc_middle::mir::{self, interpret};
 use rustc_middle::ty::codec::{RefDecodable, TyDecoder, TyEncoder};
 use rustc_middle::ty::{self, Ty, TyCtxt};
@@ -22,7 +23,7 @@ use rustc_session::Session;
 use rustc_span::hygiene::{
     ExpnId, HygieneDecodeContext, HygieneEncodeContext, SyntaxContext, SyntaxContextData,
 };
-use rustc_span::source_map::SourceMap;
+use rustc_span::source_map::{SourceMap, Spanned};
 use rustc_span::{
     BytePos, CachingSourceMapView, ExpnData, ExpnHash, Pos, RelativeBytePos, SourceFile, Span,
     SpanDecoder, SpanEncoder, StableSourceFileId, Symbol,
@@ -773,6 +774,13 @@ impl<'a, 'tcx> Decodable<CacheDecoder<'a, 'tcx>> for &'tcx [rustc_ast::InlineAsm
     }
 }
 
+impl<'a, 'tcx> Decodable<CacheDecoder<'a, 'tcx>> for &'tcx [Spanned<MonoItem<'tcx>>] {
+    #[inline]
+    fn decode(d: &mut CacheDecoder<'a, 'tcx>) -> Self {
+        RefDecodable::decode(d)
+    }
+}
+
 impl<'a, 'tcx> Decodable<CacheDecoder<'a, 'tcx>>
     for &'tcx crate::traits::specialization_graph::Graph
 {

compiler/rustc_middle/src/ty/codec.rs

@@ -12,9 +12,11 @@ use std::marker::DiscriminantKind;
 
 use rustc_data_structures::fx::FxHashMap;
 use rustc_hir::def_id::LocalDefId;
+use rustc_middle::mir::mono::MonoItem;
 use rustc_middle::ty::TyCtxt;
 use rustc_serialize::{Decodable, Encodable};
 use rustc_span::Span;
+use rustc_span::source_map::Spanned;
 use rustc_target::abi::{FieldIdx, VariantIdx};
 pub use rustc_type_ir::{TyDecoder, TyEncoder};
 
@@ -397,6 +399,15 @@ impl<'tcx, D: TyDecoder<I = TyCtxt<'tcx>>> RefDecodable<'tcx, D>
     }
 }
 
+impl<'tcx, D: TyDecoder<I = TyCtxt<'tcx>>> RefDecodable<'tcx, D> for [Spanned<MonoItem<'tcx>>] {
+    fn decode(decoder: &mut D) -> &'tcx Self {
+        decoder
+            .interner()
+            .arena
+            .alloc_from_iter((0..decoder.read_usize()).map(|_| Decodable::decode(decoder)))
+    }
+}
+
 impl<'tcx, D: TyDecoder<I = TyCtxt<'tcx>>> RefDecodable<'tcx, D>
     for ty::List<ty::BoundVariableKind>
 {

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/collector.rs

@@ -219,7 +219,7 @@ use rustc_hir::def_id::{DefId, DefIdMap, LocalDefId};
 use rustc_hir::lang_items::LangItem;
 use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrFlags;
 use rustc_middle::mir::interpret::{AllocId, ErrorHandled, GlobalAlloc, Scalar};
-use rustc_middle::mir::mono::{InstantiationMode, MonoItem};
+use rustc_middle::mir::mono::{CollectionMode, InstantiationMode, MonoItem};
 use rustc_middle::mir::visit::Visitor as MirVisitor;
 use rustc_middle::mir::{self, Location, MentionedItem, traversal};
 use rustc_middle::query::TyCtxtAt;
@@ -268,24 +268,6 @@ struct SharedState<'tcx> {
     usage_map: MTLock<UsageMap<'tcx>>,
 }
 
-/// See module-level docs on some contect for "mentioned" items.
-#[derive(Copy, Clone, Debug, PartialEq)]
-enum CollectionMode {
-    /// Collect items that are used, i.e., actually needed for codegen.
-    ///
-    /// Which items are used can depend on optimization levels, as MIR optimizations can remove
-    /// uses.
-    UsedItems,
-    /// Collect items that are mentioned. The goal of this mode is that it is independent of
-    /// optimizations: the set of "mentioned" items is computed before optimizations are run.
-    ///
-    /// The exact contents of this set are *not* a stable guarantee. (For instance, it is currently
-    /// computed after drop-elaboration. If we ever do some optimizations even in debug builds, we
-    /// might decide to run them before computing mentioned items.) The key property of this set is
-    /// that it is optimization-independent.
-    MentionedItems,
-}
-
 impl<'tcx> UsageMap<'tcx> {
     fn new() -> UsageMap<'tcx> {
         UsageMap { used_map: Default::default(), user_map: Default::default() }
@@ -447,13 +429,9 @@ fn collect_items_rec<'tcx>(
             ));
 
             rustc_data_structures::stack::ensure_sufficient_stack(|| {
-                collect_items_of_instance(
-                    tcx,
-                    instance,
-                    &mut used_items,
-                    &mut mentioned_items,
-                    mode,
-                )
+                let (used, mentioned) = tcx.items_of_instance((instance, mode));
+                used_items.extend(used);
+                mentioned_items.extend(mentioned);
             });
         }
         MonoItem::GlobalAsm(item_id) => {
@@ -1199,14 +1177,12 @@ fn assoc_fn_of_type<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId, fn_ident: Ident) ->
 /// Scans the MIR in order to find function calls, closures, and drop-glue.
 ///
 /// Anything that's found is added to `output`. Furthermore the "mentioned items" of the MIR are returned.
-#[instrument(skip(tcx, used_items, mentioned_items), level = "debug")]
+#[instrument(skip(tcx), level = "debug")]
 fn collect_items_of_instance<'tcx>(
     tcx: TyCtxt<'tcx>,
     instance: Instance<'tcx>,
-    used_items: &mut MonoItems<'tcx>,
-    mentioned_items: &mut MonoItems<'tcx>,
     mode: CollectionMode,
-) {
+) -> (MonoItems<'tcx>, MonoItems<'tcx>) {
     let body = tcx.instance_mir(instance.def);
     // Naively, in "used" collection mode, all functions get added to *both* `used_items` and
     // `mentioned_items`. Mentioned items processing will then notice that they have already been
@@ -1218,11 +1194,13 @@ fn collect_items_of_instance<'tcx>(
     // mentioned item. So instead we collect all pre-monomorphized `MentionedItem` that were already
     // added to `used_items` in a hash set, which can efficiently query in the
     // `body.mentioned_items` loop below without even having to monomorphize the item.
+    let mut used_items = Default::default();
+    let mut mentioned_items = Default::default();
     let mut used_mentioned_items = Default::default();
     let mut collector = MirUsedCollector {
         tcx,
         body,
-        used_items,
+        used_items: &mut used_items,
         used_mentioned_items: &mut used_mentioned_items,
         instance,
         visiting_call_terminator: false,
@@ -1239,7 +1217,7 @@ fn collect_items_of_instance<'tcx>(
     // them errors.
     for const_op in body.required_consts() {
         if let Some(val) = collector.eval_constant(const_op) {
-            collect_const_value(tcx, val, mentioned_items);
+            collect_const_value(tcx, val, &mut mentioned_items);
         }
     }
 
@@ -1248,9 +1226,19 @@ fn collect_items_of_instance<'tcx>(
     for item in body.mentioned_items() {
         if !collector.used_mentioned_items.contains(&item.node) {
             let item_mono = collector.monomorphize(item.node);
-            visit_mentioned_item(tcx, &item_mono, item.span, mentioned_items);
+            visit_mentioned_item(tcx, &item_mono, item.span, &mut mentioned_items);
         }
     }
+
+    (used_items, mentioned_items)
+}
+
+fn items_of_instance<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    (instance, mode): (Instance<'tcx>, CollectionMode),
+) -> (&'tcx [Spanned<MonoItem<'tcx>>], &'tcx [Spanned<MonoItem<'tcx>>]) {
+    let (used_items, mentioned_items) = collect_items_of_instance(tcx, instance, mode);
+    (tcx.arena.alloc_slice(&used_items), tcx.arena.alloc_slice(&mentioned_items))
 }
 
 /// `item` must be already monomorphized.
@@ -1623,4 +1611,5 @@ pub(crate) fn collect_crate_mono_items<'tcx>(
 
 pub(crate) fn provide(providers: &mut Providers) {
     providers.hooks.should_codegen_locally = should_codegen_locally;
+    providers.items_of_instance = items_of_instance;
 }

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 }
         })