rustc_codegen_ssa: Make upstream monomorphizations representation sparse

Upstream monomorphisations are a blessing and a curse of Zed's build performance.

On one hand, we do benefit from it, as builds with share-generics disabled are slower for us (plus we don't really want to use nightly anyways).
On the other, deserializing query results takes *a lot* of time. For some crates close to the end of our compilation pipeline, it's over 400ms per crate.

To make matters worse, I've measured a hit ratio of upstream generics. A sample of such measurement goes as follows:
```
upstream_monomorphization returned None for 28501 distinct monomorphizations.
upstream_monomorphization returned Some 2518 times. Results came from 163 distinct CrateNums.
In total, there are 619731 instantiations
```
This is horrid for us, as we're using a very small percentage of the map that we spend so much time deserializing from.

This commit tries to (rather clumsily) move us towards a sparse representation of upstream_monomorphizations. Instead of storing <DefId, (GenericArgsRef<'_>)> which is rather heavy to deserialize, we'll resort to storing Hashes of Instances. I plan to make this more foolproof, hence this commit is marked as WIP.
For one, we should probably keep the projection queries. Also, it might be worthwhile to store index pointing at entry within  exported_generics of target crate in order to remedy a potential for collisions.

This commit reduces a `touch crates/editor/src/editor.rs` scenario in Zed for me from 14.5s to 11s.

Author: osiewicz

compiler/rustc_codegen_ssa/src/back/symbol_export.rs

@@ -1,10 +1,9 @@
-use std::collections::hash_map::Entry::*;
-
 use rustc_abi::{CanonAbi, X86Call};
 use rustc_ast::expand::allocator::{
     ALLOC_ERROR_HANDLER, ALLOCATOR_METHODS, NO_ALLOC_SHIM_IS_UNSTABLE, global_fn_name,
 };
-use rustc_data_structures::unord::UnordMap;
+use rustc_data_structures::fx::{FxIndexMap, FxIndexSet};
+use rustc_hashes::Hash128;
 use rustc_hir::def::DefKind;
 use rustc_hir::def_id::{CrateNum, DefId, DefIdMap, LOCAL_CRATE, LocalDefId};
 use rustc_middle::bug;
@@ -384,86 +383,75 @@ fn exported_generic_symbols_provider_local<'tcx>(
     tcx.arena.alloc_from_iter(symbols)
 }
 
-fn upstream_monomorphizations_provider(
+fn upstream_monomorphizations_of_crate_provider(
     tcx: TyCtxt<'_>,
-    (): (),
-) -> DefIdMap<UnordMap<GenericArgsRef<'_>, CrateNum>> {
-    let cnums = tcx.crates(());
-
-    let mut instances: DefIdMap<UnordMap<_, _>> = Default::default();
-
+    cnum: CrateNum,
+) -> FxIndexSet<Hash128> {
+    use rustc_data_structures::stable_hasher::HashStable;
     let drop_in_place_fn_def_id = tcx.lang_items().drop_in_place_fn();
     let async_drop_in_place_fn_def_id = tcx.lang_items().async_drop_in_place_fn();
-
-    for &cnum in cnums.iter() {
-        for (exported_symbol, _) in tcx.exported_generic_symbols(cnum).iter() {
-            let (def_id, args) = match *exported_symbol {
-                ExportedSymbol::Generic(def_id, args) => (def_id, args),
-                ExportedSymbol::DropGlue(ty) => {
-                    if let Some(drop_in_place_fn_def_id) = drop_in_place_fn_def_id {
-                        (drop_in_place_fn_def_id, tcx.mk_args(&[ty.into()]))
-                    } else {
-                        // `drop_in_place` in place does not exist, don't try
-                        // to use it.
-                        continue;
-                    }
-                }
-                ExportedSymbol::AsyncDropGlueCtorShim(ty) => {
-                    if let Some(async_drop_in_place_fn_def_id) = async_drop_in_place_fn_def_id {
-                        (async_drop_in_place_fn_def_id, tcx.mk_args(&[ty.into()]))
-                    } else {
-                        continue;
-                    }
-                }
-                ExportedSymbol::AsyncDropGlue(def_id, ty) => (def_id, tcx.mk_args(&[ty.into()])),
-                ExportedSymbol::NonGeneric(..)
-                | ExportedSymbol::ThreadLocalShim(..)
-                | ExportedSymbol::NoDefId(..) => unreachable!("{exported_symbol:?}"),
-            };
-
-            let args_map = instances.entry(def_id).or_default();
-
-            match args_map.entry(args) {
-                Occupied(mut e) => {
-                    // If there are multiple monomorphizations available,
-                    // we select one deterministically.
-                    let other_cnum = *e.get();
-                    if tcx.stable_crate_id(other_cnum) > tcx.stable_crate_id(cnum) {
-                        e.insert(cnum);
-                    }
+    let mut instances = FxIndexSet::default();
+    for (exported_symbol, _) in tcx.exported_generic_symbols(cnum).iter() {
+        let instance = match *exported_symbol {
+            ExportedSymbol::Generic(def_id, args) => rustc_middle::ty::Instance {
+                args,
+                def: rustc_middle::ty::InstanceKind::Item(def_id),
+            },
+            ExportedSymbol::DropGlue(ty) => {
+                let Some(drop_in_place_fn_def_id) = drop_in_place_fn_def_id else {
+                    // `drop_in_place` in place does not exist, don't try
+                    // to use it.
+                    continue;
+                };
+                let args = tcx.mk_args(&[ty.clone().into()]);
+                rustc_middle::ty::Instance {
+                    args,
+                    def: rustc_middle::ty::InstanceKind::DropGlue(
+                        drop_in_place_fn_def_id,
+                        Some(ty),
+                    ),
                 }
-                Vacant(e) => {
-                    e.insert(cnum);
+            }
+            ExportedSymbol::AsyncDropGlueCtorShim(ty) => {
+                let Some(drop_in_place) = async_drop_in_place_fn_def_id else {
+                    continue;
+                };
+                let args = tcx.mk_args(&[ty.clone().into()]);
+                rustc_middle::ty::Instance {
+                    args,
+                    def: rustc_middle::ty::InstanceKind::AsyncDropGlueCtorShim(drop_in_place, ty),
                 }
             }
-        }
-    }
+            ExportedSymbol::AsyncDropGlue(def_id, ty) => rustc_middle::ty::Instance {
+                args: tcx.mk_args(&[ty.clone().into()]),
+                def: rustc_middle::ty::InstanceKind::AsyncDropGlue(def_id, ty),
+            },
+            ExportedSymbol::NonGeneric(..)
+            | ExportedSymbol::ThreadLocalShim(..)
+            | ExportedSymbol::NoDefId(..) => unreachable!("{exported_symbol:?}"),
+        };
 
+        tcx.with_stable_hashing_context(|mut tcx| {
+            let mut hasher = rustc_data_structures::stable_hasher::StableHasher::new();
+            instance.hash_stable(&mut tcx, &mut hasher);
+            instances.insert(hasher.finish());
+        });
+    }
     instances
 }
 
-fn upstream_monomorphizations_for_provider(
-    tcx: TyCtxt<'_>,
-    def_id: DefId,
-) -> Option<&UnordMap<GenericArgsRef<'_>, CrateNum>> {
-    assert!(!def_id.is_local());
-    tcx.upstream_monomorphizations(()).get(&def_id)
-}
+fn upstream_monomorphizations_provider(tcx: TyCtxt<'_>, _: ()) -> FxIndexMap<Hash128, CrateNum> {
+    let cnums = tcx.crates(());
+    let mut instances: FxIndexMap<Hash128, CrateNum> = Default::default();
 
-fn upstream_drop_glue_for_provider<'tcx>(
-    tcx: TyCtxt<'tcx>,
-    args: GenericArgsRef<'tcx>,
-) -> Option<CrateNum> {
-    let def_id = tcx.lang_items().drop_in_place_fn()?;
-    tcx.upstream_monomorphizations_for(def_id)?.get(&args).cloned()
-}
+    for cnum in cnums.iter() {
+        let hashes = tcx.upstream_monomorphizations_of_crate(*cnum);
+        for hash in hashes {
+            instances.entry(*hash).or_insert(*cnum);
+        }
+    }
 
-fn upstream_async_drop_glue_for_provider<'tcx>(
-    tcx: TyCtxt<'tcx>,
-    args: GenericArgsRef<'tcx>,
-) -> Option<CrateNum> {
-    let def_id = tcx.lang_items().async_drop_in_place_fn()?;
-    tcx.upstream_monomorphizations_for(def_id)?.get(&args).cloned()
+    instances
 }
 
 fn is_unreachable_local_definition_provider(tcx: TyCtxt<'_>, def_id: LocalDefId) -> bool {
@@ -477,12 +465,9 @@ pub(crate) fn provide(providers: &mut Providers) {
     providers.exported_generic_symbols = exported_generic_symbols_provider_local;
     providers.upstream_monomorphizations = upstream_monomorphizations_provider;
     providers.is_unreachable_local_definition = is_unreachable_local_definition_provider;
-    providers.upstream_drop_glue_for = upstream_drop_glue_for_provider;
-    providers.upstream_async_drop_glue_for = upstream_async_drop_glue_for_provider;
     providers.wasm_import_module_map = wasm_import_module_map;
     providers.extern_queries.is_reachable_non_generic = is_reachable_non_generic_provider_extern;
-    providers.extern_queries.upstream_monomorphizations_for =
-        upstream_monomorphizations_for_provider;
+    providers.upstream_monomorphizations_of_crate = upstream_monomorphizations_of_crate_provider;
 }
 
 pub(crate) fn allocator_shim_symbols(

compiler/rustc_middle/src/query/mod.rs

@@ -76,6 +76,7 @@ use rustc_data_structures::steal::Steal;
 use rustc_data_structures::svh::Svh;
 use rustc_data_structures::unord::{UnordMap, UnordSet};
 use rustc_errors::ErrorGuaranteed;
+use rustc_hashes::Hash128;
 use rustc_hir::attrs::StrippedCfgItem;
 use rustc_hir::def::{DefKind, DocLinkResMap};
 use rustc_hir::def_id::{
@@ -1960,65 +1961,15 @@ rustc_queries! {
     /// crate. Instead use the narrower `upstream_monomorphizations_for`,
     /// `upstream_drop_glue_for`, `upstream_async_drop_glue_for`, or,
     /// even better, `Instance::upstream_monomorphization()`.
-    query upstream_monomorphizations(_: ()) -> &'tcx DefIdMap<UnordMap<GenericArgsRef<'tcx>, CrateNum>> {
+    query upstream_monomorphizations(_: ()) -> &'tcx FxIndexMap<Hash128, CrateNum> {
         arena_cache
         desc { "collecting available upstream monomorphizations" }
     }
 
-    /// Returns the set of upstream monomorphizations available for the
-    /// generic function identified by the given `def_id`. The query makes
-    /// sure to make a stable selection if the same monomorphization is
-    /// available in multiple upstream crates.
-    ///
-    /// You likely want to call `Instance::upstream_monomorphization()`
-    /// instead of invoking this query directly.
-    query upstream_monomorphizations_for(def_id: DefId)
-        -> Option<&'tcx UnordMap<GenericArgsRef<'tcx>, CrateNum>>
-    {
-        desc { |tcx|
-            "collecting available upstream monomorphizations for `{}`",
-            tcx.def_path_str(def_id),
-        }
-        separate_provide_extern
-    }
-
-    /// Returns the upstream crate that exports drop-glue for the given
-    /// type (`args` is expected to be a single-item list containing the
-    /// type one wants drop-glue for).
-    ///
-    /// This is a subset of `upstream_monomorphizations_for` in order to
-    /// increase dep-tracking granularity. Otherwise adding or removing any
-    /// type with drop-glue in any upstream crate would invalidate all
-    /// functions calling drop-glue of an upstream type.
-    ///
-    /// You likely want to call `Instance::upstream_monomorphization()`
-    /// instead of invoking this query directly.
-    ///
-    /// NOTE: This query could easily be extended to also support other
-    ///       common functions that have are large set of monomorphizations
-    ///       (like `Clone::clone` for example).
-    query upstream_drop_glue_for(args: GenericArgsRef<'tcx>) -> Option<CrateNum> {
-        desc { "available upstream drop-glue for `{:?}`", args }
-    }
-
-    /// Returns the upstream crate that exports async-drop-glue for
-    /// the given type (`args` is expected to be a single-item list
-    /// containing the type one wants async-drop-glue for).
-    ///
-    /// This is a subset of `upstream_monomorphizations_for` in order
-    /// to increase dep-tracking granularity. Otherwise adding or
-    /// removing any type with async-drop-glue in any upstream crate
-    /// would invalidate all functions calling async-drop-glue of an
-    /// upstream type.
-    ///
-    /// You likely want to call `Instance::upstream_monomorphization()`
-    /// instead of invoking this query directly.
-    ///
-    /// NOTE: This query could easily be extended to also support other
-    ///       common functions that have are large set of monomorphizations
-    ///       (like `Clone::clone` for example).
-    query upstream_async_drop_glue_for(args: GenericArgsRef<'tcx>) -> Option<CrateNum> {
-        desc { "available upstream async-drop-glue for `{:?}`", args }
+    query upstream_monomorphizations_of_crate(_: CrateNum) -> &'tcx FxIndexSet<Hash128> {
+        arena_cache
+        cache_on_disk_if { true }
+        desc { "collecting available upstream monomorphizations for crate" }
     }
 
     /// Returns a list of all `extern` blocks of a crate.

compiler/rustc_middle/src/ty/instance.rs

@@ -2,7 +2,9 @@ use std::assert_matches::assert_matches;
 use std::fmt;
 
 use rustc_data_structures::fx::FxHashMap;
+use rustc_data_structures::stable_hasher::{HashStable, StableHasher};
 use rustc_errors::ErrorGuaranteed;
+use rustc_hashes::Hash128;
 use rustc_hir as hir;
 use rustc_hir::def::{CtorKind, DefKind, Namespace};
 use rustc_hir::def_id::{CrateNum, DefId};
@@ -224,18 +226,25 @@ impl<'tcx> Instance<'tcx> {
             return None;
         }
 
-        match self.def {
-            InstanceKind::Item(def) => tcx
-                .upstream_monomorphizations_for(def)
-                .and_then(|monos| monos.get(&self.args).cloned()),
-            InstanceKind::DropGlue(_, Some(_)) => tcx.upstream_drop_glue_for(self.args),
+        let defining_crates = match self.def {
+            InstanceKind::Item(_)
+            | InstanceKind::DropGlue(_, Some(_))
+            | InstanceKind::AsyncDropGlueCtorShim(_, _) => {
+                let hash: Hash128 = tcx.with_stable_hashing_context(|mut hcx| {
+                    let mut hasher = StableHasher::new();
+                    self.hash_stable(&mut hcx, &mut hasher);
+                    hasher.finish()
+                });
+                tcx.upstream_monomorphizations(()).get(&hash)
+            }
+
             InstanceKind::AsyncDropGlue(_, _) => None,
             InstanceKind::FutureDropPollShim(_, _, _) => None,
-            InstanceKind::AsyncDropGlueCtorShim(_, _) => {
-                tcx.upstream_async_drop_glue_for(self.args)
-            }
+
             _ => None,
-        }
+        }?;
+
+        Some(*defining_crates)
     }
 }