1 files changed, 121 insertions, 29 deletions

diff --git a/crates/hir_def/src/nameres.rs b/crates/hir_def/src/nameres.rs
index 50acc3f54..bd3ea9b8b 100644
--- a/crates/hir_def/src/nameres.rs
+++ b/crates/hir_def/src/nameres.rs
@@ -59,9 +59,10 @@ use std::sync::Arc;
 use base_db::{CrateId, Edition, FileId};
 use hir_expand::{diagnostics::DiagnosticSink, name::Name, InFile};
 use la_arena::Arena;
+use profile::Count;
 use rustc_hash::FxHashMap;
 use stdx::format_to;
-use syntax::ast;
+use syntax::{ast, AstNode};
 
 use crate::{
     db::DefDatabase,
@@ -74,21 +75,23 @@ use crate::{
 
 /// Contains all top-level defs from a macro-expanded crate
 #[derive(Debug, PartialEq, Eq)]
-pub struct CrateDefMap {
-    pub root: LocalModuleId,
-    pub modules: Arena<ModuleData>,
-    pub(crate) krate: CrateId,
+pub struct DefMap {
+    _c: Count<Self>,
+    parent: Option<Arc<DefMap>>,
+    root: LocalModuleId,
+    modules: Arena<ModuleData>,
+    krate: CrateId,
     /// The prelude module for this crate. This either comes from an import
     /// marked with the `prelude_import` attribute, or (in the normal case) from
     /// a dependency (`std` or `core`).
-    pub(crate) prelude: Option<ModuleId>,
-    pub(crate) extern_prelude: FxHashMap<Name, ModuleDefId>,
+    prelude: Option<ModuleId>,
+    extern_prelude: FxHashMap<Name, ModuleDefId>,
 
     edition: Edition,
     diagnostics: Vec<DefDiagnostic>,
 }
 
-impl std::ops::Index<LocalModuleId> for CrateDefMap {
+impl std::ops::Index<LocalModuleId> for DefMap {
     type Output = ModuleData;
     fn index(&self, id: LocalModuleId) -> &ModuleData {
         &self.modules[id]
@@ -109,6 +112,10 @@ pub enum ModuleOrigin {
     Inline {
         definition: AstId<ast::Module>,
     },
+    /// Pseudo-module introduced by a block scope (contains only inner items).
+    BlockExpr {
+        block: AstId<ast::BlockExpr>,
+    },
 }
 
 impl Default for ModuleOrigin {
@@ -122,7 +129,7 @@ impl ModuleOrigin {
         match self {
             ModuleOrigin::File { declaration: module, .. }
             | ModuleOrigin::Inline { definition: module, .. } => Some(*module),
-            ModuleOrigin::CrateRoot { .. } => None,
+            ModuleOrigin::CrateRoot { .. } | ModuleOrigin::BlockExpr { .. } => None,
         }
     }
 
@@ -137,7 +144,7 @@ impl ModuleOrigin {
 
     pub fn is_inline(&self) -> bool {
         match self {
-            ModuleOrigin::Inline { .. } => true,
+            ModuleOrigin::Inline { .. } | ModuleOrigin::BlockExpr { .. } => true,
             ModuleOrigin::CrateRoot { .. } | ModuleOrigin::File { .. } => false,
         }
     }
@@ -155,6 +162,9 @@ impl ModuleOrigin {
                 definition.file_id,
                 ModuleSource::Module(definition.to_node(db.upcast())),
             ),
+            ModuleOrigin::BlockExpr { block } => {
+                InFile::new(block.file_id, ModuleSource::BlockExpr(block.to_node(db.upcast())))
+            }
         }
     }
 }
@@ -169,29 +179,56 @@ pub struct ModuleData {
     pub origin: ModuleOrigin,
 }
 
-impl CrateDefMap {
-    pub(crate) fn crate_def_map_query(db: &dyn DefDatabase, krate: CrateId) -> Arc<CrateDefMap> {
+impl DefMap {
+    pub(crate) fn crate_def_map_query(db: &dyn DefDatabase, krate: CrateId) -> Arc<DefMap> {
         let _p = profile::span("crate_def_map_query").detail(|| {
             db.crate_graph()[krate].display_name.as_deref().unwrap_or_default().to_string()
         });
-        let def_map = {
-            let edition = db.crate_graph()[krate].edition;
-            let mut modules: Arena<ModuleData> = Arena::default();
-            let root = modules.alloc(ModuleData::default());
-            CrateDefMap {
-                krate,
-                edition,
-                extern_prelude: FxHashMap::default(),
-                prelude: None,
-                root,
-                modules,
-                diagnostics: Vec::new(),
-            }
-        };
-        let def_map = collector::collect_defs(db, def_map);
+        let edition = db.crate_graph()[krate].edition;
+        let def_map = DefMap::empty(krate, edition);
+        let def_map = collector::collect_defs(db, def_map, None);
         Arc::new(def_map)
     }
 
+    pub(crate) fn block_def_map_query(
+        db: &dyn DefDatabase,
+        krate: CrateId,
+        block: AstId<ast::BlockExpr>,
+    ) -> Arc<DefMap> {
+        let item_tree = db.item_tree(block.file_id);
+        let block_items = item_tree.inner_items_of_block(block.value);
+
+        let parent = parent_def_map(db, krate, block);
+
+        if block_items.is_empty() {
+            // If there are no inner items, nothing new is brought into scope, so we can just return
+            // the parent DefMap. This keeps DefMap parent chains short.
+            return parent;
+        }
+
+        let mut def_map = DefMap::empty(krate, parent.edition);
+        def_map.parent = Some(parent);
+
+        let def_map = collector::collect_defs(db, def_map, Some(block.value));
+        Arc::new(def_map)
+    }
+
+    fn empty(krate: CrateId, edition: Edition) -> DefMap {
+        let mut modules: Arena<ModuleData> = Arena::default();
+        let root = modules.alloc(ModuleData::default());
+        DefMap {
+            _c: Count::new(),
+            parent: None,
+            krate,
+            edition,
+            extern_prelude: FxHashMap::default(),
+            prelude: None,
+            root,
+            modules,
+            diagnostics: Vec::new(),
+        }
+    }
+
     pub fn add_diagnostics(
         &self,
         db: &dyn DefDatabase,
@@ -208,6 +245,26 @@ impl CrateDefMap {
             .map(|(id, _data)| id)
     }
 
+    pub fn modules(&self) -> impl Iterator<Item = (LocalModuleId, &ModuleData)> + '_ {
+        self.modules.iter()
+    }
+
+    pub fn root(&self) -> LocalModuleId {
+        self.root
+    }
+
+    pub(crate) fn krate(&self) -> CrateId {
+        self.krate
+    }
+
+    pub(crate) fn prelude(&self) -> Option<ModuleId> {
+        self.prelude
+    }
+
+    pub(crate) fn extern_prelude(&self) -> impl Iterator<Item = (&Name, &ModuleDefId)> + '_ {
+        self.extern_prelude.iter()
+    }
+
     pub(crate) fn resolve_path(
         &self,
         db: &dyn DefDatabase,
@@ -224,10 +281,15 @@ impl CrateDefMap {
     // even), as this should be a great debugging aid.
     pub fn dump(&self) -> String {
         let mut buf = String::new();
-        go(&mut buf, self, "crate", self.root);
+        let mut current_map = self;
+        while let Some(parent) = &current_map.parent {
+            go(&mut buf, current_map, "block scope", current_map.root);
+            current_map = &**parent;
+        }
+        go(&mut buf, current_map, "crate", current_map.root);
         return buf;
 
-        fn go(buf: &mut String, map: &CrateDefMap, path: &str, module: LocalModuleId) {
+        fn go(buf: &mut String, map: &DefMap, path: &str, module: LocalModuleId) {
             format_to!(buf, "{}\n", path);
 
             let mut entries: Vec<_> = map.modules[module].scope.resolutions().collect();
@@ -276,10 +338,40 @@ impl ModuleData {
     }
 }
 
+fn parent_def_map(
+    db: &dyn DefDatabase,
+    krate: CrateId,
+    block: AstId<ast::BlockExpr>,
+) -> Arc<DefMap> {
+    // FIXME: store this info in the item tree instead of reparsing here
+    let ast_id_map = db.ast_id_map(block.file_id);
+    let block_ptr = ast_id_map.get(block.value);
+    let root = match db.parse_or_expand(block.file_id) {
+        Some(it) => it,
+        None => {
+            return Arc::new(DefMap::empty(krate, Edition::Edition2018));
+        }
+    };
+    let ast = block_ptr.to_node(&root);
+
+    for ancestor in ast.syntax().ancestors().skip(1) {
+        if let Some(block_expr) = ast::BlockExpr::cast(ancestor) {
+            let ancestor_id = ast_id_map.ast_id(&block_expr);
+            let ast_id = InFile::new(block.file_id, ancestor_id);
+            let parent_map = db.block_def_map(krate, ast_id);
+            return parent_map;
+        }
+    }
+
+    // No enclosing block scope, so the parent is the crate-level DefMap.
+    db.crate_def_map(krate)
+}
+
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub enum ModuleSource {
     SourceFile(ast::SourceFile),
     Module(ast::Module),
+    BlockExpr(ast::BlockExpr),
 }
 
 mod diagnostics {