Precompute macro fragment kind

9 files changed, 113 insertions, 74 deletions

diff --git a/crates/hir_def/src/item_tree.rs b/crates/hir_def/src/item_tree.rs
index 8d13c7e04..cad8a7479 100644
--- a/crates/hir_def/src/item_tree.rs
+++ b/crates/hir_def/src/item_tree.rs
@@ -18,7 +18,7 @@ use hir_expand::{
     ast_id_map::FileAstId,
     hygiene::Hygiene,
     name::{name, AsName, Name},
-    HirFileId, InFile,
+    FragmentKind, HirFileId, InFile,
 };
 use la_arena::{Arena, Idx, RawIdx};
 use profile::Count;
@@ -656,6 +656,7 @@ pub struct MacroCall {
     /// Path to the called macro.
     pub path: Interned<ModPath>,
     pub ast_id: FileAstId<ast::MacroCall>,
+    pub fragment: FragmentKind,
 }
 
 #[derive(Debug, Clone, Eq, PartialEq)]
diff --git a/crates/hir_def/src/item_tree/lower.rs b/crates/hir_def/src/item_tree/lower.rs
index 5743b3386..fe348091d 100644
--- a/crates/hir_def/src/item_tree/lower.rs
+++ b/crates/hir_def/src/item_tree/lower.rs
@@ -624,7 +624,8 @@ impl<'a> Ctx<'a> {
     fn lower_macro_call(&mut self, m: &ast::MacroCall) -> Option<FileItemTreeId<MacroCall>> {
         let path = Interned::new(ModPath::from_src(self.db, m.path()?, &self.hygiene)?);
         let ast_id = self.source_ast_id_map.ast_id(m);
-        let res = MacroCall { path, ast_id };
+        let fragment = hir_expand::to_fragment_kind(m);
+        let res = MacroCall { path, ast_id, fragment };
         Some(id(self.data().macro_calls.alloc(res)))
     }
 
diff --git a/crates/hir_def/src/lib.rs b/crates/hir_def/src/lib.rs
index da46f16f7..e96ca953f 100644
--- a/crates/hir_def/src/lib.rs
+++ b/crates/hir_def/src/lib.rs
@@ -62,7 +62,8 @@ use hir_expand::{
     ast_id_map::FileAstId,
     eager::{expand_eager_macro, ErrorEmitted, ErrorSink},
     hygiene::Hygiene,
-    AstId, AttrId, HirFileId, InFile, MacroCallId, MacroCallKind, MacroDefId, MacroDefKind,
+    AstId, AttrId, FragmentKind, HirFileId, InFile, MacroCallId, MacroCallKind, MacroDefId,
+    MacroDefKind,
 };
 use la_arena::Idx;
 use nameres::DefMap;
@@ -652,6 +653,7 @@ impl AsMacroCall for InFile<&ast::MacroCall> {
         resolver: impl Fn(path::ModPath) -> Option<MacroDefId>,
         mut error_sink: &mut dyn FnMut(mbe::ExpandError),
     ) -> Result<Result<MacroCallId, ErrorEmitted>, UnresolvedMacro> {
+        let fragment = hir_expand::to_fragment_kind(self.value);
         let ast_id = AstId::new(self.file_id, db.ast_id_map(self.file_id).ast_id(self.value));
         let h = Hygiene::new(db.upcast(), self.file_id);
         let path = self.value.path().and_then(|path| path::ModPath::from_src(db, path, &h));
@@ -667,6 +669,7 @@ impl AsMacroCall for InFile<&ast::MacroCall> {
 
         macro_call_as_call_id(
             &AstIdWithPath::new(ast_id.file_id, ast_id.value, path),
+            fragment,
             db,
             krate,
             resolver,
@@ -695,6 +698,7 @@ pub struct UnresolvedMacro {
 
 fn macro_call_as_call_id(
     call: &AstIdWithPath<ast::MacroCall>,
+    fragment: FragmentKind,
     db: &dyn db::DefDatabase,
     krate: CrateId,
     resolver: impl Fn(path::ModPath) -> Option<MacroDefId>,
@@ -718,7 +722,11 @@ fn macro_call_as_call_id(
         .map(MacroCallId::from)
     } else {
         Ok(def
-            .as_lazy_macro(db.upcast(), krate, MacroCallKind::FnLike { ast_id: call.ast_id })
+            .as_lazy_macro(
+                db.upcast(),
+                krate,
+                MacroCallKind::FnLike { ast_id: call.ast_id, fragment },
+            )
             .into())
     };
     Ok(res)
diff --git a/crates/hir_def/src/nameres.rs b/crates/hir_def/src/nameres.rs
index 1bc72ec1f..249af6fc8 100644
--- a/crates/hir_def/src/nameres.rs
+++ b/crates/hir_def/src/nameres.rs
@@ -629,7 +629,7 @@ mod diagnostics {
                 DiagnosticKind::UnresolvedProcMacro { ast } => {
                     let mut precise_location = None;
                     let (file, ast, name) = match ast {
-                        MacroCallKind::FnLike { ast_id } => {
+                        MacroCallKind::FnLike { ast_id, .. } => {
                             let node = ast_id.to_node(db.upcast());
                             (ast_id.file_id, SyntaxNodePtr::from(AstPtr::new(&node)), None)
                         }
diff --git a/crates/hir_def/src/nameres/collector.rs b/crates/hir_def/src/nameres/collector.rs
index 05ceb1efb..e89136ed1 100644
--- a/crates/hir_def/src/nameres/collector.rs
+++ b/crates/hir_def/src/nameres/collector.rs
@@ -13,7 +13,7 @@ use hir_expand::{
     builtin_macro::find_builtin_macro,
     name::{AsName, Name},
     proc_macro::ProcMacroExpander,
-    AttrId, HirFileId, MacroCallId, MacroCallKind, MacroDefId, MacroDefKind,
+    AttrId, FragmentKind, HirFileId, MacroCallId, MacroCallKind, MacroDefId, MacroDefKind,
 };
 use hir_expand::{InFile, MacroCallLoc};
 use rustc_hash::{FxHashMap, FxHashSet};
@@ -215,7 +215,7 @@ struct MacroDirective {
 
 #[derive(Clone, Debug, Eq, PartialEq)]
 enum MacroDirectiveKind {
-    FnLike { ast_id: AstIdWithPath<ast::MacroCall> },
+    FnLike { ast_id: AstIdWithPath<ast::MacroCall>, fragment: FragmentKind },
     Derive { ast_id: AstIdWithPath<ast::Item>, derive_attr: AttrId },
 }
 
@@ -807,9 +807,10 @@ impl DefCollector<'_> {
         let mut res = ReachedFixedPoint::Yes;
         macros.retain(|directive| {
             match &directive.kind {
-                MacroDirectiveKind::FnLike { ast_id } => {
+                MacroDirectiveKind::FnLike { ast_id, fragment } => {
                     match macro_call_as_call_id(
                         ast_id,
+                        *fragment,
                         self.db,
                         self.def_map.krate,
                         |path| {
@@ -926,8 +927,9 @@ impl DefCollector<'_> {
 
         for directive in &self.unexpanded_macros {
             match &directive.kind {
-                MacroDirectiveKind::FnLike { ast_id, .. } => match macro_call_as_call_id(
+                MacroDirectiveKind::FnLike { ast_id, fragment } => match macro_call_as_call_id(
                     ast_id,
+                    *fragment,
                     self.db,
                     self.def_map.krate,
                     |path| {
@@ -1496,6 +1498,7 @@ impl ModCollector<'_, '_> {
         let mut error = None;
         match macro_call_as_call_id(
             &ast_id,
+            mac.fragment,
             self.def_collector.db,
             self.def_collector.def_map.krate,
             |path| {
@@ -1524,9 +1527,14 @@ impl ModCollector<'_, '_> {
             }
             Ok(Err(_)) => {
                 // Built-in macro failed eager expansion.
+
+                // FIXME: don't parse the file here
+                let fragment = hir_expand::to_fragment_kind(
+                    &ast_id.ast_id.to_node(self.def_collector.db.upcast()),
+                );
                 self.def_collector.def_map.diagnostics.push(DefDiagnostic::macro_error(
                     self.module_id,
-                    MacroCallKind::FnLike { ast_id: ast_id.ast_id },
+                    MacroCallKind::FnLike { ast_id: ast_id.ast_id, fragment },
                     error.unwrap().to_string(),
                 ));
                 return;
@@ -1543,7 +1551,7 @@ impl ModCollector<'_, '_> {
         self.def_collector.unexpanded_macros.push(MacroDirective {
             module_id: self.module_id,
             depth: self.macro_depth + 1,
-            kind: MacroDirectiveKind::FnLike { ast_id },
+            kind: MacroDirectiveKind::FnLike { ast_id, fragment: mac.fragment },
         });
     }
 
diff --git a/crates/hir_expand/src/builtin_macro.rs b/crates/hir_expand/src/builtin_macro.rs
index 179de61f9..0dd725ec3 100644
--- a/crates/hir_expand/src/builtin_macro.rs
+++ b/crates/hir_expand/src/builtin_macro.rs
@@ -578,6 +578,7 @@ mod tests {
                     krate,
                     kind: MacroCallKind::FnLike {
                         ast_id: AstId::new(file_id.into(), ast_id_map.ast_id(&macro_call)),
+                        fragment: FragmentKind::Expr,
                     },
                 };
 
diff --git a/crates/hir_expand/src/db.rs b/crates/hir_expand/src/db.rs
index d61f4b31a..6647e57e7 100644
--- a/crates/hir_expand/src/db.rs
+++ b/crates/hir_expand/src/db.rs
@@ -8,9 +8,7 @@ use parser::FragmentKind;
 use syntax::{
     algo::diff,
     ast::{self, NameOwner},
-    AstNode, GreenNode, Parse,
-    SyntaxKind::*,
-    SyntaxNode, SyntaxToken,
+    AstNode, GreenNode, Parse, SyntaxNode, SyntaxToken,
 };
 
 use crate::{
@@ -160,7 +158,7 @@ pub fn expand_hypothetical(
 
     let hypothetical_expansion = macro_def.expand(db, lazy_id, &tt);
 
-    let fragment_kind = to_fragment_kind(db, actual_macro_call);
+    let fragment_kind = macro_fragment_kind(db, actual_macro_call);
 
     let (node, tmap_2) =
         mbe::token_tree_to_syntax_node(&hypothetical_expansion.value, fragment_kind).ok()?;
@@ -226,7 +224,7 @@ fn parse_macro_expansion(
         None => return ExpandResult { value: None, err: result.err },
     };
 
-    let fragment_kind = to_fragment_kind(db, macro_file.macro_call_id);
+    let fragment_kind = macro_fragment_kind(db, macro_file.macro_call_id);
 
     log::debug!("expanded = {}", tt.as_debug_string());
     log::debug!("kind = {:?}", fragment_kind);
@@ -427,62 +425,15 @@ fn hygiene_frame(db: &dyn AstDatabase, file_id: HirFileId) -> Arc<HygieneFrame>
     Arc::new(HygieneFrame::new(db, file_id))
 }
 
-/// Given a `MacroCallId`, return what `FragmentKind` it belongs to.
-/// FIXME: Not completed
-fn to_fragment_kind(db: &dyn AstDatabase, id: MacroCallId) -> FragmentKind {
-    let lazy_id = match id {
-        MacroCallId::LazyMacro(id) => id,
-        MacroCallId::EagerMacro(id) => {
-            return db.lookup_intern_eager_expansion(id).fragment;
-        }
-    };
-    let syn = db.lookup_intern_macro(lazy_id).kind.node(db).value;
-
-    let parent = match syn.parent() {
-        Some(it) => it,
-        None => return FragmentKind::Statements,
-    };
-
-    match parent.kind() {
-        MACRO_ITEMS | SOURCE_FILE => FragmentKind::Items,
-        MACRO_STMTS => FragmentKind::Statements,
-        MACRO_PAT => FragmentKind::Pattern,
-        MACRO_TYPE => FragmentKind::Type,
-        ITEM_LIST => FragmentKind::Items,
-        LET_STMT => {
-            // FIXME: Handle LHS Pattern
-            FragmentKind::Expr
+fn macro_fragment_kind(db: &dyn AstDatabase, id: MacroCallId) -> FragmentKind {
+    match id {
+        MacroCallId::LazyMacro(id) => {
+            let loc: MacroCallLoc = db.lookup_intern_macro(id);
+            loc.kind.fragment_kind()
         }
-        EXPR_STMT => FragmentKind::Statements,
-        BLOCK_EXPR => FragmentKind::Statements,
-        ARG_LIST => FragmentKind::Expr,
-        TRY_EXPR => FragmentKind::Expr,
-        TUPLE_EXPR => FragmentKind::Expr,
-        PAREN_EXPR => FragmentKind::Expr,
-        ARRAY_EXPR => FragmentKind::Expr,
-        FOR_EXPR => FragmentKind::Expr,
-        PATH_EXPR => FragmentKind::Expr,
-        CLOSURE_EXPR => FragmentKind::Expr,
-        CONDITION => FragmentKind::Expr,
-        BREAK_EXPR => FragmentKind::Expr,
-        RETURN_EXPR => FragmentKind::Expr,
-        MATCH_EXPR => FragmentKind::Expr,
-        MATCH_ARM => FragmentKind::Expr,
-        MATCH_GUARD => FragmentKind::Expr,
-        RECORD_EXPR_FIELD => FragmentKind::Expr,
-        CALL_EXPR => FragmentKind::Expr,
-        INDEX_EXPR => FragmentKind::Expr,
-        METHOD_CALL_EXPR => FragmentKind::Expr,
-        FIELD_EXPR => FragmentKind::Expr,
-        AWAIT_EXPR => FragmentKind::Expr,
-        CAST_EXPR => FragmentKind::Expr,
-        REF_EXPR => FragmentKind::Expr,
-        PREFIX_EXPR => FragmentKind::Expr,
-        RANGE_EXPR => FragmentKind::Expr,
-        BIN_EXPR => FragmentKind::Expr,
-        _ => {
-            // Unknown , Just guess it is `Items`
-            FragmentKind::Items
+        MacroCallId::EagerMacro(id) => {
+            let loc: EagerCallLoc = db.lookup_intern_eager_expansion(id);
+            loc.fragment
         }
     }
 }
diff --git a/crates/hir_expand/src/eager.rs b/crates/hir_expand/src/eager.rs
index f12132f84..85491fe8b 100644
--- a/crates/hir_expand/src/eager.rs
+++ b/crates/hir_expand/src/eager.rs
@@ -175,8 +175,13 @@ fn lazy_expand(
 ) -> ExpandResult<Option<InFile<SyntaxNode>>> {
     let ast_id = db.ast_id_map(macro_call.file_id).ast_id(&macro_call.value);
 
+    let fragment = crate::to_fragment_kind(&macro_call.value);
     let id: MacroCallId = def
-        .as_lazy_macro(db, krate, MacroCallKind::FnLike { ast_id: macro_call.with_value(ast_id) })
+        .as_lazy_macro(
+            db,
+            krate,
+            MacroCallKind::FnLike { ast_id: macro_call.with_value(ast_id), fragment },
+        )
         .into();
 
     let err = db.macro_expand_error(id);
diff --git a/crates/hir_expand/src/lib.rs b/crates/hir_expand/src/lib.rs
index 0402640de..80ab3aeee 100644
--- a/crates/hir_expand/src/lib.rs
+++ b/crates/hir_expand/src/lib.rs
@@ -16,7 +16,9 @@ pub mod quote;
 pub mod eager;
 
 use either::Either;
+
 pub use mbe::{ExpandError, ExpandResult};
+pub use parser::FragmentKind;
 
 use std::hash::Hash;
 use std::sync::Arc;
@@ -290,7 +292,7 @@ pub struct MacroCallLoc {
 
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub enum MacroCallKind {
-    FnLike { ast_id: AstId<ast::MacroCall> },
+    FnLike { ast_id: AstId<ast::MacroCall>, fragment: FragmentKind },
     Derive { ast_id: AstId<ast::Item>, derive_name: String, derive_attr: AttrId },
 }
 
@@ -324,6 +326,13 @@ impl MacroCallKind {
             MacroCallKind::Derive { ast_id, .. } => Some(ast_id.to_node(db).syntax().clone()),
         }
     }
+
+    fn fragment_kind(&self) -> FragmentKind {
+        match self {
+            MacroCallKind::FnLike { fragment, .. } => *fragment,
+            MacroCallKind::Derive { .. } => FragmentKind::Items,
+        }
+    }
 }
 
 impl MacroCallId {
@@ -357,7 +366,6 @@ pub struct ExpansionInfo {
 }
 
 pub use mbe::Origin;
-use parser::FragmentKind;
 
 impl ExpansionInfo {
     pub fn call_node(&self) -> Option<InFile<SyntaxNode>> {
@@ -562,3 +570,59 @@ impl<N: AstNode> InFile<N> {
         self.with_value(self.value.syntax())
     }
 }
+
+/// Given a `MacroCallId`, return what `FragmentKind` it belongs to.
+/// FIXME: Not completed
+pub fn to_fragment_kind(call: &ast::MacroCall) -> FragmentKind {
+    use syntax::SyntaxKind::*;
+
+    let syn = call.syntax();
+
+    let parent = match syn.parent() {
+        Some(it) => it,
+        None => return FragmentKind::Statements,
+    };
+
+    match parent.kind() {
+        MACRO_ITEMS | SOURCE_FILE => FragmentKind::Items,
+        MACRO_STMTS => FragmentKind::Statements,
+        MACRO_PAT => FragmentKind::Pattern,
+        MACRO_TYPE => FragmentKind::Type,
+        ITEM_LIST => FragmentKind::Items,
+        LET_STMT => {
+            // FIXME: Handle LHS Pattern
+            FragmentKind::Expr
+        }
+        EXPR_STMT => FragmentKind::Statements,
+        BLOCK_EXPR => FragmentKind::Statements,
+        ARG_LIST => FragmentKind::Expr,
+        TRY_EXPR => FragmentKind::Expr,
+        TUPLE_EXPR => FragmentKind::Expr,
+        PAREN_EXPR => FragmentKind::Expr,
+        ARRAY_EXPR => FragmentKind::Expr,
+        FOR_EXPR => FragmentKind::Expr,
+        PATH_EXPR => FragmentKind::Expr,
+        CLOSURE_EXPR => FragmentKind::Expr,
+        CONDITION => FragmentKind::Expr,
+        BREAK_EXPR => FragmentKind::Expr,
+        RETURN_EXPR => FragmentKind::Expr,
+        MATCH_EXPR => FragmentKind::Expr,
+        MATCH_ARM => FragmentKind::Expr,
+        MATCH_GUARD => FragmentKind::Expr,
+        RECORD_EXPR_FIELD => FragmentKind::Expr,
+        CALL_EXPR => FragmentKind::Expr,
+        INDEX_EXPR => FragmentKind::Expr,
+        METHOD_CALL_EXPR => FragmentKind::Expr,
+        FIELD_EXPR => FragmentKind::Expr,
+        AWAIT_EXPR => FragmentKind::Expr,
+        CAST_EXPR => FragmentKind::Expr,
+        REF_EXPR => FragmentKind::Expr,
+        PREFIX_EXPR => FragmentKind::Expr,
+        RANGE_EXPR => FragmentKind::Expr,
+        BIN_EXPR => FragmentKind::Expr,
+        _ => {
+            // Unknown , Just guess it is `Items`
+            FragmentKind::Items
+        }
+    }
+}