rustc_codegen_ssa: Make upstream monomorphizations representation sparse

Cargo.lock

@@ -4244,6 +4244,7 @@ dependencies = [
  "rustc_expand",
  "rustc_feature",
  "rustc_fs_util",
+ "rustc_hashes",
  "rustc_hir",
  "rustc_hir_pretty",
  "rustc_incremental",

compiler/rustc_codegen_ssa/src/back/symbol_export.rs

@@ -1,8 +1,7 @@
-use std::collections::hash_map::Entry::*;
-
 use rustc_abi::{CanonAbi, X86Call};
 use rustc_ast::expand::allocator::{AllocatorKind, NO_ALLOC_SHIM_IS_UNSTABLE, global_fn_name};
-use rustc_data_structures::unord::UnordMap;
+use rustc_data_structures::fx::FxIndexMap;
+use rustc_hashes::Hash128;
 use rustc_hir::def::DefKind;
 use rustc_hir::def_id::{CrateNum, DefId, DefIdMap, LOCAL_CRATE, LocalDefId};
 use rustc_middle::bug;
@@ -395,104 +394,45 @@ fn exported_generic_symbols_provider_local<'tcx>(
     tcx.arena.alloc_from_iter(symbols)
 }
 
-fn upstream_monomorphizations_provider(
+fn is_unreachable_local_definition_provider(tcx: TyCtxt<'_>, def_id: LocalDefId) -> bool {
+    !tcx.reachable_set(()).contains(&def_id)
+}
+
+fn upstream_monomorphization_hashes_provider(
     tcx: TyCtxt<'_>,
     (): (),
-) -> DefIdMap<UnordMap<GenericArgsRef<'_>, CrateNum>> {
-    let cnums = tcx.crates(());
-
-    let mut instances: DefIdMap<UnordMap<_, _>> = Default::default();
-
-    let drop_glue_fn_def_id = tcx.lang_items().drop_glue_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_glue_fn_def_id {
-                        (drop_in_place_fn_def_id, tcx.mk_args(&[ty.into()]))
-                    } else {
-                        // `drop_glue` 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);
-                    }
-                }
-                Vacant(e) => {
-                    e.insert(cnum);
-                }
-            }
+) -> FxIndexMap<Hash128, smallvec::SmallVec<[CrateNum; 1]>> {
+    let mut map: FxIndexMap<Hash128, smallvec::SmallVec<[CrateNum; 1]>> = Default::default();
+    for &cnum in tcx.crates(()) {
+        for &hash in tcx.exported_generic_symbol_hashes(cnum) {
+            map.entry(hash).or_default().push(cnum);
         }
     }
-
-    instances
+    // Sort each candidate list by StableCrateId for determinism.
+    for candidates in map.values_mut() {
+        candidates.sort_by_key(|&cnum| tcx.stable_crate_id(cnum));
+    }
+    map
 }
 
-fn upstream_monomorphizations_for_provider(
+fn upstream_monomorphization_for_hash_provider(
     tcx: TyCtxt<'_>,
-    def_id: DefId,
-) -> Option<&UnordMap<GenericArgsRef<'_>, CrateNum>> {
-    assert!(!def_id.is_local());
-    tcx.upstream_monomorphizations(()).get(&def_id)
-}
-
-fn upstream_drop_glue_for_provider<'tcx>(
-    tcx: TyCtxt<'tcx>,
-    args: GenericArgsRef<'tcx>,
-) -> Option<CrateNum> {
-    let def_id = tcx.lang_items().drop_glue_fn()?;
-    tcx.upstream_monomorphizations_for(def_id)?.get(&args).cloned()
-}
-
-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()
-}
-
-fn is_unreachable_local_definition_provider(tcx: TyCtxt<'_>, def_id: LocalDefId) -> bool {
-    !tcx.reachable_set(()).contains(&def_id)
+    hash: Hash128,
+) -> Option<&smallvec::SmallVec<[CrateNum; 1]>> {
+    tcx.upstream_monomorphization_hashes(()).get(&hash)
 }
 
 pub(crate) fn provide(providers: &mut Providers) {
     providers.queries.reachable_non_generics = reachable_non_generics_provider;
     providers.queries.is_reachable_non_generic = is_reachable_non_generic_provider_local;
     providers.queries.exported_non_generic_symbols = exported_non_generic_symbols_provider_local;
     providers.queries.exported_generic_symbols = exported_generic_symbols_provider_local;
-    providers.queries.upstream_monomorphizations = upstream_monomorphizations_provider;
     providers.queries.is_unreachable_local_definition = is_unreachable_local_definition_provider;
-    providers.queries.upstream_drop_glue_for = upstream_drop_glue_for_provider;
-    providers.queries.upstream_async_drop_glue_for = upstream_async_drop_glue_for_provider;
+    providers.queries.upstream_monomorphization_hashes = upstream_monomorphization_hashes_provider;
+    providers.queries.upstream_monomorphization_for_hash =
+        upstream_monomorphization_for_hash_provider;
     providers.queries.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;
 }
 
 pub(crate) fn allocator_shim_symbols(

compiler/rustc_metadata/Cargo.toml

@@ -16,6 +16,7 @@ rustc_errors = { path = "../rustc_errors" }
 rustc_expand = { path = "../rustc_expand" }
 rustc_feature = { path = "../rustc_feature" }
 rustc_fs_util = { path = "../rustc_fs_util" }
+rustc_hashes = { path = "../rustc_hashes" }
 rustc_hir = { path = "../rustc_hir" }
 rustc_hir_pretty = { path = "../rustc_hir_pretty" }
 rustc_incremental = { path = "../rustc_incremental" }

compiler/rustc_metadata/src/rmeta/decoder.rs

@@ -1527,6 +1527,15 @@ impl CrateMetadata {
         tcx.arena.alloc_from_iter(self.root.exported_generic_symbols.decode((self, tcx)))
     }
 
+    fn exported_generic_symbol_hashes<'tcx>(
+        &self,
+        tcx: TyCtxt<'tcx>,
+    ) -> &'tcx [rustc_hashes::Hash128] {
+        let table: exported_symbol_hash_map::ExportedSymbolHashTableRef =
+            self.root.exported_generic_symbol_hashes.decode(&self.blob);
+        tcx.arena.alloc_from_iter(table.keys())
+    }
+
     fn get_macro(&self, tcx: TyCtxt<'_>, id: DefIndex) -> ast::MacroDef {
         match self.def_kind(id) {
             DefKind::Macro(_) => {

compiler/rustc_metadata/src/rmeta/decoder/cstore_impl.rs

@@ -405,6 +405,7 @@ provide! { tcx, def_id, other, cdata,
     }
     exported_non_generic_symbols => { cdata.exported_non_generic_symbols(tcx) }
     exported_generic_symbols => { cdata.exported_generic_symbols(tcx) }
+    exported_generic_symbol_hashes => { cdata.exported_generic_symbol_hashes(tcx) }
 
     crate_extern_paths => { cdata.source().paths().cloned().collect() }
     expn_that_defined => { cdata.get_expn_that_defined(tcx, def_id.index) }

compiler/rustc_metadata/src/rmeta/encoder.rs

@@ -701,6 +701,10 @@ impl<'a, 'tcx> EncodeContext<'a, 'tcx> {
                 )
             });
 
+        let exported_generic_symbol_hashes = stat!("exported-symbol-hashes", || {
+            self.encode_exported_generic_symbol_hashes(tcx.exported_generic_symbols(LOCAL_CRATE))
+        });
+
         // Encode the hygiene data.
         // IMPORTANT: this *must* be the last thing that we encode (other than `SourceMap`). The
         // process of encoding other items (e.g. `optimized_mir`) may cause us to load data from
@@ -768,6 +772,7 @@ impl<'a, 'tcx> EncodeContext<'a, 'tcx> {
                 stable_order_of_exportable_impls,
                 exported_non_generic_symbols,
                 exported_generic_symbols,
+                exported_generic_symbol_hashes,
                 interpret_alloc_index,
                 tables,
                 syntax_contexts,
@@ -2269,6 +2274,30 @@ impl<'a, 'tcx> EncodeContext<'a, 'tcx> {
         self.lazy_array(exported_symbols.iter().cloned())
     }
 
+    fn encode_exported_generic_symbol_hashes(
+        &mut self,
+        exported_symbols: &[(ExportedSymbol<'tcx>, SymbolExportInfo)],
+    ) -> LazyValue<exported_symbol_hash_map::ExportedSymbolHashTableRef> {
+        use exported_symbol_hash_map::{ExportedSymbolHashTable, ExportedSymbolHashTableRef};
+        use rustc_data_structures::stable_hash::{StableHash, StableHasher};
+
+        let tcx = self.tcx;
+        let mut table = ExportedSymbolHashTable::with_capacity(exported_symbols.len(), 87);
+
+        for (sym, _) in exported_symbols {
+            if let Some(instance) = sym.to_instance(tcx) {
+                let hash = tcx.with_stable_hashing_context(|mut hcx| {
+                    let mut hasher = StableHasher::new();
+                    instance.stable_hash(&mut hcx, &mut hasher);
+                    hasher.finish()
+                });
+                table.insert(&hash, &());
+            }
+        }
+
+        self.lazy(ExportedSymbolHashTableRef::Owned(table))
+    }
+
     fn encode_dylib_dependency_formats(&mut self) -> LazyArray<Option<LinkagePreference>> {
         empty_proc_macro!(self);
         let formats = self.tcx.dependency_formats(());

compiler/rustc_metadata/src/rmeta/exported_symbol_hash_map.rs

@@ -0,0 +1,87 @@
+use rustc_data_structures::owned_slice::OwnedSlice;
+use rustc_hashes::Hash128;
+use rustc_serialize::{Decodable, Decoder, Encodable, Encoder};
+
+use crate::rmeta::EncodeContext;
+use crate::rmeta::decoder::BlobDecodeContext;
+
+#[derive(Clone, Default)]
+pub(crate) struct ExportedSymbolHashConfig;
+
+impl odht::Config for ExportedSymbolHashConfig {
+    type Key = Hash128;
+    type Value = ();
+
+    type EncodedKey = [u8; 16];
+    type EncodedValue = [u8; 0];
+
+    type H = odht::UnHashFn;
+
+    #[inline]
+    fn encode_key(k: &Hash128) -> [u8; 16] {
+        k.as_u128().to_le_bytes()
+    }
+
+    #[inline]
+    fn encode_value(_v: &()) -> [u8; 0] {
+        []
+    }
+
+    #[inline]
+    fn decode_key(k: &[u8; 16]) -> Hash128 {
+        Hash128::new(u128::from_le_bytes(*k))
+    }
+
+    #[inline]
+    fn decode_value(_v: &[u8; 0]) {}
+}
+
+pub(crate) type ExportedSymbolHashTable = odht::HashTableOwned<ExportedSymbolHashConfig>;
+
+pub(crate) enum ExportedSymbolHashTableRef {
+    /// Zero-copy view into metadata bytes.
+    OwnedFromMetadata(odht::HashTable<ExportedSymbolHashConfig, OwnedSlice>),
+    /// Locally built table, used during encoding.
+    Owned(ExportedSymbolHashTable),
+}
+
+impl ExportedSymbolHashTableRef {
+    pub(crate) fn keys(&self) -> Vec<Hash128> {
+        match self {
+            Self::OwnedFromMetadata(table) => table.iter().map(|(k, ())| k).collect(),
+            Self::Owned(table) => table.iter().map(|(k, ())| k).collect(),
+        }
+    }
+}
+
+impl<'a, 'tcx> Encodable<EncodeContext<'a, 'tcx>> for ExportedSymbolHashTableRef {
+    fn encode(&self, e: &mut EncodeContext<'a, 'tcx>) {
+        let raw_bytes = match self {
+            ExportedSymbolHashTableRef::Owned(table) => table.raw_bytes(),
+            ExportedSymbolHashTableRef::OwnedFromMetadata(_) => {
+                panic!(
+                    "ExportedSymbolHashTableRef::OwnedFromMetadata variant \
+                     only exists for deserialization"
+                )
+            }
+        };
+        e.emit_usize(raw_bytes.len());
+        e.emit_raw_bytes(raw_bytes);
+    }
+}
+
+impl<'a> Decodable<BlobDecodeContext<'a>> for ExportedSymbolHashTableRef {
+    fn decode(d: &mut BlobDecodeContext<'a>) -> ExportedSymbolHashTableRef {
+        let len = d.read_usize();
+        let pos = d.position();
+        let o = d.blob().bytes().clone().slice(|blob| &blob[pos..pos + len]);
+
+        // Advance the decoder's position past the raw bytes we just sliced.
+        let _ = d.read_raw_bytes(len);
+
+        let inner = odht::HashTable::from_raw_bytes(o).unwrap_or_else(|e| {
+            panic!("decode error for ExportedSymbolHashTable: {e}");
+        });
+        ExportedSymbolHashTableRef::OwnedFromMetadata(inner)
+    }
+}

compiler/rustc_metadata/src/rmeta/mod.rs

@@ -6,6 +6,7 @@ pub(crate) use decoder::{CrateMetadata, CrateNumMap, MetadataBlob, TargetModifie
 use def_path_hash_map::DefPathHashMapRef;
 use encoder::EncodeContext;
 pub use encoder::{EncodedMetadata, encode_metadata, rendered_const};
+use exported_symbol_hash_map::ExportedSymbolHashTableRef;
 pub(crate) use parameterized::ParameterizedOverTcx;
 use rustc_abi::{FieldIdx, ReprOptions, VariantIdx};
 use rustc_ast as ast;
@@ -50,6 +51,7 @@ use crate::eii::EiiMapEncodedKeyValue;
 mod decoder;
 mod def_path_hash_map;
 mod encoder;
+mod exported_symbol_hash_map;
 mod parameterized;
 mod table;
 
@@ -277,6 +279,7 @@ pub(crate) struct CrateRoot {
     stable_order_of_exportable_impls: LazyArray<(DefIndex, usize)>,
     exported_non_generic_symbols: LazyArray<(ExportedSymbol<'static>, SymbolExportInfo)>,
     exported_generic_symbols: LazyArray<(ExportedSymbol<'static>, SymbolExportInfo)>,
+    exported_generic_symbol_hashes: LazyValue<ExportedSymbolHashTableRef>,
 
     syntax_contexts: SyntaxContextTable,
     expn_data: ExpnDataTable,

compiler/rustc_metadata/src/rmeta/parameterized.rs

@@ -73,6 +73,7 @@ trivially_parameterized_over_tcx! {
     crate::rmeta::RawDefId,
     crate::rmeta::TraitImpls,
     crate::rmeta::VariantData,
+    crate::rmeta::exported_symbol_hash_map::ExportedSymbolHashTableRef,
     rustc_abi::ReprOptions,
     rustc_ast::DelimArgs,
     rustc_hir::Attribute,

compiler/rustc_middle/src/middle/exported_symbols.rs

@@ -1,7 +1,7 @@
 use rustc_hir::def_id::{DefId, LOCAL_CRATE};
 use rustc_macros::{Decodable, Encodable, StableHash, TyDecodable, TyEncodable};
 
-use crate::ty::{self, GenericArgsRef, Ty, TyCtxt};
+use crate::ty::{self, GenericArgsRef, Instance, InstanceKind, Ty, TyCtxt};
 
 /// The SymbolExportLevel of a symbols specifies from which kinds of crates
 /// the symbol will be exported. `C` symbols will be exported from any
@@ -57,6 +57,39 @@ pub enum ExportedSymbol<'tcx> {
 }
 
 impl<'tcx> ExportedSymbol<'tcx> {
+    /// Convert to the corresponding [`Instance`], if this is a generic symbol.
+    ///
+    /// Returns `None` for `NonGeneric`, `ThreadLocalShim`, `NoDefId`, and when
+    /// required lang items are unavailable.
+    pub fn to_instance(&self, tcx: TyCtxt<'tcx>) -> Option<Instance<'tcx>> {
+        match *self {
+            ExportedSymbol::Generic(def_id, args) => {
+                Some(Instance { def: InstanceKind::Item(def_id), args })
+            }
+            ExportedSymbol::DropGlue(ty) => {
+                let def_id = tcx.lang_items().drop_glue_fn()?;
+                Some(Instance {
+                    def: InstanceKind::DropGlue(def_id, Some(ty)),
+                    args: tcx.mk_args(&[ty.into()]),
+                })
+            }
+            ExportedSymbol::AsyncDropGlueCtorShim(ty) => {
+                let def_id = tcx.lang_items().async_drop_in_place_fn()?;
+                Some(Instance {
+                    def: InstanceKind::AsyncDropGlueCtorShim(def_id, ty),
+                    args: tcx.mk_args(&[ty.into()]),
+                })
+            }
+            ExportedSymbol::AsyncDropGlue(def_id, ty) => Some(Instance {
+                def: InstanceKind::AsyncDropGlue(def_id, ty),
+                args: tcx.mk_args(&[ty.into()]),
+            }),
+            ExportedSymbol::NonGeneric(..)
+            | ExportedSymbol::ThreadLocalShim(..)
+            | ExportedSymbol::NoDefId(..) => None,
+        }
+    }
+
     /// This is the symbol name of an instance if it is instantiated in the
     /// local crate.
     pub fn symbol_name_for_local_instance(&self, tcx: TyCtxt<'tcx>) -> ty::SymbolName<'tcx> {