feat: make hightlighting linear

In https://youtu.be/qvIZZf5dmTE, we've noticed that AstIdMap does a
linear lookup when going from SyntaxNode to Id. This leads to
accidentally quadratic overall performance. Replace linear lookup with a
O(1) hashmap lookup.

Future work: don't duplicate `SyntaxNodePtr` in `AstIdMap` and switch to
"call site dependency injection" style storage (eg, store a
`HashSet<ErasedFileAstId>`).

See the explanation of the work here on YouTube
https://youtu.be/wvEgymUm7cY :-)

diff --git a/crates/hir_expand/src/ast_id_map.rs b/crates/hir_expand/src/ast_id_map.rs
index 16cf299076f295601fe704cdd015d498ecfd91ab..64387b8148207e2c1b88795dd8717d43fab0e7ed 100644 (file)
--- a/crates/hir_expand/src/ast_id_map.rs
+++ b/crates/hir_expand/src/ast_id_map.rs
@@ -14,6 +14,7 @@
 
 use la_arena::{Arena, Idx};
 use profile::Count;
+use rustc_hash::FxHashMap;
 use syntax::{ast, match_ast, AstNode, AstPtr, SyntaxNode, SyntaxNodePtr};
 
 /// `AstId` points to an AST node in a specific file.
@@ -63,6 +64,7 @@ pub fn upcast<M: AstNode>(self) -> FileAstId<M>
 #[derive(Debug, PartialEq, Eq, Default)]
 pub struct AstIdMap {
     arena: Arena<SyntaxNodePtr>,
+    map: FxHashMap<SyntaxNodePtr, ErasedFileAstId>,
     _c: Count<Self>,
 }
 
@@ -89,6 +91,7 @@ pub(crate) fn from_source(node: &SyntaxNode) -> AstIdMap {
                 }
             }
         });
+        res.map.extend(res.arena.iter().map(|(idx, ptr)| (ptr.clone(), idx)));
         res
     }
 
@@ -96,16 +99,16 @@ pub fn ast_id<N: AstNode>(&self, item: &N) -> FileAstId<N> {
         let raw = self.erased_ast_id(item.syntax());
         FileAstId { raw, _ty: PhantomData }
     }
+
     fn erased_ast_id(&self, item: &SyntaxNode) -> ErasedFileAstId {
         let ptr = SyntaxNodePtr::new(item);
-        match self.arena.iter().find(|(_id, i)| **i == ptr) {
-            Some((it, _)) => it,
-            None => panic!(
+        *self.map.get(&ptr).unwrap_or_else(|| {
+            panic!(
                 "Can't find {:?} in AstIdMap:\n{:?}",
                 item,
                 self.arena.iter().map(|(_id, i)| i).collect::<Vec<_>>(),
-            ),
-        }
+            )
+        })
     }
 
     pub fn get<N: AstNode>(&self, id: FileAstId<N>) -> AstPtr<N> {