fix: Do not replace items annotated with builtin attrs with the attr input

crates/hir/src/semantics.rs

@@ -173,6 +173,10 @@ impl<'db, DB: HirDatabase> Semantics<'db, DB> {
         self.imp.descend_node_at_offset(node, offset).find_map(N::cast)
     }
 
+    pub fn hir_file_for(&self, syntax_node: &SyntaxNode) -> HirFileId {
+        self.imp.find_file(syntax_node.clone()).file_id
+    }
+
     pub fn original_range(&self, node: &SyntaxNode) -> FileRange {
         self.imp.original_range(node)
     }

crates/hir_def/src/builtin_attr.rs

@@ -34,9 +34,6 @@ macro_rules! rustc_attr {
     };
 }
 
-/// Built-in macro-like attributes.
-pub const EXTRA_ATTRIBUTES: &[BuiltinAttribute] = &["test", "bench"];
-
 /// "Inert" built-in attributes that have a special meaning to rustc or rustdoc.
 #[rustfmt::skip]
 pub const INERT_ATTRIBUTES: &[BuiltinAttribute] = &[

crates/hir_def/src/nameres/collector.rs

@@ -1712,27 +1712,24 @@ impl ModCollector<'_, '_> {
         if path.kind == PathKind::Plain {
             if let Some(tool_module) = path.segments().first() {
                 let tool_module = tool_module.to_string();
-                if builtin_attr::TOOL_MODULES
+                let is_tool = builtin_attr::TOOL_MODULES
                     .iter()
                     .copied()
-                    .chain(self.def_collector.registered_tools.iter().map(|s| &**s))
-                    .any(|m| tool_module == *m)
-                {
+                    .chain(self.def_collector.registered_tools.iter().map(AsRef::as_ref))
+                    .any(|m| tool_module == *m);
+                if is_tool {
                     return true;
                 }
             }
 
             if let Some(name) = path.as_ident() {
                 let name = name.to_string();
-                if builtin_attr::INERT_ATTRIBUTES
+                let is_inert = builtin_attr::INERT_ATTRIBUTES
                     .iter()
-                    .chain(builtin_attr::EXTRA_ATTRIBUTES)
                     .copied()
-                    .chain(self.def_collector.registered_attrs.iter().map(|s| &**s))
-                    .any(|attr| name == *attr)
-                {
-                    return true;
-                }
+                    .chain(self.def_collector.registered_attrs.iter().map(AsRef::as_ref))
+                    .any(|attr| name == *attr);
+                return is_inert;
             }
         }
 

crates/hir_expand/src/builtin_attr.rs

@@ -17,11 +17,12 @@ macro_rules! register_builtin {
                 db: &dyn AstDatabase,
                 id: MacroCallId,
                 tt: &tt::Subtree,
+                item: &tt::Subtree,
             ) -> Result<tt::Subtree, mbe::ExpandError> {
                 let expander = match *self {
                     $( BuiltinAttrExpander::$variant => $expand, )*
                 };
-                expander(db, id, tt)
+                expander(db, id, tt, item)
             }
 
             fn find_by_name(name: &name::Name) -> Option<Self> {
@@ -61,7 +62,8 @@ pub fn find_builtin_attr(
 fn dummy_attr_expand(
     _db: &dyn AstDatabase,
     _id: MacroCallId,
-    tt: &tt::Subtree,
+    _tt: &tt::Subtree,
+    item: &tt::Subtree,
 ) -> Result<tt::Subtree, mbe::ExpandError> {
-    Ok(tt.clone())
+    Ok(item.clone())
 }

crates/hir_expand/src/db.rs

@@ -54,7 +54,12 @@ impl TokenExpander {
             TokenExpander::MacroDef { mac, .. } => mac.expand(tt),
             TokenExpander::Builtin(it) => it.expand(db, id, tt),
             // FIXME switch these to ExpandResult as well
-            TokenExpander::BuiltinAttr(it) => it.expand(db, id, tt).into(),
+            TokenExpander::BuiltinAttr(it) => match db.macro_arg(id) {
+                Some(macro_arg) => it.expand(db, id, tt, &macro_arg.0).into(),
+                None => mbe::ExpandResult::only_err(
+                    mbe::ExpandError::Other("No item argument for attribute".to_string()).into(),
+                ),
+            },
             TokenExpander::BuiltinDerive(it) => it.expand(db, id, tt).into(),
             TokenExpander::ProcMacro(_) => {
                 // We store the result in salsa db to prevent non-deterministic behavior in

crates/hir_expand/src/lib.rs

@@ -22,7 +22,7 @@ use either::Either;
 pub use mbe::{ExpandError, ExpandResult};
 pub use parser::FragmentKind;
 
-use std::{hash::Hash, sync::Arc};
+use std::{hash::Hash, iter, sync::Arc};
 
 use base_db::{impl_intern_key, salsa, CrateId, FileId, FileRange};
 use syntax::{
@@ -454,7 +454,7 @@ impl InFile<SyntaxNode> {
         self,
         db: &dyn db::AstDatabase,
     ) -> impl Iterator<Item = InFile<SyntaxNode>> + '_ {
-        std::iter::successors(Some(self), move |node| match node.value.parent() {
+        iter::successors(Some(self), move |node| match node.value.parent() {
             Some(parent) => Some(node.with_value(parent)),
             None => {
                 let parent_node = node.file_id.call_node(db)?;
@@ -562,6 +562,23 @@ impl<N: AstNode> InFile<N> {
     pub fn syntax(&self) -> InFile<&SyntaxNode> {
         self.with_value(self.value.syntax())
     }
+
+    pub fn nodes_with_attributes<'db>(
+        self,
+        db: &'db dyn db::AstDatabase,
+    ) -> impl Iterator<Item = InFile<N>> + 'db
+    where
+        N: 'db,
+    {
+        iter::successors(Some(self), move |node| {
+            let InFile { file_id, value } = node.file_id.call_node(db)?;
+            N::cast(value).map(|n| InFile::new(file_id, n))
+        })
+    }
+
+    pub fn node_with_attributes(self, db: &dyn db::AstDatabase) -> InFile<N> {
+        self.nodes_with_attributes(db).last().unwrap()
+    }
 }
 
 /// Given a `MacroCallId`, return what `FragmentKind` it belongs to.

crates/ide/src/runnables.rs

@@ -3,7 +3,7 @@ use std::fmt;
 use ast::NameOwner;
 use cfg::CfgExpr;
 use either::Either;
-use hir::{AsAssocItem, HasAttrs, HasSource, HirDisplay, Semantics};
+use hir::{AsAssocItem, HasAttrs, HasSource, HirDisplay, InFile, Semantics};
 use ide_assists::utils::test_related_attribute;
 use ide_db::{
     base_db::{FilePosition, FileRange},
@@ -232,22 +232,26 @@ fn find_related_tests(
             let functions = refs.iter().filter_map(|(range, _)| {
                 let token = file.token_at_offset(range.start()).next()?;
                 let token = sema.descend_into_macros(token);
-                token.ancestors().find_map(ast::Fn::cast)
+                token
+                    .ancestors()
+                    .find_map(ast::Fn::cast)
+                    .map(|f| hir::InFile::new(sema.hir_file_for(f.syntax()), f))
             });
 
             for fn_def in functions {
-                if let Some(runnable) = as_test_runnable(sema, &fn_def) {
+                // #[test/bench] expands to just the item causing us to lose the attribute, so recover them by going out of the attribute
+                let InFile { value: fn_def, .. } = &fn_def.node_with_attributes(sema.db);
+                if let Some(runnable) = as_test_runnable(sema, fn_def) {
                     // direct test
                     tests.insert(runnable);
-                } else if let Some(module) = parent_test_module(sema, &fn_def) {
+                } else if let Some(module) = parent_test_module(sema, fn_def) {
                     // indirect test
-                    find_related_tests_in_module(sema, &fn_def, &module, tests);
+                    find_related_tests_in_module(sema, fn_def, &module, tests);
                 }
             }
         }
     }
 }
-
 fn find_related_tests_in_module(
     sema: &Semantics<RootDatabase>,
     fn_def: &ast::Fn,
@@ -292,7 +296,8 @@ fn parent_test_module(sema: &Semantics<RootDatabase>, fn_def: &ast::Fn) -> Optio
 }
 
 pub(crate) fn runnable_fn(sema: &Semantics<RootDatabase>, def: hir::Function) -> Option<Runnable> {
-    let func = def.source(sema.db)?;
+    // #[test/bench] expands to just the item causing us to lose the attribute, so recover them by going out of the attribute
+    let func = def.source(sema.db)?.node_with_attributes(sema.db);
     let name_string = def.name(sema.db).to_string();
 
     let root = def.module(sema.db).krate().root_module(sema.db);
@@ -499,6 +504,8 @@ fn has_test_function_or_multiple_test_submodules(
         match item {
             hir::ModuleDef::Function(f) => {
                 if let Some(it) = f.source(sema.db) {
+                    // #[test/bench] expands to just the item causing us to lose the attribute, so recover them by going out of the attribute
+                    let it = it.node_with_attributes(sema.db);
                     if test_related_attribute(&it.value).is_some() {
                         return true;
                     }