Map attribute input tokens correctly

Veykril · Veykril · commit 177c70128ccc · 2021-08-21T18:13:41.000+02:00
diff --git a/crates/hir/src/semantics.rs b/crates/hir/src/semantics.rs
@@ -441,7 +441,7 @@ impl<'db> SemanticsImpl<'db> {
                                 .entry(file_id)
                                 .or_insert_with(|| file_id.expansion_info(self.db.upcast()))
                                 .as_ref()?
-                                .map_token_down(token.as_ref())?;
+                                .map_token_down(self.db.upcast(), None, token.as_ref())?;
 
                             if let Some(parent) = token.value.parent() {
                                 self.cache(find_root(&parent), token.file_id);
@@ -450,24 +450,21 @@ impl<'db> SemanticsImpl<'db> {
                             return Some(token);
                         },
                         ast::Item(item) => {
-                            match self.with_ctx(|ctx| ctx.item_to_macro_call(token.with_value(item))) {
-                                Some(call_id) => {
-                                    let file_id = call_id.as_file();
-                                    let token = self
-                                        .expansion_info_cache
-                                        .borrow_mut()
-                                        .entry(file_id)
-                                        .or_insert_with(|| file_id.expansion_info(self.db.upcast()))
-                                        .as_ref()?
-                                        .map_token_down(token.as_ref())?;
-
-                                    if let Some(parent) = token.value.parent() {
-                                        self.cache(find_root(&parent), token.file_id);
-                                    }
-
-                                    return Some(token);
+                            if let Some(call_id) = self.with_ctx(|ctx| ctx.item_to_macro_call(token.with_value(item.clone()))) {
+                                let file_id = call_id.as_file();
+                                let token = self
+                                    .expansion_info_cache
+                                    .borrow_mut()
+                                    .entry(file_id)
+                                    .or_insert_with(|| file_id.expansion_info(self.db.upcast()))
+                                    .as_ref()?
+                                    .map_token_down(self.db.upcast(), Some(item), token.as_ref())?;
+
+                                if let Some(parent) = token.value.parent() {
+                                    self.cache(find_root(&parent), token.file_id);
                                 }
-                                None => {}
+
+                                return Some(token);
                             }
                         },
                         _ => {}
@@ -479,7 +476,6 @@ impl<'db> SemanticsImpl<'db> {
         })
         .last()
         .unwrap();
-
         token.value
     }
 
diff --git a/crates/hir_def/src/attr.rs b/crates/hir_def/src/attr.rs
@@ -2,7 +2,9 @@
 
 use std::{
     convert::{TryFrom, TryInto},
-    fmt, ops,
+    fmt,
+    hash::Hash,
+    ops,
     sync::Arc,
 };
 
@@ -12,7 +14,7 @@ use either::Either;
 use hir_expand::{hygiene::Hygiene, name::AsName, AstId, InFile};
 use itertools::Itertools;
 use la_arena::ArenaMap;
-use mbe::{syntax_node_to_token_tree, DelimiterKind};
+use mbe::{syntax_node_to_token_tree, DelimiterKind, MappedSubTree};
 use smallvec::{smallvec, SmallVec};
 use syntax::{
     ast::{self, AstNode, AttrsOwner},
@@ -165,11 +167,11 @@ impl RawAttrs {
                 // Input subtree is: `(cfg, $(attr),+)`
                 // Split it up into a `cfg` subtree and the `attr` subtrees.
                 // FIXME: There should be a common API for this.
-                let mut parts = subtree.token_trees.split(
+                let mut parts = subtree.tree.token_trees.split(
                     |tt| matches!(tt, tt::TokenTree::Leaf(tt::Leaf::Punct(p)) if p.char == ','),
                 );
                 let cfg = parts.next().unwrap();
-                let cfg = Subtree { delimiter: subtree.delimiter, token_trees: cfg.to_vec() };
+                let cfg = Subtree { delimiter: subtree.tree.delimiter, token_trees: cfg.to_vec() };
                 let cfg = CfgExpr::parse(&cfg);
                 let index = attr.id;
                 let attrs = parts.filter(|a| !a.is_empty()).filter_map(|attr| {
@@ -652,14 +654,14 @@ pub enum AttrInput {
     /// `#[attr = "string"]`
     Literal(SmolStr),
     /// `#[attr(subtree)]`
-    TokenTree(Subtree),
+    TokenTree(mbe::MappedSubTree),
 }
 
 impl fmt::Display for AttrInput {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         match self {
             AttrInput::Literal(lit) => write!(f, " = \"{}\"", lit.escape_debug()),
-            AttrInput::TokenTree(subtree) => subtree.fmt(f),
+            AttrInput::TokenTree(subtree) => subtree.tree.fmt(f),
         }
     }
 }
@@ -679,7 +681,8 @@ impl Attr {
             };
             Some(Interned::new(AttrInput::Literal(value)))
         } else if let Some(tt) = ast.token_tree() {
-            Some(Interned::new(AttrInput::TokenTree(syntax_node_to_token_tree(tt.syntax()).0)))
+            let (tree, map) = syntax_node_to_token_tree(tt.syntax());
+            Some(Interned::new(AttrInput::TokenTree(MappedSubTree { tree, map })))
         } else {
             None
         };
@@ -712,6 +715,7 @@ impl Attr {
             Some(AttrInput::TokenTree(args)) => {
                 let mut counter = 0;
                 let paths = args
+                    .tree
                     .token_trees
                     .iter()
                     .group_by(move |tt| {
@@ -756,7 +760,7 @@ pub struct AttrQuery<'a> {
 impl<'a> AttrQuery<'a> {
     pub fn tt_values(self) -> impl Iterator<Item = &'a Subtree> {
         self.attrs().filter_map(|attr| match attr.input.as_deref()? {
-            AttrInput::TokenTree(it) => Some(it),
+            AttrInput::TokenTree(it) => Some(&it.tree),
             _ => None,
         })
     }
diff --git a/crates/hir_def/src/lib.rs b/crates/hir_def/src/lib.rs
@@ -786,13 +786,13 @@ fn attr_macro_as_call_id(
         .ok_or_else(|| UnresolvedMacro { path: item_attr.path.clone() })?;
     let mut arg = match &macro_attr.input {
         Some(input) => match &**input {
-            attr::AttrInput::Literal(_) => tt::Subtree::default(),
+            attr::AttrInput::Literal(_) => Default::default(),
             attr::AttrInput::TokenTree(tt) => tt.clone(),
         },
-        None => tt::Subtree::default(),
+        None => Default::default(),
     };
     // The parentheses are always disposed here.
-    arg.delimiter = None;
+    arg.tree.delimiter = None;
 
     let res = def.as_lazy_macro(
         db.upcast(),
diff --git a/crates/hir_def/src/nameres/collector.rs b/crates/hir_def/src/nameres/collector.rs
@@ -289,7 +289,7 @@ impl DefCollector<'_> {
                     || *attr_name == hir_expand::name![register_tool]
                 {
                     match attr.input.as_deref() {
-                        Some(AttrInput::TokenTree(subtree)) => match &*subtree.token_trees {
+                        Some(AttrInput::TokenTree(subtree)) => match &*subtree.tree.token_trees {
                             [tt::TokenTree::Leaf(tt::Leaf::Ident(name))] => name.as_name(),
                             _ => continue,
                         },
diff --git a/crates/hir_expand/src/db.rs b/crates/hir_expand/src/db.rs
@@ -223,7 +223,7 @@ fn parse_macro_expansion(
         Ok(it) => it,
         Err(err) => {
             log::debug!(
-                "failed to parse expanstion to {:?} = {}",
+                "failed to parse expansion to {:?} = {}",
                 fragment_kind,
                 tt.as_debug_string()
             );
@@ -386,11 +386,15 @@ fn expand_proc_macro(db: &dyn AstDatabase, id: MacroCallId) -> ExpandResult<tt::
     };
 
     let attr_arg = match &loc.kind {
-        MacroCallKind::Attr { attr_args, .. } => Some(attr_args),
+        MacroCallKind::Attr { attr_args, .. } => {
+            let mut attr_args = attr_args.tree.clone();
+            mbe::Shift::new(&macro_arg.0).shift_all(&mut attr_args);
+            Some(attr_args)
+        }
         _ => None,
     };
 
-    expander.expand(db, loc.krate, &macro_arg.0, attr_arg)
+    expander.expand(db, loc.krate, &macro_arg.0, attr_arg.as_ref())
 }
 
 fn is_self_replicating(from: &SyntaxNode, to: &SyntaxNode) -> bool {
diff --git a/crates/hir_expand/src/lib.rs b/crates/hir_expand/src/lib.rs
@@ -27,7 +27,7 @@ use std::{hash::Hash, iter, sync::Arc};
 use base_db::{impl_intern_key, salsa, CrateId, FileId, FileRange};
 use syntax::{
     algo::skip_trivia_token,
-    ast::{self, AstNode},
+    ast::{self, AstNode, AttrsOwner},
     Direction, SyntaxNode, SyntaxToken, TextRange, TextSize,
 };
 
@@ -36,6 +36,7 @@ use crate::{
     builtin_attr::BuiltinAttrExpander,
     builtin_derive::BuiltinDeriveExpander,
     builtin_macro::{BuiltinFnLikeExpander, EagerExpander},
+    db::TokenExpander,
     proc_macro::ProcMacroExpander,
 };
 
@@ -132,6 +133,17 @@ impl HirFileId {
                     };
                     Some(InFile::new(id.file_id, def_tt))
                 });
+                let def_or_attr_input = def.or_else(|| match loc.kind {
+                    MacroCallKind::Attr { ast_id, invoc_attr_index, .. } => {
+                        let tt = ast_id
+                            .to_node(db)
+                            .attrs()
+                            .nth(invoc_attr_index as usize)
+                            .and_then(|attr| attr.token_tree())?;
+                        Some(InFile::new(ast_id.file_id, tt))
+                    }
+                    _ => None,
+                });
 
                 let macro_def = db.macro_def(loc.def)?;
                 let (parse, exp_map) = db.parse_macro_expansion(macro_file).value?;
@@ -140,7 +152,8 @@ impl HirFileId {
                 Some(ExpansionInfo {
                     expanded: InFile::new(self, parse.syntax_node()),
                     arg: InFile::new(loc.kind.file_id(), arg_tt),
-                    def,
+                    attr_input_or_mac_def: def_or_attr_input,
+                    macro_arg_shift: mbe::Shift::new(&macro_arg.0),
                     macro_arg,
                     macro_def,
                     exp_map,
@@ -270,7 +283,7 @@ pub enum MacroCallKind {
     Attr {
         ast_id: AstId<ast::Item>,
         attr_name: String,
-        attr_args: tt::Subtree,
+        attr_args: mbe::MappedSubTree,
         /// Syntactical index of the invoking `#[attribute]`.
         ///
         /// Outer attributes are counted first, then inner attributes. This does not support
@@ -335,11 +348,12 @@ impl MacroCallId {
 pub struct ExpansionInfo {
     expanded: InFile<SyntaxNode>,
     arg: InFile<SyntaxNode>,
-    /// The `macro_rules!` arguments.
-    def: Option<InFile<ast::TokenTree>>,
+    /// The `macro_rules!` arguments or attribute input.
+    attr_input_or_mac_def: Option<InFile<ast::TokenTree>>,
 
-    macro_def: Arc<db::TokenExpander>,
+    macro_def: Arc<TokenExpander>,
     macro_arg: Arc<(tt::Subtree, mbe::TokenMap)>,
+    macro_arg_shift: mbe::Shift,
     exp_map: Arc<mbe::TokenMap>,
 }
 
@@ -350,11 +364,53 @@ impl ExpansionInfo {
         Some(self.arg.with_value(self.arg.value.parent()?))
     }
 
-    pub fn map_token_down(&self, token: InFile<&SyntaxToken>) -> Option<InFile<SyntaxToken>> {
+    pub fn map_token_down(
+        &self,
+        db: &dyn db::AstDatabase,
+        item: Option<ast::Item>,
+        token: InFile<&SyntaxToken>,
+    ) -> Option<InFile<SyntaxToken>> {
         assert_eq!(token.file_id, self.arg.file_id);
-        let range = token.value.text_range().checked_sub(self.arg.value.text_range().start())?;
-        let token_id = self.macro_arg.1.token_by_range(range)?;
-        let token_id = self.macro_def.map_id_down(token_id);
+        let token_id = if let Some(item) = item {
+            let call_id = match self.expanded.file_id.0 {
+                HirFileIdRepr::FileId(_) => return None,
+                HirFileIdRepr::MacroFile(macro_file) => macro_file.macro_call_id,
+            };
+            let loc = db.lookup_intern_macro(call_id);
+
+            let token_range = token.value.text_range();
+            match &loc.kind {
+                MacroCallKind::Attr { attr_args, invoc_attr_index, .. } => {
+                    let attr = item.attrs().nth(*invoc_attr_index as usize)?;
+                    match attr.token_tree() {
+                        Some(token_tree)
+                            if token_tree.syntax().text_range().contains_range(token_range) =>
+                        {
+                            let attr_input_start =
+                                token_tree.left_delimiter_token()?.text_range().start();
+                            let range = token.value.text_range().checked_sub(attr_input_start)?;
+                            let token_id =
+                                self.macro_arg_shift.shift(attr_args.map.token_by_range(range)?);
+                            Some(token_id)
+                        }
+                        _ => None,
+                    }
+                }
+                _ => None,
+            }
+        } else {
+            None
+        };
+
+        let token_id = match token_id {
+            Some(token_id) => token_id,
+            None => {
+                let range =
+                    token.value.text_range().checked_sub(self.arg.value.text_range().start())?;
+                let token_id = self.macro_arg.1.token_by_range(range)?;
+                self.macro_def.map_id_down(token_id)
+            }
+        };
 
         let range = self.exp_map.range_by_token(token_id, token.value.kind())?;
 
@@ -365,20 +421,36 @@ impl ExpansionInfo {
 
     pub fn map_token_up(
         &self,
+        db: &dyn db::AstDatabase,
         token: InFile<&SyntaxToken>,
     ) -> Option<(InFile<SyntaxToken>, Origin)> {
         let token_id = self.exp_map.token_by_range(token.value.text_range())?;
+        let (mut token_id, origin) = self.macro_def.map_id_up(token_id);
 
-        let (token_id, origin) = self.macro_def.map_id_up(token_id);
-        let (token_map, tt) = match origin {
-            mbe::Origin::Call => (&self.macro_arg.1, self.arg.clone()),
-            mbe::Origin::Def => match (&*self.macro_def, self.def.as_ref()) {
-                (
-                    db::TokenExpander::MacroRules { def_site_token_map, .. }
-                    | db::TokenExpander::MacroDef { def_site_token_map, .. },
-                    Some(tt),
-                ) => (def_site_token_map, tt.syntax().cloned()),
-                _ => panic!("`Origin::Def` used with non-`macro_rules!` macro"),
+        let call_id = match self.expanded.file_id.0 {
+            HirFileIdRepr::FileId(_) => return None,
+            HirFileIdRepr::MacroFile(macro_file) => macro_file.macro_call_id,
+        };
+        let loc = db.lookup_intern_macro(call_id);
+
+        let (token_map, tt) = match &loc.kind {
+            MacroCallKind::Attr { attr_args, .. } => match self.macro_arg_shift.unshift(token_id) {
+                Some(unshifted) => {
+                    token_id = unshifted;
+                    (&attr_args.map, self.attr_input_or_mac_def.clone()?.syntax().cloned())
+                }
+                None => (&self.macro_arg.1, self.arg.clone()),
+            },
+            _ => match origin {
+                mbe::Origin::Call => (&self.macro_arg.1, self.arg.clone()),
+                mbe::Origin::Def => match (&*self.macro_def, self.attr_input_or_mac_def.as_ref()) {
+                    (
+                        TokenExpander::MacroRules { def_site_token_map, .. }
+                        | TokenExpander::MacroDef { def_site_token_map, .. },
+                        Some(tt),
+                    ) => (def_site_token_map, tt.syntax().cloned()),
+                    _ => panic!("`Origin::Def` used with non-`macro_rules!` macro"),
+                },
             },
         };
 
@@ -532,7 +604,7 @@ fn ascend_call_token(
     expansion: &ExpansionInfo,
     token: InFile<SyntaxToken>,
 ) -> Option<InFile<SyntaxToken>> {
-    let (mapped, origin) = expansion.map_token_up(token.as_ref())?;
+    let (mapped, origin) = expansion.map_token_up(db, token.as_ref())?;
     if origin != Origin::Call {
         return None;
     }
diff --git a/crates/mbe/src/lib.rs b/crates/mbe/src/lib.rs
diff --git a/crates/mbe/src/token_map.rs b/crates/mbe/src/token_map.rs

Original file line number	Diff line number	Diff line change
`@@ -289,7 +289,7 @@ impl DefCollector<'_> {`
`289`	`289`	`\|\| *attr_name == hir_expand::name![register_tool]`
`290`	`290`	`{`
`291`	`291`	`match attr.input.as_deref() {`
`292`		`- Some(AttrInput::TokenTree(subtree)) => match &*subtree.token_trees {`
	`292`	`+ Some(AttrInput::TokenTree(subtree)) => match &*subtree.tree.token_trees {`
`293`	`293`	`[tt::TokenTree::Leaf(tt::Leaf::Ident(name))] => name.as_name(),`
`294`	`294`	`_ => continue,`
`295`	`295`	`},`