EII collection queries

EII collection queries

Author: jdonszelmann

compiler/rustc_interface/src/passes.rs

@@ -1061,6 +1061,9 @@ fn run_required_analyses(tcx: TyCtxt<'_>) {
         parallel!(
             {
                 sess.time("looking_for_entry_point", || tcx.ensure_ok().entry_fn(()));
+                sess.time("check_externally_implementable_items", || {
+                    tcx.ensure_ok().check_externally_implementable_items(())
+                });
 
                 sess.time("looking_for_derive_registrar", || {
                     tcx.ensure_ok().proc_macro_decls_static(())

compiler/rustc_metadata/src/eii.rs

@@ -0,0 +1,48 @@
+use rustc_data_structures::fx::FxIndexMap;
+use rustc_hir::attrs::{AttributeKind, EiiDecl, EiiImpl};
+use rustc_hir::def_id::DefId;
+use rustc_hir::find_attr;
+use rustc_middle::query::LocalCrate;
+use rustc_middle::ty::TyCtxt;
+
+// basically the map below but flattened out
+pub(crate) type EiiMapEncodedKeyValue = (DefId, (EiiDecl, Vec<(DefId, EiiImpl)>));
+
+pub(crate) type EiiMap = FxIndexMap<
+    DefId, // the defid of the macro that declared the eii
+    (
+        // the corresponding declaration
+        EiiDecl,
+        // all the given implementations, indexed by defid.
+        // We expect there to be only one, but collect them all to give errors if there are more
+        // (or if there are none) in the final crate we build.
+        FxIndexMap<DefId, EiiImpl>,
+    ),
+>;
+
+pub(crate) fn collect<'tcx>(tcx: TyCtxt<'tcx>, LocalCrate: LocalCrate) -> EiiMap {
+    let mut eiis = EiiMap::default();
+
+    // iterate over all items in the current crate
+    // FIXME(speed up)
+    for id in tcx.hir_crate_items(()).definitions() {
+        for i in
+            find_attr!(tcx.get_all_attrs(id), AttributeKind::EiiImpls(e) => e).into_iter().flatten()
+        {
+            eiis.entry(i.eii_macro)
+                .or_insert_with(|| {
+                    // find the decl for this one if it wasn't in yet (maybe it's from the local crate? not very useful but not illegal)
+                    (find_attr!(tcx.get_all_attrs(i.eii_macro), AttributeKind::EiiExternTarget(d) => *d).unwrap(), Default::default())
+                }).1.insert(id.into(), *i);
+        }
+
+        // if we find a new declaration, add it to the list without a known implementation
+        if let Some(decl) =
+            find_attr!(tcx.get_all_attrs(id), AttributeKind::EiiExternTarget(d) => *d)
+        {
+            eiis.entry(id.into()).or_insert((decl, Default::default()));
+        }
+    }
+
+    eiis
+}

compiler/rustc_metadata/src/lib.rs

@@ -16,6 +16,7 @@
 pub use rmeta::provide;
 
 mod dependency_format;
+mod eii;
 mod foreign_modules;
 mod native_libs;
 mod rmeta;

compiler/rustc_metadata/src/rmeta/decoder.rs

@@ -39,6 +39,7 @@ use rustc_span::{
 use tracing::debug;
 
 use crate::creader::CStore;
+use crate::eii::EiiMapEncodedKeyValue;
 use crate::rmeta::table::IsDefault;
 use crate::rmeta::*;
 
@@ -1487,6 +1488,13 @@ impl<'a> CrateMetadataRef<'a> {
         )
     }
 
+    fn get_externally_implementable_items(
+        self,
+        tcx: TyCtxt<'_>,
+    ) -> impl Iterator<Item = EiiMapEncodedKeyValue> {
+        self.root.externally_implementable_items.decode((self, tcx))
+    }
+
     fn get_missing_lang_items<'tcx>(self, tcx: TyCtxt<'tcx>) -> &'tcx [LangItem] {
         tcx.arena.alloc_from_iter(self.root.lang_items_missing.decode((self, tcx)))
     }

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

@@ -25,7 +25,7 @@ use super::{Decodable, DecodeIterator};
 use crate::creader::{CStore, LoadedMacro};
 use crate::rmeta::AttrFlags;
 use crate::rmeta::table::IsDefault;
-use crate::{foreign_modules, native_libs};
+use crate::{eii, foreign_modules, native_libs};
 
 trait ProcessQueryValue<'tcx, T> {
     fn process_decoded(self, _tcx: TyCtxt<'tcx>, _err: impl Fn() -> !) -> T;
@@ -330,9 +330,22 @@ provide! { tcx, def_id, other, cdata,
     is_private_dep => { cdata.private_dep }
     is_panic_runtime => { cdata.root.panic_runtime }
     is_compiler_builtins => { cdata.root.compiler_builtins }
+
+    // FIXME: to be replaced with externally_implementable_items below
     has_global_allocator => { cdata.root.has_global_allocator }
+    // FIXME: to be replaced with externally_implementable_items below
     has_alloc_error_handler => { cdata.root.has_alloc_error_handler }
+    // FIXME: to be replaced with externally_implementable_items below
     has_panic_handler => { cdata.root.has_panic_handler }
+
+    externally_implementable_items => {
+        cdata.get_externally_implementable_items(tcx)
+            .map(|(decl_did, (decl, impls))| (
+                decl_did,
+                (decl, impls.into_iter().collect())
+            )).collect()
+    }
+
     is_profiler_runtime => { cdata.root.profiler_runtime }
     required_panic_strategy => { cdata.root.required_panic_strategy }
     panic_in_drop_strategy => { cdata.root.panic_in_drop_strategy }
@@ -430,6 +443,7 @@ pub(in crate::rmeta) fn provide(providers: &mut Providers) {
         },
         native_libraries: native_libs::collect,
         foreign_modules: foreign_modules::collect,
+        externally_implementable_items: eii::collect,
 
         // Returns a map from a sufficiently visible external item (i.e., an
         // external item that is visible from at least one local module) to a

compiler/rustc_metadata/src/rmeta/encoder.rs

@@ -35,6 +35,7 @@ use rustc_span::{
 };
 use tracing::{debug, instrument, trace};
 
+use crate::eii::EiiMapEncodedKeyValue;
 use crate::errors::{FailCreateFileEncoder, FailWriteFile};
 use crate::rmeta::*;
 
@@ -620,6 +621,9 @@ impl<'a, 'tcx> EncodeContext<'a, 'tcx> {
         // We have already encoded some things. Get their combined size from the current position.
         stats.push(("preamble", self.position()));
 
+        let externally_implementable_items = stat!("externally-implementable-items", || self
+            .encode_externally_implementable_items());
+
         let (crate_deps, dylib_dependency_formats) =
             stat!("dep", || (self.encode_crate_deps(), self.encode_dylib_dependency_formats()));
 
@@ -738,6 +742,7 @@ impl<'a, 'tcx> EncodeContext<'a, 'tcx> {
                     attrs,
                     sym::default_lib_allocator,
                 ),
+                externally_implementable_items,
                 proc_macro_data,
                 debugger_visualizers,
                 compiler_builtins: ast::attr::contains_name(attrs, sym::compiler_builtins),
@@ -1649,6 +1654,15 @@ impl<'a, 'tcx> EncodeContext<'a, 'tcx> {
         }
     }
 
+    fn encode_externally_implementable_items(&mut self) -> LazyArray<EiiMapEncodedKeyValue> {
+        empty_proc_macro!(self);
+        let externally_implementable_items = self.tcx.externally_implementable_items(LOCAL_CRATE);
+
+        self.lazy_array(externally_implementable_items.iter().map(|(decl_did, (decl, impls))| {
+            (*decl_did, (decl.clone(), impls.iter().map(|(impl_did, i)| (*impl_did, *i)).collect()))
+        }))
+    }
+
     #[instrument(level = "trace", skip(self))]
     fn encode_info_for_adt(&mut self, local_def_id: LocalDefId) {
         let def_id = local_def_id.to_def_id();

compiler/rustc_metadata/src/rmeta/mod.rs

@@ -44,6 +44,7 @@ use table::TableBuilder;
 use {rustc_ast as ast, rustc_hir as hir};
 
 use crate::creader::CrateMetadataRef;
+use crate::eii::EiiMapEncodedKeyValue;
 
 mod decoder;
 mod def_path_hash_map;
@@ -250,6 +251,7 @@ pub(crate) struct CrateRoot {
     has_alloc_error_handler: bool,
     has_panic_handler: bool,
     has_default_lib_allocator: bool,
+    externally_implementable_items: LazyArray<EiiMapEncodedKeyValue>,
 
     crate_deps: LazyArray<CrateDep>,
     dylib_dependency_formats: LazyArray<Option<LinkagePreference>>,

compiler/rustc_metadata/src/rmeta/parameterized.rs

@@ -87,6 +87,8 @@ trivially_parameterized_over_tcx! {
     rustc_hir::Safety,
     rustc_hir::Stability,
     rustc_hir::attrs::Deprecation,
+    rustc_hir::attrs::EiiDecl,
+    rustc_hir::attrs::EiiImpl,
     rustc_hir::attrs::StrippedCfgItem<rustc_hir::def_id::DefIndex>,
     rustc_hir::def::DefKind,
     rustc_hir::def::DocLinkResMap,

compiler/rustc_middle/src/query/erase.rs

@@ -296,6 +296,8 @@ trivial! {
     rustc_ast::expand::allocator::AllocatorKind,
     rustc_hir::DefaultBodyStability,
     rustc_hir::attrs::Deprecation,
+    rustc_hir::attrs::EiiDecl,
+    rustc_hir::attrs::EiiImpl,
     rustc_data_structures::svh::Svh,
     rustc_errors::ErrorGuaranteed,
     rustc_hir::Constness,

compiler/rustc_middle/src/query/mod.rs

@@ -76,7 +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_hir::attrs::StrippedCfgItem;
+use rustc_hir::attrs::{EiiDecl, EiiImpl, StrippedCfgItem};
 use rustc_hir::def::{DefKind, DocLinkResMap};
 use rustc_hir::def_id::{
     CrateNum, DefId, DefIdMap, LocalDefId, LocalDefIdMap, LocalDefIdSet, LocalModDefId,
@@ -2753,6 +2753,17 @@ rustc_queries! {
         desc { |tcx| "checking what set of sanitizers are enabled on `{}`", tcx.def_path_str(key) }
         feedable
     }
+
+    query check_externally_implementable_items(_: ()) {
+        desc { "check externally implementable items" }
+    }
+
+    /// Returns a list of all `externally implementable items` crate.
+    query externally_implementable_items(_: CrateNum) -> &'tcx FxIndexMap<DefId, (EiiDecl, FxIndexMap<DefId, EiiImpl>)> {
+        arena_cache
+        desc { "looking up the externally implementable items of a crate" }
+        separate_provide_extern
+    }
 }
 
 rustc_with_all_queries! { define_callbacks! }