Derive block attributes from block item tree

Author: Veykril

crates/hir-def/src/attr.rs

@@ -431,12 +431,10 @@ impl AttrsWithOwner {
                         .item_tree(db)
                         .raw_attrs(AttrOwner::ModItem(definition_tree_id.value.into()))
                         .clone(),
-                    ModuleOrigin::BlockExpr { block } => RawAttrs::from_attrs_owner(
-                        db.upcast(),
-                        InFile::new(block.file_id, block.to_node(db.upcast()))
-                            .as_ref()
-                            .map(|it| it as &dyn ast::HasAttrs),
-                    ),
+                    ModuleOrigin::BlockExpr { id, block } => {
+                        let tree = db.block_item_tree_query(id);
+                        tree.raw_attrs(AttrOwner::TopLevel).clone()
+                    }
                 }
             }
             AttrDefId::FieldId(it) => {

crates/hir-def/src/body/lower.rs

@@ -1100,7 +1100,9 @@ impl ExprCollector<'_> {
                 ast::Stmt::ExprStmt(es) => matches!(es.expr(), Some(ast::Expr::MacroExpr(_))),
                 _ => false,
             });
-            statement_has_item || matches!(block.tail_expr(), Some(ast::Expr::MacroExpr(_)))
+            statement_has_item
+                || matches!(block.tail_expr(), Some(ast::Expr::MacroExpr(_)))
+                || (block.may_carry_attributes() && block.attrs().next().is_some())
         };
 
         let block_id = if block_has_items {

crates/hir-def/src/db.rs

@@ -82,6 +82,12 @@ pub trait DefDatabase: InternDatabase + ExpandDatabase + Upcast<dyn ExpandDataba
     #[salsa::invoke(ItemTree::file_item_tree_query)]
     fn file_item_tree(&self, file_id: HirFileId) -> Arc<ItemTree>;
 
+    #[salsa::invoke(ItemTree::block_item_tree_query)]
+    // FIXME: Investigate memory usage increase if this were not transparent
+    // Also make sure to `shrink_to_fit` if you do
+    #[salsa::transparent]
+    fn block_item_tree_query(&self, block_id: BlockId) -> Arc<ItemTree>;
+
     #[salsa::invoke(crate_def_map_wait)]
     #[salsa::transparent]
     fn crate_def_map(&self, krate: CrateId) -> Arc<DefMap>;

crates/hir-def/src/item_tree.rs

@@ -143,6 +143,14 @@ impl ItemTree {
         Arc::new(item_tree)
     }
 
+    pub(crate) fn block_item_tree_query(db: &dyn DefDatabase, block: BlockId) -> Arc<ItemTree> {
+        let loc = db.lookup_intern_block(block);
+        let block = loc.ast_id.to_node(db.upcast());
+
+        let ctx = lower::Ctx::new(db, loc.ast_id.file_id);
+        Arc::new(ctx.lower_block(&block))
+    }
+
     /// Returns an iterator over all items located at the top level of the `HirFileId` this
     /// `ItemTree` was created from.
     pub fn top_level_items(&self) -> &[ModItem] {
@@ -178,13 +186,6 @@ impl ItemTree {
         self.data.get_or_insert_with(Box::default)
     }
 
-    fn block_item_tree(db: &dyn DefDatabase, block: BlockId) -> Arc<ItemTree> {
-        let loc = db.lookup_intern_block(block);
-        let block = loc.ast_id.to_node(db.upcast());
-        let ctx = lower::Ctx::new(db, loc.ast_id.file_id);
-        Arc::new(ctx.lower_block(&block))
-    }
-
     fn shrink_to_fit(&mut self) {
         if let Some(data) = &mut self.data {
             let ItemTreeData {
@@ -382,7 +383,7 @@ impl TreeId {
 
     pub(crate) fn item_tree(&self, db: &dyn DefDatabase) -> Arc<ItemTree> {
         match self.block {
-            Some(block) => ItemTree::block_item_tree(db, block),
+            Some(block) => ItemTree::block_item_tree_query(db, block),
             None => db.file_item_tree(self.file),
         }
     }

crates/hir-def/src/item_tree/lower.rs

@@ -77,6 +77,9 @@ impl<'a> Ctx<'a> {
     }
 
     pub(super) fn lower_block(mut self, block: &ast::BlockExpr) -> ItemTree {
+        self.tree
+            .attrs
+            .insert(AttrOwner::TopLevel, RawAttrs::new(self.db.upcast(), block, self.hygiene()));
         self.tree.top_level = block
             .statements()
             .filter_map(|stmt| match stmt {

crates/hir-def/src/nameres.rs

@@ -227,6 +227,7 @@ pub enum ModuleOrigin {
     },
     /// Pseudo-module introduced by a block scope (contains only inner items).
     BlockExpr {
+        id: BlockId,
         block: AstId<ast::BlockExpr>,
     },
 }
@@ -269,7 +270,7 @@ impl ModuleOrigin {
                 definition.file_id,
                 ModuleSource::Module(definition.to_node(db.upcast())),
             ),
-            ModuleOrigin::BlockExpr { block } => {
+            ModuleOrigin::BlockExpr { block, .. } => {
                 InFile::new(block.file_id, ModuleSource::BlockExpr(block.to_node(db.upcast())))
             }
         }
@@ -325,8 +326,10 @@ impl DefMap {
         // modules declared by blocks with items. At the moment, we don't use
         // this visibility for anything outside IDE, so that's probably OK.
         let visibility = Visibility::Module(ModuleId { krate, local_id, block: None });
-        let module_data =
-            ModuleData::new(ModuleOrigin::BlockExpr { block: block.ast_id }, visibility);
+        let module_data = ModuleData::new(
+            ModuleOrigin::BlockExpr { block: block.ast_id, id: block_id },
+            visibility,
+        );
 
         let mut def_map = DefMap::empty(krate, parent_map.data.edition, module_data);
         def_map.data = parent_map.data.clone();
@@ -643,7 +646,7 @@ impl ModuleData {
                 definition.into()
             }
             ModuleOrigin::Inline { definition, .. } => definition.file_id,
-            ModuleOrigin::BlockExpr { block } => block.file_id,
+            ModuleOrigin::BlockExpr { block, .. } => block.file_id,
         }
     }
 

crates/hir-expand/src/attrs.rs

@@ -342,14 +342,7 @@ fn inner_attributes(
             ast::Impl(it) => it.assoc_item_list()?.syntax().clone(),
             ast::Module(it) => it.item_list()?.syntax().clone(),
             ast::BlockExpr(it) => {
-                use syntax::SyntaxKind::{BLOCK_EXPR , EXPR_STMT};
-                // Block expressions accept outer and inner attributes, but only when they are the outer
-                // expression of an expression statement or the final expression of another block expression.
-                let may_carry_attributes = matches!(
-                    it.syntax().parent().map(|it| it.kind()),
-                     Some(BLOCK_EXPR | EXPR_STMT)
-                );
-                if !may_carry_attributes {
+                if !it.may_carry_attributes() {
                     return None
                 }
                 syntax.clone()

crates/syntax/src/ast/node_ext.rs

@@ -61,6 +61,14 @@ impl ast::BlockExpr {
     pub fn tail_expr(&self) -> Option<ast::Expr> {
         self.stmt_list()?.tail_expr()
     }
+    /// Block expressions accept outer and inner attributes, but only when they are the outer
+    /// expression of an expression statement or the final expression of another block expression.
+    pub fn may_carry_attributes(&self) -> bool {
+        matches!(
+            self.syntax().parent().map(|it| it.kind()),
+            Some(SyntaxKind::BLOCK_EXPR | SyntaxKind::EXPR_STMT)
+        )
+    }
 }
 
 #[derive(Debug, PartialEq, Eq, Clone)]