fix: Reimplement mapping out of test/bench attributes for runnables

crates/hir_expand/src/lib.rs

@@ -701,6 +701,16 @@ 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.

crates/ide/src/runnables.rs

@@ -236,7 +236,9 @@ fn find_related_tests(
                 .filter_map(|token| token.ancestors().find_map(ast::Fn::cast))
                 .map(|f| hir::InFile::new(sema.hir_file_for(f.syntax()), f));
 
-            for InFile { value: ref fn_def, .. } in functions {
+            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);
                 if let Some(runnable) = as_test_runnable(sema, fn_def) {
                     // direct test
                     tests.insert(runnable);
@@ -292,7 +294,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)?.map_out_of_test_attr(sema.db);
     let name_string = def.name(sema.db).to_string();
 
     let root = def.module(sema.db).krate().root_module(sema.db);
@@ -501,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.map_out_of_test_attr(sema.db);
                     if test_related_attribute(&it.value).is_some() {
                         return true;
                     }