Restructure syntax element highlighting

crates/ide/src/syntax_highlighting.rs

@@ -309,6 +309,7 @@ fn traverse(
                     match (token.kind(), parent.kind()) {
                         (T![ident], NAME | NAME_REF) => parent.into(),
                         (T![self] | T![super] | T![crate], NAME_REF) => parent.into(),
+                        (INT_NUMBER, NAME_REF) => parent.into(),
                         _ => token.into(),
                     }
                 }

crates/ide/src/syntax_highlighting/highlight.rs

@@ -8,7 +8,7 @@ use ide_db::{
 };
 use rustc_hash::FxHashMap;
 use syntax::{
-    ast, AstNode, AstToken, NodeOrToken, SyntaxElement,
+    ast, match_ast, AstNode, AstToken, NodeOrToken, SyntaxElement,
     SyntaxKind::{self, *},
     SyntaxNode, SyntaxToken, T,
 };
@@ -25,73 +25,38 @@ pub(super) fn element(
     syntactic_name_ref_highlighting: bool,
     element: SyntaxElement,
 ) -> Option<(Highlight, Option<u64>)> {
-    let mut binding_hash = None;
+    match element {
-    let highlight: Highlight = match element.kind() {
+        NodeOrToken::Node(it) => {
-        FN => {
+            node(sema, krate, bindings_shadow_count, syntactic_name_ref_highlighting, it)
-            bindings_shadow_count.clear();
-            return None;
         }
-        // Highlight definitions depending on the "type" of the definition.
+        NodeOrToken::Token(it) => Some((token(sema, krate, it)?, None)),
-        NAME => {
-            let name = element.into_node().and_then(ast::Name::cast).unwrap();
-            highlight_name(sema, bindings_shadow_count, &mut binding_hash, krate, name)
     }
-        // Highlight references like the definitions they resolve to
-        NAME_REF if element.ancestors().any(|it| it.kind() == ATTR) => {
-            // FIXME: We highlight paths in attributes slightly differently to work around this module
-            // currently not knowing about tool attributes and rustc builtin attributes as
-            // we do not want to resolve those to functions that may be defined in scope.
-            let name_ref = element.into_node().and_then(ast::NameRef::cast).unwrap();
-            highlight_name_ref_in_attr(sema, name_ref)
 }
-        NAME_REF => {
+
-            let name_ref = element.into_node().and_then(ast::NameRef::cast).unwrap();
+fn token(
-            highlight_name_ref(
+    sema: &Semantics<RootDatabase>,
-                sema,
+    krate: Option<hir::Crate>,
-                krate,
+    token: SyntaxToken,
-                bindings_shadow_count,
+) -> Option<Highlight> {
-                &mut binding_hash,
+    let highlight: Highlight = if let Some(comment) = ast::Comment::cast(token.clone()) {
-                syntactic_name_ref_highlighting,
-                name_ref,
-            )
-        }
-        // Simple token-based highlighting
-        COMMENT => {
-            let comment = element.into_token().and_then(ast::Comment::cast)?;
         let h = HlTag::Comment;
         match comment.kind().doc {
             Some(_) => h | HlMod::Documentation,
             None => h.into(),
         }
-        }
+    } else {
+        match token.kind() {
             STRING | BYTE_STRING => HlTag::StringLiteral.into(),
-        ATTR => HlTag::Attribute.into(),
+            INT_NUMBER if token.ancestors().nth(1).map_or(false, |it| it.kind() == FIELD_EXPR) => {
-        INT_NUMBER if element.ancestors().nth(1).map_or(false, |it| it.kind() == FIELD_EXPR) => {
                 SymbolKind::Field.into()
             }
             INT_NUMBER | FLOAT_NUMBER => HlTag::NumericLiteral.into(),
             BYTE => HlTag::ByteLiteral.into(),
             CHAR => HlTag::CharLiteral.into(),
-        QUESTION => HlTag::Operator(HlOperator::Other) | HlMod::ControlFlow,
+            T![?] => HlTag::Operator(HlOperator::Other) | HlMod::ControlFlow,
-        LIFETIME => {
+            IDENT if parent_matches::<ast::TokenTree>(&token) => {
-            let lifetime = element.into_node().and_then(ast::Lifetime::cast).unwrap();
+                if let Some(attr) = token.ancestors().nth(2).and_then(ast::Attr::cast) {
-
+                    match try_resolve_derive_input_at(sema, &attr, &token) {
-            match NameClass::classify_lifetime(sema, &lifetime) {
-                Some(NameClass::Definition(def)) => {
-                    highlight_def(sema, krate, def) | HlMod::Definition
-                }
-                None => match NameRefClass::classify_lifetime(sema, &lifetime) {
-                    Some(NameRefClass::Definition(def)) => highlight_def(sema, krate, def),
-                    _ => SymbolKind::LifetimeParam.into(),
-                },
-                _ => Highlight::from(SymbolKind::LifetimeParam) | HlMod::Definition,
-            }
-        }
-        IDENT if parent_matches::<ast::TokenTree>(&element) => {
-            if let Some((attr, token)) =
-                element.ancestors().nth(2).and_then(ast::Attr::cast).zip(element.as_token())
-            {
-                match try_resolve_derive_input_at(sema, &attr, token) {
                         Some(makro) => highlight_def(sema, krate, Definition::Macro(makro)),
                         None => HlTag::None.into(),
                     }
@@ -100,10 +65,10 @@ pub(super) fn element(
                 }
             }
             p if p.is_punct() => match p {
-            T![&] if parent_matches::<ast::BinExpr>(&element) => HlOperator::Bitwise.into(),
+                T![&] if parent_matches::<ast::BinExpr>(&token) => HlOperator::Bitwise.into(),
                 T![&] => {
                     let h = HlTag::Operator(HlOperator::Other).into();
-                let is_unsafe = element
+                    let is_unsafe = token
                         .parent()
                         .and_then(ast::RefExpr::cast)
                         .map_or(false, |ref_expr| sema.is_unsafe_ref_expr(&ref_expr));
@@ -113,13 +78,15 @@ pub(super) fn element(
                         h
                     }
                 }
-            T![::] | T![->] | T![=>] | T![..] | T![=] | T![@] | T![.] => HlOperator::Other.into(),
+                T![::] | T![->] | T![=>] | T![..] | T![=] | T![@] | T![.] => {
-            T![!] if parent_matches::<ast::MacroCall>(&element) => SymbolKind::Macro.into(),
+                    HlOperator::Other.into()
-            T![!] if parent_matches::<ast::NeverType>(&element) => HlTag::BuiltinType.into(),
+                }
-            T![!] if parent_matches::<ast::PrefixExpr>(&element) => HlOperator::Logical.into(),
+                T![!] if parent_matches::<ast::MacroCall>(&token) => SymbolKind::Macro.into(),
-            T![*] if parent_matches::<ast::PtrType>(&element) => HlTag::Keyword.into(),
+                T![!] if parent_matches::<ast::NeverType>(&token) => HlTag::BuiltinType.into(),
-            T![*] if parent_matches::<ast::PrefixExpr>(&element) => {
+                T![!] if parent_matches::<ast::PrefixExpr>(&token) => HlOperator::Logical.into(),
-                let prefix_expr = element.parent().and_then(ast::PrefixExpr::cast)?;
+                T![*] if parent_matches::<ast::PtrType>(&token) => HlTag::Keyword.into(),
+                T![*] if parent_matches::<ast::PrefixExpr>(&token) => {
+                    let prefix_expr = token.parent().and_then(ast::PrefixExpr::cast)?;
 
                     let expr = prefix_expr.expr()?;
                     let ty = sema.type_of_expr(&expr)?.original;
@@ -131,8 +98,8 @@ pub(super) fn element(
                         HlPunct::Other.into()
                     }
                 }
-            T![-] if parent_matches::<ast::PrefixExpr>(&element) => {
+                T![-] if parent_matches::<ast::PrefixExpr>(&token) => {
-                let prefix_expr = element.parent().and_then(ast::PrefixExpr::cast)?;
+                    let prefix_expr = token.parent().and_then(ast::PrefixExpr::cast)?;
 
                     let expr = prefix_expr.expr()?;
                     match expr {
@@ -141,32 +108,32 @@ pub(super) fn element(
                     }
                     .into()
                 }
-            _ if parent_matches::<ast::PrefixExpr>(&element) => HlOperator::Other.into(),
+                _ if parent_matches::<ast::PrefixExpr>(&token) => HlOperator::Other.into(),
-            T![+] | T![-] | T![*] | T![/] if parent_matches::<ast::BinExpr>(&element) => {
+                T![+] | T![-] | T![*] | T![/] if parent_matches::<ast::BinExpr>(&token) => {
                     HlOperator::Arithmetic.into()
                 }
-            T![+=] | T![-=] | T![*=] | T![/=] if parent_matches::<ast::BinExpr>(&element) => {
+                T![+=] | T![-=] | T![*=] | T![/=] if parent_matches::<ast::BinExpr>(&token) => {
                     Highlight::from(HlOperator::Arithmetic) | HlMod::Mutable
                 }
-            T![|] | T![&] | T![!] | T![^] if parent_matches::<ast::BinExpr>(&element) => {
+                T![|] | T![&] | T![!] | T![^] if parent_matches::<ast::BinExpr>(&token) => {
                     HlOperator::Bitwise.into()
                 }
-            T![|=] | T![&=] | T![^=] if parent_matches::<ast::BinExpr>(&element) => {
+                T![|=] | T![&=] | T![^=] if parent_matches::<ast::BinExpr>(&token) => {
                     Highlight::from(HlOperator::Bitwise) | HlMod::Mutable
                 }
-            T![&&] | T![||] if parent_matches::<ast::BinExpr>(&element) => {
+                T![&&] | T![||] if parent_matches::<ast::BinExpr>(&token) => {
                     HlOperator::Logical.into()
                 }
                 T![>] | T![<] | T![==] | T![>=] | T![<=] | T![!=]
-                if parent_matches::<ast::BinExpr>(&element) =>
+                    if parent_matches::<ast::BinExpr>(&token) =>
                 {
                     HlOperator::Comparison.into()
                 }
-            _ if parent_matches::<ast::BinExpr>(&element) => HlOperator::Other.into(),
+                _ if parent_matches::<ast::BinExpr>(&token) => HlOperator::Other.into(),
-            _ if parent_matches::<ast::RangeExpr>(&element) => HlOperator::Other.into(),
+                _ if parent_matches::<ast::RangeExpr>(&token) => HlOperator::Other.into(),
-            _ if parent_matches::<ast::RangePat>(&element) => HlOperator::Other.into(),
+                _ if parent_matches::<ast::RangePat>(&token) => HlOperator::Other.into(),
-            _ if parent_matches::<ast::RestPat>(&element) => HlOperator::Other.into(),
+                _ if parent_matches::<ast::RestPat>(&token) => HlOperator::Other.into(),
-            _ if parent_matches::<ast::Attr>(&element) => HlTag::Attribute.into(),
+                _ if parent_matches::<ast::Attr>(&token) => HlTag::Attribute.into(),
                 kind => match kind {
                     T!['['] | T![']'] => HlPunct::Bracket,
                     T!['{'] | T!['}'] => HlPunct::Brace,
@@ -180,7 +147,6 @@ pub(super) fn element(
                 }
                 .into(),
             },
-
             k if k.is_keyword() => {
                 let h = Highlight::new(HlTag::Keyword);
                 match k {
@@ -195,32 +161,84 @@ pub(super) fn element(
                     | T![return]
                     | T![while]
                     | T![yield] => h | HlMod::ControlFlow,
-                T![for] if !is_child_of_impl(&element) => h | HlMod::ControlFlow,
+                    T![for] if !is_child_of_impl(&token) => h | HlMod::ControlFlow,
                     T![unsafe] => h | HlMod::Unsafe,
                     T![true] | T![false] => HlTag::BoolLiteral.into(),
                     // self is handled as either a Name or NameRef already
                     T![self] => return None,
-                T![ref] => element
+                    T![ref] => token
                         .parent()
                         .and_then(ast::IdentPat::cast)
                         .and_then(|ident_pat| {
-                        if sema.is_unsafe_ident_pat(&ident_pat) {
+                            (sema.is_unsafe_ident_pat(&ident_pat)).then(|| HlMod::Unsafe)
-                            Some(HlMod::Unsafe)
-                        } else {
-                            None
-                        }
                         })
-                    .map(|modifier| h | modifier)
+                        .map_or(h, |modifier| h | modifier),
-                    .unwrap_or(h),
                     T![async] => h | HlMod::Async,
                     _ => h,
                 }
             }
-
             _ => return None,
+        }
     };
+    Some(highlight)
+}
 
-    return Some((highlight, binding_hash));
+fn node(
+    sema: &Semantics<RootDatabase>,
+    krate: Option<hir::Crate>,
+    bindings_shadow_count: &mut FxHashMap<hir::Name, u32>,
+    syntactic_name_ref_highlighting: bool,
+    node: SyntaxNode,
+) -> Option<(Highlight, Option<u64>)> {
+    let mut binding_hash = None;
+    let highlight = match_ast! {
+        match node {
+            ast::Fn(__) => {
+                bindings_shadow_count.clear();
+                return None;
+            },
+            ast::Attr(__) => {
+                HlTag::Attribute.into()
+            },
+            // Highlight definitions depending on the "type" of the definition.
+            ast::Name(name) => {
+                highlight_name(sema, bindings_shadow_count, &mut binding_hash, krate, name)
+            },
+            // Highlight references like the definitions they resolve to
+            ast::NameRef(name_ref) => {
+                if node.ancestors().any(|it| it.kind() == ATTR) {
+
+                    // FIXME: We highlight paths in attributes slightly differently to work around this module
+                    // currently not knowing about tool attributes and rustc builtin attributes as
+                    // we do not want to resolve those to functions that may be defined in scope.
+                    highlight_name_ref_in_attr(sema, name_ref)
+                } else {
+                    highlight_name_ref(
+                        sema,
+                        krate,
+                        bindings_shadow_count,
+                        &mut binding_hash,
+                        syntactic_name_ref_highlighting,
+                        name_ref,
+                    )
+                }
+            },
+            ast::Lifetime(lifetime) => {
+                match NameClass::classify_lifetime(sema, &lifetime) {
+                    Some(NameClass::Definition(def)) => {
+                        highlight_def(sema, krate, def) | HlMod::Definition
+                    }
+                    None => match NameRefClass::classify_lifetime(sema, &lifetime) {
+                        Some(NameRefClass::Definition(def)) => highlight_def(sema, krate, def),
+                        _ => SymbolKind::LifetimeParam.into(),
+                    },
+                    _ => Highlight::from(SymbolKind::LifetimeParam) | HlMod::Definition,
+                }
+            },
+            _ => return None,
+        }
+    };
+    Some((highlight, binding_hash))
 }
 
 fn highlight_name_ref_in_attr(sema: &Semantics<RootDatabase>, name_ref: ast::NameRef) -> Highlight {
@@ -715,12 +733,12 @@ fn parents_match(mut node: NodeOrToken<SyntaxNode, SyntaxToken>, mut kinds: &[Sy
 }
 
 #[inline]
-fn parent_matches<N: AstNode>(element: &SyntaxElement) -> bool {
+fn parent_matches<N: AstNode>(token: &SyntaxToken) -> bool {
-    element.parent().map_or(false, |it| N::can_cast(it.kind()))
+    token.parent().map_or(false, |it| N::can_cast(it.kind()))
 }
 
-fn is_child_of_impl(element: &SyntaxElement) -> bool {
+fn is_child_of_impl(token: &SyntaxToken) -> bool {
-    match element.parent() {
+    match token.parent() {
         Some(e) => e.kind() == IMPL,
         _ => false,
     }

crates/ide/src/syntax_highlighting/inject.rs

@@ -12,11 +12,10 @@ use syntax::{
 
 use crate::{
     doc_links::{doc_attributes, extract_definitions_from_docs, resolve_doc_path_for_def},
+    syntax_highlighting::{highlights::Highlights, injector::Injector},
     Analysis, HlMod, HlRange, HlTag, RootDatabase,
 };
 
-use super::{highlights::Highlights, injector::Injector};
-
 pub(super) fn ra_fixture(
     hl: &mut Highlights,
     sema: &Semantics<RootDatabase>,

crates/ide/src/syntax_highlighting/injector.rs

@@ -17,9 +17,11 @@ impl Injector {
         assert_eq!(len, source_range.len());
         self.add_impl(text, Some(source_range.start()));
     }
+
     pub(super) fn add_unmapped(&mut self, text: &str) {
         self.add_impl(text, None);
     }
+
     fn add_impl(&mut self, text: &str, source: Option<TextSize>) {
         let len = TextSize::of(text);
         let target_range = TextRange::at(TextSize::of(&self.buf), len);
@@ -30,6 +32,7 @@ impl Injector {
     pub(super) fn text(&self) -> &str {
         &self.buf
     }
+
     pub(super) fn map_range_up(&self, range: TextRange) -> impl Iterator<Item = TextRange> + '_ {
         equal_range_by(&self.ranges, |&(r, _)| TextRange::ordering(r, range)).filter_map(move |i| {
             let (target_range, delta) = self.ranges[i];