1 files changed, 190 insertions, 52 deletions

diff --git a/crates/hir_def/src/nameres.rs b/crates/hir_def/src/nameres.rs
index 50acc3f54..003d668ca 100644
--- a/crates/hir_def/src/nameres.rs
+++ b/crates/hir_def/src/nameres.rs
@@ -59,6 +59,7 @@ 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;
@@ -69,26 +70,45 @@ use crate::{
     nameres::{diagnostics::DefDiagnostic, path_resolution::ResolveMode},
     path::ModPath,
     per_ns::PerNs,
-    AstId, LocalModuleId, ModuleDefId, ModuleId,
+    AstId, BlockId, BlockLoc, LocalModuleId, ModuleDefId, ModuleId,
 };
 
-/// Contains all top-level defs from a macro-expanded crate
+/// Contains the results of (early) name resolution.
+///
+/// A `DefMap` stores the module tree and the definitions that are in scope in every module after
+/// item-level macros have been expanded.
+///
+/// Every crate has a primary `DefMap` whose root is the crate's main file (`main.rs`/`lib.rs`),
+/// computed by the `crate_def_map` query. Additionally, every block expression introduces the
+/// opportunity to write arbitrary item and module hierarchies, and thus gets its own `DefMap` that
+/// is computed by the `block_def_map` query.
 #[derive(Debug, PartialEq, Eq)]
-pub struct CrateDefMap {
-    pub root: LocalModuleId,
-    pub modules: Arena<ModuleData>,
-    pub(crate) krate: CrateId,
+pub struct DefMap {
+    _c: Count<Self>,
+    block: Option<BlockInfo>,
+    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 {
+/// For `DefMap`s computed for a block expression, this stores its location in the parent map.
+#[derive(Debug, PartialEq, Eq, Clone, Copy)]
+struct BlockInfo {
+    /// The `BlockId` this `DefMap` was created from.
+    block: BlockId,
+    /// The containing module.
+    parent: ModuleId,
+}
+
+impl std::ops::Index<LocalModuleId> for DefMap {
     type Output = ModuleData;
     fn index(&self, id: LocalModuleId) -> &ModuleData {
         &self.modules[id]
@@ -109,6 +129,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 +146,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 +161,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 +179,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 +196,54 @@ 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,
+        block_id: BlockId,
+    ) -> Option<Arc<DefMap>> {
+        let block: BlockLoc = db.lookup_intern_block(block_id);
+
+        let item_tree = db.item_tree(block.ast_id.file_id);
+        if item_tree.inner_items_of_block(block.ast_id.value).is_empty() {
+            return None;
+        }
+
+        let block_info = BlockInfo { block: block_id, parent: block.module };
+
+        let parent_map = block.module.def_map(db);
+        let mut def_map = DefMap::empty(block.module.krate, parent_map.edition);
+        def_map.block = Some(block_info);
+
+        let def_map = collector::collect_defs(db, def_map, Some(block.ast_id));
+        Some(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(),
+            block: None,
+            krate,
+            edition,
+            extern_prelude: FxHashMap::default(),
+            prelude: None,
+            root,
+            modules,
+            diagnostics: Vec::new(),
+        }
+    }
+
     pub fn add_diagnostics(
         &self,
         db: &dyn DefDatabase,
@@ -208,6 +260,46 @@ 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 block_id(&self) -> Option<BlockId> {
+        self.block.as_ref().map(|block| block.block)
+    }
+
+    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 fn module_id(&self, local_id: LocalModuleId) -> ModuleId {
+        let block = self.block.as_ref().map(|b| b.block);
+        ModuleId { krate: self.krate, local_id, block }
+    }
+
+    pub(crate) fn crate_root(&self, db: &dyn DefDatabase) -> ModuleId {
+        self.with_ancestor_maps(db, self.root, &mut |def_map, _module| {
+            if def_map.block.is_none() {
+                Some(def_map.module_id(def_map.root))
+            } else {
+                None
+            }
+        })
+        .expect("DefMap chain without root")
+    }
+
     pub(crate) fn resolve_path(
         &self,
         db: &dyn DefDatabase,
@@ -220,37 +312,68 @@ impl CrateDefMap {
         (res.resolved_def, res.segment_index)
     }
 
+    /// Ascends the `DefMap` hierarchy and calls `f` with every `DefMap` and containing module.
+    ///
+    /// If `f` returns `Some(val)`, iteration is stopped and `Some(val)` is returned. If `f` returns
+    /// `None`, iteration continues.
+    pub fn with_ancestor_maps<T>(
+        &self,
+        db: &dyn DefDatabase,
+        local_mod: LocalModuleId,
+        f: &mut dyn FnMut(&DefMap, LocalModuleId) -> Option<T>,
+    ) -> Option<T> {
+        if let Some(it) = f(self, local_mod) {
+            return Some(it);
+        }
+        let mut block = self.block;
+        while let Some(block_info) = block {
+            let parent = block_info.parent.def_map(db);
+            if let Some(it) = f(&parent, block_info.parent.local_id) {
+                return Some(it);
+            }
+            block = parent.block;
+        }
+
+        None
+    }
+
+    /// If this `DefMap` is for a block expression, returns the module containing the block (which
+    /// might again be a block, or a module inside a block).
+    pub fn parent(&self) -> Option<ModuleId> {
+        Some(self.block?.parent)
+    }
+
+    /// Returns the module containing `local_mod`, either the parent `mod`, or the module containing
+    /// the block, if `self` corresponds to a block expression.
+    pub fn containing_module(&self, local_mod: LocalModuleId) -> Option<ModuleId> {
+        match &self[local_mod].parent {
+            Some(parent) => Some(self.module_id(*parent)),
+            None => match &self.block {
+                Some(block) => Some(block.parent),
+                None => None,
+            },
+        }
+    }
+
     // FIXME: this can use some more human-readable format (ideally, an IR
     // even), as this should be a great debugging aid.
-    pub fn dump(&self) -> String {
+    pub fn dump(&self, db: &dyn DefDatabase) -> String {
         let mut buf = String::new();
-        go(&mut buf, self, "crate", self.root);
+        let mut arc;
+        let mut current_map = self;
+        while let Some(block) = &current_map.block {
+            go(&mut buf, current_map, "block scope", current_map.root);
+            buf.push('\n');
+            arc = block.parent.def_map(db);
+            current_map = &*arc;
+        }
+        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();
-            entries.sort_by_key(|(name, _)| name.clone());
-
-            for (name, def) in entries {
-                format_to!(buf, "{}:", name.map_or("_".to_string(), |name| name.to_string()));
-
-                if def.types.is_some() {
-                    buf.push_str(" t");
-                }
-                if def.values.is_some() {
-                    buf.push_str(" v");
-                }
-                if def.macros.is_some() {
-                    buf.push_str(" m");
-                }
-                if def.is_none() {
-                    buf.push_str(" _");
-                }
-
-                buf.push('\n');
-            }
+            map.modules[module].scope.dump(buf);
 
             for (name, child) in map.modules[module].children.iter() {
                 let path = format!("{}::{}", path, name);
@@ -280,6 +403,7 @@ impl ModuleData {
 pub enum ModuleSource {
     SourceFile(ast::SourceFile),
     Module(ast::Module),
+    BlockExpr(ast::BlockExpr),
 }
 
 mod diagnostics {
@@ -305,6 +429,8 @@ mod diagnostics {
 
         UnresolvedProcMacro { ast: MacroCallKind },
 
+        UnresolvedMacroCall { ast: AstId<ast::MacroCall> },
+
         MacroError { ast: MacroCallKind, message: String },
     }
 
@@ -365,6 +491,13 @@ mod diagnostics {
             Self { in_module: container, kind: DiagnosticKind::MacroError { ast, message } }
         }
 
+        pub(super) fn unresolved_macro_call(
+            container: LocalModuleId,
+            ast: AstId<ast::MacroCall>,
+        ) -> Self {
+            Self { in_module: container, kind: DiagnosticKind::UnresolvedMacroCall { ast } }
+        }
+
         pub(super) fn add_to(
             &self,
             db: &dyn DefDatabase,
@@ -477,6 +610,11 @@ mod diagnostics {
                     });
                 }
 
+                DiagnosticKind::UnresolvedMacroCall { ast } => {
+                    let node = ast.to_node(db.upcast());
+                    sink.push(UnresolvedMacroCall { file: ast.file_id, node: AstPtr::new(&node) });
+                }
+
                 DiagnosticKind::MacroError { ast, message } => {
                     let (file, ast) = match ast {
                         MacroCallKind::FnLike(ast) => {