Skip test/bench attr expansion in resolution instead of collection

6 files changed, 47 insertions, 37 deletions

diff --git a/crates/hir_def/src/builtin_attr.rs b/crates/hir_def/src/builtin_attr.rs
index 15fb904e6f3..6cd185ceeb2 100644
--- a/crates/hir_def/src/builtin_attr.rs
+++ b/crates/hir_def/src/builtin_attr.rs
@@ -34,11 +34,6 @@ macro_rules! rustc_attr {
     };
 }
 
-// FIXME: We shouldn't special case these at all, but as of now expanding attributes severely degrades
-// user experience due to lacking support.
-/// 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] = &[
diff --git a/crates/hir_def/src/lib.rs b/crates/hir_def/src/lib.rs
index 5ddef48495b..1abda2d66aa 100644
--- a/crates/hir_def/src/lib.rs
+++ b/crates/hir_def/src/lib.rs
@@ -778,21 +778,18 @@ fn attr_macro_as_call_id(
     krate: CrateId,
     resolver: impl Fn(path::ModPath) -> Option<MacroDefId>,
 ) -> Result<MacroCallId, UnresolvedMacro> {
-    let def: MacroDefId = resolver(item_attr.path.clone())
+    let attr_path = &item_attr.path;
+
+    let def = resolver(attr_path.clone())
         .filter(MacroDefId::is_attribute)
-        .ok_or_else(|| UnresolvedMacro { path: item_attr.path.clone() })?;
-    let last_segment = item_attr
-        .path
-        .segments()
-        .last()
-        .ok_or_else(|| UnresolvedMacro { path: item_attr.path.clone() })?;
-    let mut arg = match &macro_attr.input {
-        Some(input) => match &**input {
-            attr::AttrInput::Literal(_) => Default::default(),
-            attr::AttrInput::TokenTree(tt, map) => (tt.clone(), map.clone()),
-        },
-        None => Default::default(),
+        .ok_or_else(|| UnresolvedMacro { path: attr_path.clone() })?;
+    let last_segment =
+        attr_path.segments().last().ok_or_else(|| UnresolvedMacro { path: attr_path.clone() })?;
+    let mut arg = match macro_attr.input.as_deref() {
+        Some(attr::AttrInput::TokenTree(tt, map)) => (tt.clone(), map.clone()),
+        _ => Default::default(),
     };
+
     // The parentheses are always disposed here.
     arg.0.delimiter = None;
 
diff --git a/crates/hir_def/src/nameres/collector.rs b/crates/hir_def/src/nameres/collector.rs
index b486beea7dc..d1e54cd46d5 100644
--- a/crates/hir_def/src/nameres/collector.rs
+++ b/crates/hir_def/src/nameres/collector.rs
@@ -9,7 +9,7 @@ use base_db::{CrateId, Edition, FileId, ProcMacroId};
 use cfg::{CfgExpr, CfgOptions};
 use hir_expand::{
     ast_id_map::FileAstId,
-    builtin_attr_macro::find_builtin_attr,
+    builtin_attr_macro::{find_builtin_attr, is_builtin_test_or_bench_attr},
     builtin_derive_macro::find_builtin_derive,
     builtin_fn_macro::find_builtin_macro,
     name::{name, AsName, Name},
@@ -1142,7 +1142,30 @@ impl DefCollector<'_> {
                     ) {
                         Ok(call_id) => {
                             let loc: MacroCallLoc = self.db.lookup_intern_macro_call(call_id);
-                            if let MacroDefKind::ProcMacro(exp, ..) = &loc.def.kind {
+
+                            // Skip #[test]/#[bench] expansion, which would merely result in more memory usage
+                            // due to duplicating functions into macro expansions
+                            if is_builtin_test_or_bench_attr(loc.def) {
+                                let file_id = ast_id.ast_id.file_id;
+                                let item_tree = self.db.file_item_tree(file_id);
+                                let mod_dir = self.mod_dirs[&directive.module_id].clone();
+                                self.skip_attrs.insert(InFile::new(file_id, *mod_item), attr.id);
+                                ModCollector {
+                                    def_collector: &mut *self,
+                                    macro_depth: directive.depth,
+                                    module_id: directive.module_id,
+                                    tree_id: TreeId::new(file_id, None),
+                                    item_tree: &item_tree,
+                                    mod_dir,
+                                }
+                                .collect(&[*mod_item]);
+
+                                // Remove the original directive since we resolved it.
+                                res = ReachedFixedPoint::No;
+                                return false;
+                            }
+
+                            if let MacroDefKind::ProcMacro(exp, ..) = loc.def.kind {
                                 if exp.is_dummy() {
                                     // Proc macros that cannot be expanded are treated as not
                                     // resolved, in order to fall back later.
@@ -1774,7 +1797,6 @@ impl ModCollector<'_, '_> {
                 let name = name.to_smol_str();
                 let is_inert = builtin_attr::INERT_ATTRIBUTES
                     .iter()
-                    .chain(builtin_attr::EXTRA_ATTRIBUTES)
                     .copied()
                     .chain(self.def_collector.registered_attrs.iter().map(AsRef::as_ref))
                     .any(|attr| name == *attr);
diff --git a/crates/hir_expand/src/builtin_attr_macro.rs b/crates/hir_expand/src/builtin_attr_macro.rs
index 2e461864a0b..ec587daf9ba 100644
--- a/crates/hir_expand/src/builtin_attr_macro.rs
+++ b/crates/hir_expand/src/builtin_attr_macro.rs
@@ -46,6 +46,16 @@ register_builtin! {
     (test_case, TestCase) => dummy_attr_expand
 }
 
+pub fn is_builtin_test_or_bench_attr(makro: MacroDefId) -> bool {
+    match makro.kind {
+        MacroDefKind::BuiltInAttr(expander, ..) => {
+            BuiltinAttrExpander::find_by_name(&name!(test)) == Some(expander)
+                || BuiltinAttrExpander::find_by_name(&name!(bench)) == Some(expander)
+        }
+        _ => false,
+    }
+}
+
 pub fn find_builtin_attr(
     ident: &name::Name,
     krate: CrateId,
diff --git a/crates/hir_expand/src/lib.rs b/crates/hir_expand/src/lib.rs
index 671fcb0ba8b..fdb639f55d2 100644
--- a/crates/hir_expand/src/lib.rs
+++ b/crates/hir_expand/src/lib.rs
@@ -701,16 +701,6 @@ impl<N: AstNode> InFile<N> {
     }
 }
 
-impl InFile<ast::Fn> {
-    pub fn map_out_of_test_attr(self, db: &dyn db::AstDatabase) -> InFile<ast::Fn> {
-        (|| {
-            let InFile { file_id, value } = self.file_id.call_node(db)?;
-            ast::Fn::cast(value).map(|n| InFile::new(file_id, n))
-        })()
-        .unwrap_or(self)
-    }
-}
-
 /// In Rust, macros expand token trees to token trees. When we want to turn a
 /// token tree into an AST node, we need to figure out what kind of AST node we
 /// want: something like `foo` can be a type, an expression, or a pattern.
diff --git a/crates/ide/src/runnables.rs b/crates/ide/src/runnables.rs
index 00b60a39026..21130e06075 100644
--- a/crates/ide/src/runnables.rs
+++ b/crates/ide/src/runnables.rs
@@ -237,8 +237,7 @@ fn find_related_tests(
                 .map(|f| hir::InFile::new(sema.hir_file_for(f.syntax()), f));
 
             for fn_def in functions {
-                // #[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.map_out_of_test_attr(sema.db);
+                let InFile { value: fn_def, .. } = &fn_def;
                 if let Some(runnable) = as_test_runnable(sema, fn_def) {
                     // direct test
                     tests.insert(runnable);
@@ -294,8 +293,7 @@ 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> {
-    // #[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)?.map_out_of_test_attr(sema.db);
+    let func = def.source(sema.db)?;
     let name_string = def.name(sema.db).to_string();
 
     let root = def.module(sema.db).krate().root_module(sema.db);
@@ -504,8 +502,6 @@ 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.map_out_of_test_attr(sema.db);
                     if test_related_attribute(&it.value).is_some() {
                         return true;
                     }