From 278e7c33115a448eb7bbe6420404331c1820f1af Mon Sep 17 00:00:00 2001
From: Aleksey Kladov <aleksey.kladov@gmail.com>
Date: Sun, 5 Dec 2021 17:19:48 +0300
Subject: [PATCH] more frugal map

---
 Cargo.lock                          |  2 +-
 crates/hir_expand/Cargo.toml        |  2 +-
 crates/hir_expand/src/ast_id_map.rs | 41 +++++++++++++++++++----------
 3 files changed, 29 insertions(+), 16 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 4d22ebdc181..173bbb801dd 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -525,11 +525,11 @@ dependencies = [
  "cfg",
  "cov-mark",
  "either",
+ "hashbrown",
  "itertools",
  "la-arena",
  "limit",
  "mbe",
- "once_cell",
  "profile",
  "rustc-hash",
  "syntax",
diff --git a/crates/hir_expand/Cargo.toml b/crates/hir_expand/Cargo.toml
index 7a610365306..039861f7f7b 100644
--- a/crates/hir_expand/Cargo.toml
+++ b/crates/hir_expand/Cargo.toml
@@ -16,7 +16,7 @@ either = "1.5.3"
 rustc-hash = "1.0.0"
 la-arena = { version = "0.3.0", path = "../../lib/arena" }
 itertools = "0.10.0"
-once_cell = "1"
+hashbrown = { version = "0.11", features = ["inline-more"], default-features = false }
 
 base_db = { path = "../base_db", version = "0.0.0" }
 cfg = { path = "../cfg", 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 9db2bc64101..43c8a3db5c5 100644
--- a/crates/hir_expand/src/ast_id_map.rs
+++ b/crates/hir_expand/src/ast_id_map.rs
@@ -8,13 +8,13 @@
 use std::{
     any::type_name,
     fmt,
-    hash::{Hash, Hasher},
+    hash::{BuildHasher, BuildHasherDefault, Hash, Hasher},
     marker::PhantomData,
 };
 
 use la_arena::{Arena, Idx};
 use profile::Count;
-use rustc_hash::FxHashMap;
+use rustc_hash::FxHasher;
 use syntax::{ast, match_ast, AstNode, AstPtr, SyntaxNode, SyntaxNodePtr};
 
 /// `AstId` points to an AST node in a specific file.
@@ -63,11 +63,10 @@ pub fn upcast<M: AstNode>(self) -> FileAstId<M>
 /// Maps items' `SyntaxNode`s to `ErasedFileAstId`s and back.
 #[derive(Default)]
 pub struct AstIdMap {
+    /// Maps stable id to unstable ptr.
     arena: Arena<SyntaxNodePtr>,
-    /// Reversed mapping lazily derived from [`self.arena`].
-    ///
-    /// FIXE: Do not store `SyntaxNodePtr` twice.
-    map: once_cell::sync::OnceCell<FxHashMap<SyntaxNodePtr, ErasedFileAstId>>,
+    /// Reverse: map ptr to id.
+    map: hashbrown::HashMap<Idx<SyntaxNodePtr>, (), ()>,
     _c: Count<Self>,
 }
 
@@ -107,26 +106,34 @@ pub(crate) fn from_source(node: &SyntaxNode) -> AstIdMap {
                 }
             }
         });
+        res.map = hashbrown::HashMap::with_capacity_and_hasher(res.arena.len(), ());
+        for (idx, ptr) in res.arena.iter() {
+            let hash = hash_ptr(ptr);
+            match res.map.raw_entry_mut().from_hash(hash, |idx2| *idx2 == idx) {
+                hashbrown::hash_map::RawEntryMut::Occupied(_) => unreachable!(),
+                hashbrown::hash_map::RawEntryMut::Vacant(entry) => {
+                    entry.insert_with_hasher(hash, idx, (), |&idx| hash_ptr(&res.arena[idx]));
+                }
+            }
+        }
         res
     }
-    fn map(&self) -> &FxHashMap<SyntaxNodePtr, ErasedFileAstId> {
-        self.map.get_or_init(|| self.arena.iter().map(|(idx, ptr)| (ptr.clone(), idx)).collect())
-    }
 
     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);
-        *self.map().get(&ptr).unwrap_or_else(|| {
-            panic!(
+        let hash = hash_ptr(&ptr);
+        match self.map.raw_entry().from_hash(hash, |&idx| self.arena[idx] == ptr) {
+            Some((&idx, &())) => idx,
+            None => 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> {
@@ -138,6 +145,12 @@ fn alloc(&mut self, item: &SyntaxNode) -> ErasedFileAstId {
     }
 }
 
+fn hash_ptr(ptr: &SyntaxNodePtr) -> u64 {
+    let mut hasher = BuildHasherDefault::<FxHasher>::default().build_hasher();
+    ptr.hash(&mut hasher);
+    hasher.finish()
+}
+
 /// 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
-- 
2.44.0