1 files changed, 17 insertions, 7 deletions

diff --git a/crates/hir_expand/src/ast_id_map.rs b/crates/hir_expand/src/ast_id_map.rs
index 2401b0cc5..16cf29907 100644
--- a/crates/hir_expand/src/ast_id_map.rs
+++ b/crates/hir_expand/src/ast_id_map.rs
@@ -13,7 +13,8 @@ use std::{
 };
 
 use la_arena::{Arena, Idx};
-use syntax::{ast, AstNode, AstPtr, SyntaxNode, SyntaxNodePtr};
+use profile::Count;
+use syntax::{ast, match_ast, AstNode, AstPtr, SyntaxNode, SyntaxNodePtr};
 
 /// `AstId` points to an AST node in a specific file.
 pub struct FileAstId<N: AstNode> {
@@ -62,22 +63,31 @@ type ErasedFileAstId = Idx<SyntaxNodePtr>;
 #[derive(Debug, PartialEq, Eq, Default)]
 pub struct AstIdMap {
     arena: Arena<SyntaxNodePtr>,
+    _c: Count<Self>,
 }
 
 impl AstIdMap {
     pub(crate) fn from_source(node: &SyntaxNode) -> AstIdMap {
         assert!(node.parent().is_none());
-        let mut res = AstIdMap { arena: Arena::default() };
+        let mut res = AstIdMap::default();
         // By walking the tree in breadth-first order we make sure that parents
         // get lower ids then children. That is, adding a new child does not
         // change parent's id. This means that, say, adding a new function to a
         // trait does not change ids of top-level items, which helps caching.
-        bdfs(node, |it| match ast::Item::cast(it) {
-            Some(module_item) => {
-                res.alloc(module_item.syntax());
-                true
+        bdfs(node, |it| {
+            match_ast! {
+                match it {
+                    ast::Item(module_item) => {
+                        res.alloc(module_item.syntax());
+                        true
+                    },
+                    ast::BlockExpr(block) => {
+                        res.alloc(block.syntax());
+                        true
+                    },
+                    _ => false,
+                }
             }
-            None => false,
         });
         res
     }