3 files changed, 117 insertions, 68 deletions

diff --git a/crates/hir_def/src/body/tests.rs b/crates/hir_def/src/body/tests.rs
index 2e5d0a01e..da60072ce 100644
--- a/crates/hir_def/src/body/tests.rs
+++ b/crates/hir_def/src/body/tests.rs
@@ -1,7 +1,10 @@
-use base_db::{fixture::WithFixture, SourceDatabase};
+mod block;
+
+use base_db::{fixture::WithFixture, FilePosition, SourceDatabase};
+use expect_test::Expect;
 use test_utils::mark;
 
-use crate::{test_db::TestDB, ModuleDefId};
+use crate::{test_db::TestDB, BlockId, ModuleDefId};
 
 use super::*;
 
@@ -31,6 +34,115 @@ fn check_diagnostics(ra_fixture: &str) {
     db.check_diagnostics();
 }
 
+fn block_def_map_at(ra_fixture: &str) -> Arc<DefMap> {
+    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);
+        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;
+            }
+        }
+    }
+}
+
+fn block_at_pos(db: &dyn DefDatabase, def_map: &DefMap, position: FilePosition) -> Option<BlockId> {
+    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();
+
+                // Find the smallest (innermost) function containing the cursor.
+                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"))
+}
+
+fn check_at(ra_fixture: &str, expect: Expect) {
+    let def_map = block_def_map_at(ra_fixture);
+    let actual = def_map.dump();
+    expect.assert_eq(&actual);
+}
+
 #[test]
 fn your_stack_belongs_to_me() {
     mark::check!(your_stack_belongs_to_me);
diff --git a/crates/hir_def/src/nameres/tests/block.rs b/crates/hir_def/src/body/tests/block.rs
index 6cc659513..6b1ed2555 100644
--- a/crates/hir_def/src/nameres/tests/block.rs
+++ b/crates/hir_def/src/body/tests/block.rs
@@ -1,4 +1,5 @@
 use super::*;
+use expect_test::expect;
 
 #[test]
 fn inner_item_smoke() {
diff --git a/crates/hir_def/src/nameres/tests.rs b/crates/hir_def/src/nameres/tests.rs
index b36d0b59b..723481c36 100644
--- a/crates/hir_def/src/nameres/tests.rs
+++ b/crates/hir_def/src/nameres/tests.rs
@@ -4,16 +4,14 @@ mod macros;
 mod mod_resolution;
 mod diagnostics;
 mod primitives;
-mod block;
 
 use std::sync::Arc;
 
-use base_db::{fixture::WithFixture, FilePosition, SourceDatabase};
+use base_db::{fixture::WithFixture, SourceDatabase};
 use expect_test::{expect, Expect};
-use syntax::AstNode;
 use test_utils::mark;
 
-use crate::{db::DefDatabase, nameres::*, test_db::TestDB, Lookup};
+use crate::{db::DefDatabase, nameres::*, test_db::TestDB};
 
 fn compute_crate_def_map(ra_fixture: &str) -> Arc<DefMap> {
     let db = TestDB::with_files(ra_fixture);
@@ -21,74 +19,12 @@ fn compute_crate_def_map(ra_fixture: &str) -> Arc<DefMap> {
     db.crate_def_map(krate)
 }
 
-fn compute_block_def_map(ra_fixture: &str) -> Arc<DefMap> {
-    let (db, position) = TestDB::with_position(ra_fixture);
-
-    // FIXME: perhaps we should make this use body lowering tests instead?
-
-    let module = db.module_for_file(position.file_id);
-    let mut def_map = db.crate_def_map(module.krate);
-    while let Some(new_def_map) = descend_def_map_at_position(&db, position, def_map.clone()) {
-        def_map = new_def_map;
-    }
-
-    // FIXME: select the right module, not the root
-
-    def_map
-}
-
-fn descend_def_map_at_position(
-    db: &dyn DefDatabase,
-    position: FilePosition,
-    def_map: Arc<DefMap>,
-) -> Option<Arc<DefMap>> {
-    for (local_id, module_data) in def_map.modules() {
-        let mod_def = module_data.origin.definition_source(db);
-        let ast_map = db.ast_id_map(mod_def.file_id);
-        let item_tree = db.item_tree(mod_def.file_id);
-        let root = db.parse_or_expand(mod_def.file_id).unwrap();
-        for item in module_data.scope.declarations() {
-            match item {
-                ModuleDefId::FunctionId(it) => {
-                    // Technically blocks can be inside any type (due to arrays and const generics),
-                    // and also in const/static initializers. For tests we only really care about
-                    // functions though.
-
-                    let ast = ast_map.get(item_tree[it.lookup(db).id.value].ast_id).to_node(&root);
-
-                    if ast.syntax().text_range().contains(position.offset) {
-                        // Cursor inside function, descend into its body's DefMap.
-                        // Note that we don't handle block *expressions* inside function bodies.
-                        let ast_map = db.ast_id_map(position.file_id.into());
-                        let ast_id = ast_map.ast_id(&ast.body().unwrap());
-                        let block = BlockLoc {
-                            ast_id: InFile::new(position.file_id.into(), ast_id),
-                            module: def_map.module_id(local_id),
-                        };
-                        let block_id = db.intern_block(block);
-                        return Some(db.block_def_map(block_id));
-                    }
-                }
-                _ => continue,
-            }
-        }
-    }
-
-    None
-}
-
 fn check(ra_fixture: &str, expect: Expect) {
     let def_map = compute_crate_def_map(ra_fixture);
     let actual = def_map.dump();
     expect.assert_eq(&actual);
 }
 
-fn check_at(ra_fixture: &str, expect: Expect) {
-    let def_map = compute_block_def_map(ra_fixture);
-    let actual = def_map.dump();
-    expect.assert_eq(&actual);
-}
-
 #[test]
 fn crate_def_map_smoke_test() {
     check(