4 files changed, 30 insertions, 12 deletions

diff --git a/crates/hir_expand/Cargo.toml b/crates/hir_expand/Cargo.toml
index b535a3d4f..1b1c7da65 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.1.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/db.rs b/crates/hir_expand/src/db.rs
index c62086390..467516eb7 100644
--- a/crates/hir_expand/src/db.rs
+++ b/crates/hir_expand/src/db.rs
@@ -118,7 +118,7 @@ pub fn expand_hypothetical(
         parse_macro_with_arg(db, macro_file, Some(std::sync::Arc::new((tt, tmap_1)))).value?;
     let token_id = macro_def.0.map_id_down(token_id);
     let range = tmap_2.range_by_token(token_id)?.by_kind(token_to_map.kind())?;
-    let token = syntax::algo::find_covering_element(&node.syntax_node(), range).into_token()?;
+    let token = node.syntax_node().covering_element(range).into_token()?;
     Some((node.syntax_node(), token))
 }
 
diff --git a/crates/hir_expand/src/lib.rs b/crates/hir_expand/src/lib.rs
index 3fa1b1d77..e388ddacc 100644
--- a/crates/hir_expand/src/lib.rs
+++ b/crates/hir_expand/src/lib.rs
@@ -22,7 +22,7 @@ use std::sync::Arc;
 
 use base_db::{impl_intern_key, salsa, CrateId, FileId, FileRange};
 use syntax::{
-    algo::{self, skip_trivia_token},
+    algo::skip_trivia_token,
     ast::{self, AstNode},
     Direction, SyntaxNode, SyntaxToken, TextRange, TextSize,
 };
@@ -335,7 +335,7 @@ impl ExpansionInfo {
 
         let range = self.exp_map.range_by_token(token_id)?.by_kind(token.value.kind())?;
 
-        let token = algo::find_covering_element(&self.expanded.value, range).into_token()?;
+        let token = self.expanded.value.covering_element(range).into_token()?;
 
         Some(self.expanded.with_value(token))
     }
@@ -360,8 +360,8 @@ impl ExpansionInfo {
         };
 
         let range = token_map.range_by_token(token_id)?.by_kind(token.value.kind())?;
-        let token = algo::find_covering_element(&tt.value, range + tt.value.text_range().start())
-            .into_token()?;
+        let token =
+            tt.value.covering_element(range + tt.value.text_range().start()).into_token()?;
         Some((tt.with_value(token), origin))
     }
 }