Resolve tests per file instead of per crate in test explorer

8 files changed, 142 insertions, 40 deletions

diff --git a/crates/ide/src/lib.rs b/crates/ide/src/lib.rs
index 7a6c6ec41d0..ad48d803895 100644
--- a/crates/ide/src/lib.rs
+++ b/crates/ide/src/lib.rs
@@ -353,6 +353,10 @@ impl Analysis {
         self.with_db(|db| test_explorer::discover_tests_in_crate(db, crate_id))
     }
 
+    pub fn discover_tests_in_file(&self, file_id: FileId) -> Cancellable<Vec<TestItem>> {
+        self.with_db(|db| test_explorer::discover_tests_in_file(db, file_id))
+    }
+
     /// Renders the crate graph to GraphViz "dot" syntax.
     pub fn view_crate_graph(&self, full: bool) -> Cancellable<Result<String, String>> {
         self.with_db(|db| view_crate_graph::view_crate_graph(db, full))
diff --git a/crates/ide/src/test_explorer.rs b/crates/ide/src/test_explorer.rs
index ca471399703..99e24308607 100644
--- a/crates/ide/src/test_explorer.rs
+++ b/crates/ide/src/test_explorer.rs
@@ -7,7 +7,7 @@ use ide_db::{
 };
 use syntax::TextRange;
 
-use crate::{navigation_target::ToNav, runnables::runnable_fn, Runnable, TryToNav};
+use crate::{runnables::runnable_fn, NavigationTarget, Runnable, TryToNav};
 
 #[derive(Debug)]
 pub enum TestItemKind {
@@ -56,7 +56,12 @@ fn find_crate_by_id(crate_graph: &CrateGraph, crate_id: &str) -> Option<CrateId>
     })
 }
 
-fn discover_tests_in_module(db: &RootDatabase, module: Module, prefix_id: String) -> Vec<TestItem> {
+fn discover_tests_in_module(
+    db: &RootDatabase,
+    module: Module,
+    prefix_id: String,
+    only_in_this_file: bool,
+) -> Vec<TestItem> {
     let sema = Semantics::new(db);
 
     let mut r = vec![];
@@ -64,9 +69,9 @@ fn discover_tests_in_module(db: &RootDatabase, module: Module, prefix_id: String
         let module_name =
             c.name(db).as_ref().and_then(|n| n.as_str()).unwrap_or("[mod without name]").to_owned();
         let module_id = format!("{prefix_id}::{module_name}");
-        let module_children = discover_tests_in_module(db, c, module_id.clone());
+        let module_children = discover_tests_in_module(db, c, module_id.clone(), only_in_this_file);
         if !module_children.is_empty() {
-            let nav = c.to_nav(db).call_site;
+            let nav = NavigationTarget::from_module_to_decl(sema.db, c).call_site;
             r.push(TestItem {
                 id: module_id,
                 kind: TestItemKind::Module,
@@ -76,7 +81,9 @@ fn discover_tests_in_module(db: &RootDatabase, module: Module, prefix_id: String
                 text_range: Some(nav.focus_or_full_range()),
                 runnable: None,
             });
-            r.extend(module_children);
+            if !only_in_this_file || c.is_inline(db) {
+                r.extend(module_children);
+            }
         }
     }
     for def in module.declarations(db) {
@@ -112,6 +119,55 @@ pub(crate) fn discover_tests_in_crate_by_test_id(
     discover_tests_in_crate(db, crate_id)
 }
 
+pub(crate) fn discover_tests_in_file(db: &RootDatabase, file_id: FileId) -> Vec<TestItem> {
+    let sema = Semantics::new(db);
+
+    let Some(module) = sema.file_to_module_def(file_id) else { return vec![] };
+    let Some((mut tests, id)) = find_module_id_and_test_parents(&sema, module) else {
+        return vec![];
+    };
+    tests.extend(discover_tests_in_module(db, module, id, true));
+    tests
+}
+
+fn find_module_id_and_test_parents(
+    sema: &Semantics<'_, RootDatabase>,
+    module: Module,
+) -> Option<(Vec<TestItem>, String)> {
+    let Some(parent) = module.parent(sema.db) else {
+        let name = module.krate().display_name(sema.db)?.to_string();
+        return Some((
+            vec![TestItem {
+                id: name.clone(),
+                kind: TestItemKind::Crate(module.krate().into()),
+                label: name.clone(),
+                parent: None,
+                file: None,
+                text_range: None,
+                runnable: None,
+            }],
+            name,
+        ));
+    };
+    let (mut r, mut id) = find_module_id_and_test_parents(sema, parent)?;
+    let parent = Some(id.clone());
+    id += "::";
+    let module_name = &module.name(sema.db);
+    let module_name = module_name.as_ref().and_then(|n| n.as_str()).unwrap_or("[mod without name]");
+    id += module_name;
+    let nav = NavigationTarget::from_module_to_decl(sema.db, module).call_site;
+    r.push(TestItem {
+        id: id.clone(),
+        kind: TestItemKind::Module,
+        label: module_name.to_owned(),
+        parent,
+        file: Some(nav.file_id),
+        text_range: Some(nav.focus_or_full_range()),
+        runnable: None,
+    });
+    Some((r, id))
+}
+
 pub(crate) fn discover_tests_in_crate(db: &RootDatabase, crate_id: CrateId) -> Vec<TestItem> {
     let crate_graph = db.crate_graph();
     if !crate_graph[crate_id].origin.is_local() {
@@ -133,6 +189,6 @@ pub(crate) fn discover_tests_in_crate(db: &RootDatabase, crate_id: CrateId) -> V
         text_range: None,
         runnable: None,
     }];
-    r.extend(discover_tests_in_module(db, module, crate_test_id));
+    r.extend(discover_tests_in_module(db, module, crate_test_id, false));
     r
 }
diff --git a/crates/rust-analyzer/src/handlers/request.rs b/crates/rust-analyzer/src/handlers/request.rs
index 4d4103a7ad7..77692ed3ae7 100644
--- a/crates/rust-analyzer/src/handlers/request.rs
+++ b/crates/rust-analyzer/src/handlers/request.rs
@@ -238,9 +238,12 @@ pub(crate) fn handle_discover_test(
     let (tests, scope) = match params.test_id {
         Some(id) => {
             let crate_id = id.split_once("::").map(|it| it.0).unwrap_or(&id);
-            (snap.analysis.discover_tests_in_crate_by_test_id(crate_id)?, vec![crate_id.to_owned()])
+            (
+                snap.analysis.discover_tests_in_crate_by_test_id(crate_id)?,
+                Some(vec![crate_id.to_owned()]),
+            )
         }
-        None => (snap.analysis.discover_test_roots()?, vec![]),
+        None => (snap.analysis.discover_test_roots()?, None),
     };
     for t in &tests {
         hack_recover_crate_name::insert_name(t.id.clone());
@@ -254,6 +257,7 @@ pub(crate) fn handle_discover_test(
             })
             .collect(),
         scope,
+        scope_file: None,
     })
 }
 
diff --git a/crates/rust-analyzer/src/lsp/ext.rs b/crates/rust-analyzer/src/lsp/ext.rs
index 5d3d75efcba..eac982f1b27 100644
--- a/crates/rust-analyzer/src/lsp/ext.rs
+++ b/crates/rust-analyzer/src/lsp/ext.rs
@@ -194,7 +194,8 @@ pub struct TestItem {
 #[serde(rename_all = "camelCase")]
 pub struct DiscoverTestResults {
     pub tests: Vec<TestItem>,
-    pub scope: Vec<String>,
+    pub scope: Option<Vec<String>>,
+    pub scope_file: Option<Vec<TextDocumentIdentifier>>,
 }
 
 pub enum DiscoverTest {}
diff --git a/crates/rust-analyzer/src/main_loop.rs b/crates/rust-analyzer/src/main_loop.rs
index 666bf2920e4..38df3235125 100644
--- a/crates/rust-analyzer/src/main_loop.rs
+++ b/crates/rust-analyzer/src/main_loop.rs
@@ -9,9 +9,8 @@ use std::{
 use always_assert::always;
 use crossbeam_channel::{never, select, Receiver};
 use ide_db::base_db::{SourceDatabase, SourceDatabaseExt, VfsPath};
-use itertools::Itertools;
 use lsp_server::{Connection, Notification, Request};
-use lsp_types::notification::Notification as _;
+use lsp_types::{notification::Notification as _, TextDocumentIdentifier};
 use stdx::thread::ThreadIntent;
 use vfs::FileId;
 
@@ -533,22 +532,14 @@ impl GlobalState {
             let snapshot = self.snapshot();
             move || {
                 let tests = subscriptions
-                    .into_iter()
-                    .filter_map(|f| snapshot.analysis.crates_for(f).ok())
-                    .flatten()
-                    .unique()
-                    .filter_map(|c| snapshot.analysis.discover_tests_in_crate(c).ok())
+                    .iter()
+                    .copied()
+                    .filter_map(|f| snapshot.analysis.discover_tests_in_file(f).ok())
                     .flatten()
                     .collect::<Vec<_>>();
                 for t in &tests {
                     hack_recover_crate_name::insert_name(t.id.clone());
                 }
-                let scope = tests
-                    .iter()
-                    .filter_map(|t| Some(t.id.split_once("::")?.0))
-                    .unique()
-                    .map(|it| it.to_owned())
-                    .collect();
                 Task::DiscoverTest(lsp_ext::DiscoverTestResults {
                     tests: tests
                         .into_iter()
@@ -557,7 +548,13 @@ impl GlobalState {
                             to_proto::test_item(&snapshot, t, line_index.as_ref())
                         })
                         .collect(),
-                    scope,
+                    scope: None,
+                    scope_file: Some(
+                        subscriptions
+                            .into_iter()
+                            .map(|f| TextDocumentIdentifier { uri: to_proto::url(&snapshot, f) })
+                            .collect(),
+                    ),
                 })
             }
         });
diff --git a/docs/dev/lsp-extensions.md b/docs/dev/lsp-extensions.md
index 8db66687aae..939b1819c7e 100644
--- a/docs/dev/lsp-extensions.md
+++ b/docs/dev/lsp-extensions.md
@@ -1,5 +1,5 @@
 <!---
-lsp/ext.rs hash: d5febcbf63650753
+lsp/ext.rs hash: 223f48a89a5126a0
 
 If you need to change the above hash to make the test pass, please check if you
 need to adjust this doc as well and ping this issue:
@@ -440,7 +440,11 @@ interface DiscoverTestResults {
     // For each test which its id is in this list, the response
     // contains all tests that are children of this test, and
     // client should remove old tests not included in the response.
-    scope: string[];
+    scope: string[] | undefined;
+    // For each file which its uri is in this list, the response
+    // contains all tests that are located in this file, and
+    // client should remove old tests not included in the response.
+    scopeFile: lc.TextDocumentIdentifier[] | undefined;    
 }
 ```
 
diff --git a/editors/code/src/lsp_ext.ts b/editors/code/src/lsp_ext.ts
index ca8106371b0..9a7a4aae959 100644
--- a/editors/code/src/lsp_ext.ts
+++ b/editors/code/src/lsp_ext.ts
@@ -83,7 +83,11 @@ export type TestItem = {
     range?: lc.Range | undefined;
     runnable?: Runnable | undefined;
 };
-export type DiscoverTestResults = { tests: TestItem[]; scope: string[] };
+export type DiscoverTestResults = {
+    tests: TestItem[];
+    scope: string[] | undefined;
+    scopeFile: lc.TextDocumentIdentifier[] | undefined;
+};
 export type TestState =
     | { tag: "failed"; message: string }
     | { tag: "passed" }
diff --git a/editors/code/src/test_explorer.ts b/editors/code/src/test_explorer.ts
index ae5351a1858..de41d2a57ec 100644
--- a/editors/code/src/test_explorer.ts
+++ b/editors/code/src/test_explorer.ts
@@ -12,6 +12,7 @@ export const prepareTestExplorer = (
 ) => {
     let currentTestRun: vscode.TestRun | undefined;
     let idToTestMap: Map<string, vscode.TestItem> = new Map();
+    const fileToTestMap: Map<string, vscode.TestItem[]> = new Map();
     const idToRunnableMap: Map<string, ra.Runnable> = new Map();
 
     testController.createRunProfile(
@@ -59,6 +60,18 @@ export const prepareTestExplorer = (
         false,
     );
 
+    const deleteTest = (item: vscode.TestItem, parentList: vscode.TestItemCollection) => {
+        parentList.delete(item.id);
+        idToTestMap.delete(item.id);
+        idToRunnableMap.delete(item.id);
+        if (item.uri) {
+            fileToTestMap.set(
+                item.uri.toString(),
+                fileToTestMap.get(item.uri.toString())!.filter((t) => t.id !== item.id),
+            );
+        }
+    };
+
     const addTest = (item: ra.TestItem) => {
         const parentList = item.parent
             ? idToTestMap.get(item.parent)!.children
@@ -76,7 +89,7 @@ export const prepareTestExplorer = (
                 oldTest.range = range;
                 return;
             }
-            parentList.delete(item.id);
+            deleteTest(oldTest, parentList);
         }
         const iconToVscodeMap = {
             package: "package",
@@ -91,6 +104,12 @@ export const prepareTestExplorer = (
         test.range = range;
         test.canResolveChildren = item.canResolveChildren;
         idToTestMap.set(item.id, test);
+        if (uri) {
+            if (!fileToTestMap.has(uri.toString())) {
+                fileToTestMap.set(uri.toString(), []);
+            }
+            fileToTestMap.get(uri.toString())!.push(test);
+        }
         if (item.runnable) {
             idToRunnableMap.set(item.id, item.runnable);
         }
@@ -98,7 +117,7 @@ export const prepareTestExplorer = (
     };
 
     const addTestGroup = (testsAndScope: ra.DiscoverTestResults) => {
-        const { tests, scope } = testsAndScope;
+        const { tests, scope, scopeFile } = testsAndScope;
         const testSet: Set<string> = new Set();
         for (const test of tests) {
             addTest(test);
@@ -107,25 +126,38 @@ export const prepareTestExplorer = (
         // FIXME(hack_recover_crate_name): We eagerly resolve every test if we got a lazy top level response (detected
         // by checking that `scope` is empty). This is not a good thing and wastes cpu and memory unnecessarily, so we
         // should remove it.
-        if (scope.length === 0) {
+        if (!scope && !scopeFile) {
             for (const test of tests) {
                 void testController.resolveHandler!(idToTestMap.get(test.id));
             }
         }
-        if (!scope) {
-            return;
+        if (scope) {
+            const recursivelyRemove = (tests: vscode.TestItemCollection) => {
+                for (const [_, test] of tests) {
+                    if (!testSet.has(test.id)) {
+                        deleteTest(test, tests);
+                    } else {
+                        recursivelyRemove(test.children);
+                    }
+                }
+            };
+            for (const root of scope) {
+                recursivelyRemove(idToTestMap.get(root)!.children);
+            }
         }
-        const recursivelyRemove = (tests: vscode.TestItemCollection) => {
-            for (const [testId, _] of tests) {
-                if (!testSet.has(testId)) {
-                    tests.delete(testId);
-                } else {
-                    recursivelyRemove(tests.get(testId)!.children);
+        if (scopeFile) {
+            const removeByFile = (file: vscode.Uri) => {
+                const testsToBeRemoved = (fileToTestMap.get(file.toString()) || []).filter(
+                    (t) => !testSet.has(t.id),
+                );
+                for (const test of testsToBeRemoved) {
+                    const parentList = test.parent?.children || testController.items;
+                    deleteTest(test, parentList);
                 }
+            };
+            for (const file of scopeFile) {
+                removeByFile(vscode.Uri.parse(file.uri));
             }
-        };
-        for (const root of scope) {
-            recursivelyRemove(idToTestMap.get(root)!.children);
         }
     };