fix: multi-token mapping aware goto definition

1 files changed, 56 insertions, 39 deletions

diff --git a/crates/ide/src/goto_definition.rs b/crates/ide/src/goto_definition.rs
index 645b405e5e8..986ef4cf9b3 100644
--- a/crates/ide/src/goto_definition.rs
+++ b/crates/ide/src/goto_definition.rs
@@ -39,48 +39,65 @@ pub(crate) fn goto_definition(
             kind if kind.is_trivia() => 0,
             _ => 1,
         })?;
-    let token = sema.descend_into_macros(original_token.clone());
-    let parent = token.parent()?;
-    if let Some(_) = ast::Comment::cast(token.clone()) {
-        let (attributes, def) = doc_attributes(&sema, &parent)?;
-        let (docs, doc_mapping) = attributes.docs_with_rangemap(db)?;
-        let (_, link, ns) =
-            extract_definitions_from_docs(&docs).into_iter().find(|&(range, ..)| {
-                doc_mapping.map(range).map_or(false, |InFile { file_id, value: range }| {
-                    file_id == position.file_id.into() && range.contains(position.offset)
-                })
-            })?;
-        let nav = resolve_doc_path_for_def(db, def, &link, ns)?.try_to_nav(db)?;
-        return Some(RangeInfo::new(original_token.text_range(), vec![nav]));
-    }
-
-    let navs = match_ast! {
-        match parent {
-            ast::NameRef(name_ref) => {
-                reference_definition(&sema, Either::Right(&name_ref))
-            },
-            ast::Name(name) => {
-                match NameClass::classify(&sema, &name)? {
-                    NameClass::Definition(def) | NameClass::ConstReference(def) => {
-                        try_find_trait_item_definition(sema.db, &def).unwrap_or_else(|| def_to_nav(sema.db, def))
-                    }
-                    NameClass::PatFieldShorthand { local_def, field_ref } => {
-                        local_and_field_to_nav(sema.db, local_def, field_ref)
+    let tokens = sema.descend_into_macros_many(original_token.clone());
+    let navs = tokens
+        .clone()
+        .into_iter()
+        .filter_map(|token| {
+            let parent = token.parent()?;
+            if let Some(_) = ast::Comment::cast(token.clone()) {
+                let (attributes, def) = doc_attributes(&sema, &parent)?;
+                let (docs, doc_mapping) = attributes.docs_with_rangemap(db)?;
+                let (_, link, ns) =
+                    extract_definitions_from_docs(&docs).into_iter().find(|&(range, ..)| {
+                        doc_mapping.map(range).map_or(false, |InFile { file_id, value: range }| {
+                            file_id == position.file_id.into() && range.contains(position.offset)
+                        })
+                    })?;
+                let nav = resolve_doc_path_for_def(db, def, &link, ns)?.try_to_nav(db)?;
+                return Some(nav);
+            }
+            None
+        })
+        .collect::<Vec<NavigationTarget>>();
+    if navs.len() > 0 {
+        return Some(RangeInfo::new(original_token.text_range(), navs));
+    }
+    let navs = tokens
+        .into_iter()
+        .filter_map(|token| {
+            let parent = token.parent()?;
+            let navs = match_ast! {
+                match parent {
+                    ast::NameRef(name_ref) => {
+                        reference_definition(&sema, Either::Right(&name_ref))
                     },
-                }
-            },
-            ast::Lifetime(lt) => if let Some(name_class) = NameClass::classify_lifetime(&sema, &lt) {
-                match name_class {
-                    NameClass::Definition(def) => def_to_nav(sema.db, def),
+                    ast::Name(name) => {
+                        match NameClass::classify(&sema, &name)? {
+                            NameClass::Definition(def) | NameClass::ConstReference(def) => {
+                                try_find_trait_item_definition(sema.db, &def).unwrap_or_else(|| def_to_nav(sema.db, def))
+                            }
+                            NameClass::PatFieldShorthand { local_def, field_ref } => {
+                                local_and_field_to_nav(sema.db, local_def, field_ref)
+                            },
+                        }
+                    },
+                    ast::Lifetime(lt) => if let Some(name_class) = NameClass::classify_lifetime(&sema, &lt) {
+                        match name_class {
+                            NameClass::Definition(def) => def_to_nav(sema.db, def),
+                            _ => return None,
+                        }
+                    } else {
+                        reference_definition(&sema, Either::Left(&lt))
+                    },
+                    ast::TokenTree(tt) => try_lookup_include_path_or_derive(&sema, tt, token, position.file_id)?,
                     _ => return None,
                 }
-            } else {
-                reference_definition(&sema, Either::Left(&lt))
-            },
-            ast::TokenTree(tt) => try_lookup_include_path_or_derive(&sema, tt, token, position.file_id)?,
-            _ => return None,
-        }
-    };
+            };
+            Some(navs)
+        })
+        .flatten()
+        .collect::<Vec<NavigationTarget>>();
 
     Some(RangeInfo::new(original_token.text_range(), navs))
 }