Map tokens from include! expansion to the included file

Author: Jonas Schievink

crates/hir-expand/src/builtin_fn_macro.rs

@@ -3,7 +3,7 @@
 use base_db::{AnchoredPath, Edition, FileId};
 use cfg::CfgExpr;
 use either::Either;
-use mbe::{parse_exprs_with_sep, parse_to_token_tree};
+use mbe::{parse_exprs_with_sep, parse_to_token_tree, TokenMap};
 use syntax::{
     ast::{self, AstToken},
     SmolStr,
@@ -67,7 +67,7 @@ macro_rules! register_builtin {
 pub struct ExpandedEager {
     pub(crate) subtree: tt::Subtree,
     /// The included file ID of the include macro.
-    pub(crate) included_file: Option<FileId>,
+    pub(crate) included_file: Option<(FileId, TokenMap)>,
 }
 
 impl ExpandedEager {
@@ -566,14 +566,14 @@ fn include_expand(
         let path = parse_string(tt)?;
         let file_id = relative_file(db, arg_id, &path, false)?;
 
-        let subtree =
-            parse_to_token_tree(&db.file_text(file_id)).ok_or(mbe::ExpandError::ConversionError)?.0;
-        Ok((subtree, file_id))
+        let (subtree, map) =
+            parse_to_token_tree(&db.file_text(file_id)).ok_or(mbe::ExpandError::ConversionError)?;
+        Ok((subtree, map, file_id))
     })();
 
     match res {
-        Ok((subtree, file_id)) => {
-            ExpandResult::ok(ExpandedEager { subtree, included_file: Some(file_id) })
+        Ok((subtree, map, file_id)) => {
+            ExpandResult::ok(ExpandedEager { subtree, included_file: Some((file_id, map)) })
         }
         Err(e) => ExpandResult::with_err(
             ExpandedEager { subtree: tt::Subtree::empty(), included_file: None },

crates/hir-expand/src/db.rs

@@ -13,10 +13,11 @@ use syntax::{
 };
 
 use crate::{
-    ast_id_map::AstIdMap, builtin_attr_macro::pseudo_derive_attr_expansion, fixup,
-    hygiene::HygieneFrame, tt, BuiltinAttrExpander, BuiltinDeriveExpander, BuiltinFnLikeExpander,
-    ExpandError, ExpandResult, ExpandTo, HirFileId, HirFileIdRepr, MacroCallId, MacroCallKind,
-    MacroCallLoc, MacroDefId, MacroDefKind, MacroFile, ProcMacroExpander,
+    ast_id_map::AstIdMap, builtin_attr_macro::pseudo_derive_attr_expansion,
+    builtin_fn_macro::EagerExpander, fixup, hygiene::HygieneFrame, tt, BuiltinAttrExpander,
+    BuiltinDeriveExpander, BuiltinFnLikeExpander, ExpandError, ExpandResult, ExpandTo, HirFileId,
+    HirFileIdRepr, MacroCallId, MacroCallKind, MacroCallLoc, MacroDefId, MacroDefKind, MacroFile,
+    ProcMacroExpander,
 };
 
 /// Total limit on the number of tokens produced by any macro invocation.
@@ -33,6 +34,8 @@ pub enum TokenExpander {
     DeclarativeMacro { mac: mbe::DeclarativeMacro, def_site_token_map: mbe::TokenMap },
     /// Stuff like `line!` and `file!`.
     Builtin(BuiltinFnLikeExpander),
+    /// Built-in eagerly expanded fn-like macros (`include!`, `concat!`, etc.)
+    BuiltinEager(EagerExpander),
     /// `global_allocator` and such.
     BuiltinAttr(BuiltinAttrExpander),
     /// `derive(Copy)` and such.
@@ -51,6 +54,9 @@ impl TokenExpander {
         match self {
             TokenExpander::DeclarativeMacro { mac, .. } => mac.expand(tt).map_err(Into::into),
             TokenExpander::Builtin(it) => it.expand(db, id, tt).map_err(Into::into),
+            TokenExpander::BuiltinEager(it) => {
+                it.expand(db, id, tt).map_err(Into::into).map(|res| res.subtree)
+            }
             TokenExpander::BuiltinAttr(it) => it.expand(db, id, tt),
             TokenExpander::BuiltinDerive(it) => it.expand(db, id, tt),
             TokenExpander::ProcMacro(_) => {
@@ -66,6 +72,7 @@ impl TokenExpander {
         match self {
             TokenExpander::DeclarativeMacro { mac, .. } => mac.map_id_down(id),
             TokenExpander::Builtin(..)
+            | TokenExpander::BuiltinEager(..)
             | TokenExpander::BuiltinAttr(..)
             | TokenExpander::BuiltinDerive(..)
             | TokenExpander::ProcMacro(..) => id,
@@ -76,6 +83,7 @@ impl TokenExpander {
         match self {
             TokenExpander::DeclarativeMacro { mac, .. } => mac.map_id_up(id),
             TokenExpander::Builtin(..)
+            | TokenExpander::BuiltinEager(..)
             | TokenExpander::BuiltinAttr(..)
             | TokenExpander::BuiltinDerive(..)
             | TokenExpander::ProcMacro(..) => (id, mbe::Origin::Call),
@@ -412,10 +420,8 @@ fn macro_def(
         MacroDefKind::BuiltInDerive(expander, _) => {
             Ok(Arc::new(TokenExpander::BuiltinDerive(expander)))
         }
-        MacroDefKind::BuiltInEager(..) => {
-            // FIXME: Return a random error here just to make the types align.
-            // This obviously should do something real instead.
-            Err(mbe::ParseError::UnexpectedToken("unexpected eager macro".into()))
+        MacroDefKind::BuiltInEager(expander, ..) => {
+            Ok(Arc::new(TokenExpander::BuiltinEager(expander)))
         }
         MacroDefKind::ProcMacro(expander, ..) => Ok(Arc::new(TokenExpander::ProcMacro(expander))),
     }

crates/hir-expand/src/lib.rs

@@ -20,6 +20,7 @@ pub mod mod_path;
 pub mod attrs;
 mod fixup;
 
+use mbe::TokenMap;
 pub use mbe::{Origin, ValueResult};
 
 use ::tt::token_id as tt;
@@ -139,7 +140,7 @@ pub enum MacroDefKind {
 struct EagerCallInfo {
     /// NOTE: This can be *either* the expansion result, *or* the argument to the eager macro!
     arg_or_expansion: Arc<tt::Subtree>,
-    included_file: Option<FileId>,
+    included_file: Option<(FileId, TokenMap)>,
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
@@ -206,7 +207,7 @@ impl HirFileId {
                 HirFileIdRepr::MacroFile(MacroFile { macro_call_id }) => {
                     let loc: MacroCallLoc = db.lookup_intern_macro_call(macro_call_id);
                     file_id = match loc.eager {
-                        Some(EagerCallInfo { included_file: Some(file), .. }) => file.into(),
+                        Some(EagerCallInfo { included_file: Some((file, _)), .. }) => file.into(),
                         _ => loc.kind.file_id(),
                     };
                 }
@@ -319,7 +320,7 @@ impl HirFileId {
         match self.macro_file() {
             Some(macro_file) => {
                 let loc: MacroCallLoc = db.lookup_intern_macro_call(macro_file.macro_call_id);
-                matches!(loc.eager, Some(EagerCallInfo { included_file: Some(_), .. }))
+                matches!(loc.eager, Some(EagerCallInfo { included_file: Some(..), .. }))
             }
             _ => false,
         }
@@ -677,6 +678,16 @@ impl ExpansionInfo {
         let call_id = self.expanded.file_id.macro_file()?.macro_call_id;
         let loc = db.lookup_intern_macro_call(call_id);
 
+        if let Some((file, map)) = loc.eager.and_then(|e| e.included_file) {
+            // Special case: map tokens from `include!` expansions to the included file
+            let range = map.first_range_by_token(token_id, token.value.kind())?;
+            let source = db.parse(file);
+
+            let token = source.syntax_node().covering_element(range).into_token()?;
+
+            return Some((InFile::new(file.into(), token), Origin::Call));
+        }
+
         // Attributes are a bit special for us, they have two inputs, the input tokentree and the annotated item.
         let (token_map, tt) = match &loc.kind {
             MacroCallKind::Attr { attr_args, is_derive: true, .. } => {

crates/ide/src/goto_definition.rs

@@ -833,8 +833,7 @@ fn test() {
 #[rustc_builtin_macro]
 macro_rules! include {}
 
-  include!("foo.rs");
-//^^^^^^^^^^^^^^^^^^^
+include!("foo.rs");
 
 fn f() {
     foo$0();
@@ -846,6 +845,7 @@ mod confuse_index {
 
 //- /foo.rs
 fn foo() {}
+ //^^^
         "#,
         );
     }