4 files changed, 93 insertions, 13 deletions

diff --git a/crates/ra_hir/src/source_binder.rs b/crates/ra_hir/src/source_binder.rs
index c42ceabdf..797f90d50 100644
--- a/crates/ra_hir/src/source_binder.rs
+++ b/crates/ra_hir/src/source_binder.rs
@@ -131,6 +131,7 @@ pub struct ReferenceDescriptor {
 }
 
 pub struct Expansion {
+    macro_file_kind: MacroFileKind,
     macro_call_id: MacroCallId,
 }
 
@@ -145,7 +146,7 @@ impl Expansion {
     }
 
     pub fn file_id(&self) -> HirFileId {
-        self.macro_call_id.as_file(MacroFileKind::Items)
+        self.macro_call_id.as_file(self.macro_file_kind)
     }
 }
 
@@ -439,7 +440,10 @@ impl SourceAnalyzer {
             db.ast_id_map(macro_call.file_id).ast_id(macro_call.value),
         );
         let macro_call_loc = MacroCallLoc { def, ast_id };
-        Some(Expansion { macro_call_id: db.intern_macro(macro_call_loc) })
+        Some(Expansion {
+            macro_call_id: db.intern_macro(macro_call_loc),
+            macro_file_kind: to_macro_file_kind(macro_call.value),
+        })
     }
 
     #[cfg(test)]
@@ -538,3 +542,35 @@ fn adjust(
         })
         .map(|(_ptr, scope)| *scope)
 }
+
+/// Given a `ast::MacroCall`, return what `MacroKindFile` it belongs to.
+/// FIXME: Not completed  
+fn to_macro_file_kind(macro_call: &ast::MacroCall) -> MacroFileKind {
+    let syn = macro_call.syntax();
+    let parent = match syn.parent() {
+        Some(it) => it,
+        None => {
+            // FIXME:
+            // If it is root, which means the parent HirFile
+            // MacroKindFile must be non-items
+            // return expr now.
+            return MacroFileKind::Expr;
+        }
+    };
+
+    match parent.kind() {
+        MACRO_ITEMS | SOURCE_FILE => MacroFileKind::Items,
+        LET_STMT => {
+            // FIXME: Handle Pattern
+            MacroFileKind::Expr
+        }
+        EXPR_STMT => MacroFileKind::Statements,
+        BLOCK => MacroFileKind::Statements,
+        ARG_LIST => MacroFileKind::Expr,
+        TRY_EXPR => MacroFileKind::Expr,
+        _ => {
+            // Unknown , Just guess it is `Items`
+            MacroFileKind::Items
+        }
+    }
+}
diff --git a/crates/ra_hir_expand/src/db.rs b/crates/ra_hir_expand/src/db.rs
index 3c11c8a22..e1d93a8ef 100644
--- a/crates/ra_hir_expand/src/db.rs
+++ b/crates/ra_hir_expand/src/db.rs
@@ -151,6 +151,7 @@ pub(crate) fn parse_macro(
     let fragment_kind = match macro_file.macro_file_kind {
         MacroFileKind::Items => FragmentKind::Items,
         MacroFileKind::Expr => FragmentKind::Expr,
+        MacroFileKind::Statements => FragmentKind::Statements,
     };
     let (parse, rev_token_map) = mbe::token_tree_to_syntax_node(&tt, fragment_kind).ok()?;
     Some((parse, Arc::new(rev_token_map)))
diff --git a/crates/ra_hir_expand/src/lib.rs b/crates/ra_hir_expand/src/lib.rs
index 1389f64ce..126d12fbb 100644
--- a/crates/ra_hir_expand/src/lib.rs
+++ b/crates/ra_hir_expand/src/lib.rs
@@ -109,6 +109,7 @@ pub struct MacroFile {
 pub enum MacroFileKind {
     Items,
     Expr,
+    Statements,
 }
 
 /// `MacroCallId` identifies a particular macro invocation, like
diff --git a/crates/ra_ide_api/src/expand_macro.rs b/crates/ra_ide_api/src/expand_macro.rs
index 7f39262dc..7dbf33a16 100644
--- a/crates/ra_ide_api/src/expand_macro.rs
+++ b/crates/ra_ide_api/src/expand_macro.rs
@@ -84,24 +84,19 @@ fn insert_whitespaces(syn: SyntaxNode) -> String {
         };
 
         res += &match token.kind() {
-            k @ _
-                if (k.is_keyword() || k.is_literal() || k == IDENT)
-                    && is_next(|it| !it.is_punct(), true) =>
-            {
+            k @ _ if is_text(k) && is_next(|it| !it.is_punct(), true) => {
                 token.text().to_string() + " "
             }
             L_CURLY if is_next(|it| it != R_CURLY, true) => {
                 indent += 1;
-                format!(" {{\n{}", "  ".repeat(indent))
+                let leading_space = if is_last(|it| is_text(it), false) { " " } else { "" };
+                format!("{}{{\n{}", leading_space, "  ".repeat(indent))
             }
             R_CURLY if is_last(|it| it != L_CURLY, true) => {
                 indent = indent.checked_sub(1).unwrap_or(0);
-                format!("\n}}{}", "  ".repeat(indent))
-            }
-            R_CURLY => {
-                indent = indent.checked_sub(1).unwrap_or(0);
-                format!("}}\n{}", "  ".repeat(indent))
+                format!("\n{}}}", "  ".repeat(indent))
             }
+            R_CURLY => format!("}}\n{}", "  ".repeat(indent)),
             T![;] => format!(";\n{}", "  ".repeat(indent)),
             T![->] => " -> ".to_string(),
             T![=] => " = ".to_string(),
@@ -112,7 +107,11 @@ fn insert_whitespaces(syn: SyntaxNode) -> String {
         last = Some(token.kind());
     }
 
-    res
+    return res;
+
+    fn is_text(k: SyntaxKind) -> bool {
+        k.is_keyword() || k.is_literal() || k == IDENT
+    }
 }
 
 #[cfg(test)]
@@ -175,4 +174,47 @@ fn some_thing() -> u32 {
 }
 "###);
     }
+
+    #[test]
+    fn macro_expand_match_ast() {
+        let res = check_expand_macro(
+            r#"
+        //- /lib.rs
+        macro_rules! match_ast {
+            (match $node:ident { $($tt:tt)* }) => { match_ast!(match ($node) { $($tt)* }) };
+        
+            (match ($node:expr) {
+                $( ast::$ast:ident($it:ident) => $res:block, )*
+                _ => $catch_all:expr $(,)?
+            }) => {{
+                $( if let Some($it) = ast::$ast::cast($node.clone()) $res else )*
+                { $catch_all }
+            }};
+        }        
+
+        fn main() {
+            mat<|>ch_ast! {
+                match container {
+                    ast::TraitDef(it) => {},
+                    ast::ImplBlock(it) => {},
+                    _ => { continue },
+                }
+            }
+        }
+        "#,
+        );
+
+        assert_eq!(res.name, "match_ast");
+        assert_snapshot!(res.expansion, @r###"
+{
+  if let Some(it) = ast::TraitDef::cast(container.clone()){}
+  else if let Some(it) = ast::ImplBlock::cast(container.clone()){}
+  else {
+    {
+      continue
+    }
+  }
+}
+"###);
+    }
 }