feat: Support for GOTO def from *inside* files included with include! macro

7 files changed, 163 insertions, 12 deletions

diff --git a/Cargo.lock b/Cargo.lock
index a743d1c870a..b626a16947b 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -483,6 +483,7 @@ dependencies = [
  "profile",
  "rustc-hash",
  "smallvec",
+ "span",
  "stdx",
  "syntax",
  "triomphe",
diff --git a/crates/hir-def/src/db.rs b/crates/hir-def/src/db.rs
index c9789ceb207..f773fe4c6ab 100644
--- a/crates/hir-def/src/db.rs
+++ b/crates/hir-def/src/db.rs
@@ -1,9 +1,10 @@
 //! Defines database & queries for name resolution.
-use base_db::{salsa, CrateId, SourceDatabase, Upcast};
+use base_db::{salsa, CrateId, FileId, SourceDatabase, Upcast};
 use either::Either;
 use hir_expand::{db::ExpandDatabase, HirFileId, MacroDefId};
 use intern::Interned;
 use la_arena::ArenaMap;
+use span::MacroCallId;
 use syntax::{ast, AstPtr};
 use triomphe::Arc;
 
@@ -234,6 +235,22 @@ pub trait DefDatabase: InternDatabase + ExpandDatabase + Upcast<dyn ExpandDataba
     fn crate_notable_traits(&self, krate: CrateId) -> Option<Arc<[TraitId]>>;
 
     fn crate_supports_no_std(&self, crate_id: CrateId) -> bool;
+
+    fn include_macro_invoc(&self, crate_id: CrateId) -> Vec<(MacroCallId, FileId)>;
+}
+
+// return: macro call id and include file id
+fn include_macro_invoc(db: &dyn DefDatabase, krate: CrateId) -> Vec<(MacroCallId, FileId)> {
+    db.crate_def_map(krate)
+        .modules
+        .values()
+        .flat_map(|m| m.scope.iter_macro_invoc())
+        .filter_map(|invoc| {
+            db.lookup_intern_macro_call(*invoc.1)
+                .include_file_id(db.upcast(), *invoc.1)
+                .map(|x| (*invoc.1, x))
+        })
+        .collect()
 }
 
 fn crate_def_map_wait(db: &dyn DefDatabase, krate: CrateId) -> Arc<DefMap> {
diff --git a/crates/hir-def/src/item_scope.rs b/crates/hir-def/src/item_scope.rs
index 168ee4acffb..6237ea7353f 100644
--- a/crates/hir-def/src/item_scope.rs
+++ b/crates/hir-def/src/item_scope.rs
@@ -336,6 +336,12 @@ impl ItemScope {
     pub(crate) fn macro_invoc(&self, call: AstId<ast::MacroCall>) -> Option<MacroCallId> {
         self.macro_invocations.get(&call).copied()
     }
+
+    pub(crate) fn iter_macro_invoc(
+        &self,
+    ) -> impl Iterator<Item = (&AstId<ast::MacroCall>, &MacroCallId)> {
+        self.macro_invocations.iter()
+    }
 }
 
 impl ItemScope {
diff --git a/crates/hir-expand/src/lib.rs b/crates/hir-expand/src/lib.rs
index bd216ccca82..b7d59791b24 100644
--- a/crates/hir-expand/src/lib.rs
+++ b/crates/hir-expand/src/lib.rs
@@ -520,6 +520,24 @@ impl MacroCallLoc {
             }
         }
     }
+
+    pub fn include_file_id(
+        &self,
+        db: &dyn ExpandDatabase,
+        macro_call_id: MacroCallId,
+    ) -> Option<FileId> {
+        if self.def.is_include() {
+            if let Some(eager) = &self.eager {
+                if let Ok(it) =
+                    builtin_fn_macro::include_input_to_file_id(db, macro_call_id, &eager.arg)
+                {
+                    return Some(it);
+                }
+            }
+        }
+
+        None
+    }
 }
 
 impl MacroCallKind {
diff --git a/crates/hir/Cargo.toml b/crates/hir/Cargo.toml
index e4e4bcea610..668a14cd55e 100644
--- a/crates/hir/Cargo.toml
+++ b/crates/hir/Cargo.toml
@@ -30,9 +30,10 @@ profile.workspace = true
 stdx.workspace = true
 syntax.workspace = true
 tt.workspace = true
+span.workspace = true
 
 [features]
 in-rust-tree = []
 
 [lints]
-workspace = true
\ No newline at end of file
+workspace = true
diff --git a/crates/hir/src/semantics.rs b/crates/hir/src/semantics.rs
index fdb94a6d5a7..46f3997620a 100644
--- a/crates/hir/src/semantics.rs
+++ b/crates/hir/src/semantics.rs
@@ -20,11 +20,12 @@ use hir_def::{
 };
 use hir_expand::{
     attrs::collect_attrs, db::ExpandDatabase, files::InRealFile, name::AsName, ExpansionInfo,
-    InMacroFile, MacroCallId, MacroFileId, MacroFileIdExt,
+    HirFileIdExt, InMacroFile, MacroCallId, MacroFileId, MacroFileIdExt,
 };
 use itertools::Itertools;
 use rustc_hash::{FxHashMap, FxHashSet};
 use smallvec::{smallvec, SmallVec};
+use span::Span;
 use stdx::TupleExt;
 use syntax::{
     algo::skip_trivia_token,
@@ -607,22 +608,97 @@ impl<'db> SemanticsImpl<'db> {
         res
     }
 
-    fn descend_into_macros_impl(
+    // return:
+    // SourceAnalyzer(file_id that original call include!)
+    // macro file id
+    // token in include! macro mapped from token in params
+    // span for the mapped token
+    fn is_from_include_file(
         &self,
         token: SyntaxToken,
+    ) -> Option<(SourceAnalyzer, HirFileId, SyntaxToken, Span)> {
+        let parent = token.parent()?;
+        let file_id = self.find_file(&parent).file_id.file_id()?;
+
+        // iterate related crates and find all include! invocations that include_file_id matches
+        for (invoc, _) in self
+            .db
+            .relevant_crates(file_id)
+            .iter()
+            .flat_map(|krate| self.db.include_macro_invoc(*krate))
+            .filter(|(_, include_file_id)| *include_file_id == file_id)
+        {
+            // find file_id which original calls include!
+            let Some(callnode) = invoc.as_file().original_call_node(self.db.upcast()) else {
+                continue;
+            };
+
+            // call .parse to avoid panic in .find_file
+            let _ = self.parse(callnode.file_id);
+            let Some(sa) = self.analyze_no_infer(&callnode.value) else { continue };
+
+            let expinfo = invoc.as_macro_file().expansion_info(self.db.upcast());
+            {
+                let InMacroFile { file_id, value } = expinfo.expanded();
+                self.cache(value, file_id.into());
+            }
+
+            // map token to the corresponding span in include! macro file
+            let Some((_, span)) =
+                expinfo.exp_map.iter().find(|(_, x)| x.range == token.text_range())
+            else {
+                continue;
+            };
+
+            // get mapped token in the include! macro file
+            let Some(InMacroFile { file_id: _, value: mapped_tokens }) =
+                expinfo.map_range_down(span)
+            else {
+                continue;
+            };
+
+            // if we find one, then return
+            if let Some(t) = mapped_tokens.into_iter().next() {
+                return Some((sa, invoc.as_file(), t, span));
+            };
+        }
+
+        None
+    }
+
+    fn descend_into_macros_impl(
+        &self,
+        mut token: SyntaxToken,
         f: &mut dyn FnMut(InFile<SyntaxToken>) -> ControlFlow<()>,
     ) {
         let _p = profile::span("descend_into_macros");
+
+        let mut include_macro_file_id_and_span = None;
+
         let sa = match token.parent().and_then(|parent| self.analyze_no_infer(&parent)) {
             Some(it) => it,
-            None => return,
+            None => {
+                // if we cannot find a source analyzer for this token, then we try to find out whether this file is included from other file
+                let Some((it, macro_file_id, mapped_token, s)) = self.is_from_include_file(token)
+                else {
+                    return;
+                };
+
+                include_macro_file_id_and_span = Some((macro_file_id, s));
+                token = mapped_token;
+                it
+            }
         };
 
-        let span = match sa.file_id.file_id() {
-            Some(file_id) => self.db.real_span_map(file_id).span_for_range(token.text_range()),
-            None => {
-                stdx::never!();
-                return;
+        let span = if let Some((_, s)) = include_macro_file_id_and_span {
+            s
+        } else {
+            match sa.file_id.file_id() {
+                Some(file_id) => self.db.real_span_map(file_id).span_for_range(token.text_range()),
+                None => {
+                    stdx::never!();
+                    return;
+                }
             }
         };
 
@@ -630,6 +706,13 @@ impl<'db> SemanticsImpl<'db> {
         let mut mcache = self.macro_call_cache.borrow_mut();
         let def_map = sa.resolver.def_map();
 
+        let mut stack: Vec<(_, SmallVec<[_; 2]>)> =
+            if let Some((macro_file_id, _)) = include_macro_file_id_and_span {
+                vec![(macro_file_id, smallvec![token])]
+            } else {
+                vec![(sa.file_id, smallvec![token])]
+            };
+
         let mut process_expansion_for_token = |stack: &mut Vec<_>, macro_file| {
             let expansion_info = cache
                 .entry(macro_file)
@@ -651,8 +734,6 @@ impl<'db> SemanticsImpl<'db> {
             res
         };
 
-        let mut stack: Vec<(_, SmallVec<[_; 2]>)> = vec![(sa.file_id, smallvec![token])];
-
         while let Some((file_id, mut tokens)) = stack.pop() {
             while let Some(token) = tokens.pop() {
                 let was_not_remapped = (|| {
diff --git a/crates/ide/src/goto_definition.rs b/crates/ide/src/goto_definition.rs
index d64295bdd69..073e0036182 100644
--- a/crates/ide/src/goto_definition.rs
+++ b/crates/ide/src/goto_definition.rs
@@ -226,6 +226,7 @@ mod tests {
             .map(|(FileRange { file_id, range }, _)| FileRange { file_id, range })
             .sorted_by_key(cmp)
             .collect::<Vec<_>>();
+
         assert_eq!(expected, navs);
     }
 
@@ -237,6 +238,32 @@ mod tests {
     }
 
     #[test]
+    fn goto_def_in_included_file() {
+        check(
+            r#"
+//- /main.rs
+#[rustc_builtin_macro]
+macro_rules! include {}
+
+include!("a.rs");
+
+fn main() {
+    foo();
+}
+
+//- /a.rs
+fn func_in_include() {
+ //^^^^^^^^^^^^^^^
+}
+
+fn foo() {
+    func_in_include$0();
+}
+"#,
+        );
+    }
+
+    #[test]
     fn goto_def_if_items_same_name() {
         check(
             r#"