From 19f1ff5c7063dac772c1d324d11661393a3b5f7e Mon Sep 17 00:00:00 2001
From: Lukas Wirth <lukastw97@gmail.com>
Date: Mon, 3 Jan 2022 16:00:45 +0100
Subject: [PATCH] give `resolve_derive_ident` a more robust api

---
 crates/hir/src/semantics.rs                   | 37 +++++++++++--------
 .../ide/src/syntax_highlighting/highlight.rs  | 16 +++++---
 crates/ide_db/src/defs.rs                     |  8 +++-
 3 files changed, 39 insertions(+), 22 deletions(-)

diff --git a/crates/hir/src/semantics.rs b/crates/hir/src/semantics.rs
index 33cbf74f5e4..b685d260815 100644
--- a/crates/hir/src/semantics.rs
+++ b/crates/hir/src/semantics.rs
@@ -356,8 +356,12 @@ impl<'db, DB: HirDatabase> Semantics<'db, DB> {
         self.imp.resolve_bind_pat_to_const(pat)
     }
 
-    pub fn resolve_derive_ident(&self, ident: &ast::Ident) -> Option<PathResolution> {
-        self.imp.resolve_derive_ident(ident)
+    pub fn resolve_derive_ident(
+        &self,
+        derive: &ast::Attr,
+        ident: &ast::Ident,
+    ) -> Option<PathResolution> {
+        self.imp.resolve_derive_ident(derive, ident)
     }
 
     // FIXME: use this instead?
@@ -900,23 +904,26 @@ impl<'db> SemanticsImpl<'db> {
         self.analyze(pat.syntax()).resolve_bind_pat_to_const(self.db, pat)
     }
 
-    fn resolve_derive_ident(&self, ident: &ast::Ident) -> Option<PathResolution> {
+    fn resolve_derive_ident(
+        &self,
+        derive: &ast::Attr,
+        ident: &ast::Ident,
+    ) -> Option<PathResolution> {
+        debug_assert!(ident.syntax().parent().and_then(ast::TokenTree::cast).is_some());
+        debug_assert!(ident.syntax().ancestors().any(|anc| anc == *derive.syntax()));
         // derive macros are always at depth 2, tokentree -> meta -> attribute
         let syntax = ident.syntax();
-        let attr = syntax.ancestors().nth(2).and_then(ast::Attr::cast)?;
 
-        let tt = attr.token_tree()?;
-        if !tt.syntax().text_range().contains_range(ident.syntax().text_range()) {
-            return None;
-        }
-
-        let file = self.find_file(attr.syntax());
-        let adt = attr.syntax().parent().and_then(ast::Adt::cast)?;
+        let tt = derive.token_tree()?;
+        let file = self.find_file(derive.syntax());
+        let adt = derive.syntax().parent().and_then(ast::Adt::cast)?;
 
         let res = self.with_ctx(|ctx| {
-            let attr_def = ctx.attr_to_def(file.with_value(attr.clone()))?;
-            let derives = ctx
-                .attr_to_derive_macro_call(file.with_value(&adt), file.with_value(attr.clone()))?;
+            let attr_def = ctx.attr_to_def(file.with_value(derive.clone()))?;
+            let derives = ctx.attr_to_derive_macro_call(
+                file.with_value(&adt),
+                file.with_value(derive.clone()),
+            )?;
 
             let mut derive_paths = attr_def.parse_path_comma_token_tree()?;
 
@@ -951,7 +958,7 @@ impl<'db> SemanticsImpl<'db> {
         match res {
             Either::Left(path) => resolve_hir_path(
                 self.db,
-                &self.scope(attr.syntax()).resolver,
+                &self.scope(derive.syntax()).resolver,
                 &Path::from_known_path(path, []),
             )
             .filter(|res| matches!(res, PathResolution::Def(ModuleDef::Module(_)))),
diff --git a/crates/ide/src/syntax_highlighting/highlight.rs b/crates/ide/src/syntax_highlighting/highlight.rs
index f82ccaa1c76..688f181d916 100644
--- a/crates/ide/src/syntax_highlighting/highlight.rs
+++ b/crates/ide/src/syntax_highlighting/highlight.rs
@@ -39,11 +39,17 @@ pub(super) fn token(
         INT_NUMBER | FLOAT_NUMBER => HlTag::NumericLiteral.into(),
         BYTE => HlTag::ByteLiteral.into(),
         CHAR => HlTag::CharLiteral.into(),
-        IDENT if parent_matches::<ast::TokenTree>(&token) => {
-            match sema.resolve_derive_ident(&ast::Ident::cast(token).unwrap()) {
-                Some(res) => highlight_def(sema, krate, Definition::from(res)),
-                None => HlTag::None.into(),
-            }
+        IDENT => {
+            let tt = ast::TokenTree::cast(token.parent()?)?;
+            let ident = ast::Ident::cast(token)?;
+            // from this point on we are inside a token tree, this only happens for identifiers
+            // that were not mapped down into macro invocations
+            (|| {
+                let attr = tt.parent_meta()?.parent_attr()?;
+                let res = sema.resolve_derive_ident(&attr, &ident)?;
+                Some(highlight_def(sema, krate, Definition::from(res)))
+            })()
+            .unwrap_or_else(|| HlTag::None.into())
         }
         p if p.is_punct() => punctuation(sema, token, p),
         k if k.is_keyword() => keyword(sema, token, k)?,
diff --git a/crates/ide_db/src/defs.rs b/crates/ide_db/src/defs.rs
index d69fd9bd217..1501c4eda53 100644
--- a/crates/ide_db/src/defs.rs
+++ b/crates/ide_db/src/defs.rs
@@ -55,8 +55,12 @@ impl Definition {
             let attr = ast::TokenTree::cast(parent.clone())
                 .and_then(|tt| tt.parent_meta())
                 .and_then(|meta| meta.parent_attr());
-            if let Some(_) = attr {
-                return sema.resolve_derive_ident(&ident).map(Into::into).into_iter().collect();
+            if let Some(attr) = attr {
+                return sema
+                    .resolve_derive_ident(&attr, &ident)
+                    .map(Into::into)
+                    .into_iter()
+                    .collect();
             }
         }
         Self::from_node(sema, &parent)