4 files changed, 84 insertions, 63 deletions

diff --git a/crates/hir_expand/src/builtin_macro.rs b/crates/hir_expand/src/builtin_macro.rs
index 0142a06ed..af9802144 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
+        }
+    }
+}