Merge #8796

8796: internal: rewrite `#[derive]` removal to be based on AST (take 2) r=jonas-schievink a=jonas-schievink

Second attempt of https://github.com/rust-analyzer/rust-analyzer/pull/8443, this uses syntactical attribute offsets in `hir_expand`, and changes `attr.rs` to make those easy to derive.

This will make it easy to add similar attribute removal for attribute macros, unblocking them.

Co-authored-by: Jonas Schievink <jonasschievink@gmail.com>

8 files changed, 205 insertions, 175 deletions

diff --git a/crates/hir_def/src/attr.rs b/crates/hir_def/src/attr.rs
index a2479016e57..aadd4e44ae7 100644
--- a/crates/hir_def/src/attr.rs
+++ b/crates/hir_def/src/attr.rs
@@ -9,7 +9,7 @@ use std::{
 use base_db::CrateId;
 use cfg::{CfgExpr, CfgOptions};
 use either::Either;
-use hir_expand::{hygiene::Hygiene, name::AsName, AstId, AttrId, InFile};
+use hir_expand::{hygiene::Hygiene, name::AsName, AstId, InFile};
 use itertools::Itertools;
 use la_arena::ArenaMap;
 use mbe::ast_to_token_tree;
@@ -101,17 +101,13 @@ impl RawAttrs {
         hygiene: &Hygiene,
     ) -> Self {
         let entries = collect_attrs(owner)
-            .enumerate()
-            .flat_map(|(i, attr)| {
-                let index = AttrId(i as u32);
-                match attr {
-                    Either::Left(attr) => Attr::from_src(db, attr, hygiene, index),
-                    Either::Right(comment) => comment.doc_comment().map(|doc| Attr {
-                        id: index,
-                        input: Some(AttrInput::Literal(SmolStr::new(doc))),
-                        path: Interned::new(ModPath::from(hir_expand::name!(doc))),
-                    }),
-                }
+            .flat_map(|(id, attr)| match attr {
+                Either::Left(attr) => Attr::from_src(db, attr, hygiene, id),
+                Either::Right(comment) => comment.doc_comment().map(|doc| Attr {
+                    id,
+                    input: Some(AttrInput::Literal(SmolStr::new(doc))),
+                    path: Interned::new(ModPath::from(hir_expand::name!(doc))),
+                }),
             })
             .collect::<Arc<_>>();
 
@@ -124,6 +120,7 @@ impl RawAttrs {
     }
 
     pub(crate) fn merge(&self, other: Self) -> Self {
+        // FIXME: This needs to fixup `AttrId`s
         match (&self.entries, &other.entries) {
             (None, None) => Self::EMPTY,
             (Some(entries), None) | (None, Some(entries)) => {
@@ -375,39 +372,26 @@ impl AttrsWithOwner {
 
                 let def_map = module.def_map(db);
                 let mod_data = &def_map[module.local_id];
-                let attrs = match mod_data.declaration_source(db) {
+                match mod_data.declaration_source(db) {
                     Some(it) => {
-                        let mut attrs: Vec<_> = collect_attrs(&it.value as &dyn ast::AttrsOwner)
-                            .map(|attr| InFile::new(it.file_id, attr))
-                            .collect();
+                        let mut map = AttrSourceMap::new(InFile::new(it.file_id, &it.value));
                         if let InFile { file_id, value: ModuleSource::SourceFile(file) } =
                             mod_data.definition_source(db)
                         {
-                            attrs.extend(
-                                collect_attrs(&file as &dyn ast::AttrsOwner)
-                                    .map(|attr| InFile::new(file_id, attr)),
-                            )
+                            map.merge(AttrSourceMap::new(InFile::new(file_id, &file)));
                         }
-                        attrs
+                        return map;
                     }
                     None => {
                         let InFile { file_id, value } = mod_data.definition_source(db);
-                        match &value {
-                            ModuleSource::SourceFile(file) => {
-                                collect_attrs(file as &dyn ast::AttrsOwner)
-                            }
-                            ModuleSource::Module(module) => {
-                                collect_attrs(module as &dyn ast::AttrsOwner)
-                            }
-                            ModuleSource::BlockExpr(block) => {
-                                collect_attrs(block as &dyn ast::AttrsOwner)
-                            }
-                        }
-                        .map(|attr| InFile::new(file_id, attr))
-                        .collect()
+                        let attrs_owner = match &value {
+                            ModuleSource::SourceFile(file) => file as &dyn ast::AttrsOwner,
+                            ModuleSource::Module(module) => module as &dyn ast::AttrsOwner,
+                            ModuleSource::BlockExpr(block) => block as &dyn ast::AttrsOwner,
+                        };
+                        return AttrSourceMap::new(InFile::new(file_id, attrs_owner));
                     }
-                };
-                return AttrSourceMap { attrs };
+                }
             }
             AttrDefId::FieldId(id) => {
                 let map = db.fields_attrs_source_map(id.parent);
@@ -462,11 +446,7 @@ impl AttrsWithOwner {
             },
         };
 
-        AttrSourceMap {
-            attrs: collect_attrs(&owner.value)
-                .map(|attr| InFile::new(owner.file_id, attr))
-                .collect(),
-        }
+        AttrSourceMap::new(owner.as_ref().map(|node| node as &dyn AttrsOwner))
     }
 
     pub fn docs_with_rangemap(
@@ -518,7 +498,7 @@ impl AttrsWithOwner {
         if buf.is_empty() {
             None
         } else {
-            Some((Documentation(buf), DocsRangeMap { mapping, source: self.source_map(db).attrs }))
+            Some((Documentation(buf), DocsRangeMap { mapping, source_map: self.source_map(db) }))
         }
     }
 }
@@ -559,27 +539,59 @@ fn inner_attributes(
 }
 
 pub struct AttrSourceMap {
-    attrs: Vec<InFile<Either<ast::Attr, ast::Comment>>>,
+    attrs: Vec<InFile<ast::Attr>>,
+    doc_comments: Vec<InFile<ast::Comment>>,
 }
 
 impl AttrSourceMap {
+    fn new(owner: InFile<&dyn ast::AttrsOwner>) -> Self {
+        let mut attrs = Vec::new();
+        let mut doc_comments = Vec::new();
+        for (_, attr) in collect_attrs(owner.value) {
+            match attr {
+                Either::Left(attr) => attrs.push(owner.with_value(attr)),
+                Either::Right(comment) => doc_comments.push(owner.with_value(comment)),
+            }
+        }
+
+        Self { attrs, doc_comments }
+    }
+
+    fn merge(&mut self, other: Self) {
+        self.attrs.extend(other.attrs);
+        self.doc_comments.extend(other.doc_comments);
+    }
+
     /// Maps the lowered `Attr` back to its original syntax node.
     ///
     /// `attr` must come from the `owner` used for AttrSourceMap
     ///
     /// Note that the returned syntax node might be a `#[cfg_attr]`, or a doc comment, instead of
     /// the attribute represented by `Attr`.
-    pub fn source_of(&self, attr: &Attr) -> InFile<&Either<ast::Attr, ast::Comment>> {
-        self.attrs
-            .get(attr.id.0 as usize)
-            .unwrap_or_else(|| panic!("cannot find `Attr` at index {:?}", attr.id))
-            .as_ref()
+    pub fn source_of(&self, attr: &Attr) -> InFile<Either<ast::Attr, ast::Comment>> {
+        self.source_of_id(attr.id)
+    }
+
+    fn source_of_id(&self, id: AttrId) -> InFile<Either<ast::Attr, ast::Comment>> {
+        if id.is_doc_comment {
+            self.doc_comments
+                .get(id.ast_index as usize)
+                .unwrap_or_else(|| panic!("cannot find doc comment at index {:?}", id))
+                .clone()
+                .map(|attr| Either::Right(attr))
+        } else {
+            self.attrs
+                .get(id.ast_index as usize)
+                .unwrap_or_else(|| panic!("cannot find `Attr` at index {:?}", id))
+                .clone()
+                .map(|attr| Either::Left(attr))
+        }
     }
 }
 
 /// A struct to map text ranges from [`Documentation`] back to TextRanges in the syntax tree.
 pub struct DocsRangeMap {
-    source: Vec<InFile<Either<ast::Attr, ast::Comment>>>,
+    source_map: AttrSourceMap,
     // (docstring-line-range, attr_index, attr-string-range)
     // a mapping from the text range of a line of the [`Documentation`] to the attribute index and
     // the original (untrimmed) syntax doc line
@@ -596,7 +608,7 @@ impl DocsRangeMap {
 
         let relative_range = range - line_docs_range.start();
 
-        let &InFile { file_id, value: ref source } = &self.source[idx.0 as usize];
+        let &InFile { file_id, value: ref source } = &self.source_map.source_of_id(idx);
         match source {
             Either::Left(_) => None, // FIXME, figure out a nice way to handle doc attributes here
             // as well as for whats done in syntax highlight doc injection
@@ -615,6 +627,12 @@ impl DocsRangeMap {
     }
 }
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub(crate) struct AttrId {
+    is_doc_comment: bool,
+    pub(crate) ast_index: u32,
+}
+
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct Attr {
     pub(crate) id: AttrId,
@@ -749,22 +767,32 @@ fn attrs_from_item_tree<N: ItemTreeNode>(id: ItemTreeId<N>, db: &dyn DefDatabase
 
 fn collect_attrs(
     owner: &dyn ast::AttrsOwner,
-) -> impl Iterator<Item = Either<ast::Attr, ast::Comment>> {
+) -> impl Iterator<Item = (AttrId, Either<ast::Attr, ast::Comment>)> {
     let (inner_attrs, inner_docs) = inner_attributes(owner.syntax())
         .map_or((None, None), |(attrs, docs)| (Some(attrs), Some(docs)));
 
     let outer_attrs = owner.attrs().filter(|attr| attr.kind().is_outer());
-    let attrs = outer_attrs
-        .chain(inner_attrs.into_iter().flatten())
-        .map(|attr| (attr.syntax().text_range().start(), Either::Left(attr)));
+    let attrs =
+        outer_attrs.chain(inner_attrs.into_iter().flatten()).enumerate().map(|(idx, attr)| {
+            (
+                AttrId { ast_index: idx as u32, is_doc_comment: false },
+                attr.syntax().text_range().start(),
+                Either::Left(attr),
+            )
+        });
 
     let outer_docs =
         ast::CommentIter::from_syntax_node(owner.syntax()).filter(ast::Comment::is_outer);
-    let docs = outer_docs
-        .chain(inner_docs.into_iter().flatten())
-        .map(|docs_text| (docs_text.syntax().text_range().start(), Either::Right(docs_text)));
+    let docs =
+        outer_docs.chain(inner_docs.into_iter().flatten()).enumerate().map(|(idx, docs_text)| {
+            (
+                AttrId { ast_index: idx as u32, is_doc_comment: true },
+                docs_text.syntax().text_range().start(),
+                Either::Right(docs_text),
+            )
+        });
     // sort here by syntax node offset because the source can have doc attributes and doc strings be interleaved
-    docs.chain(attrs).sorted_by_key(|&(offset, _)| offset).map(|(_, attr)| attr)
+    docs.chain(attrs).sorted_by_key(|&(_, offset, _)| offset).map(|(id, _, attr)| (id, attr))
 }
 
 pub(crate) fn variants_attrs_source_map(
diff --git a/crates/hir_def/src/lib.rs b/crates/hir_def/src/lib.rs
index e96ca953ff7..a82ea5957e9 100644
--- a/crates/hir_def/src/lib.rs
+++ b/crates/hir_def/src/lib.rs
@@ -62,14 +62,14 @@ use hir_expand::{
     ast_id_map::FileAstId,
     eager::{expand_eager_macro, ErrorEmitted, ErrorSink},
     hygiene::Hygiene,
-    AstId, AttrId, FragmentKind, HirFileId, InFile, MacroCallId, MacroCallKind, MacroDefId,
-    MacroDefKind,
+    AstId, FragmentKind, HirFileId, InFile, MacroCallId, MacroCallKind, MacroDefId, MacroDefKind,
 };
 use la_arena::Idx;
 use nameres::DefMap;
 use path::ModPath;
 use syntax::ast;
 
+use crate::attr::AttrId;
 use crate::builtin_type::BuiltinType;
 use item_tree::{
     Const, Enum, Function, Impl, ItemTreeId, ItemTreeNode, ModItem, Static, Struct, Trait,
@@ -753,7 +753,7 @@ fn derive_macro_as_call_id(
             MacroCallKind::Derive {
                 ast_id: item_attr.ast_id,
                 derive_name: last_segment.to_string(),
-                derive_attr,
+                derive_attr_index: derive_attr.ast_index,
             },
         )
         .into();
diff --git a/crates/hir_def/src/nameres/collector.rs b/crates/hir_def/src/nameres/collector.rs
index e89136ed1dc..adfb78c942f 100644
--- a/crates/hir_def/src/nameres/collector.rs
+++ b/crates/hir_def/src/nameres/collector.rs
@@ -13,14 +13,14 @@ use hir_expand::{
     builtin_macro::find_builtin_macro,
     name::{AsName, Name},
     proc_macro::ProcMacroExpander,
-    AttrId, FragmentKind, HirFileId, MacroCallId, MacroCallKind, MacroDefId, MacroDefKind,
+    FragmentKind, HirFileId, MacroCallId, MacroCallKind, MacroDefId, MacroDefKind,
 };
 use hir_expand::{InFile, MacroCallLoc};
 use rustc_hash::{FxHashMap, FxHashSet};
 use syntax::ast;
 
 use crate::{
-    attr::Attrs,
+    attr::{AttrId, Attrs},
     db::DefDatabase,
     derive_macro_as_call_id,
     intern::Interned,
diff --git a/crates/hir_expand/src/builtin_derive.rs b/crates/hir_expand/src/builtin_derive.rs
index 537c0302866..b6a6d602f7f 100644
--- a/crates/hir_expand/src/builtin_derive.rs
+++ b/crates/hir_expand/src/builtin_derive.rs
@@ -269,7 +269,7 @@ mod tests {
     use expect_test::{expect, Expect};
     use name::AsName;
 
-    use crate::{test_db::TestDB, AstId, AttrId, MacroCallId, MacroCallKind, MacroCallLoc};
+    use crate::{test_db::TestDB, AstId, MacroCallId, MacroCallKind, MacroCallLoc};
 
     use super::*;
 
@@ -320,7 +320,7 @@ $0
             kind: MacroCallKind::Derive {
                 ast_id,
                 derive_name: name.to_string(),
-                derive_attr: AttrId(0),
+                derive_attr_index: 0,
             },
         };
 
diff --git a/crates/hir_expand/src/db.rs b/crates/hir_expand/src/db.rs
index 6647e57e7dc..9fa419fcf14 100644
--- a/crates/hir_expand/src/db.rs
+++ b/crates/hir_expand/src/db.rs
@@ -12,9 +12,9 @@ use syntax::{
 };
 
 use crate::{
-    ast_id_map::AstIdMap, hygiene::HygieneFrame, BuiltinDeriveExpander, BuiltinFnLikeExpander,
-    EagerCallLoc, EagerMacroId, HirFileId, HirFileIdRepr, LazyMacroId, MacroCallId, MacroCallLoc,
-    MacroDefId, MacroDefKind, MacroFile, ProcMacroExpander,
+    ast_id_map::AstIdMap, hygiene::HygieneFrame, input::process_macro_input, BuiltinDeriveExpander,
+    BuiltinFnLikeExpander, EagerCallLoc, EagerMacroId, HirFileId, HirFileIdRepr, LazyMacroId,
+    MacroCallId, MacroCallLoc, MacroDefId, MacroDefKind, MacroFile, ProcMacroExpander,
 };
 
 /// Total limit on the number of tokens produced by any macro invocation.
@@ -281,6 +281,7 @@ fn macro_arg_text(db: &dyn AstDatabase, id: MacroCallId) -> Option<GreenNode> {
     };
     let loc = db.lookup_intern_macro(id);
     let arg = loc.kind.arg(db)?;
+    let arg = process_macro_input(db, arg, id);
     Some(arg.green().into())
 }
 
diff --git a/crates/hir_expand/src/input.rs b/crates/hir_expand/src/input.rs
new file mode 100644
index 00000000000..112216859cf
--- /dev/null
+++ b/crates/hir_expand/src/input.rs
@@ -0,0 +1,94 @@
+//! Macro input conditioning.
+
+use syntax::{
+    ast::{self, AttrsOwner},
+    AstNode, SyntaxNode,
+};
+
+use crate::{
+    db::AstDatabase,
+    name::{name, AsName},
+    LazyMacroId, MacroCallKind, MacroCallLoc,
+};
+
+pub(crate) fn process_macro_input(
+    db: &dyn AstDatabase,
+    node: SyntaxNode,
+    id: LazyMacroId,
+) -> SyntaxNode {
+    let loc: MacroCallLoc = db.lookup_intern_macro(id);
+
+    match loc.kind {
+        MacroCallKind::FnLike { .. } => node,
+        MacroCallKind::Derive { derive_attr_index, .. } => {
+            let item = match ast::Item::cast(node.clone()) {
+                Some(item) => item,
+                None => return node,
+            };
+
+            remove_derives_up_to(item, derive_attr_index as usize).syntax().clone()
+        }
+    }
+}
+
+/// Removes `#[derive]` attributes from `item`, up to `attr_index`.
+fn remove_derives_up_to(item: ast::Item, attr_index: usize) -> ast::Item {
+    let item = item.clone_for_update();
+    for attr in item.attrs().take(attr_index + 1) {
+        if let Some(name) =
+            attr.path().and_then(|path| path.as_single_segment()).and_then(|seg| seg.name_ref())
+        {
+            if name.as_name() == name![derive] {
+                attr.syntax().detach();
+            }
+        }
+    }
+    item
+}
+
+#[cfg(test)]
+mod tests {
+    use base_db::fixture::WithFixture;
+    use base_db::SourceDatabase;
+    use expect_test::{expect, Expect};
+
+    use crate::test_db::TestDB;
+
+    use super::*;
+
+    fn test_remove_derives_up_to(attr: usize, ra_fixture: &str, expect: Expect) {
+        let (db, file_id) = TestDB::with_single_file(&ra_fixture);
+        let parsed = db.parse(file_id);
+
+        let mut items: Vec<_> =
+            parsed.syntax_node().descendants().filter_map(ast::Item::cast).collect();
+        assert_eq!(items.len(), 1);
+
+        let item = remove_derives_up_to(items.pop().unwrap(), attr);
+        expect.assert_eq(&item.to_string());
+    }
+
+    #[test]
+    fn remove_derive() {
+        test_remove_derives_up_to(
+            2,
+            r#"
+#[allow(unused)]
+#[derive(Copy)]
+#[derive(Hello)]
+#[derive(Clone)]
+struct A {
+    bar: u32
+}
+        "#,
+            expect![[r#"
+#[allow(unused)]
+
+
+#[derive(Clone)]
+struct A {
+    bar: u32
+}"#]],
+        );
+    }
+}
diff --git a/crates/hir_expand/src/lib.rs b/crates/hir_expand/src/lib.rs
index 80ab3aeee7e..5df11856e90 100644
--- a/crates/hir_expand/src/lib.rs
+++ b/crates/hir_expand/src/lib.rs
@@ -14,6 +14,7 @@ pub mod builtin_macro;
 pub mod proc_macro;
 pub mod quote;
 pub mod eager;
+mod input;
 
 use either::Either;
 
@@ -292,13 +293,21 @@ pub struct MacroCallLoc {
 
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub enum MacroCallKind {
-    FnLike { ast_id: AstId<ast::MacroCall>, fragment: FragmentKind },
-    Derive { ast_id: AstId<ast::Item>, derive_name: String, derive_attr: AttrId },
+    FnLike {
+        ast_id: AstId<ast::MacroCall>,
+        fragment: FragmentKind,
+    },
+    Derive {
+        ast_id: AstId<ast::Item>,
+        derive_name: String,
+        /// Syntactical index of the invoking `#[derive]` attribute.
+        ///
+        /// Outer attributes are counted first, then inner attributes. This does not support
+        /// out-of-line modules, which may have attributes spread across 2 files!
+        derive_attr_index: u32,
+    },
 }
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
-pub struct AttrId(pub u32);
-
 impl MacroCallKind {
     fn file_id(&self) -> HirFileId {
         match self {
diff --git a/crates/hir_expand/src/proc_macro.rs b/crates/hir_expand/src/proc_macro.rs
index 75e95081609..d5643393ae2 100644
--- a/crates/hir_expand/src/proc_macro.rs
+++ b/crates/hir_expand/src/proc_macro.rs
@@ -2,7 +2,6 @@
 
 use crate::db::AstDatabase;
 use base_db::{CrateId, ProcMacroId};
-use tt::buffer::{Cursor, TokenBuffer};
 
 #[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
 pub struct ProcMacroExpander {
@@ -44,9 +43,6 @@ impl ProcMacroExpander {
                     .clone()
                     .ok_or_else(|| err!("No derive macro found."))?;
 
-                let tt = remove_derive_attrs(tt)
-                    .ok_or_else(|| err!("Fail to remove derive for custom derive"))?;
-
                 // Proc macros have access to the environment variables of the invoking crate.
                 let env = &krate_graph[calling_crate].env;
 
@@ -56,101 +52,3 @@ impl ProcMacroExpander {
         }
     }
 }
-
-fn eat_punct(cursor: &mut Cursor, c: char) -> bool {
-    if let Some(tt::buffer::TokenTreeRef::Leaf(tt::Leaf::Punct(punct), _)) = cursor.token_tree() {
-        if punct.char == c {
-            *cursor = cursor.bump();
-            return true;
-        }
-    }
-    false
-}
-
-fn eat_subtree(cursor: &mut Cursor, kind: tt::DelimiterKind) -> bool {
-    if let Some(tt::buffer::TokenTreeRef::Subtree(subtree, _)) = cursor.token_tree() {
-        if Some(kind) == subtree.delimiter_kind() {
-            *cursor = cursor.bump_subtree();
-            return true;
-        }
-    }
-    false
-}
-
-fn eat_ident(cursor: &mut Cursor, t: &str) -> bool {
-    if let Some(tt::buffer::TokenTreeRef::Leaf(tt::Leaf::Ident(ident), _)) = cursor.token_tree() {
-        if t == ident.text.as_str() {
-            *cursor = cursor.bump();
-            return true;
-        }
-    }
-    false
-}
-
-fn remove_derive_attrs(tt: &tt::Subtree) -> Option<tt::Subtree> {
-    let buffer = TokenBuffer::from_tokens(&tt.token_trees);
-    let mut p = buffer.begin();
-    let mut result = tt::Subtree::default();
-
-    while !p.eof() {
-        let curr = p;
-
-        if eat_punct(&mut p, '#') {
-            eat_punct(&mut p, '!');
-            let parent = p;
-            if eat_subtree(&mut p, tt::DelimiterKind::Bracket) {
-                if eat_ident(&mut p, "derive") {
-                    p = parent.bump();
-                    continue;
-                }
-            }
-        }
-
-        result.token_trees.push(curr.token_tree()?.cloned());
-        p = curr.bump();
-    }
-
-    Some(result)
-}
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-    use test_utils::assert_eq_text;
-
-    #[test]
-    fn test_remove_derive_attrs() {
-        let tt = mbe::parse_to_token_tree(
-            r#"
-    #[allow(unused)]
-    #[derive(Copy)]
-    #[derive(Hello)]
-    struct A {
-        bar: u32
-    }
-"#,
-        )
-        .unwrap()
-        .0;
-        let result = format!("{:#?}", remove_derive_attrs(&tt).unwrap());
-
-        assert_eq_text!(
-            r#"
-SUBTREE $
-  PUNCH   # [alone] 0
-  SUBTREE [] 1
-    IDENT   allow 2
-    SUBTREE () 3
-      IDENT   unused 4
-  IDENT   struct 15
-  IDENT   A 16
-  SUBTREE {} 17
-    IDENT   bar 18
-    PUNCH   : [alone] 19
-    IDENT   u32 20
-"#
-            .trim(),
-            &result
-        );
-    }
-}