4 files changed, 208 insertions, 111 deletions

diff --git a/crates/ra_ide/src/expand.rs b/crates/ra_ide/src/expand.rs
index b82259a3d..831438c09 100644
--- a/crates/ra_ide/src/expand.rs
+++ b/crates/ra_ide/src/expand.rs
@@ -79,6 +79,14 @@ pub(crate) fn descend_into_macros(
     let source_analyzer =
         hir::SourceAnalyzer::new(db, src.with_value(src.value.parent()).as_ref(), None);
 
+    descend_into_macros_with_analyzer(db, &source_analyzer, src)
+}
+
+pub(crate) fn descend_into_macros_with_analyzer(
+    db: &RootDatabase,
+    source_analyzer: &hir::SourceAnalyzer,
+    src: InFile<SyntaxToken>,
+) -> InFile<SyntaxToken> {
     successors(Some(src), |token| {
         let macro_call = token.value.ancestors().find_map(ast::MacroCall::cast)?;
         let tt = macro_call.token_tree()?;
diff --git a/crates/ra_ide/src/snapshots/highlighting.html b/crates/ra_ide/src/snapshots/highlighting.html
index 1d130544f..1cc55e78b 100644
--- a/crates/ra_ide/src/snapshots/highlighting.html
+++ b/crates/ra_ide/src/snapshots/highlighting.html
@@ -34,6 +34,16 @@ pre                 { color: #DCDCCC; background: #3F3F3F; font-size: 22px; padd
     <span class="function">foo</span>::&lt;<span class="type.builtin">i32</span>&gt;();
 }
 
+<span class="macro">macro_rules</span><span class="macro">!</span> def_fn {
+    ($($tt:tt)*) =&gt; {$($tt)*}
+}
+
+<span class="macro">def_fn</span><span class="macro">!</span>{
+    <span class="keyword">fn</span> <span class="function">bar</span>() -&gt; <span class="type.builtin">u32</span> {
+        <span class="literal.numeric">100</span>
+    }
+}
+
 <span class="comment">// comment</span>
 <span class="keyword">fn</span> <span class="function">main</span>() {
     <span class="macro">println</span><span class="macro">!</span>(<span class="string">"Hello, {}!"</span>, <span class="literal.numeric">92</span>);
diff --git a/crates/ra_ide/src/snapshots/rainbow_highlighting.html b/crates/ra_ide/src/snapshots/rainbow_highlighting.html
index d90ee8540..918fd4b97 100644
--- a/crates/ra_ide/src/snapshots/rainbow_highlighting.html
+++ b/crates/ra_ide/src/snapshots/rainbow_highlighting.html
@@ -24,14 +24,14 @@ pre                 { color: #DCDCCC; background: #3F3F3F; font-size: 22px; padd
 .keyword\.control  { color: #F0DFAF; font-weight: bold; }
 </style>
 <pre><code><span class="keyword">fn</span> <span class="function">main</span>() {
-    <span class="keyword">let</span> <span class="variable" data-binding-hash="8723171760279909834" style="color: hsl(307,91%,75%);">hello</span> = <span class="string">"hello"</span>;
-    <span class="keyword">let</span> <span class="variable" data-binding-hash="14702933417323009544" style="color: hsl(108,90%,49%);">x</span> = <span class="variable" data-binding-hash="8723171760279909834" style="color: hsl(307,91%,75%);">hello</span>.to_string();
-    <span class="keyword">let</span> <span class="variable" data-binding-hash="5443150872754369068" style="color: hsl(215,43%,43%);">y</span> = <span class="variable" data-binding-hash="8723171760279909834" style="color: hsl(307,91%,75%);">hello</span>.to_string();
+    <span class="keyword">let</span> <span class="variable" data-binding-hash="2217585909179791122" style="color: hsl(280,74%,48%);">hello</span> = <span class="string">"hello"</span>;
+    <span class="keyword">let</span> <span class="variable" data-binding-hash="4303609361109701698" style="color: hsl(242,75%,88%);">x</span> = <span class="variable" data-binding-hash="2217585909179791122" style="color: hsl(280,74%,48%);">hello</span>.to_string();
+    <span class="keyword">let</span> <span class="variable" data-binding-hash="13865792086344377029" style="color: hsl(340,64%,86%);">y</span> = <span class="variable" data-binding-hash="2217585909179791122" style="color: hsl(280,74%,48%);">hello</span>.to_string();
 
-    <span class="keyword">let</span> <span class="variable" data-binding-hash="17358108296605513516" style="color: hsl(331,46%,60%);">x</span> = <span class="string">"other color please!"</span>;
-    <span class="keyword">let</span> <span class="variable" data-binding-hash="2073121142529774969" style="color: hsl(320,43%,74%);">y</span> = <span class="variable" data-binding-hash="17358108296605513516" style="color: hsl(331,46%,60%);">x</span>.to_string();
+    <span class="keyword">let</span> <span class="variable" data-binding-hash="7011301204224269512" style="color: hsl(198,45%,40%);">x</span> = <span class="string">"other color please!"</span>;
+    <span class="keyword">let</span> <span class="variable" data-binding-hash="12461245066629867975" style="color: hsl(132,91%,68%);">y</span> = <span class="variable" data-binding-hash="7011301204224269512" style="color: hsl(198,45%,40%);">x</span>.to_string();
 }
 
 <span class="keyword">fn</span> <span class="function">bar</span>() {
-    <span class="keyword">let</span> <span class="keyword">mut</span> <span class="variable.mut" data-binding-hash="8723171760279909834" style="color: hsl(307,91%,75%);">hello</span> = <span class="string">"hello"</span>;
+    <span class="keyword">let</span> <span class="keyword">mut</span> <span class="variable.mut" data-binding-hash="2217585909179791122" style="color: hsl(280,74%,48%);">hello</span> = <span class="string">"hello"</span>;
 }</code></pre>
\ No newline at end of file
diff --git a/crates/ra_ide/src/syntax_highlighting.rs b/crates/ra_ide/src/syntax_highlighting.rs
index 0411977b9..530b984fc 100644
--- a/crates/ra_ide/src/syntax_highlighting.rs
+++ b/crates/ra_ide/src/syntax_highlighting.rs
@@ -1,14 +1,18 @@
 //! FIXME: write short doc here
 
-use rustc_hash::{FxHashMap, FxHashSet};
+use rustc_hash::FxHashMap;
 
-use hir::{InFile, Name, SourceBinder};
+use hir::{HirFileId, InFile, Name, SourceAnalyzer, SourceBinder};
 use ra_db::SourceDatabase;
 use ra_prof::profile;
-use ra_syntax::{ast, AstNode, Direction, SyntaxElement, SyntaxKind, SyntaxKind::*, TextRange, T};
+use ra_syntax::{
+    ast, AstNode, Direction, SyntaxElement, SyntaxKind, SyntaxKind::*, SyntaxToken, TextRange,
+    WalkEvent, T,
+};
 
 use crate::{
     db::RootDatabase,
+    expand::descend_into_macros_with_analyzer,
     references::{
         classify_name, classify_name_ref,
         NameKind::{self, *},
@@ -72,121 +76,186 @@ pub(crate) fn highlight(db: &RootDatabase, file_id: FileId) -> Vec<HighlightedRa
     let parse = db.parse(file_id);
     let root = parse.tree().syntax().clone();
 
-    fn calc_binding_hash(file_id: FileId, name: &Name, shadow_count: u32) -> u64 {
-        fn hash<T: std::hash::Hash + std::fmt::Debug>(x: T) -> u64 {
-            use std::{collections::hash_map::DefaultHasher, hash::Hasher};
+    let mut sb = SourceBinder::new(db);
+    let mut bindings_shadow_count: FxHashMap<Name, u32> = FxHashMap::default();
+    let mut res = Vec::new();
+    let analyzer = sb.analyze(InFile::new(file_id.into(), &root), None);
 
-            let mut hasher = DefaultHasher::new();
-            x.hash(&mut hasher);
-            hasher.finish()
+    let mut in_macro_call = None;
+
+    for event in root.preorder_with_tokens() {
+        match event {
+            WalkEvent::Enter(node) => match node.kind() {
+                MACRO_CALL => {
+                    in_macro_call = Some(node.clone());
+                    if let Some(range) = highlight_macro(InFile::new(file_id.into(), node)) {
+                        res.push(HighlightedRange { range, tag: tags::MACRO, binding_hash: None });
+                    }
+                }
+                _ if in_macro_call.is_some() => {
+                    if let Some(token) = node.as_token() {
+                        if let Some((tag, binding_hash)) = highlight_token_tree(
+                            db,
+                            &mut sb,
+                            &analyzer,
+                            &mut bindings_shadow_count,
+                            InFile::new(file_id.into(), token.clone()),
+                        ) {
+                            res.push(HighlightedRange {
+                                range: node.text_range(),
+                                tag,
+                                binding_hash,
+                            });
+                        }
+                    }
+                }
+                _ => {
+                    if let Some((tag, binding_hash)) = highlight_node(
+                        db,
+                        &mut sb,
+                        &mut bindings_shadow_count,
+                        InFile::new(file_id.into(), node.clone()),
+                    ) {
+                        res.push(HighlightedRange { range: node.text_range(), tag, binding_hash });
+                    }
+                }
+            },
+            WalkEvent::Leave(node) => {
+                if let Some(m) = in_macro_call.as_ref() {
+                    if *m == node {
+                        in_macro_call = None;
+                    }
+                }
+            }
         }
+    }
 
-        hash((file_id, name, shadow_count))
+    res
+}
+
+fn highlight_macro(node: InFile<SyntaxElement>) -> Option<TextRange> {
+    let macro_call = ast::MacroCall::cast(node.value.as_node()?.clone())?;
+    let path = macro_call.path()?;
+    let name_ref = path.segment()?.name_ref()?;
+
+    let range_start = name_ref.syntax().text_range().start();
+    let mut range_end = name_ref.syntax().text_range().end();
+    for sibling in path.syntax().siblings_with_tokens(Direction::Next) {
+        match sibling.kind() {
+            T![!] | IDENT => range_end = sibling.text_range().end(),
+            _ => (),
+        }
     }
 
-    let mut sb = SourceBinder::new(db);
+    Some(TextRange::from_to(range_start, range_end))
+}
 
-    // Visited nodes to handle highlighting priorities
-    // FIXME: retain only ranges here
-    let mut highlighted: FxHashSet<SyntaxElement> = FxHashSet::default();
-    let mut bindings_shadow_count: FxHashMap<Name, u32> = FxHashMap::default();
+fn highlight_token_tree(
+    db: &RootDatabase,
+    sb: &mut SourceBinder<RootDatabase>,
+    analyzer: &SourceAnalyzer,
+    bindings_shadow_count: &mut FxHashMap<Name, u32>,
+    token: InFile<SyntaxToken>,
+) -> Option<(&'static str, Option<u64>)> {
+    if token.value.parent().kind() != TOKEN_TREE {
+        return None;
+    }
+    let token = descend_into_macros_with_analyzer(db, analyzer, token);
+    let expanded = {
+        let parent = token.value.parent();
+        // We only care Name and Name_ref
+        match (token.value.kind(), parent.kind()) {
+            (IDENT, NAME) | (IDENT, NAME_REF) => token.with_value(parent.into()),
+            _ => token.map(|it| it.into()),
+        }
+    };
 
-    let mut res = Vec::new();
-    for node in root.descendants_with_tokens() {
-        if highlighted.contains(&node) {
-            continue;
+    highlight_node(db, sb, bindings_shadow_count, expanded)
+}
+
+fn highlight_node(
+    db: &RootDatabase,
+    sb: &mut SourceBinder<RootDatabase>,
+    bindings_shadow_count: &mut FxHashMap<Name, u32>,
+    node: InFile<SyntaxElement>,
+) -> Option<(&'static str, Option<u64>)> {
+    let mut binding_hash = None;
+    let tag = match node.value.kind() {
+        FN_DEF => {
+            bindings_shadow_count.clear();
+            return None;
         }
-        let mut binding_hash = None;
-        let tag = match node.kind() {
-            FN_DEF => {
-                bindings_shadow_count.clear();
-                continue;
-            }
-            COMMENT => tags::LITERAL_COMMENT,
-            STRING | RAW_STRING | RAW_BYTE_STRING | BYTE_STRING => tags::LITERAL_STRING,
-            ATTR => tags::LITERAL_ATTRIBUTE,
-            // Special-case field init shorthand
-            NAME_REF if node.parent().and_then(ast::RecordField::cast).is_some() => tags::FIELD,
-            NAME_REF if node.ancestors().any(|it| it.kind() == ATTR) => continue,
-            NAME_REF => {
-                let name_ref = node.as_node().cloned().and_then(ast::NameRef::cast).unwrap();
-                let name_kind = classify_name_ref(&mut sb, InFile::new(file_id.into(), &name_ref))
-                    .map(|d| d.kind);
-                match name_kind {
-                    Some(name_kind) => {
-                        if let Local(local) = &name_kind {
-                            if let Some(name) = local.name(db) {
-                                let shadow_count =
-                                    bindings_shadow_count.entry(name.clone()).or_default();
-                                binding_hash =
-                                    Some(calc_binding_hash(file_id, &name, *shadow_count))
-                            }
-                        };
-
-                        highlight_name(db, name_kind)
-                    }
-                    _ => continue,
-                }
-            }
-            NAME => {
-                let name = node.as_node().cloned().and_then(ast::Name::cast).unwrap();
-                let name_kind =
-                    classify_name(&mut sb, InFile::new(file_id.into(), &name)).map(|d| d.kind);
-
-                if let Some(Local(local)) = &name_kind {
-                    if let Some(name) = local.name(db) {
-                        let shadow_count = bindings_shadow_count.entry(name.clone()).or_default();
-                        *shadow_count += 1;
-                        binding_hash = Some(calc_binding_hash(file_id, &name, *shadow_count))
-                    }
-                };
-
-                match name_kind {
-                    Some(name_kind) => highlight_name(db, name_kind),
-                    None => name.syntax().parent().map_or(tags::FUNCTION, |x| match x.kind() {
-                        STRUCT_DEF | ENUM_DEF | TRAIT_DEF | TYPE_ALIAS_DEF => tags::TYPE,
-                        TYPE_PARAM => tags::TYPE_PARAM,
-                        RECORD_FIELD_DEF => tags::FIELD,
-                        _ => tags::FUNCTION,
-                    }),
+        COMMENT => tags::LITERAL_COMMENT,
+        STRING | RAW_STRING | RAW_BYTE_STRING | BYTE_STRING => tags::LITERAL_STRING,
+        ATTR => tags::LITERAL_ATTRIBUTE,
+        // Special-case field init shorthand
+        NAME_REF if node.value.parent().and_then(ast::RecordField::cast).is_some() => tags::FIELD,
+        NAME_REF if node.value.ancestors().any(|it| it.kind() == ATTR) => return None,
+        NAME_REF => {
+            let name_ref = node.value.as_node().cloned().and_then(ast::NameRef::cast).unwrap();
+            let name_kind = classify_name_ref(sb, node.with_value(&name_ref)).map(|d| d.kind);
+            match name_kind {
+                Some(name_kind) => {
+                    if let Local(local) = &name_kind {
+                        if let Some(name) = local.name(db) {
+                            let shadow_count =
+                                bindings_shadow_count.entry(name.clone()).or_default();
+                            binding_hash =
+                                Some(calc_binding_hash(node.file_id, &name, *shadow_count))
+                        }
+                    };
+
+                    highlight_name(db, name_kind)
                 }
+                _ => return None,
             }
-            INT_NUMBER | FLOAT_NUMBER => tags::LITERAL_NUMERIC,
-            BYTE => tags::LITERAL_BYTE,
-            CHAR => tags::LITERAL_CHAR,
-            LIFETIME => tags::TYPE_LIFETIME,
-            T![unsafe] => tags::KEYWORD_UNSAFE,
-            k if is_control_keyword(k) => tags::KEYWORD_CONTROL,
-            k if k.is_keyword() => tags::KEYWORD,
-            _ => {
-                if let Some(macro_call) = node.as_node().cloned().and_then(ast::MacroCall::cast) {
-                    if let Some(path) = macro_call.path() {
-                        if let Some(segment) = path.segment() {
-                            if let Some(name_ref) = segment.name_ref() {
-                                highlighted.insert(name_ref.syntax().clone().into());
-                                let range_start = name_ref.syntax().text_range().start();
-                                let mut range_end = name_ref.syntax().text_range().end();
-                                for sibling in path.syntax().siblings_with_tokens(Direction::Next) {
-                                    match sibling.kind() {
-                                        T![!] | IDENT => range_end = sibling.text_range().end(),
-                                        _ => (),
-                                    }
-                                }
-                                res.push(HighlightedRange {
-                                    range: TextRange::from_to(range_start, range_end),
-                                    tag: tags::MACRO,
-                                    binding_hash: None,
-                                })
-                            }
-                        }
-                    }
+        }
+        NAME => {
+            let name = node.value.as_node().cloned().and_then(ast::Name::cast).unwrap();
+            let name_kind = classify_name(sb, node.with_value(&name)).map(|d| d.kind);
+
+            if let Some(Local(local)) = &name_kind {
+                if let Some(name) = local.name(db) {
+                    let shadow_count = bindings_shadow_count.entry(name.clone()).or_default();
+                    *shadow_count += 1;
+                    binding_hash = Some(calc_binding_hash(node.file_id, &name, *shadow_count))
                 }
-                continue;
+            };
+
+            match name_kind {
+                Some(name_kind) => highlight_name(db, name_kind),
+                None => name.syntax().parent().map_or(tags::FUNCTION, |x| match x.kind() {
+                    STRUCT_DEF | ENUM_DEF | TRAIT_DEF | TYPE_ALIAS_DEF => tags::TYPE,
+                    TYPE_PARAM => tags::TYPE_PARAM,
+                    RECORD_FIELD_DEF => tags::FIELD,
+                    _ => tags::FUNCTION,
+                }),
             }
-        };
-        res.push(HighlightedRange { range: node.text_range(), tag, binding_hash })
+        }
+        INT_NUMBER | FLOAT_NUMBER => tags::LITERAL_NUMERIC,
+        BYTE => tags::LITERAL_BYTE,
+        CHAR => tags::LITERAL_CHAR,
+        LIFETIME => tags::TYPE_LIFETIME,
+        T![unsafe] => tags::KEYWORD_UNSAFE,
+        k if is_control_keyword(k) => tags::KEYWORD_CONTROL,
+        k if k.is_keyword() => tags::KEYWORD,
+
+        _ => return None,
+    };
+
+    return Some((tag, binding_hash));
+
+    fn calc_binding_hash(file_id: HirFileId, name: &Name, shadow_count: u32) -> u64 {
+        fn hash<T: std::hash::Hash + std::fmt::Debug>(x: T) -> u64 {
+            use std::{collections::hash_map::DefaultHasher, hash::Hasher};
+
+            let mut hasher = DefaultHasher::new();
+            x.hash(&mut hasher);
+            hasher.finish()
+        }
+
+        hash((file_id, name, shadow_count))
     }
-    res
 }
 
 pub(crate) fn highlight_as_html(db: &RootDatabase, file_id: FileId, rainbow: bool) -> String {
@@ -331,6 +400,16 @@ fn foo<T>() -> T {
     foo::<i32>();
 }
 
+macro_rules! def_fn {
+    ($($tt:tt)*) => {$($tt)*}
+}
+
+def_fn!{
+    fn bar() -> u32 {
+        100
+    }
+}
+
 // comment
 fn main() {
     println!("Hello, {}!", 92);