2 files changed, 99 insertions, 104 deletions

diff --git a/crates/hir_def/src/body/tests.rs b/crates/hir_def/src/body/tests.rs
index a92134ba7..bb43569d7 100644
--- a/crates/hir_def/src/body/tests.rs
+++ b/crates/hir_def/src/body/tests.rs
@@ -1,10 +1,10 @@
 mod block;
 
-use base_db::{fixture::WithFixture, FilePosition, SourceDatabase};
+use base_db::{fixture::WithFixture, SourceDatabase};
 use expect_test::Expect;
 use test_utils::mark;
 
-use crate::{test_db::TestDB, BlockId, ModuleDefId};
+use crate::{test_db::TestDB, ModuleDefId};
 
 use super::*;
 
@@ -37,104 +37,8 @@ fn check_diagnostics(ra_fixture: &str) {
 fn block_def_map_at(ra_fixture: &str) -> String {
     let (db, position) = crate::test_db::TestDB::with_position(ra_fixture);
 
-    let krate = db.crate_graph().iter().next().unwrap();
-    let def_map = db.crate_def_map(krate);
-
-    let mut block =
-        block_at_pos(&db, &def_map, position).expect("couldn't find enclosing function or block");
-    loop {
-        let def_map = db.block_def_map(block).unwrap_or_else(|| def_map.clone());
-        let new_block = block_at_pos(&db, &def_map, position);
-        match new_block {
-            Some(new_block) => {
-                assert_ne!(block, new_block);
-                block = new_block;
-            }
-            None => {
-                return def_map.dump(&db);
-            }
-        }
-    }
-}
-
-fn block_at_pos(db: &dyn DefDatabase, def_map: &DefMap, position: FilePosition) -> Option<BlockId> {
-    // Find the smallest (innermost) function containing the cursor.
-    let mut size = None;
-    let mut fn_def = None;
-    for (_, module) in def_map.modules() {
-        let file_id = module.definition_source(db).file_id;
-        if file_id != position.file_id.into() {
-            continue;
-        }
-        let root = db.parse_or_expand(file_id).unwrap();
-        let ast_map = db.ast_id_map(file_id);
-        let item_tree = db.item_tree(file_id);
-        for decl in module.scope.declarations() {
-            if let ModuleDefId::FunctionId(it) = decl {
-                let ast = ast_map.get(item_tree[it.lookup(db).id.value].ast_id).to_node(&root);
-                let range = ast.syntax().text_range();
-
-                if !range.contains(position.offset) {
-                    continue;
-                }
-
-                let new_size = match size {
-                    None => range.len(),
-                    Some(size) => {
-                        if range.len() < size {
-                            range.len()
-                        } else {
-                            size
-                        }
-                    }
-                };
-                if size != Some(new_size) {
-                    size = Some(new_size);
-                    fn_def = Some(it);
-                }
-            }
-        }
-    }
-
-    let (body, source_map) = db.body_with_source_map(fn_def?.into());
-
-    // Now find the smallest encompassing block expression in the function body.
-    let mut size = None;
-    let mut block_id = None;
-    for (expr_id, expr) in body.exprs.iter() {
-        if let Expr::Block { id, .. } = expr {
-            if let Ok(ast) = source_map.expr_syntax(expr_id) {
-                if ast.file_id != position.file_id.into() {
-                    continue;
-                }
-
-                let root = db.parse_or_expand(ast.file_id).unwrap();
-                let ast = ast.value.to_node(&root);
-                let range = ast.syntax().text_range();
-
-                if !range.contains(position.offset) {
-                    continue;
-                }
-
-                let new_size = match size {
-                    None => range.len(),
-                    Some(size) => {
-                        if range.len() < size {
-                            range.len()
-                        } else {
-                            size
-                        }
-                    }
-                };
-                if size != Some(new_size) {
-                    size = Some(new_size);
-                    block_id = Some(*id);
-                }
-            }
-        }
-    }
-
-    Some(block_id.expect("can't find block containing cursor"))
+    let module = db.module_at_position(position);
+    module.def_map(&db).dump(&db)
 }
 
 fn check_at(ra_fixture: &str, expect: Expect) {
diff --git a/crates/hir_def/src/test_db.rs b/crates/hir_def/src/test_db.rs
index 6665d902d..eda982c85 100644
--- a/crates/hir_def/src/test_db.rs
+++ b/crates/hir_def/src/test_db.rs
@@ -5,17 +5,17 @@ use std::{
     sync::{Arc, Mutex},
 };
 
-use base_db::{salsa, CrateId, FileId, FileLoader, FileLoaderDelegate, Upcast};
+use base_db::{salsa, CrateId, FileId, FileLoader, FileLoaderDelegate, FilePosition, Upcast};
 use base_db::{AnchoredPath, SourceDatabase};
-use hir_expand::db::AstDatabase;
 use hir_expand::diagnostics::Diagnostic;
 use hir_expand::diagnostics::DiagnosticSinkBuilder;
+use hir_expand::{db::AstDatabase, InFile};
 use rustc_hash::FxHashMap;
 use rustc_hash::FxHashSet;
-use syntax::{TextRange, TextSize};
+use syntax::{algo, ast, AstNode, TextRange, TextSize};
 use test_utils::extract_annotations;
 
-use crate::{db::DefDatabase, ModuleDefId, ModuleId};
+use crate::{db::DefDatabase, nameres::DefMap, Lookup, ModuleDefId, ModuleId};
 
 #[salsa::database(
     base_db::SourceDatabaseExtStorage,
@@ -84,6 +84,97 @@ impl TestDB {
         panic!("Can't find module for file")
     }
 
+    pub(crate) fn module_at_position(&self, position: FilePosition) -> ModuleId {
+        let file_module = self.module_for_file(position.file_id);
+        let mut def_map = file_module.def_map(self);
+
+        def_map = match self.block_at_position(&def_map, position) {
+            Some(it) => it,
+            None => return file_module,
+        };
+        loop {
+            let new_map = self.block_at_position(&def_map, position);
+            match new_map {
+                Some(new_block) if !Arc::ptr_eq(&new_block, &def_map) => {
+                    def_map = new_block;
+                }
+                _ => {
+                    // FIXME: handle `mod` inside block expression
+                    return def_map.module_id(def_map.root());
+                }
+            }
+        }
+    }
+
+    fn block_at_position(&self, def_map: &DefMap, position: FilePosition) -> Option<Arc<DefMap>> {
+        // Find the smallest (innermost) function in `def_map` containing the cursor.
+        let mut size = None;
+        let mut fn_def = None;
+        for (_, module) in def_map.modules() {
+            let file_id = module.definition_source(self).file_id;
+            if file_id != position.file_id.into() {
+                continue;
+            }
+            let root = self.parse_or_expand(file_id).unwrap();
+            let ast_map = self.ast_id_map(file_id);
+            let item_tree = self.item_tree(file_id);
+            for decl in module.scope.declarations() {
+                if let ModuleDefId::FunctionId(it) = decl {
+                    let ast =
+                        ast_map.get(item_tree[it.lookup(self).id.value].ast_id).to_node(&root);
+                    let range = ast.syntax().text_range();
+
+                    if !range.contains(position.offset) {
+                        continue;
+                    }
+
+                    let new_size = match size {
+                        None => range.len(),
+                        Some(size) => {
+                            if range.len() < size {
+                                range.len()
+                            } else {
+                                size
+                            }
+                        }
+                    };
+                    if size != Some(new_size) {
+                        size = Some(new_size);
+                        fn_def = Some(it);
+                    }
+                }
+            }
+        }
+
+        // Find the innermost block expression that has a `DefMap`.
+        let def_with_body = fn_def?.into();
+        let (_, source_map) = self.body_with_source_map(def_with_body);
+        let scopes = self.expr_scopes(def_with_body);
+        let root = self.parse(position.file_id);
+
+        let scope_iter = algo::ancestors_at_offset(&root.syntax_node(), position.offset)
+            .filter_map(|node| {
+                let block = ast::BlockExpr::cast(node)?;
+                let expr = ast::Expr::from(block);
+                let expr_id = source_map.node_expr(InFile::new(position.file_id.into(), &expr))?;
+                let scope = scopes.scope_for(expr_id).unwrap();
+                Some(scope)
+            });
+
+        for scope in scope_iter {
+            let containing_blocks =
+                scopes.scope_chain(Some(scope)).filter_map(|scope| scopes.block(scope));
+
+            for block in containing_blocks {
+                if let Some(def_map) = self.block_def_map(block) {
+                    return Some(def_map);
+                }
+            }
+        }
+
+        None
+    }
+
     pub(crate) fn log(&self, f: impl FnOnce()) -> Vec<salsa::Event> {
         *self.events.lock().unwrap() = Some(Vec::new());
         f();