2 files changed, 231 insertions, 242 deletions

diff --git a/crates/ra_ide/src/syntax_highlighting.rs b/crates/ra_ide/src/syntax_highlighting.rs
index e8ca7d652..796f0e545 100644
--- a/crates/ra_ide/src/syntax_highlighting.rs
+++ b/crates/ra_ide/src/syntax_highlighting.rs
@@ -1,7 +1,9 @@
-//! FIXME: write short doc here
+//! Implements syntax highlighting.
 
 mod tags;
 mod html;
+#[cfg(test)]
+mod tests;
 
 use hir::{Name, Semantics};
 use ra_ide_db::{
@@ -10,16 +12,14 @@ use ra_ide_db::{
 };
 use ra_prof::profile;
 use ra_syntax::{
-    ast, AstNode, Direction, NodeOrToken, SyntaxElement, SyntaxKind, SyntaxKind::*, SyntaxToken,
-    TextRange, WalkEvent, T,
+    ast, AstNode, Direction, NodeOrToken, SyntaxElement, SyntaxKind::*, TextRange, WalkEvent, T,
 };
 use rustc_hash::FxHashMap;
 
 use crate::{references::classify_name_ref, FileId};
 
-pub use tags::{Highlight, HighlightModifier, HighlightModifiers, HighlightTag};
-
 pub(crate) use html::highlight_as_html;
+pub use tags::{Highlight, HighlightModifier, HighlightModifiers, HighlightTag};
 
 #[derive(Debug)]
 pub struct HighlightedRange {
@@ -28,21 +28,6 @@ pub struct HighlightedRange {
     pub binding_hash: Option<u64>,
 }
 
-fn is_control_keyword(kind: SyntaxKind) -> bool {
-    match kind {
-        T![for]
-        | T![loop]
-        | T![while]
-        | T![continue]
-        | T![break]
-        | T![if]
-        | T![else]
-        | T![match]
-        | T![return] => true,
-        _ => false,
-    }
-}
-
 pub(crate) fn highlight(
     db: &RootDatabase,
     file_id: FileId,
@@ -71,20 +56,24 @@ pub(crate) fn highlight(
 
     let mut current_macro_call: Option<ast::MacroCall> = None;
 
+    // Walk all nodes, keeping track of whether we are inside a macro or not.
+    // If in macro, expand it first and highlight the expanded code.
     for event in root.preorder_with_tokens() {
         let event_range = match &event {
             WalkEvent::Enter(it) => it.text_range(),
             WalkEvent::Leave(it) => it.text_range(),
         };
 
-        if event_range.intersection(&range_to_highlight).is_none() {
+        // Element outside of the viewport, no need to highlight
+        if range_to_highlight.intersection(&event_range).is_none() {
             continue;
         }
 
+        // Track "inside macro" state
         match event.clone().map(|it| it.into_node().and_then(ast::MacroCall::cast)) {
             WalkEvent::Enter(Some(mc)) => {
                 current_macro_call = Some(mc.clone());
-                if let Some(range) = highlight_macro(&mc) {
+                if let Some(range) = macro_call_range(&mc) {
                     res.push(HighlightedRange {
                         range,
                         highlight: HighlightTag::Macro.into(),
@@ -101,37 +90,40 @@ pub(crate) fn highlight(
             _ => (),
         }
 
-        let node = match event {
+        let element = match event {
             WalkEvent::Enter(it) => it,
             WalkEvent::Leave(_) => continue,
         };
+        let range = element.text_range();
 
-        if current_macro_call.is_some() {
-            if let Some(token) = node.into_token() {
-                if let Some((highlight, binding_hash)) =
-                    highlight_token_tree(&sema, &mut bindings_shadow_count, token.clone())
-                {
-                    res.push(HighlightedRange {
-                        range: token.text_range(),
-                        highlight,
-                        binding_hash,
-                    });
-                }
+        let element_to_highlight = if current_macro_call.is_some() {
+            // Inside a macro -- expand it first
+            let token = match element.into_token() {
+                Some(it) if it.parent().kind() == TOKEN_TREE => it,
+                _ => continue,
+            };
+            let token = sema.descend_into_macros(token.clone());
+            let parent = token.parent();
+            // We only care Name and Name_ref
+            match (token.kind(), parent.kind()) {
+                (IDENT, NAME) | (IDENT, NAME_REF) => parent.into(),
+                _ => token.into(),
             }
-            continue;
-        }
+        } else {
+            element
+        };
 
         if let Some((highlight, binding_hash)) =
-            highlight_node(&sema, &mut bindings_shadow_count, node.clone())
+            highlight_element(&sema, &mut bindings_shadow_count, element_to_highlight)
         {
-            res.push(HighlightedRange { range: node.text_range(), highlight, binding_hash });
+            res.push(HighlightedRange { range, highlight, binding_hash });
         }
     }
 
     res
 }
 
-fn highlight_macro(macro_call: &ast::MacroCall) -> Option<TextRange> {
+fn macro_call_range(macro_call: &ast::MacroCall) -> Option<TextRange> {
     let path = macro_call.path()?;
     let name_ref = path.segment()?.name_ref()?;
 
@@ -147,67 +139,22 @@ fn highlight_macro(macro_call: &ast::MacroCall) -> Option<TextRange> {
     Some(TextRange::from_to(range_start, range_end))
 }
 
-fn highlight_token_tree(
-    sema: &Semantics<RootDatabase>,
-    bindings_shadow_count: &mut FxHashMap<Name, u32>,
-    token: SyntaxToken,
-) -> Option<(Highlight, Option<u64>)> {
-    if token.parent().kind() != TOKEN_TREE {
-        return None;
-    }
-    let token = sema.descend_into_macros(token.clone());
-    let expanded = {
-        let parent = token.parent();
-        // We only care Name and Name_ref
-        match (token.kind(), parent.kind()) {
-            (IDENT, NAME) | (IDENT, NAME_REF) => parent.into(),
-            _ => token.into(),
-        }
-    };
-
-    highlight_node(sema, bindings_shadow_count, expanded)
-}
-
-fn highlight_node(
+fn highlight_element(
     sema: &Semantics<RootDatabase>,
     bindings_shadow_count: &mut FxHashMap<Name, u32>,
-    node: SyntaxElement,
+    element: SyntaxElement,
 ) -> Option<(Highlight, Option<u64>)> {
     let db = sema.db;
     let mut binding_hash = None;
-    let highlight: Highlight = match node.kind() {
+    let highlight: Highlight = match element.kind() {
         FN_DEF => {
             bindings_shadow_count.clear();
             return None;
         }
-        COMMENT => HighlightTag::Comment.into(),
-        STRING | RAW_STRING | RAW_BYTE_STRING | BYTE_STRING => HighlightTag::LiteralString.into(),
-        ATTR => HighlightTag::Attribute.into(),
-        // Special-case field init shorthand
-        NAME_REF if node.parent().and_then(ast::RecordField::cast).is_some() => {
-            HighlightTag::Field.into()
-        }
-        NAME_REF if node.ancestors().any(|it| it.kind() == ATTR) => return None,
-        NAME_REF => {
-            let name_ref = node.as_node().cloned().and_then(ast::NameRef::cast).unwrap();
-            let name_kind = classify_name_ref(sema, &name_ref);
-            match name_kind {
-                Some(name_kind) => {
-                    if let NameDefinition::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(&name, *shadow_count))
-                        }
-                    };
-
-                    highlight_name(db, name_kind)
-                }
-                _ => return None,
-            }
-        }
+
+        // Highlight definitions depending on the "type" of the definition.
         NAME => {
-            let name = node.as_node().cloned().and_then(ast::Name::cast).unwrap();
+            let name = element.into_node().and_then(ast::Name::cast).unwrap();
             let name_kind = classify_name(sema, &name);
 
             if let Some(NameDefinition::Local(local)) = &name_kind {
@@ -220,25 +167,56 @@ fn highlight_node(
 
             match name_kind {
                 Some(name_kind) => highlight_name(db, name_kind),
-                None => name.syntax().parent().map_or(HighlightTag::Function.into(), |x| {
-                    match x.kind() {
-                        STRUCT_DEF | ENUM_DEF | TRAIT_DEF | TYPE_ALIAS_DEF => {
-                            HighlightTag::Type.into()
-                        }
-                        TYPE_PARAM => HighlightTag::TypeParam.into(),
-                        RECORD_FIELD_DEF => HighlightTag::Field.into(),
-                        _ => HighlightTag::Function.into(),
-                    }
-                }),
+                None => highlight_name_by_syntax(name),
             }
         }
+
+        // Highlight references like the definitions they resolve to
+
+        // Special-case field init shorthand
+        NAME_REF if element.parent().and_then(ast::RecordField::cast).is_some() => {
+            HighlightTag::Field.into()
+        }
+        NAME_REF if element.ancestors().any(|it| it.kind() == ATTR) => return None,
+        NAME_REF => {
+            let name_ref = element.into_node().and_then(ast::NameRef::cast).unwrap();
+            let name_kind = classify_name_ref(sema, &name_ref)?;
+
+            if let NameDefinition::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(&name, *shadow_count))
+                }
+            };
+
+            highlight_name(db, name_kind)
+        }
+
+        // Simple token-based highlighting
+        COMMENT => HighlightTag::Comment.into(),
+        STRING | RAW_STRING | RAW_BYTE_STRING | BYTE_STRING => HighlightTag::LiteralString.into(),
+        ATTR => HighlightTag::Attribute.into(),
         INT_NUMBER | FLOAT_NUMBER => HighlightTag::LiteralNumeric.into(),
         BYTE => HighlightTag::LiteralByte.into(),
         CHAR => HighlightTag::LiteralChar.into(),
         LIFETIME => HighlightTag::TypeLifetime.into(),
-        T![unsafe] => HighlightTag::Keyword | HighlightModifier::Unsafe,
-        k if is_control_keyword(k) => HighlightTag::Keyword | HighlightModifier::Control,
-        k if k.is_keyword() => HighlightTag::Keyword.into(),
+
+        k if k.is_keyword() => {
+            let h = Highlight::new(HighlightTag::Keyword);
+            match k {
+                T![break]
+                | T![continue]
+                | T![else]
+                | T![for]
+                | T![if]
+                | T![loop]
+                | T![match]
+                | T![return]
+                | T![while] => h | HighlightModifier::Control,
+                T![unsafe] => h | HighlightModifier::Unsafe,
+                _ => h,
+            }
+        }
 
         _ => return None,
     };
@@ -262,17 +240,19 @@ fn highlight_name(db: &RootDatabase, def: NameDefinition) -> Highlight {
     match def {
         NameDefinition::Macro(_) => HighlightTag::Macro,
         NameDefinition::StructField(_) => HighlightTag::Field,
-        NameDefinition::ModuleDef(hir::ModuleDef::Module(_)) => HighlightTag::Module,
-        NameDefinition::ModuleDef(hir::ModuleDef::Function(_)) => HighlightTag::Function,
-        NameDefinition::ModuleDef(hir::ModuleDef::Adt(_)) => HighlightTag::Type,
-        NameDefinition::ModuleDef(hir::ModuleDef::EnumVariant(_)) => HighlightTag::Constant,
-        NameDefinition::ModuleDef(hir::ModuleDef::Const(_)) => HighlightTag::Constant,
-        NameDefinition::ModuleDef(hir::ModuleDef::Static(_)) => HighlightTag::Constant,
-        NameDefinition::ModuleDef(hir::ModuleDef::Trait(_)) => HighlightTag::Type,
-        NameDefinition::ModuleDef(hir::ModuleDef::TypeAlias(_)) => HighlightTag::Type,
-        NameDefinition::ModuleDef(hir::ModuleDef::BuiltinType(_)) => {
-            return HighlightTag::Type | HighlightModifier::Builtin
-        }
+        NameDefinition::ModuleDef(def) => match def {
+            hir::ModuleDef::Module(_) => HighlightTag::Module,
+            hir::ModuleDef::Function(_) => HighlightTag::Function,
+            hir::ModuleDef::Adt(_) => HighlightTag::Type,
+            hir::ModuleDef::EnumVariant(_) => HighlightTag::Constant,
+            hir::ModuleDef::Const(_) => HighlightTag::Constant,
+            hir::ModuleDef::Static(_) => HighlightTag::Constant,
+            hir::ModuleDef::Trait(_) => HighlightTag::Type,
+            hir::ModuleDef::TypeAlias(_) => HighlightTag::Type,
+            hir::ModuleDef::BuiltinType(_) => {
+                return HighlightTag::Type | HighlightModifier::Builtin
+            }
+        },
         NameDefinition::SelfType(_) => HighlightTag::TypeSelf,
         NameDefinition::TypeParam(_) => HighlightTag::TypeParam,
         NameDefinition::Local(local) => {
@@ -286,136 +266,18 @@ fn highlight_name(db: &RootDatabase, def: NameDefinition) -> Highlight {
     .into()
 }
 
-#[cfg(test)]
-mod tests {
-    use std::fs;
-
-    use test_utils::{assert_eq_text, project_dir, read_text};
+fn highlight_name_by_syntax(name: ast::Name) -> Highlight {
+    let default = HighlightTag::Function.into();
 
-    use crate::{
-        mock_analysis::{single_file, MockAnalysis},
-        FileRange, TextRange,
+    let parent = match name.syntax().parent() {
+        Some(it) => it,
+        _ => return default,
     };
 
-    #[test]
-    fn test_highlighting() {
-        let (analysis, file_id) = single_file(
-            r#"
-#[derive(Clone, Debug)]
-struct Foo {
-    pub x: i32,
-    pub y: i32,
-}
-
-fn foo<T>() -> T {
-    unimplemented!();
-    foo::<i32>();
-}
-
-macro_rules! def_fn {
-    ($($tt:tt)*) => {$($tt)*}
-}
-
-def_fn!{
-    fn bar() -> u32 {
-        100
-    }
-}
-
-// comment
-fn main() {
-    println!("Hello, {}!", 92);
-
-    let mut vec = Vec::new();
-    if true {
-        let x = 92;
-        vec.push(Foo { x, y: 1 });
-    }
-    unsafe { vec.set_len(0); }
-
-    let mut x = 42;
-    let y = &mut x;
-    let z = &y;
-
-    y;
-}
-
-enum E<X> {
-    V(X)
-}
-
-impl<X> E<X> {
-    fn new<T>() -> E<T> {}
-}
-"#
-            .trim(),
-        );
-        let dst_file = project_dir().join("crates/ra_ide/src/snapshots/highlighting.html");
-        let actual_html = &analysis.highlight_as_html(file_id, false).unwrap();
-        let expected_html = &read_text(&dst_file);
-        fs::write(dst_file, &actual_html).unwrap();
-        assert_eq_text!(expected_html, actual_html);
-    }
-
-    #[test]
-    fn test_rainbow_highlighting() {
-        let (analysis, file_id) = single_file(
-            r#"
-fn main() {
-    let hello = "hello";
-    let x = hello.to_string();
-    let y = hello.to_string();
-
-    let x = "other color please!";
-    let y = x.to_string();
-}
-
-fn bar() {
-    let mut hello = "hello";
-}
-"#
-            .trim(),
-        );
-        let dst_file = project_dir().join("crates/ra_ide/src/snapshots/rainbow_highlighting.html");
-        let actual_html = &analysis.highlight_as_html(file_id, true).unwrap();
-        let expected_html = &read_text(&dst_file);
-        fs::write(dst_file, &actual_html).unwrap();
-        assert_eq_text!(expected_html, actual_html);
-    }
-
-    #[test]
-    fn accidentally_quadratic() {
-        let file = project_dir().join("crates/ra_syntax/test_data/accidentally_quadratic");
-        let src = fs::read_to_string(file).unwrap();
-
-        let mut mock = MockAnalysis::new();
-        let file_id = mock.add_file("/main.rs", &src);
-        let host = mock.analysis_host();
-
-        // let t = std::time::Instant::now();
-        let _ = host.analysis().highlight(file_id).unwrap();
-        // eprintln!("elapsed: {:?}", t.elapsed());
-    }
-
-    #[test]
-    fn test_ranges() {
-        let (analysis, file_id) = single_file(
-            r#"
-            #[derive(Clone, Debug)]
-            struct Foo {
-                pub x: i32,
-                pub y: i32,
-            }"#,
-        );
-
-        // The "x"
-        let highlights = &analysis
-            .highlight_range(FileRange {
-                file_id,
-                range: TextRange::offset_len(82.into(), 1.into()),
-            })
-            .unwrap();
-
-        assert_eq!(&highlights[0].highlight.to_string(), "field");
+    match parent.kind() {
+        STRUCT_DEF | ENUM_DEF | TRAIT_DEF | TYPE_ALIAS_DEF => HighlightTag::Type.into(),
+        TYPE_PARAM => HighlightTag::TypeParam.into(),
+        RECORD_FIELD_DEF => HighlightTag::Field.into(),
+        _ => default,
     }
 }
diff --git a/crates/ra_ide/src/syntax_highlighting/tests.rs b/crates/ra_ide/src/syntax_highlighting/tests.rs
new file mode 100644
index 000000000..ff23d4ac5
--- /dev/null
+++ b/crates/ra_ide/src/syntax_highlighting/tests.rs
@@ -0,0 +1,127 @@
+use std::fs;
+
+use test_utils::{assert_eq_text, project_dir, read_text};
+
+use crate::{
+    mock_analysis::{single_file, MockAnalysis},
+    FileRange, TextRange,
+};
+
+#[test]
+fn test_highlighting() {
+    let (analysis, file_id) = single_file(
+        r#"
+#[derive(Clone, Debug)]
+struct Foo {
+    pub x: i32,
+    pub y: i32,
+}
+
+fn foo<T>() -> T {
+    unimplemented!();
+    foo::<i32>();
+}
+
+macro_rules! def_fn {
+    ($($tt:tt)*) => {$($tt)*}
+}
+
+def_fn!{
+    fn bar() -> u32 {
+        100
+    }
+}
+
+// comment
+fn main() {
+    println!("Hello, {}!", 92);
+
+    let mut vec = Vec::new();
+    if true {
+        let x = 92;
+        vec.push(Foo { x, y: 1 });
+    }
+    unsafe { vec.set_len(0); }
+
+    let mut x = 42;
+    let y = &mut x;
+    let z = &y;
+
+    y;
+}
+
+enum E<X> {
+    V(X)
+}
+
+impl<X> E<X> {
+    fn new<T>() -> E<T> {}
+}
+"#
+        .trim(),
+    );
+    let dst_file = project_dir().join("crates/ra_ide/src/snapshots/highlighting.html");
+    let actual_html = &analysis.highlight_as_html(file_id, false).unwrap();
+    let expected_html = &read_text(&dst_file);
+    fs::write(dst_file, &actual_html).unwrap();
+    assert_eq_text!(expected_html, actual_html);
+}
+
+#[test]
+fn test_rainbow_highlighting() {
+    let (analysis, file_id) = single_file(
+        r#"
+fn main() {
+    let hello = "hello";
+    let x = hello.to_string();
+    let y = hello.to_string();
+
+    let x = "other color please!";
+    let y = x.to_string();
+}
+
+fn bar() {
+    let mut hello = "hello";
+}
+"#
+        .trim(),
+    );
+    let dst_file = project_dir().join("crates/ra_ide/src/snapshots/rainbow_highlighting.html");
+    let actual_html = &analysis.highlight_as_html(file_id, true).unwrap();
+    let expected_html = &read_text(&dst_file);
+    fs::write(dst_file, &actual_html).unwrap();
+    assert_eq_text!(expected_html, actual_html);
+}
+
+#[test]
+fn accidentally_quadratic() {
+    let file = project_dir().join("crates/ra_syntax/test_data/accidentally_quadratic");
+    let src = fs::read_to_string(file).unwrap();
+
+    let mut mock = MockAnalysis::new();
+    let file_id = mock.add_file("/main.rs", &src);
+    let host = mock.analysis_host();
+
+    // let t = std::time::Instant::now();
+    let _ = host.analysis().highlight(file_id).unwrap();
+    // eprintln!("elapsed: {:?}", t.elapsed());
+}
+
+#[test]
+fn test_ranges() {
+    let (analysis, file_id) = single_file(
+        r#"
+            #[derive(Clone, Debug)]
+            struct Foo {
+                pub x: i32,
+                pub y: i32,
+            }"#,
+    );
+
+    // The "x"
+    let highlights = &analysis
+        .highlight_range(FileRange { file_id, range: TextRange::offset_len(82.into(), 1.into()) })
+        .unwrap();
+
+    assert_eq!(&highlights[0].highlight.to_string(), "field");
+}