Fix eager token mapping panics

crates/hir-expand/src/db.rs

@@ -308,7 +308,7 @@ fn macro_arg(
 ) -> Option<Arc<(tt::Subtree, mbe::TokenMap, fixup::SyntaxFixupUndoInfo)>> {
     let loc = db.lookup_intern_macro_call(id);
 
-    if let Some(EagerCallInfo { arg, arg_id: Some(_), error: _ }) = loc.eager.as_deref() {
+    if let Some(EagerCallInfo { arg, arg_id: _, error: _ }) = loc.eager.as_deref() {
         return Some(Arc::new((arg.0.clone(), arg.1.clone(), Default::default())));
     }
 
@@ -371,7 +371,7 @@ fn censor_for_macro_input(loc: &MacroCallLoc, node: &SyntaxNode) -> FxHashSet<Sy
 
 fn macro_arg_text(db: &dyn ExpandDatabase, id: MacroCallId) -> Option<GreenNode> {
     let loc = db.lookup_intern_macro_call(id);
-    let arg = loc.kind.arg(db)?;
+    let arg = loc.kind.arg(db)?.value;
     if matches!(loc.kind, MacroCallKind::FnLike { .. }) {
         let first = arg.first_child_or_token().map_or(T![.], |it| it.kind());
         let last = arg.last_child_or_token().map_or(T![.], |it| it.kind());
@@ -446,8 +446,14 @@ fn macro_def(
 fn macro_expand(db: &dyn ExpandDatabase, id: MacroCallId) -> ExpandResult<Arc<tt::Subtree>> {
     let _p = profile::span("macro_expand");
     let loc = db.lookup_intern_macro_call(id);
+
+    // This might look a bit odd, but we do not expand the inputs to eager macros here.
+    // Eager macros inputs are expanded, well, eagerly when we collect the macro calls.
+    // That kind of expansion uses the ast id map of an eager macros input though which goes through
+    // the HirFileId machinery. As eager macro inputs are assigned a macro file id that query
+    // will end up going through here again, whereas we want to just want to inspect the raw input.
+    // As such we just return the input subtree here.
     if let Some(EagerCallInfo { arg, arg_id: None, error }) = loc.eager.as_deref() {
-        // This is an input expansion for an eager macro. These are already pre-expanded
         return ExpandResult { value: Arc::new(arg.0.clone()), err: error.clone() };
     }
 

crates/hir-expand/src/eager.rs

@@ -67,26 +67,37 @@ pub fn expand_eager_macro_input(
         })),
         kind: MacroCallKind::FnLike { ast_id: call_id, expand_to: ExpandTo::Expr },
     });
-    let arg_as_expr = match db.macro_arg_text(arg_id) {
-        Some(it) => it,
-        None => {
-            return Ok(ExpandResult {
-                value: None,
-                err: Some(ExpandError::other("invalid token tree")),
-            })
-        }
+
+    let ExpandResult { value: expanded_eager_input, err } = {
+        let arg_as_expr = match db.macro_arg_text(arg_id) {
+            Some(it) => it,
+            None => {
+                return Ok(ExpandResult {
+                    value: None,
+                    err: Some(ExpandError::other("invalid token tree")),
+                })
+            }
+        };
+
+        eager_macro_recur(
+            db,
+            &Hygiene::new(db, macro_call.file_id),
+            InFile::new(arg_id.as_file(), SyntaxNode::new_root(arg_as_expr)),
+            krate,
+            resolver,
+        )?
     };
-    let ExpandResult { value: expanded_eager_input, err } = eager_macro_recur(
-        db,
-        &Hygiene::new(db, macro_call.file_id),
-        InFile::new(arg_id.as_file(), SyntaxNode::new_root(arg_as_expr)),
-        krate,
-        resolver,
-    )?;
+
     let Some(expanded_eager_input) = expanded_eager_input else {
         return Ok(ExpandResult { value: None, err });
     };
-    let (mut subtree, token_map) = mbe::syntax_node_to_token_tree(&expanded_eager_input);
+    // FIXME: This token map is pointless, it points into the expanded eager syntax node, but that
+    // node doesn't exist outside this function so we can't use this tokenmap.
+    // Ideally we'd need to patch the tokenmap of the pre-expanded input and then put that here
+    // or even better, forego expanding into a SyntaxNode altogether and instead construct a subtree
+    // in place! But that is kind of difficult.
+    let (mut subtree, _token_map) = mbe::syntax_node_to_token_tree(&expanded_eager_input);
+    let token_map = Default::default();
     subtree.delimiter = crate::tt::Delimiter::unspecified();
 
     let loc = MacroCallLoc {

crates/hir-expand/src/hygiene.rs

@@ -149,16 +149,12 @@ impl HygieneInfo {
                     token_id = unshifted;
                     (&attr_args.1, self.attr_input_or_mac_def_start?)
                 }
-                None => (
-                    &self.macro_arg.1,
-                    InFile::new(loc.kind.file_id(), loc.kind.arg(db)?.text_range().start()),
-                ),
+                None => (&self.macro_arg.1, loc.kind.arg(db)?.map(|it| it.text_range().start())),
             },
             _ => match origin {
-                mbe::Origin::Call => (
-                    &self.macro_arg.1,
-                    InFile::new(loc.kind.file_id(), loc.kind.arg(db)?.text_range().start()),
-                ),
+                mbe::Origin::Call => {
+                    (&self.macro_arg.1, loc.kind.arg(db)?.map(|it| it.text_range().start()))
+                }
                 mbe::Origin::Def => match (&*self.macro_def, &self.attr_input_or_mac_def_start) {
                     (TokenExpander::DeclarativeMacro { def_site_token_map, .. }, Some(tt)) => {
                         (def_site_token_map, *tt)

crates/hir-expand/src/lib.rs

@@ -154,8 +154,9 @@ pub enum MacroDefKind {
 struct EagerCallInfo {
     /// NOTE: This can be *either* the expansion result, *or* the argument to the eager macro!
     arg: Arc<(tt::Subtree, TokenMap)>,
-    /// call id of the eager macro's input file. If this is none, macro call containing this call info
-    /// is an eager macro's input, otherwise it is its output.
+    /// Call id of the eager macro's input file (this is the macro file for its fully expanded input).
+    /// If this is none, `arg` contains the pre-expanded input, otherwise arg contains the
+    /// post-expanded input.
     arg_id: Option<MacroCallId>,
     error: Option<ExpandError>,
 }
@@ -270,53 +271,7 @@ impl HirFileId {
     /// Return expansion information if it is a macro-expansion file
     pub fn expansion_info(self, db: &dyn db::ExpandDatabase) -> Option<ExpansionInfo> {
         let macro_file = self.macro_file()?;
-        let loc: MacroCallLoc = db.lookup_intern_macro_call(macro_file.macro_call_id);
-
-        let arg_tt = loc.kind.arg(db)?;
-
-        let macro_def = db.macro_def(loc.def).ok()?;
-        let (parse, exp_map) = db.parse_macro_expansion(macro_file).value;
-        let macro_arg = db.macro_arg(macro_file.macro_call_id).unwrap_or_else(|| {
-            Arc::new((
-                tt::Subtree { delimiter: tt::Delimiter::UNSPECIFIED, token_trees: Vec::new() },
-                Default::default(),
-                Default::default(),
-            ))
-        });
-
-        let def = loc.def.ast_id().left().and_then(|id| {
-            let def_tt = match id.to_node(db) {
-                ast::Macro::MacroRules(mac) => mac.token_tree()?,
-                ast::Macro::MacroDef(_) if matches!(*macro_def, TokenExpander::BuiltinAttr(_)) => {
-                    return None
-                }
-                ast::Macro::MacroDef(mac) => mac.body()?,
-            };
-            Some(InFile::new(id.file_id, def_tt))
-        });
-        let attr_input_or_mac_def = def.or_else(|| match loc.kind {
-            MacroCallKind::Attr { ast_id, invoc_attr_index, .. } => {
-                // FIXME: handle `cfg_attr`
-                let tt = ast_id
-                    .to_node(db)
-                    .doc_comments_and_attrs()
-                    .nth(invoc_attr_index.ast_index())
-                    .and_then(Either::left)?
-                    .token_tree()?;
-                Some(InFile::new(ast_id.file_id, tt))
-            }
-            _ => None,
-        });
-
-        Some(ExpansionInfo {
-            expanded: InFile::new(self, parse.syntax_node()),
-            arg: InFile::new(loc.kind.file_id(), arg_tt),
-            attr_input_or_mac_def,
-            macro_arg_shift: mbe::Shift::new(&macro_arg.0),
-            macro_arg,
-            macro_def,
-            exp_map,
-        })
+        ExpansionInfo::new(db, macro_file)
     }
 
     pub fn as_builtin_derive_attr_node(
@@ -603,13 +558,18 @@ impl MacroCallKind {
         FileRange { range, file_id }
     }
 
-    fn arg(&self, db: &dyn db::ExpandDatabase) -> Option<SyntaxNode> {
+    fn arg(&self, db: &dyn db::ExpandDatabase) -> Option<InFile<SyntaxNode>> {
         match self {
-            MacroCallKind::FnLike { ast_id, .. } => {
-                Some(ast_id.to_node(db).token_tree()?.syntax().clone())
+            MacroCallKind::FnLike { ast_id, .. } => ast_id
+                .to_in_file_node(db)
+                .map(|it| Some(it.token_tree()?.syntax().clone()))
+                .transpose(),
+            MacroCallKind::Derive { ast_id, .. } => {
+                Some(ast_id.to_in_file_node(db).syntax().cloned())
+            }
+            MacroCallKind::Attr { ast_id, .. } => {
+                Some(ast_id.to_in_file_node(db).syntax().cloned())
             }
-            MacroCallKind::Derive { ast_id, .. } => Some(ast_id.to_node(db).syntax().clone()),
-            MacroCallKind::Attr { ast_id, .. } => Some(ast_id.to_node(db).syntax().clone()),
         }
     }
 }
@@ -627,7 +587,7 @@ impl MacroCallId {
 /// ExpansionInfo mainly describes how to map text range between src and expanded macro
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct ExpansionInfo {
-    expanded: InFile<SyntaxNode>,
+    expanded: InMacroFile<SyntaxNode>,
     /// The argument TokenTree or item for attributes
     arg: InFile<SyntaxNode>,
     /// The `macro_rules!` or attribute input.
@@ -643,7 +603,7 @@ pub struct ExpansionInfo {
 
 impl ExpansionInfo {
     pub fn expanded(&self) -> InFile<SyntaxNode> {
-        self.expanded.clone()
+        self.expanded.clone().into()
     }
 
     pub fn call_node(&self) -> Option<InFile<SyntaxNode>> {
@@ -674,7 +634,7 @@ impl ExpansionInfo {
         let token_id_in_attr_input = if let Some(item) = item {
             // check if we are mapping down in an attribute input
             // this is a special case as attributes can have two inputs
-            let call_id = self.expanded.file_id.macro_file()?.macro_call_id;
+            let call_id = self.expanded.file_id.macro_call_id;
             let loc = db.lookup_intern_macro_call(call_id);
 
             let token_range = token.value.text_range();
@@ -720,7 +680,7 @@ impl ExpansionInfo {
                 let relative_range =
                     token.value.text_range().checked_sub(self.arg.value.text_range().start())?;
                 let token_id = self.macro_arg.1.token_by_range(relative_range)?;
-                // conditionally shift the id by a declaratives macro definition
+                // conditionally shift the id by a declarative macro definition
                 self.macro_def.map_id_down(token_id)
             }
         };
@@ -730,7 +690,7 @@ impl ExpansionInfo {
             .ranges_by_token(token_id, token.value.kind())
             .flat_map(move |range| self.expanded.value.covering_element(range).into_token());
 
-        Some(tokens.map(move |token| self.expanded.with_value(token)))
+        Some(tokens.map(move |token| InFile::new(self.expanded.file_id.into(), token)))
     }
 
     /// Map a token up out of the expansion it resides in into the arguments of the macro call of the expansion.
@@ -739,12 +699,13 @@ impl ExpansionInfo {
         db: &dyn db::ExpandDatabase,
         token: InFile<&SyntaxToken>,
     ) -> Option<(InFile<SyntaxToken>, Origin)> {
+        assert_eq!(token.file_id, self.expanded.file_id.into());
         // Fetch the id through its text range,
         let token_id = self.exp_map.token_by_range(token.value.text_range())?;
         // conditionally unshifting the id to accommodate for macro-rules def site
         let (mut token_id, origin) = self.macro_def.map_id_up(token_id);
 
-        let call_id = self.expanded.file_id.macro_file()?.macro_call_id;
+        let call_id = self.expanded.file_id.macro_call_id;
         let loc = db.lookup_intern_macro_call(call_id);
 
         // Special case: map tokens from `include!` expansions to the included file
@@ -794,6 +755,63 @@ impl ExpansionInfo {
             tt.value.covering_element(range + tt.value.text_range().start()).into_token()?;
         Some((tt.with_value(token), origin))
     }
+
+    fn new(db: &dyn db::ExpandDatabase, macro_file: MacroFile) -> Option<ExpansionInfo> {
+        let loc: MacroCallLoc = db.lookup_intern_macro_call(macro_file.macro_call_id);
+
+        let arg_tt = loc.kind.arg(db)?;
+
+        let macro_def = db.macro_def(loc.def).ok()?;
+        let (parse, exp_map) = db.parse_macro_expansion(macro_file).value;
+        let expanded = InMacroFile { file_id: macro_file, value: parse.syntax_node() };
+
+        let macro_arg = db
+            .macro_arg(match loc.eager.as_deref() {
+                Some(&EagerCallInfo { arg_id: Some(_), .. }) => return None,
+                _ => macro_file.macro_call_id,
+            })
+            .unwrap_or_else(|| {
+                Arc::new((
+                    tt::Subtree { delimiter: tt::Delimiter::UNSPECIFIED, token_trees: Vec::new() },
+                    Default::default(),
+                    Default::default(),
+                ))
+            });
+
+        let def = loc.def.ast_id().left().and_then(|id| {
+            let def_tt = match id.to_node(db) {
+                ast::Macro::MacroRules(mac) => mac.token_tree()?,
+                ast::Macro::MacroDef(_) if matches!(*macro_def, TokenExpander::BuiltinAttr(_)) => {
+                    return None
+                }
+                ast::Macro::MacroDef(mac) => mac.body()?,
+            };
+            Some(InFile::new(id.file_id, def_tt))
+        });
+        let attr_input_or_mac_def = def.or_else(|| match loc.kind {
+            MacroCallKind::Attr { ast_id, invoc_attr_index, .. } => {
+                // FIXME: handle `cfg_attr`
+                let tt = ast_id
+                    .to_node(db)
+                    .doc_comments_and_attrs()
+                    .nth(invoc_attr_index.ast_index())
+                    .and_then(Either::left)?
+                    .token_tree()?;
+                Some(InFile::new(ast_id.file_id, tt))
+            }
+            _ => None,
+        });
+
+        Some(ExpansionInfo {
+            expanded,
+            arg: arg_tt,
+            attr_input_or_mac_def,
+            macro_arg_shift: mbe::Shift::new(&macro_arg.0),
+            macro_arg,
+            macro_def,
+            exp_map,
+        })
+    }
 }
 
 /// `AstId` points to an AST node in any file.
@@ -805,6 +823,9 @@ impl<N: AstIdNode> AstId<N> {
     pub fn to_node(&self, db: &dyn db::ExpandDatabase) -> N {
         self.to_ptr(db).to_node(&db.parse_or_expand(self.file_id))
     }
+    pub fn to_in_file_node(&self, db: &dyn db::ExpandDatabase) -> InFile<N> {
+        InFile::new(self.file_id, self.to_ptr(db).to_node(&db.parse_or_expand(self.file_id)))
+    }
     pub fn to_ptr(&self, db: &dyn db::ExpandDatabase) -> AstPtr<N> {
         db.ast_id_map(self.file_id).get(self.value)
     }
@@ -820,6 +841,7 @@ impl ErasedAstId {
         db.ast_id_map(self.file_id).get_raw(self.value)
     }
 }
+
 /// `InFile<T>` stores a value of `T` inside a particular file/syntax tree.
 ///
 /// Typical usages are:
@@ -1038,6 +1060,18 @@ impl InFile<SyntaxToken> {
     }
 }
 
+#[derive(Debug, PartialEq, Eq, Clone, Copy, Hash)]
+pub struct InMacroFile<T> {
+    pub file_id: MacroFile,
+    pub value: T,
+}
+
+impl<T> From<InMacroFile<T>> for InFile<T> {
+    fn from(macro_file: InMacroFile<T>) -> Self {
+        InFile { file_id: macro_file.file_id.into(), value: macro_file.value }
+    }
+}
+
 fn ascend_node_border_tokens(
     db: &dyn db::ExpandDatabase,
     InFile { file_id, value: node }: InFile<&SyntaxNode>,

crates/ide/src/syntax_highlighting.rs

@@ -265,10 +265,14 @@ fn traverse(
 
         // set macro and attribute highlighting states
         match event.clone() {
-            Enter(NodeOrToken::Node(node)) if ast::TokenTree::can_cast(node.kind()) => {
+            Enter(NodeOrToken::Node(node))
+                if current_macro.is_none() && ast::TokenTree::can_cast(node.kind()) =>
+            {
                 tt_level += 1;
             }
-            Leave(NodeOrToken::Node(node)) if ast::TokenTree::can_cast(node.kind()) => {
+            Leave(NodeOrToken::Node(node))
+                if current_macro.is_none() && ast::TokenTree::can_cast(node.kind()) =>
+            {
                 tt_level -= 1;
             }
             Enter(NodeOrToken::Node(node)) if ast::Attr::can_cast(node.kind()) => {
@@ -387,7 +391,7 @@ fn traverse(
             };
         let descended_element = if in_macro {
             // Attempt to descend tokens into macro-calls.
-            match element {
+            let res = match element {
                 NodeOrToken::Token(token) if token.kind() != COMMENT => {
                     let token = match attr_or_derive_item {
                         Some(AttrOrDerive::Attr(_)) => {
@@ -412,7 +416,8 @@ fn traverse(
                     }
                 }
                 e => e,
-            }
+            };
+            res
         } else {
             element
         };