Merge #9970

9970: feat: Implement attribute input token mapping, fix attribute item token mapping r=Veykril a=Veykril

![image](https://user-images.githubusercontent.com/3757771/130328577-4c1ad72c-51b1-47c3-8d3d-3242ec44a355.png)

The token mapping for items with attributes got overwritten partially by the attributes non-item input, since attributes have two different inputs, the item and the direct input both.
This PR gives attributes a second TokenMap for its direct input. We now shift all normal input IDs by the item input maximum(we maybe wanna swap this see below) similar to what we do for macro-rules/def. For mapping down we then have to figure out whether we are inside the direct attribute input or its item input to pick the appropriate mapping which can be done with some token range comparisons.

Fixes #9867

Co-authored-by: Lukas Wirth <[email protected]>

Author: bors[bot]

crates/hir/src/semantics.rs

@@ -474,7 +474,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);
@@ -483,24 +483,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);
                             }
                         },
                         _ => {}
@@ -512,7 +509,6 @@ impl<'db> SemanticsImpl<'db> {
         })
         .last()
         .unwrap();
-
         token.value
     }
 

crates/hir_def/src/attr.rs

@@ -2,7 +2,9 @@
 
 use std::{
     convert::{TryFrom, TryInto},
-    fmt, ops,
+    fmt,
+    hash::Hash,
+    ops,
     sync::Arc,
 };
 
@@ -158,7 +160,7 @@ impl RawAttrs {
                 }
 
                 let subtree = match attr.input.as_deref() {
-                    Some(AttrInput::TokenTree(it)) => it,
+                    Some(AttrInput::TokenTree(it, _)) => it,
                     _ => return smallvec![attr.clone()],
                 };
 
@@ -258,7 +260,7 @@ impl Attrs {
     pub fn docs(&self) -> Option<Documentation> {
         let docs = self.by_key("doc").attrs().flat_map(|attr| match attr.input.as_deref()? {
             AttrInput::Literal(s) => Some(s),
-            AttrInput::TokenTree(_) => None,
+            AttrInput::TokenTree(..) => None,
         });
         let indent = docs
             .clone()
@@ -463,7 +465,7 @@ impl AttrsWithOwner {
         // FIXME: code duplication in `docs` above
         let docs = self.by_key("doc").attrs().flat_map(|attr| match attr.input.as_deref()? {
             AttrInput::Literal(s) => Some((s, attr.id)),
-            AttrInput::TokenTree(_) => None,
+            AttrInput::TokenTree(..) => None,
         });
         let indent = docs
             .clone()
@@ -652,14 +654,14 @@ pub enum AttrInput {
     /// `#[attr = "string"]`
     Literal(SmolStr),
     /// `#[attr(subtree)]`
-    TokenTree(Subtree),
+    TokenTree(tt::Subtree, mbe::TokenMap),
 }
 
 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.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(tree, map)))
         } else {
             None
         };
@@ -709,7 +712,7 @@ impl Attr {
         }
 
         match self.input.as_deref() {
-            Some(AttrInput::TokenTree(args)) => {
+            Some(AttrInput::TokenTree(args, _)) => {
                 let mut counter = 0;
                 let paths = args
                     .token_trees
@@ -756,7 +759,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),
             _ => None,
         })
     }

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::TokenTree(tt) => tt.clone(),
+            attr::AttrInput::Literal(_) => Default::default(),
+            attr::AttrInput::TokenTree(tt, map) => (tt.clone(), map.clone()),
         },
-        None => tt::Subtree::default(),
+        None => Default::default(),
     };
     // The parentheses are always disposed here.
-    arg.delimiter = None;
+    arg.0.delimiter = None;
 
     let res = def.as_lazy_macro(
         db.upcast(),

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.token_trees {
                             [tt::TokenTree::Leaf(tt::Leaf::Ident(name))] => name.as_name(),
                             _ => continue,
                         },

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.0.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 {

crates/hir_expand/src/hygiene.rs

@@ -9,12 +9,15 @@ use db::TokenExpander;
 use either::Either;
 use mbe::Origin;
 use parser::SyntaxKind;
-use syntax::{ast, AstNode, SyntaxNode, TextRange, TextSize};
+use syntax::{
+    ast::{self, AttrsOwner},
+    AstNode, SyntaxNode, TextRange, TextSize,
+};
 
 use crate::{
     db::{self, AstDatabase},
     name::{AsName, Name},
-    HirFileId, HirFileIdRepr, InFile, MacroCallLoc, MacroDefKind, MacroFile,
+    HirFileId, HirFileIdRepr, InFile, MacroCallKind, MacroCallLoc, MacroDefKind, MacroFile,
 };
 
 #[derive(Clone, Debug)]
@@ -121,11 +124,12 @@ impl HygieneFrames {
 #[derive(Debug, Clone, PartialEq, Eq)]
 struct HygieneInfo {
     file: MacroFile,
-    /// The `macro_rules!` arguments.
-    def_start: Option<InFile<TextSize>>,
+    /// The start offset of the `macro_rules!` arguments or attribute input.
+    attr_input_or_mac_def_start: Option<InFile<TextSize>>,
 
     macro_def: Arc<TokenExpander>,
     macro_arg: Arc<(tt::Subtree, mbe::TokenMap)>,
+    macro_arg_shift: mbe::Shift,
     exp_map: Arc<mbe::TokenMap>,
 }
 
@@ -136,22 +140,34 @@ impl HygieneInfo {
         token: TextRange,
     ) -> Option<(InFile<TextRange>, Origin)> {
         let token_id = self.exp_map.token_by_range(token)?;
+        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 => {
-                let call_id = self.file.macro_call_id;
-                let loc: MacroCallLoc = db.lookup_intern_macro(call_id);
-                let arg_start = loc.kind.arg(db)?.text_range().start();
-                (&self.macro_arg.1, InFile::new(loc.kind.file_id(), arg_start))
-            }
-            mbe::Origin::Def => match (&*self.macro_def, self.def_start) {
-                (
-                    TokenExpander::MacroDef { def_site_token_map, .. }
-                    | TokenExpander::MacroRules { def_site_token_map, .. },
-                    Some(tt),
-                ) => (def_site_token_map, tt),
-                _ => panic!("`Origin::Def` used with non-`macro_rules!` macro"),
+        let loc = db.lookup_intern_macro(self.file.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.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()),
+                ),
+            },
+            _ => match origin {
+                mbe::Origin::Call => (
+                    &self.macro_arg.1,
+                    InFile::new(loc.kind.file_id(), loc.kind.arg(db)?.text_range().start()),
+                ),
+                mbe::Origin::Def => match (&*self.macro_def, &self.attr_input_or_mac_def_start) {
+                    (
+                        TokenExpander::MacroDef { def_site_token_map, .. }
+                        | TokenExpander::MacroRules { def_site_token_map, .. },
+                        Some(tt),
+                    ) => (def_site_token_map, *tt),
+                    _ => panic!("`Origin::Def` used with non-`macro_rules!` macro"),
+                },
             },
         };
 
@@ -165,19 +181,34 @@ fn make_hygiene_info(
     macro_file: MacroFile,
     loc: &MacroCallLoc,
 ) -> Option<HygieneInfo> {
-    let def_offset = loc.def.ast_id().left().and_then(|id| {
+    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()?.syntax().text_range().start(),
-            ast::Macro::MacroDef(mac) => mac.body()?.syntax().text_range().start(),
+            ast::Macro::MacroRules(mac) => mac.token_tree()?,
+            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, .. } => {
+            let tt = ast_id.to_node(db).attrs().nth(invoc_attr_index as usize)?.token_tree()?;
+            Some(InFile::new(ast_id.file_id, tt))
+        }
+        _ => None,
+    });
 
     let macro_def = db.macro_def(loc.def)?;
     let (_, exp_map) = db.parse_macro_expansion(macro_file).value?;
     let macro_arg = db.macro_arg(macro_file.macro_call_id)?;
 
-    Some(HygieneInfo { file: macro_file, def_start: def_offset, macro_arg, macro_def, exp_map })
+    Some(HygieneInfo {
+        file: macro_file,
+        attr_input_or_mac_def_start: attr_input_or_mac_def
+            .map(|it| it.map(|tt| tt.syntax().text_range().start())),
+        macro_arg_shift: mbe::Shift::new(&macro_arg.0),
+        macro_arg,
+        macro_def,
+        exp_map,
+    })
 }
 
 impl HygieneFrame {
@@ -214,7 +245,7 @@ impl HygieneFrame {
             Some(it) => it,
         };
 
-        let def_site = info.def_start.map(|it| db.hygiene_frame(it.file_id));
+        let def_site = info.attr_input_or_mac_def_start.map(|it| db.hygiene_frame(it.file_id));
         let call_site = Some(db.hygiene_frame(calling_file));
 
         HygieneFrame { expansion: Some(info), local_inner, krate, call_site, def_site }