3 files changed, 27 insertions, 9 deletions

diff --git a/crates/hir_expand/Cargo.toml b/crates/hir_expand/Cargo.toml
index b535a3d4f..5271110d2 100644
--- a/crates/hir_expand/Cargo.toml
+++ b/crates/hir_expand/Cargo.toml
@@ -13,7 +13,7 @@ doctest = false
 log = "0.4.8"
 either = "1.5.3"
 rustc-hash = "1.0.0"
-la-arena = "0.1.0"
+la-arena = { version = "0.2.0", path = "../../lib/arena" }
 
 base_db = { path = "../base_db", version = "0.0.0" }
 syntax = { path = "../syntax", version = "0.0.0" }
diff --git a/crates/hir_expand/src/ast_id_map.rs b/crates/hir_expand/src/ast_id_map.rs
index f4f6e11fd..2401b0cc5 100644
--- a/crates/hir_expand/src/ast_id_map.rs
+++ b/crates/hir_expand/src/ast_id_map.rs
@@ -72,10 +72,12 @@ impl AstIdMap {
         // 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.
-        bfs(node, |it| {
-            if let Some(module_item) = ast::Item::cast(it) {
+        bdfs(node, |it| match ast::Item::cast(it) {
+            Some(module_item) => {
                 res.alloc(module_item.syntax());
+                true
             }
+            None => false,
         });
         res
     }
@@ -105,14 +107,30 @@ impl AstIdMap {
     }
 }
 
-/// Walks the subtree in bfs order, calling `f` for each node.
-fn bfs(node: &SyntaxNode, mut f: impl FnMut(SyntaxNode)) {
+/// Walks the subtree in bdfs order, calling `f` for each node. What is bdfs
+/// order? It is a mix of breadth-first and depth first orders. Nodes for which
+/// `f` returns true are visited breadth-first, all the other nodes are explored
+/// depth-first.
+///
+/// In other words, the size of the bfs queue is bound by the number of "true"
+/// nodes.
+fn bdfs(node: &SyntaxNode, mut f: impl FnMut(SyntaxNode) -> bool) {
     let mut curr_layer = vec![node.clone()];
     let mut next_layer = vec![];
     while !curr_layer.is_empty() {
         curr_layer.drain(..).for_each(|node| {
-            next_layer.extend(node.children());
-            f(node);
+            let mut preorder = node.preorder();
+            while let Some(event) = preorder.next() {
+                match event {
+                    syntax::WalkEvent::Enter(node) => {
+                        if f(node.clone()) {
+                            next_layer.extend(node.children());
+                            preorder.skip_subtree();
+                        }
+                    }
+                    syntax::WalkEvent::Leave(_) => {}
+                }
+            }
         });
         std::mem::swap(&mut curr_layer, &mut next_layer);
     }
diff --git a/crates/hir_expand/src/proc_macro.rs b/crates/hir_expand/src/proc_macro.rs
index 1923daca5..75e950816 100644
--- a/crates/hir_expand/src/proc_macro.rs
+++ b/crates/hir_expand/src/proc_macro.rs
@@ -135,7 +135,6 @@ mod tests {
         let result = format!("{:#?}", remove_derive_attrs(&tt).unwrap());
 
         assert_eq_text!(
-            &result,
             r#"
 SUBTREE $
   PUNCH   # [alone] 0
@@ -150,7 +149,8 @@ SUBTREE $
     PUNCH   : [alone] 19
     IDENT   u32 20
 "#
-            .trim()
+            .trim(),
+            &result
         );
     }
 }