1 files changed, 62 insertions, 48 deletions

diff --git a/crates/ide/src/runnables.rs b/crates/ide/src/runnables.rs
index 557563d7e..f5ee7de86 100644
--- a/crates/ide/src/runnables.rs
+++ b/crates/ide/src/runnables.rs
@@ -3,7 +3,7 @@ use std::fmt;
 use assists::utils::test_related_attribute;
 use cfg::CfgExpr;
 use hir::{AsAssocItem, HasAttrs, HasSource, Semantics};
-use ide_db::RootDatabase;
+use ide_db::{defs::Definition, RootDatabase};
 use itertools::Itertools;
 use syntax::{
     ast::{self, AstNode, AttrsOwner, ModuleItemOwner},
@@ -96,21 +96,26 @@ impl Runnable {
 pub(crate) fn runnables(db: &RootDatabase, file_id: FileId) -> Vec<Runnable> {
     let sema = Semantics::new(db);
     let source_file = sema.parse(file_id);
-    source_file.syntax().descendants().filter_map(|i| runnable(&sema, i)).collect()
-}
-
-pub(crate) fn runnable(sema: &Semantics<RootDatabase>, item: SyntaxNode) -> Option<Runnable> {
-    let runnable_item = match_ast! {
-        match (item.clone()) {
-            ast::Fn(func) => {
-                let def = sema.to_def(&func)?;
-                runnable_fn(sema, def)
-            },
-            ast::Module(it) => runnable_mod(sema, it),
-            _ => None,
-        }
-    };
-    runnable_item.or_else(|| runnable_doctest(sema, item))
+    source_file
+        .syntax()
+        .descendants()
+        .filter_map(|item| {
+            let runnable = match_ast! {
+                match item {
+                    ast::Fn(func) => {
+                        let def = sema.to_def(&func)?;
+                        runnable_fn(&sema, def)
+                    },
+                    ast::Module(it) => runnable_mod(&sema, it),
+                    _ => None,
+                }
+            };
+            runnable.or_else(|| match doc_owner_to_def(&sema, item)? {
+                Definition::ModuleDef(def) => module_def_doctest(&sema, def),
+                _ => None,
+            })
+        })
+        .collect()
 }
 
 pub(crate) fn runnable_fn(sema: &Semantics<RootDatabase>, def: hir::Function) -> Option<Runnable> {
@@ -145,20 +150,49 @@ pub(crate) fn runnable_fn(sema: &Semantics<RootDatabase>, def: hir::Function) ->
     Some(Runnable { nav, kind, cfg })
 }
 
-fn runnable_doctest(sema: &Semantics<RootDatabase>, item: SyntaxNode) -> Option<Runnable> {
-    match_ast! {
+pub(crate) fn runnable_mod(
+    sema: &Semantics<RootDatabase>,
+    module: ast::Module,
+) -> Option<Runnable> {
+    if !has_test_function_or_multiple_test_submodules(&module) {
+        return None;
+    }
+    let module_def = sema.to_def(&module)?;
+
+    let path = module_def
+        .path_to_root(sema.db)
+        .into_iter()
+        .rev()
+        .filter_map(|it| it.name(sema.db))
+        .join("::");
+
+    let def = sema.to_def(&module)?;
+    let attrs = def.attrs(sema.db);
+    let cfg = attrs.cfg();
+    let nav = module_def.to_nav(sema.db);
+    Some(Runnable { nav, kind: RunnableKind::TestMod { path }, cfg })
+}
+
+// FIXME: figure out a proper API here.
+pub(crate) fn doc_owner_to_def(
+    sema: &Semantics<RootDatabase>,
+    item: SyntaxNode,
+) -> Option<Definition> {
+    let res: hir::ModuleDef = match_ast! {
         match item {
-            ast::Fn(it) => module_def_doctest(sema, sema.to_def(&it)?.into()),
-            ast::Struct(it) => module_def_doctest(sema, sema.to_def(&it)?.into()),
-            ast::Enum(it) => module_def_doctest(sema, sema.to_def(&it)?.into()),
-            ast::Union(it) => module_def_doctest(sema, sema.to_def(&it)?.into()),
-            ast::Trait(it) => module_def_doctest(sema, sema.to_def(&it)?.into()),
-            ast::Const(it) => module_def_doctest(sema, sema.to_def(&it)?.into()),
-            ast::Static(it) => module_def_doctest(sema, sema.to_def(&it)?.into()),
-            ast::TypeAlias(it) => module_def_doctest(sema, sema.to_def(&it)?.into()),
-            _ => None,
+            ast::SourceFile(it) => sema.scope(&item).module()?.into(),
+            ast::Fn(it) => sema.to_def(&it)?.into(),
+            ast::Struct(it) => sema.to_def(&it)?.into(),
+            ast::Enum(it) => sema.to_def(&it)?.into(),
+            ast::Union(it) => sema.to_def(&it)?.into(),
+            ast::Trait(it) => sema.to_def(&it)?.into(),
+            ast::Const(it) => sema.to_def(&it)?.into(),
+            ast::Static(it) => sema.to_def(&it)?.into(),
+            ast::TypeAlias(it) => sema.to_def(&it)?.into(),
+            _ => return None,
         }
-    }
+    };
+    Some(Definition::ModuleDef(res))
 }
 
 fn module_def_doctest(sema: &Semantics<RootDatabase>, def: hir::ModuleDef) -> Option<Runnable> {
@@ -253,26 +287,6 @@ fn has_runnable_doc_test(attrs: &hir::Attrs) -> bool {
     })
 }
 
-fn runnable_mod(sema: &Semantics<RootDatabase>, module: ast::Module) -> Option<Runnable> {
-    if !has_test_function_or_multiple_test_submodules(&module) {
-        return None;
-    }
-    let module_def = sema.to_def(&module)?;
-
-    let path = module_def
-        .path_to_root(sema.db)
-        .into_iter()
-        .rev()
-        .filter_map(|it| it.name(sema.db))
-        .join("::");
-
-    let def = sema.to_def(&module)?;
-    let attrs = def.attrs(sema.db);
-    let cfg = attrs.cfg();
-    let nav = module_def.to_nav(sema.db);
-    Some(Runnable { nav, kind: RunnableKind::TestMod { path }, cfg })
-}
-
 // We could create runnables for modules with number_of_test_submodules > 0,
 // but that bloats the runnables for no real benefit, since all tests can be run by the submodule already
 fn has_test_function_or_multiple_test_submodules(module: &ast::Module) -> bool {