resolve/expand: Cache intermediate results of #[derive] expansion

petrochenkov · petrochenkov · commit fbf1bec48228 · 2021-04-04T17:51:41.000+03:00
diff --git a/compiler/rustc_builtin_macros/src/derive.rs b/compiler/rustc_builtin_macros/src/derive.rs
@@ -1,6 +1,6 @@
 use crate::cfg_eval::cfg_eval;
 
-use rustc_ast::{self as ast, token, ItemKind, MetaItemKind, NestedMetaItem, StmtKind};
+use rustc_ast::{self as ast, attr, token, ItemKind, MetaItemKind, NestedMetaItem, StmtKind};
 use rustc_errors::{struct_span_err, Applicability};
 use rustc_expand::base::{Annotatable, ExpandResult, ExtCtxt, Indeterminate, MultiItemModifier};
 use rustc_feature::AttributeTemplate;
@@ -26,32 +26,39 @@ impl MultiItemModifier for Expander {
             return ExpandResult::Ready(vec![item]);
         }
 
-        let template =
-            AttributeTemplate { list: Some("Trait1, Trait2, ..."), ..Default::default() };
-        let attr = ecx.attribute(meta_item.clone());
-        validate_attr::check_builtin_attribute(&sess.parse_sess, &attr, sym::derive, template);
+        let result =
+            ecx.resolver.resolve_derives(ecx.current_expansion.id, ecx.force_mode, &|| {
+                let template =
+                    AttributeTemplate { list: Some("Trait1, Trait2, ..."), ..Default::default() };
+                let attr = attr::mk_attr_outer(meta_item.clone());
+                validate_attr::check_builtin_attribute(
+                    &sess.parse_sess,
+                    &attr,
+                    sym::derive,
+                    template,
+                );
 
-        let derives: Vec<_> = attr
-            .meta_item_list()
-            .unwrap_or_default()
-            .into_iter()
-            .filter_map(|nested_meta| match nested_meta {
-                NestedMetaItem::MetaItem(meta) => Some(meta),
-                NestedMetaItem::Literal(lit) => {
-                    // Reject `#[derive("Debug")]`.
-                    report_unexpected_literal(sess, &lit);
-                    None
-                }
-            })
-            .map(|meta| {
-                // Reject `#[derive(Debug = "value", Debug(abc))]`, but recover the paths.
-                report_path_args(sess, &meta);
-                meta.path
-            })
-            .collect();
+                attr.meta_item_list()
+                    .unwrap_or_default()
+                    .into_iter()
+                    .filter_map(|nested_meta| match nested_meta {
+                        NestedMetaItem::MetaItem(meta) => Some(meta),
+                        NestedMetaItem::Literal(lit) => {
+                            // Reject `#[derive("Debug")]`.
+                            report_unexpected_literal(sess, &lit);
+                            None
+                        }
+                    })
+                    .map(|meta| {
+                        // Reject `#[derive(Debug = "value", Debug(abc))]`, but recover the paths.
+                        report_path_args(sess, &meta);
+                        meta.path
+                    })
+                    .map(|path| (path, None))
+                    .collect()
+            });
 
-        // FIXME: Try to cache intermediate results to avoid collecting same paths multiple times.
-        match ecx.resolver.resolve_derives(ecx.current_expansion.id, derives, ecx.force_mode) {
+        match result {
             Ok(()) => ExpandResult::Ready(cfg_eval(ecx, item)),
             Err(Indeterminate) => ExpandResult::Retry(item),
         }
diff --git a/compiler/rustc_expand/src/base.rs b/compiler/rustc_expand/src/base.rs
@@ -868,6 +868,8 @@ impl SyntaxExtension {
 /// Error type that denotes indeterminacy.
 pub struct Indeterminate;
 
+pub type DeriveResolutions = Vec<(ast::Path, Option<Lrc<SyntaxExtension>>)>;
+
 pub trait ResolverExpand {
     fn next_node_id(&mut self) -> NodeId;
 
@@ -904,15 +906,12 @@ pub trait ResolverExpand {
     fn resolve_derives(
         &mut self,
         expn_id: ExpnId,
-        derives: Vec<ast::Path>,
         force: bool,
+        derive_paths: &dyn Fn() -> DeriveResolutions,
     ) -> Result<(), Indeterminate>;
     /// Take resolutions for paths inside the `#[derive(...)]` attribute with the given `ExpnId`
     /// back from resolver.
-    fn take_derive_resolutions(
-        &mut self,
-        expn_id: ExpnId,
-    ) -> Option<Vec<(Lrc<SyntaxExtension>, ast::Path)>>;
+    fn take_derive_resolutions(&mut self, expn_id: ExpnId) -> Option<DeriveResolutions>;
     /// Path resolution logic for `#[cfg_accessible(path)]`.
     fn cfg_accessible(&mut self, expn_id: ExpnId, path: &ast::Path) -> Result<bool, Indeterminate>;
 }
diff --git a/compiler/rustc_expand/src/expand.rs b/compiler/rustc_expand/src/expand.rs
@@ -515,7 +515,7 @@ impl<'a, 'b> MacroExpander<'a, 'b> {
                             invocations.reserve(derives.len());
                             derives
                                 .into_iter()
-                                .map(|(_exts, path)| {
+                                .map(|(path, _exts)| {
                                     // FIXME: Consider using the derive resolutions (`_exts`)
                                     // instead of enqueuing the derives to be resolved again later.
                                     let expn_id = ExpnId::fresh(None);
diff --git a/compiler/rustc_resolve/src/lib.rs b/compiler/rustc_resolve/src/lib.rs
@@ -37,7 +37,7 @@ use rustc_data_structures::fx::{FxHashMap, FxHashSet, FxIndexMap};
 use rustc_data_structures::ptr_key::PtrKey;
 use rustc_data_structures::sync::Lrc;
 use rustc_errors::{struct_span_err, Applicability, DiagnosticBuilder};
-use rustc_expand::base::{SyntaxExtension, SyntaxExtensionKind};
+use rustc_expand::base::{DeriveResolutions, SyntaxExtension, SyntaxExtensionKind};
 use rustc_hir::def::Namespace::*;
 use rustc_hir::def::{self, CtorOf, DefKind, NonMacroAttrKind, PartialRes};
 use rustc_hir::def_id::{CrateNum, DefId, DefIdMap, LocalDefId, CRATE_DEF_INDEX};
@@ -851,6 +851,12 @@ enum BuiltinMacroState {
     AlreadySeen(Span),
 }
 
+struct DeriveData {
+    resolutions: DeriveResolutions,
+    helper_attrs: Vec<Ident>,
+    has_derive_copy: bool,
+}
+
 /// The main resolver class.
 ///
 /// This is the visitor that walks the whole crate.
@@ -973,8 +979,9 @@ pub struct Resolver<'a> {
     output_macro_rules_scopes: FxHashMap<ExpnId, MacroRulesScopeRef<'a>>,
     /// Helper attributes that are in scope for the given expansion.
     helper_attrs: FxHashMap<ExpnId, Vec<Ident>>,
-    /// Resolutions for paths inside the `#[derive(...)]` attribute with the given `ExpnId`.
-    derive_resolutions: FxHashMap<ExpnId, Vec<(Lrc<SyntaxExtension>, ast::Path)>>,
+    /// Ready or in-progress results of resolving paths inside the `#[derive(...)]` attribute
+    /// with the given `ExpnId`.
+    derive_data: FxHashMap<ExpnId, DeriveData>,
 
     /// Avoid duplicated errors for "name already defined".
     name_already_seen: FxHashMap<Symbol, Span>,
@@ -1310,7 +1317,7 @@ impl<'a> Resolver<'a> {
             invocation_parent_scopes: Default::default(),
             output_macro_rules_scopes: Default::default(),
             helper_attrs: Default::default(),
-            derive_resolutions: Default::default(),
+            derive_data: Default::default(),
             local_macro_def_scopes: FxHashMap::default(),
             name_already_seen: FxHashMap::default(),
             potentially_unused_imports: Vec::new(),
diff --git a/compiler/rustc_resolve/src/macros.rs b/compiler/rustc_resolve/src/macros.rs
@@ -4,7 +4,7 @@
 use crate::imports::ImportResolver;
 use crate::Namespace::*;
 use crate::{AmbiguityError, AmbiguityErrorMisc, AmbiguityKind, BuiltinMacroState, Determinacy};
-use crate::{CrateLint, ParentScope, ResolutionError, Resolver, Scope, ScopeSet, Weak};
+use crate::{CrateLint, DeriveData, ParentScope, ResolutionError, Resolver, Scope, ScopeSet, Weak};
 use crate::{ModuleKind, ModuleOrUniformRoot, NameBinding, PathResult, Segment, ToNameBinding};
 use rustc_ast::{self as ast, Inline, ItemKind, ModKind, NodeId};
 use rustc_ast_lowering::ResolverAstLowering;
@@ -14,8 +14,8 @@ use rustc_data_structures::fx::FxHashSet;
 use rustc_data_structures::ptr_key::PtrKey;
 use rustc_data_structures::sync::Lrc;
 use rustc_errors::struct_span_err;
-use rustc_expand::base::Annotatable;
-use rustc_expand::base::{Indeterminate, ResolverExpand, SyntaxExtension, SyntaxExtensionKind};
+use rustc_expand::base::{Annotatable, DeriveResolutions, Indeterminate, ResolverExpand};
+use rustc_expand::base::{SyntaxExtension, SyntaxExtensionKind};
 use rustc_expand::compile_declarative_macro;
 use rustc_expand::expand::{AstFragment, Invocation, InvocationKind, SupportsMacroExpansion};
 use rustc_feature::is_builtin_attr_name;
@@ -359,58 +359,69 @@ impl<'a> ResolverExpand for Resolver<'a> {
     fn resolve_derives(
         &mut self,
         expn_id: ExpnId,
-        derives: Vec<ast::Path>,
         force: bool,
+        derive_paths: &dyn Fn() -> DeriveResolutions,
     ) -> Result<(), Indeterminate> {
         // Block expansion of the container until we resolve all derives in it.
         // This is required for two reasons:
         // - Derive helper attributes are in scope for the item to which the `#[derive]`
         //   is applied, so they have to be produced by the container's expansion rather
         //   than by individual derives.
         // - Derives in the container need to know whether one of them is a built-in `Copy`.
-        // FIXME: Try to cache intermediate results to avoid resolving same derives multiple times.
+        // Temporarily take the data to avoid borrow checker conflicts.
+        let mut derive_data = mem::take(&mut self.derive_data);
+        let entry = derive_data.entry(expn_id).or_insert_with(|| DeriveData {
+            resolutions: derive_paths(),
+            helper_attrs: Vec::new(),
+            has_derive_copy: false,
+        });
         let parent_scope = self.invocation_parent_scopes[&expn_id];
-        let mut exts = Vec::new();
-        let mut helper_attrs = Vec::new();
-        let mut has_derive_copy = false;
-        for path in derives {
-            exts.push((
-                match self.resolve_macro_path(
-                    &path,
-                    Some(MacroKind::Derive),
-                    &parent_scope,
-                    true,
-                    force,
-                ) {
-                    Ok((Some(ext), _)) => {
-                        let span =
-                            path.segments.last().unwrap().ident.span.normalize_to_macros_2_0();
-                        helper_attrs
-                            .extend(ext.helper_attrs.iter().map(|name| Ident::new(*name, span)));
-                        has_derive_copy |= ext.builtin_name == Some(sym::Copy);
-                        ext
-                    }
-                    Ok(_) | Err(Determinacy::Determined) => self.dummy_ext(MacroKind::Derive),
-                    Err(Determinacy::Undetermined) => return Err(Indeterminate),
-                },
-                path,
-            ))
+        for (path, opt_ext) in &mut entry.resolutions {
+            if opt_ext.is_none() {
+                *opt_ext = Some(
+                    match self.resolve_macro_path(
+                        &path,
+                        Some(MacroKind::Derive),
+                        &parent_scope,
+                        true,
+                        force,
+                    ) {
+                        Ok((Some(ext), _)) => {
+                            if !ext.helper_attrs.is_empty() {
+                                let last_seg = path.segments.last().unwrap();
+                                let span = last_seg.ident.span.normalize_to_macros_2_0();
+                                entry.helper_attrs.extend(
+                                    ext.helper_attrs.iter().map(|name| Ident::new(*name, span)),
+                                );
+                            }
+                            entry.has_derive_copy |= ext.builtin_name == Some(sym::Copy);
+                            ext
+                        }
+                        Ok(_) | Err(Determinacy::Determined) => self.dummy_ext(MacroKind::Derive),
+                        Err(Determinacy::Undetermined) => {
+                            assert!(self.derive_data.is_empty());
+                            self.derive_data = derive_data;
+                            return Err(Indeterminate);
+                        }
+                    },
+                );
+            }
         }
-        self.derive_resolutions.insert(expn_id, exts);
-        self.helper_attrs.insert(expn_id, helper_attrs);
+        // If we get to here, then `derive_data` for the given `expn_id` will only be accessed by
+        // `take_derive_resolutions` later, so we can steal `helper_attrs` instead of cloning them.
+        self.helper_attrs.insert(expn_id, mem::take(&mut entry.helper_attrs));
         // Mark this derive as having `Copy` either if it has `Copy` itself or if its parent derive
         // has `Copy`, to support cases like `#[derive(Clone, Copy)] #[derive(Debug)]`.
-        if has_derive_copy || self.has_derive_copy(parent_scope.expansion) {
+        if entry.has_derive_copy || self.has_derive_copy(parent_scope.expansion) {
             self.containers_deriving_copy.insert(expn_id);
         }
+        assert!(self.derive_data.is_empty());
+        self.derive_data = derive_data;
         Ok(())
     }
 
-    fn take_derive_resolutions(
-        &mut self,
-        expn_id: ExpnId,
-    ) -> Option<Vec<(Lrc<SyntaxExtension>, ast::Path)>> {
-        self.derive_resolutions.remove(&expn_id)
+    fn take_derive_resolutions(&mut self, expn_id: ExpnId) -> Option<DeriveResolutions> {
+        self.derive_data.remove(&expn_id).map(|data| data.resolutions)
     }
 
     // The function that implements the resolution logic of `#[cfg_accessible(path)]`.
