Try caching macro calls more aggressively

crates/hir/src/semantics/source_to_def.rs

@@ -98,22 +98,26 @@ use hir_def::{
     FieldId, FunctionId, GenericDefId, GenericParamId, ImplId, LifetimeParamId, MacroId, ModuleId,
     StaticId, StructId, TraitAliasId, TraitId, TypeAliasId, TypeParamId, UnionId, UseId, VariantId,
 };
-use hir_expand::{attrs::AttrId, name::AsName, HirFileId, HirFileIdExt, MacroCallId};
+use hir_expand::{
+    attrs::AttrId, name::AsName, ExpansionInfo, HirFileId, HirFileIdExt, MacroCallId,
+};
 use rustc_hash::FxHashMap;
 use smallvec::SmallVec;
+use span::MacroFileId;
 use stdx::impl_from;
 use syntax::{
     ast::{self, HasName},
-    AstNode, SyntaxNode,
+    AstNode, AstPtr, SyntaxNode,
 };
 
 use crate::{db::HirDatabase, InFile};
 
 pub(super) type SourceToDefCache = FxHashMap<(ChildContainer, HirFileId), DynMap>;
 
-pub(super) struct SourceToDefCtx<'a, 'b> {
-    pub(super) db: &'b dyn HirDatabase,
-    pub(super) dynmap_cache: &'a mut SourceToDefCache,
+pub(super) struct SourceToDefCtx<'a, 'dyn_cache> {
+    pub(super) db: &'a dyn HirDatabase,
+    pub(super) dynmap_cache: &'dyn_cache mut SourceToDefCache,
+    pub(super) expansion_info_cache: &'a mut FxHashMap<MacroFileId, ExpansionInfo>,
 }
 
 impl SourceToDefCtx<'_, '_> {
@@ -135,19 +139,21 @@ impl SourceToDefCtx<'_, '_> {
         mods
     }
 
-    pub(super) fn module_to_def(&mut self, src: InFile<ast::Module>) -> Option<ModuleId> {
+    pub(super) fn module_to_def(&mut self, src: InFile<&ast::Module>) -> Option<ModuleId> {
         let _p = tracing::span!(tracing::Level::INFO, "module_to_def").entered();
-        let parent_declaration = src
-            .syntax()
-            .ancestors_with_macros(self.db.upcast())
-            .find_map(|it| it.map(Either::<ast::Module, ast::BlockExpr>::cast).transpose())
+        let parent_declaration = self
+            .ancestors_with_macros(src.syntax_ref(), |_, ancestor| {
+                ancestor.map(Either::<ast::Module, ast::BlockExpr>::cast).transpose()
+            })
             .map(|it| it.transpose());
 
         let parent_module = match parent_declaration {
             Some(Either::Right(parent_block)) => self
-                .block_to_def(parent_block)
+                .block_to_def(parent_block.as_ref())
                 .map(|block| self.db.block_def_map(block).root_module_id()),
-            Some(Either::Left(parent_declaration)) => self.module_to_def(parent_declaration),
+            Some(Either::Left(parent_declaration)) => {
+                self.module_to_def(parent_declaration.as_ref())
+            }
             None => {
                 let file_id = src.file_id.original_file(self.db.upcast());
                 self.file_to_def(file_id).first().copied()
@@ -160,73 +166,79 @@ impl SourceToDefCtx<'_, '_> {
         Some(def_map.module_id(child_id))
     }
 
-    pub(super) fn source_file_to_def(&self, src: InFile<ast::SourceFile>) -> Option<ModuleId> {
+    pub(super) fn source_file_to_def(&self, src: InFile<&ast::SourceFile>) -> Option<ModuleId> {
         let _p = tracing::span!(tracing::Level::INFO, "source_file_to_def").entered();
         let file_id = src.file_id.original_file(self.db.upcast());
         self.file_to_def(file_id).first().copied()
     }
 
-    pub(super) fn trait_to_def(&mut self, src: InFile<ast::Trait>) -> Option<TraitId> {
+    pub(super) fn trait_to_def(&mut self, src: InFile<&ast::Trait>) -> Option<TraitId> {
         self.to_def(src, keys::TRAIT)
     }
     pub(super) fn trait_alias_to_def(
         &mut self,
-        src: InFile<ast::TraitAlias>,
+        src: InFile<&ast::TraitAlias>,
     ) -> Option<TraitAliasId> {
         self.to_def(src, keys::TRAIT_ALIAS)
     }
-    pub(super) fn impl_to_def(&mut self, src: InFile<ast::Impl>) -> Option<ImplId> {
+    pub(super) fn impl_to_def(&mut self, src: InFile<&ast::Impl>) -> Option<ImplId> {
         self.to_def(src, keys::IMPL)
     }
-    pub(super) fn fn_to_def(&mut self, src: InFile<ast::Fn>) -> Option<FunctionId> {
+    pub(super) fn fn_to_def(&mut self, src: InFile<&ast::Fn>) -> Option<FunctionId> {
         self.to_def(src, keys::FUNCTION)
     }
-    pub(super) fn struct_to_def(&mut self, src: InFile<ast::Struct>) -> Option<StructId> {
+    pub(super) fn struct_to_def(&mut self, src: InFile<&ast::Struct>) -> Option<StructId> {
         self.to_def(src, keys::STRUCT)
     }
-    pub(super) fn enum_to_def(&mut self, src: InFile<ast::Enum>) -> Option<EnumId> {
+    pub(super) fn enum_to_def(&mut self, src: InFile<&ast::Enum>) -> Option<EnumId> {
         self.to_def(src, keys::ENUM)
     }
-    pub(super) fn union_to_def(&mut self, src: InFile<ast::Union>) -> Option<UnionId> {
+    pub(super) fn union_to_def(&mut self, src: InFile<&ast::Union>) -> Option<UnionId> {
         self.to_def(src, keys::UNION)
     }
-    pub(super) fn static_to_def(&mut self, src: InFile<ast::Static>) -> Option<StaticId> {
+    pub(super) fn static_to_def(&mut self, src: InFile<&ast::Static>) -> Option<StaticId> {
         self.to_def(src, keys::STATIC)
     }
-    pub(super) fn const_to_def(&mut self, src: InFile<ast::Const>) -> Option<ConstId> {
+    pub(super) fn const_to_def(&mut self, src: InFile<&ast::Const>) -> Option<ConstId> {
         self.to_def(src, keys::CONST)
     }
-    pub(super) fn type_alias_to_def(&mut self, src: InFile<ast::TypeAlias>) -> Option<TypeAliasId> {
+    pub(super) fn type_alias_to_def(
+        &mut self,
+        src: InFile<&ast::TypeAlias>,
+    ) -> Option<TypeAliasId> {
         self.to_def(src, keys::TYPE_ALIAS)
     }
-    pub(super) fn record_field_to_def(&mut self, src: InFile<ast::RecordField>) -> Option<FieldId> {
+    pub(super) fn record_field_to_def(
+        &mut self,
+        src: InFile<&ast::RecordField>,
+    ) -> Option<FieldId> {
         self.to_def(src, keys::RECORD_FIELD)
     }
-    pub(super) fn tuple_field_to_def(&mut self, src: InFile<ast::TupleField>) -> Option<FieldId> {
+    pub(super) fn tuple_field_to_def(&mut self, src: InFile<&ast::TupleField>) -> Option<FieldId> {
         self.to_def(src, keys::TUPLE_FIELD)
     }
-    pub(super) fn block_to_def(&mut self, src: InFile<ast::BlockExpr>) -> Option<BlockId> {
+    pub(super) fn block_to_def(&mut self, src: InFile<&ast::BlockExpr>) -> Option<BlockId> {
         self.to_def(src, keys::BLOCK)
     }
     pub(super) fn enum_variant_to_def(
         &mut self,
-        src: InFile<ast::Variant>,
+        src: InFile<&ast::Variant>,
     ) -> Option<EnumVariantId> {
         self.to_def(src, keys::ENUM_VARIANT)
     }
     pub(super) fn extern_crate_to_def(
         &mut self,
-        src: InFile<ast::ExternCrate>,
+        src: InFile<&ast::ExternCrate>,
     ) -> Option<ExternCrateId> {
         self.to_def(src, keys::EXTERN_CRATE)
     }
     #[allow(dead_code)]
-    pub(super) fn use_to_def(&mut self, src: InFile<ast::Use>) -> Option<UseId> {
+    pub(super) fn use_to_def(&mut self, src: InFile<&ast::Use>) -> Option<UseId> {
         self.to_def(src, keys::USE)
     }
     pub(super) fn adt_to_def(
         &mut self,
-        InFile { file_id, value }: InFile<ast::Adt>,
+        InFile { file_id, value }: InFile<&ast::Adt>,
     ) -> Option<AdtId> {
         match value {
             ast::Adt::Enum(it) => self.enum_to_def(InFile::new(file_id, it)).map(AdtId::EnumId),
@@ -238,11 +250,11 @@ impl SourceToDefCtx<'_, '_> {
     }
     pub(super) fn bind_pat_to_def(
         &mut self,
-        src: InFile<ast::IdentPat>,
+        src: InFile<&ast::IdentPat>,
     ) -> Option<(DefWithBodyId, BindingId)> {
-        let container = self.find_pat_or_label_container(src.syntax())?;
+        let container = self.find_pat_or_label_container(src.syntax_ref())?;
         let (body, source_map) = self.db.body_with_source_map(container);
-        let src = src.map(ast::Pat::from);
+        let src = src.cloned().map(ast::Pat::from);
         let pat_id = source_map.node_pat(src.as_ref())?;
         // the pattern could resolve to a constant, verify that that is not the case
         if let crate::Pat::Bind { id, .. } = body[pat_id] {
@@ -253,25 +265,33 @@ impl SourceToDefCtx<'_, '_> {
     }
     pub(super) fn self_param_to_def(
         &mut self,
-        src: InFile<ast::SelfParam>,
+        src: InFile<&ast::SelfParam>,
     ) -> Option<(DefWithBodyId, BindingId)> {
-        let container = self.find_pat_or_label_container(src.syntax())?;
+        let container = self.find_pat_or_label_container(src.syntax_ref())?;
         let body = self.db.body(container);
         Some((container, body.self_param?))
     }
     pub(super) fn label_to_def(
         &mut self,
-        src: InFile<ast::Label>,
+        src: InFile<&ast::Label>,
     ) -> Option<(DefWithBodyId, LabelId)> {
-        let container = self.find_pat_or_label_container(src.syntax())?;
+        let container = self.find_pat_or_label_container(src.syntax_ref())?;
         let (_body, source_map) = self.db.body_with_source_map(container);
-        let label_id = source_map.node_label(src.as_ref())?;
+        let label_id = source_map.node_label(src)?;
         Some((container, label_id))
     }
 
-    pub(super) fn item_to_macro_call(&mut self, src: InFile<ast::Item>) -> Option<MacroCallId> {
-        let map = self.dyn_map(src.as_ref())?;
-        map[keys::ATTR_MACRO_CALL].get(&src.value).copied()
+    pub(super) fn item_to_macro_call(&mut self, src: InFile<&ast::Item>) -> Option<MacroCallId> {
+        let map = self.dyn_map(src)?;
+        map[keys::ATTR_MACRO_CALL].get(&AstPtr::new(src.value)).copied()
+    }
+
+    pub(super) fn macro_call_to_macro_call(
+        &mut self,
+        src: InFile<&ast::MacroCall>,
+    ) -> Option<MacroCallId> {
+        let map = self.dyn_map(src)?;
+        map[keys::MACRO_CALL].get(&AstPtr::new(src.value)).copied()
     }
 
     /// (AttrId, derive attribute call id, derive call ids)
@@ -282,7 +302,7 @@ impl SourceToDefCtx<'_, '_> {
     ) -> Option<(AttrId, MacroCallId, &[Option<MacroCallId>])> {
         let map = self.dyn_map(item)?;
         map[keys::DERIVE_MACRO_CALL]
-            .get(&src.value)
+            .get(&AstPtr::new(&src.value))
             .map(|&(attr_id, call_id, ref ids)| (attr_id, call_id, &**ids))
     }
 
@@ -292,10 +312,10 @@ impl SourceToDefCtx<'_, '_> {
 
     fn to_def<Ast: AstNode + 'static, ID: Copy + 'static>(
         &mut self,
-        src: InFile<Ast>,
+        src: InFile<&Ast>,
         key: Key<Ast, ID>,
     ) -> Option<ID> {
-        self.dyn_map(src.as_ref())?[key].get(&src.value).copied()
+        self.dyn_map(src)?[key].get(&AstPtr::new(src.value)).copied()
     }
 
     fn dyn_map<Ast: AstNode + 'static>(&mut self, src: InFile<&Ast>) -> Option<&DynMap> {
@@ -310,33 +330,42 @@ impl SourceToDefCtx<'_, '_> {
             .or_insert_with(|| container.child_by_source(db, file_id))
     }
 
-    pub(super) fn type_param_to_def(&mut self, src: InFile<ast::TypeParam>) -> Option<TypeParamId> {
-        let container: ChildContainer = self.find_generic_param_container(src.syntax())?.into();
+    pub(super) fn type_param_to_def(
+        &mut self,
+        src: InFile<&ast::TypeParam>,
+    ) -> Option<TypeParamId> {
+        let container: ChildContainer = self.find_generic_param_container(src.syntax_ref())?.into();
         let dyn_map = self.cache_for(container, src.file_id);
-        dyn_map[keys::TYPE_PARAM].get(&src.value).copied().map(TypeParamId::from_unchecked)
+        dyn_map[keys::TYPE_PARAM]
+            .get(&AstPtr::new(src.value))
+            .copied()
+            .map(TypeParamId::from_unchecked)
     }
 
     pub(super) fn lifetime_param_to_def(
         &mut self,
-        src: InFile<ast::LifetimeParam>,
+        src: InFile<&ast::LifetimeParam>,
     ) -> Option<LifetimeParamId> {
-        let container: ChildContainer = self.find_generic_param_container(src.syntax())?.into();
+        let container: ChildContainer = self.find_generic_param_container(src.syntax_ref())?.into();
         let dyn_map = self.cache_for(container, src.file_id);
-        dyn_map[keys::LIFETIME_PARAM].get(&src.value).copied()
+        dyn_map[keys::LIFETIME_PARAM].get(&AstPtr::new(src.value)).copied()
     }
 
     pub(super) fn const_param_to_def(
         &mut self,
-        src: InFile<ast::ConstParam>,
+        src: InFile<&ast::ConstParam>,
     ) -> Option<ConstParamId> {
-        let container: ChildContainer = self.find_generic_param_container(src.syntax())?.into();
+        let container: ChildContainer = self.find_generic_param_container(src.syntax_ref())?.into();
         let dyn_map = self.cache_for(container, src.file_id);
-        dyn_map[keys::CONST_PARAM].get(&src.value).copied().map(ConstParamId::from_unchecked)
+        dyn_map[keys::CONST_PARAM]
+            .get(&AstPtr::new(src.value))
+            .copied()
+            .map(ConstParamId::from_unchecked)
     }
 
     pub(super) fn generic_param_to_def(
         &mut self,
-        InFile { file_id, value }: InFile<ast::GenericParam>,
+        InFile { file_id, value }: InFile<&ast::GenericParam>,
     ) -> Option<GenericParamId> {
         match value {
             ast::GenericParam::ConstParam(it) => {
@@ -351,34 +380,111 @@ impl SourceToDefCtx<'_, '_> {
         }
     }
 
-    pub(super) fn macro_to_def(&mut self, src: InFile<ast::Macro>) -> Option<MacroId> {
-        self.dyn_map(src.as_ref()).and_then(|it| match &src.value {
+    pub(super) fn macro_to_def(&mut self, src: InFile<&ast::Macro>) -> Option<MacroId> {
+        self.dyn_map(src).and_then(|it| match src.value {
             ast::Macro::MacroRules(value) => {
-                it[keys::MACRO_RULES].get(value).copied().map(MacroId::from)
+                it[keys::MACRO_RULES].get(&AstPtr::new(value)).copied().map(MacroId::from)
+            }
+            ast::Macro::MacroDef(value) => {
+                it[keys::MACRO2].get(&AstPtr::new(value)).copied().map(MacroId::from)
             }
-            ast::Macro::MacroDef(value) => it[keys::MACRO2].get(value).copied().map(MacroId::from),
         })
     }
 
-    pub(super) fn proc_macro_to_def(&mut self, src: InFile<ast::Fn>) -> Option<MacroId> {
-        self.dyn_map(src.as_ref())
-            .and_then(|it| it[keys::PROC_MACRO].get(&src.value).copied().map(MacroId::from))
+    pub(super) fn proc_macro_to_def(&mut self, src: InFile<&ast::Fn>) -> Option<MacroId> {
+        self.dyn_map(src).and_then(|it| {
+            it[keys::PROC_MACRO].get(&AstPtr::new(src.value)).copied().map(MacroId::from)
+        })
     }
 
     pub(super) fn find_container(&mut self, src: InFile<&SyntaxNode>) -> Option<ChildContainer> {
-        for container in src.ancestors_with_macros(self.db.upcast()) {
-            if let Some(res) = self.container_to_def(container) {
-                return Some(res);
-            }
+        let def =
+            self.ancestors_with_macros(src, |this, container| this.container_to_def(container));
+        if let Some(def) = def {
+            return Some(def);
         }
 
         let def = self.file_to_def(src.file_id.original_file(self.db.upcast())).first().copied()?;
         Some(def.into())
     }
 
+    /// Skips the attributed item that caused the macro invocation we are climbing up
+    fn ancestors_with_macros<T>(
+        &mut self,
+        node: InFile<&SyntaxNode>,
+        mut cb: impl FnMut(&mut Self, InFile<SyntaxNode>) -> Option<T>,
+    ) -> Option<T> {
+        use hir_expand::MacroFileIdExt;
+        let parent = |this: &mut Self, node: InFile<&SyntaxNode>| match node.value.parent() {
+            Some(parent) => Some(node.with_value(parent)),
+            None => {
+                let macro_file = node.file_id.macro_file()?;
+
+                let expansion_info = this
+                    .expansion_info_cache
+                    .entry(macro_file)
+                    .or_insert_with(|| macro_file.expansion_info(this.db.upcast()));
+
+                expansion_info.call_node().map(|node| node?.parent()).transpose()
+            }
+        };
+        let mut node = node.cloned();
+        while let Some(parent) = parent(self, node.as_ref()) {
+            if let Some(res) = cb(self, parent.clone()) {
+                return Some(res);
+            }
+            node = parent;
+        }
+        None
+    }
+
+    fn find_generic_param_container(&mut self, src: InFile<&SyntaxNode>) -> Option<GenericDefId> {
+        self.ancestors_with_macros(src, |this, InFile { file_id, value }| {
+            let item = ast::Item::cast(value)?;
+            match &item {
+                ast::Item::Fn(it) => this.fn_to_def(InFile::new(file_id, it)).map(Into::into),
+                ast::Item::Struct(it) => {
+                    this.struct_to_def(InFile::new(file_id, it)).map(Into::into)
+                }
+                ast::Item::Enum(it) => this.enum_to_def(InFile::new(file_id, it)).map(Into::into),
+                ast::Item::Trait(it) => this.trait_to_def(InFile::new(file_id, it)).map(Into::into),
+                ast::Item::TraitAlias(it) => {
+                    this.trait_alias_to_def(InFile::new(file_id, it)).map(Into::into)
+                }
+                ast::Item::TypeAlias(it) => {
+                    this.type_alias_to_def(InFile::new(file_id, it)).map(Into::into)
+                }
+                ast::Item::Impl(it) => this.impl_to_def(InFile::new(file_id, it)).map(Into::into),
+                _ => None,
+            }
+        })
+    }
+
+    fn find_pat_or_label_container(&mut self, src: InFile<&SyntaxNode>) -> Option<DefWithBodyId> {
+        self.ancestors_with_macros(src, |this, InFile { file_id, value }| {
+            let item = match ast::Item::cast(value.clone()) {
+                Some(it) => it,
+                None => {
+                    let variant = ast::Variant::cast(value.clone())?;
+                    return this
+                        .enum_variant_to_def(InFile::new(file_id, &variant))
+                        .map(Into::into);
+                }
+            };
+            match &item {
+                ast::Item::Fn(it) => this.fn_to_def(InFile::new(file_id, it)).map(Into::into),
+                ast::Item::Const(it) => this.const_to_def(InFile::new(file_id, it)).map(Into::into),
+                ast::Item::Static(it) => {
+                    this.static_to_def(InFile::new(file_id, it)).map(Into::into)
+                }
+                _ => None,
+            }
+        })
+    }
+
     fn container_to_def(&mut self, container: InFile<SyntaxNode>) -> Option<ChildContainer> {
         let cont = if let Some(item) = ast::Item::cast(container.value.clone()) {
-            match item {
+            match &item {
                 ast::Item::Module(it) => self.module_to_def(container.with_value(it))?.into(),
                 ast::Item::Trait(it) => self.trait_to_def(container.with_value(it))?.into(),
                 ast::Item::TraitAlias(it) => {
@@ -413,63 +519,11 @@ impl SourceToDefCtx<'_, '_> {
             }
         } else {
             let it = ast::Variant::cast(container.value)?;
-            let def = self.enum_variant_to_def(InFile::new(container.file_id, it))?;
+            let def = self.enum_variant_to_def(InFile::new(container.file_id, &it))?;
             DefWithBodyId::from(def).into()
         };
         Some(cont)
     }
-
-    fn find_generic_param_container(&mut self, src: InFile<&SyntaxNode>) -> Option<GenericDefId> {
-        let ancestors = src.ancestors_with_macros(self.db.upcast());
-        for InFile { file_id, value } in ancestors {
-            let item = match ast::Item::cast(value) {
-                Some(it) => it,
-                None => continue,
-            };
-            let res: GenericDefId = match item {
-                ast::Item::Fn(it) => self.fn_to_def(InFile::new(file_id, it))?.into(),
-                ast::Item::Struct(it) => self.struct_to_def(InFile::new(file_id, it))?.into(),
-                ast::Item::Union(it) => self.union_to_def(InFile::new(file_id, it))?.into(),
-                ast::Item::Enum(it) => self.enum_to_def(InFile::new(file_id, it))?.into(),
-                ast::Item::Trait(it) => self.trait_to_def(InFile::new(file_id, it))?.into(),
-                ast::Item::TraitAlias(it) => {
-                    self.trait_alias_to_def(InFile::new(file_id, it))?.into()
-                }
-                ast::Item::TypeAlias(it) => {
-                    self.type_alias_to_def(InFile::new(file_id, it))?.into()
-                }
-                ast::Item::Impl(it) => self.impl_to_def(InFile::new(file_id, it))?.into(),
-                _ => continue,
-            };
-            return Some(res);
-        }
-        None
-    }
-
-    fn find_pat_or_label_container(&mut self, src: InFile<&SyntaxNode>) -> Option<DefWithBodyId> {
-        let ancestors = src.ancestors_with_macros(self.db.upcast());
-        for InFile { file_id, value } in ancestors {
-            let item = match ast::Item::cast(value.clone()) {
-                Some(it) => it,
-                None => {
-                    if let Some(variant) = ast::Variant::cast(value.clone()) {
-                        return self
-                            .enum_variant_to_def(InFile::new(file_id, variant))
-                            .map(Into::into);
-                    }
-                    continue;
-                }
-            };
-            let res: DefWithBodyId = match item {
-                ast::Item::Const(it) => self.const_to_def(InFile::new(file_id, it))?.into(),
-                ast::Item::Static(it) => self.static_to_def(InFile::new(file_id, it))?.into(),
-                ast::Item::Fn(it) => self.fn_to_def(InFile::new(file_id, it))?.into(),
-                _ => continue,
-            };
-            return Some(res);
-        }
-        None
-    }
 }
 
 #[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
@@ -501,6 +555,7 @@ impl_from! {
 
 impl ChildContainer {
     fn child_by_source(self, db: &dyn HirDatabase, file_id: HirFileId) -> DynMap {
+        let _p = tracing::span!(tracing::Level::INFO, "ChildContainer::child_by_source").entered();
         let db = db.upcast();
         match self {
             ChildContainer::DefWithBodyId(it) => it.child_by_source(db, file_id),