Merge pull request #20527 from ChayimFriedman2/cache-next-solver

perf: Cache trait solving across queries in the same revision

7 files changed, 140 insertions, 66 deletions

diff --git a/src/tools/rust-analyzer/crates/base-db/src/lib.rs b/src/tools/rust-analyzer/crates/base-db/src/lib.rs
index dbf949c470c..14544acc11b 100644
--- a/src/tools/rust-analyzer/crates/base-db/src/lib.rs
+++ b/src/tools/rust-analyzer/crates/base-db/src/lib.rs
@@ -7,7 +7,12 @@ pub use salsa_macros;
 mod change;
 mod input;
 
-use std::{cell::RefCell, hash::BuildHasherDefault, panic, sync::Once};
+use std::{
+    cell::RefCell,
+    hash::BuildHasherDefault,
+    panic,
+    sync::{Once, atomic::AtomicUsize},
+};
 
 pub use crate::{
     change::FileChange,
@@ -328,6 +333,27 @@ pub trait SourceDatabase: salsa::Database {
 
     #[doc(hidden)]
     fn crates_map(&self) -> Arc<CratesMap>;
+
+    fn nonce_and_revision(&self) -> (Nonce, salsa::Revision);
+}
+
+static NEXT_NONCE: AtomicUsize = AtomicUsize::new(0);
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub struct Nonce(usize);
+
+impl Default for Nonce {
+    #[inline]
+    fn default() -> Self {
+        Nonce::new()
+    }
+}
+
+impl Nonce {
+    #[inline]
+    pub fn new() -> Nonce {
+        Nonce(NEXT_NONCE.fetch_add(1, std::sync::atomic::Ordering::SeqCst))
+    }
 }
 
 /// Crate related data shared by the whole workspace.
diff --git a/src/tools/rust-analyzer/crates/hir-def/src/test_db.rs b/src/tools/rust-analyzer/crates/hir-def/src/test_db.rs
index e30a5b65a1f..1e2f354f975 100644
--- a/src/tools/rust-analyzer/crates/hir-def/src/test_db.rs
+++ b/src/tools/rust-analyzer/crates/hir-def/src/test_db.rs
@@ -3,7 +3,7 @@
 use std::{fmt, panic, sync::Mutex};
 
 use base_db::{
-    Crate, CrateGraphBuilder, CratesMap, FileSourceRootInput, FileText, RootQueryDb,
+    Crate, CrateGraphBuilder, CratesMap, FileSourceRootInput, FileText, Nonce, RootQueryDb,
     SourceDatabase, SourceRoot, SourceRootId, SourceRootInput,
 };
 use hir_expand::{InFile, files::FilePosition};
@@ -20,12 +20,12 @@ use crate::{
 };
 
 #[salsa_macros::db]
-#[derive(Clone)]
 pub(crate) struct TestDB {
     storage: salsa::Storage<Self>,
     files: Arc<base_db::Files>,
     crates_map: Arc<CratesMap>,
     events: Arc<Mutex<Option<Vec<salsa::Event>>>>,
+    nonce: Nonce,
 }
 
 impl Default for TestDB {
@@ -44,6 +44,7 @@ impl Default for TestDB {
             events,
             files: Default::default(),
             crates_map: Default::default(),
+            nonce: Nonce::new(),
         };
         this.set_expand_proc_attr_macros_with_durability(true, Durability::HIGH);
         // This needs to be here otherwise `CrateGraphBuilder` panics.
@@ -53,6 +54,18 @@ impl Default for TestDB {
     }
 }
 
+impl Clone for TestDB {
+    fn clone(&self) -> Self {
+        Self {
+            storage: self.storage.clone(),
+            files: self.files.clone(),
+            crates_map: self.crates_map.clone(),
+            events: self.events.clone(),
+            nonce: Nonce::new(),
+        }
+    }
+}
+
 #[salsa_macros::db]
 impl salsa::Database for TestDB {}
 
@@ -117,6 +130,10 @@ impl SourceDatabase for TestDB {
     fn crates_map(&self) -> Arc<CratesMap> {
         self.crates_map.clone()
     }
+
+    fn nonce_and_revision(&self) -> (Nonce, salsa::Revision) {
+        (self.nonce, salsa::plumbing::ZalsaDatabase::zalsa(self).current_revision())
+    }
 }
 
 impl TestDB {
diff --git a/src/tools/rust-analyzer/crates/hir-ty/src/infer.rs b/src/tools/rust-analyzer/crates/hir-ty/src/infer.rs
index d778cc0e30e..71b33a0e903 100644
--- a/src/tools/rust-analyzer/crates/hir-ty/src/infer.rs
+++ b/src/tools/rust-analyzer/crates/hir-ty/src/infer.rs
@@ -88,54 +88,52 @@ pub(crate) use closure::{CaptureKind, CapturedItem, CapturedItemWithoutTy};
 
 /// The entry point of type inference.
 pub(crate) fn infer_query(db: &dyn HirDatabase, def: DefWithBodyId) -> Arc<InferenceResult> {
-    crate::next_solver::with_new_cache(|| {
-        let _p = tracing::info_span!("infer_query").entered();
-        let resolver = def.resolver(db);
-        let body = db.body(def);
-        let mut ctx = InferenceContext::new(db, def, &body, resolver);
-
-        match def {
-            DefWithBodyId::FunctionId(f) => {
-                ctx.collect_fn(f);
-            }
-            DefWithBodyId::ConstId(c) => ctx.collect_const(c, &db.const_signature(c)),
-            DefWithBodyId::StaticId(s) => ctx.collect_static(&db.static_signature(s)),
-            DefWithBodyId::VariantId(v) => {
-                ctx.return_ty = TyBuilder::builtin(
-                    match db.enum_signature(v.lookup(db).parent).variant_body_type() {
-                        hir_def::layout::IntegerType::Pointer(signed) => match signed {
-                            true => BuiltinType::Int(BuiltinInt::Isize),
-                            false => BuiltinType::Uint(BuiltinUint::Usize),
-                        },
-                        hir_def::layout::IntegerType::Fixed(size, signed) => match signed {
-                            true => BuiltinType::Int(match size {
-                                Integer::I8 => BuiltinInt::I8,
-                                Integer::I16 => BuiltinInt::I16,
-                                Integer::I32 => BuiltinInt::I32,
-                                Integer::I64 => BuiltinInt::I64,
-                                Integer::I128 => BuiltinInt::I128,
-                            }),
-                            false => BuiltinType::Uint(match size {
-                                Integer::I8 => BuiltinUint::U8,
-                                Integer::I16 => BuiltinUint::U16,
-                                Integer::I32 => BuiltinUint::U32,
-                                Integer::I64 => BuiltinUint::U64,
-                                Integer::I128 => BuiltinUint::U128,
-                            }),
-                        },
+    let _p = tracing::info_span!("infer_query").entered();
+    let resolver = def.resolver(db);
+    let body = db.body(def);
+    let mut ctx = InferenceContext::new(db, def, &body, resolver);
+
+    match def {
+        DefWithBodyId::FunctionId(f) => {
+            ctx.collect_fn(f);
+        }
+        DefWithBodyId::ConstId(c) => ctx.collect_const(c, &db.const_signature(c)),
+        DefWithBodyId::StaticId(s) => ctx.collect_static(&db.static_signature(s)),
+        DefWithBodyId::VariantId(v) => {
+            ctx.return_ty = TyBuilder::builtin(
+                match db.enum_signature(v.lookup(db).parent).variant_body_type() {
+                    hir_def::layout::IntegerType::Pointer(signed) => match signed {
+                        true => BuiltinType::Int(BuiltinInt::Isize),
+                        false => BuiltinType::Uint(BuiltinUint::Usize),
                     },
-                );
-            }
+                    hir_def::layout::IntegerType::Fixed(size, signed) => match signed {
+                        true => BuiltinType::Int(match size {
+                            Integer::I8 => BuiltinInt::I8,
+                            Integer::I16 => BuiltinInt::I16,
+                            Integer::I32 => BuiltinInt::I32,
+                            Integer::I64 => BuiltinInt::I64,
+                            Integer::I128 => BuiltinInt::I128,
+                        }),
+                        false => BuiltinType::Uint(match size {
+                            Integer::I8 => BuiltinUint::U8,
+                            Integer::I16 => BuiltinUint::U16,
+                            Integer::I32 => BuiltinUint::U32,
+                            Integer::I64 => BuiltinUint::U64,
+                            Integer::I128 => BuiltinUint::U128,
+                        }),
+                    },
+                },
+            );
         }
+    }
 
-        ctx.infer_body();
+    ctx.infer_body();
 
-        ctx.infer_mut_body();
+    ctx.infer_mut_body();
 
-        ctx.infer_closures();
+    ctx.infer_closures();
 
-        Arc::new(ctx.resolve_all())
-    })
+    Arc::new(ctx.resolve_all())
 }
 
 pub(crate) fn infer_cycle_result(_: &dyn HirDatabase, _: DefWithBodyId) -> Arc<InferenceResult> {
diff --git a/src/tools/rust-analyzer/crates/hir-ty/src/mir/lower.rs b/src/tools/rust-analyzer/crates/hir-ty/src/mir/lower.rs
index 052be11e433..8c03ca939e3 100644
--- a/src/tools/rust-analyzer/crates/hir-ty/src/mir/lower.rs
+++ b/src/tools/rust-analyzer/crates/hir-ty/src/mir/lower.rs
@@ -2168,9 +2168,7 @@ pub fn mir_body_query(db: &dyn HirDatabase, def: DefWithBodyId) -> Result<Arc<Mi
     let _p = tracing::info_span!("mir_body_query", ?detail).entered();
     let body = db.body(def);
     let infer = db.infer(def);
-    let mut result = crate::next_solver::with_new_cache(|| {
-        lower_to_mir(db, def, &body, &infer, body.body_expr)
-    })?;
+    let mut result = lower_to_mir(db, def, &body, &infer, body.body_expr)?;
     result.shrink_to_fit();
     Ok(Arc::new(result))
 }
diff --git a/src/tools/rust-analyzer/crates/hir-ty/src/next_solver/interner.rs b/src/tools/rust-analyzer/crates/hir-ty/src/next_solver/interner.rs
index 11c79ff742d..6e1a5e96455 100644
--- a/src/tools/rust-analyzer/crates/hir-ty/src/next_solver/interner.rs
+++ b/src/tools/rust-analyzer/crates/hir-ty/src/next_solver/interner.rs
@@ -2148,37 +2148,48 @@ TrivialTypeTraversalImpls! {
     Placeholder<BoundVar>,
 }
 
-pub(crate) use tls_cache::with_new_cache;
 mod tls_cache {
     use crate::db::HirDatabase;
 
     use super::DbInterner;
+    use base_db::Nonce;
     use rustc_type_ir::search_graph::GlobalCache;
+    use salsa::Revision;
     use std::cell::RefCell;
 
-    scoped_tls::scoped_thread_local!(static GLOBAL_CACHE: RefCell<rustc_type_ir::search_graph::GlobalCache<DbInterner<'static>>>);
+    struct Cache {
+        cache: GlobalCache<DbInterner<'static>>,
+        revision: Revision,
+        db_nonce: Nonce,
+    }
 
-    pub(crate) fn with_new_cache<T>(f: impl FnOnce() -> T) -> T {
-        GLOBAL_CACHE.set(&RefCell::new(GlobalCache::default()), f)
+    thread_local! {
+        static GLOBAL_CACHE: RefCell<Option<Cache>> = const { RefCell::new(None) };
     }
 
     pub(super) fn with_cache<'db, T>(
-        _db: &'db dyn HirDatabase,
+        db: &'db dyn HirDatabase,
         f: impl FnOnce(&mut GlobalCache<DbInterner<'db>>) -> T,
     ) -> T {
-        // SAFETY: No idea
-        let call = move |slot: &RefCell<_>| {
+        GLOBAL_CACHE.with_borrow_mut(|handle| {
+            let (db_nonce, revision) = db.nonce_and_revision();
+            let handle = match handle {
+                Some(handle) => {
+                    if handle.revision != revision || db_nonce != handle.db_nonce {
+                        *handle = Cache { cache: GlobalCache::default(), revision, db_nonce };
+                    }
+                    handle
+                }
+                None => handle.insert(Cache { cache: GlobalCache::default(), revision, db_nonce }),
+            };
+
+            // SAFETY: No idea
             f(unsafe {
                 std::mem::transmute::<
                     &mut GlobalCache<DbInterner<'static>>,
                     &mut GlobalCache<DbInterner<'db>>,
-                >(&mut *slot.borrow_mut())
+                >(&mut handle.cache)
             })
-        };
-        if GLOBAL_CACHE.is_set() {
-            GLOBAL_CACHE.with(call)
-        } else {
-            GLOBAL_CACHE.set(&RefCell::new(GlobalCache::default()), || GLOBAL_CACHE.with(call))
-        }
+        })
     }
 }
diff --git a/src/tools/rust-analyzer/crates/hir-ty/src/test_db.rs b/src/tools/rust-analyzer/crates/hir-ty/src/test_db.rs
index 775136dc0cb..2a92aa52e0c 100644
--- a/src/tools/rust-analyzer/crates/hir-ty/src/test_db.rs
+++ b/src/tools/rust-analyzer/crates/hir-ty/src/test_db.rs
@@ -3,8 +3,8 @@
 use std::{fmt, panic, sync::Mutex};
 
 use base_db::{
-    CrateGraphBuilder, CratesMap, FileSourceRootInput, FileText, RootQueryDb, SourceDatabase,
-    SourceRoot, SourceRootId, SourceRootInput,
+    CrateGraphBuilder, CratesMap, FileSourceRootInput, FileText, Nonce, RootQueryDb,
+    SourceDatabase, SourceRoot, SourceRootId, SourceRootInput,
 };
 
 use hir_def::{ModuleId, db::DefDatabase, nameres::crate_def_map};
@@ -17,12 +17,12 @@ use test_utils::extract_annotations;
 use triomphe::Arc;
 
 #[salsa_macros::db]
-#[derive(Clone)]
 pub(crate) struct TestDB {
     storage: salsa::Storage<Self>,
     files: Arc<base_db::Files>,
     crates_map: Arc<CratesMap>,
     events: Arc<Mutex<Option<Vec<salsa::Event>>>>,
+    nonce: Nonce,
 }
 
 impl Default for TestDB {
@@ -41,6 +41,7 @@ impl Default for TestDB {
             events,
             files: Default::default(),
             crates_map: Default::default(),
+            nonce: Nonce::new(),
         };
         this.set_expand_proc_attr_macros_with_durability(true, Durability::HIGH);
         // This needs to be here otherwise `CrateGraphBuilder` panics.
@@ -50,6 +51,18 @@ impl Default for TestDB {
     }
 }
 
+impl Clone for TestDB {
+    fn clone(&self) -> Self {
+        Self {
+            storage: self.storage.clone(),
+            files: self.files.clone(),
+            crates_map: self.crates_map.clone(),
+            events: self.events.clone(),
+            nonce: Nonce::new(),
+        }
+    }
+}
+
 impl fmt::Debug for TestDB {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         f.debug_struct("TestDB").finish()
@@ -109,6 +122,10 @@ impl SourceDatabase for TestDB {
     fn crates_map(&self) -> Arc<CratesMap> {
         self.crates_map.clone()
     }
+
+    fn nonce_and_revision(&self) -> (Nonce, salsa::Revision) {
+        (self.nonce, salsa::plumbing::ZalsaDatabase::zalsa(self).current_revision())
+    }
 }
 
 #[salsa_macros::db]
diff --git a/src/tools/rust-analyzer/crates/ide-db/src/lib.rs b/src/tools/rust-analyzer/crates/ide-db/src/lib.rs
index 9d2474d91da..44bccd86d87 100644
--- a/src/tools/rust-analyzer/crates/ide-db/src/lib.rs
+++ b/src/tools/rust-analyzer/crates/ide-db/src/lib.rs
@@ -51,7 +51,7 @@ use salsa::Durability;
 use std::{fmt, mem::ManuallyDrop};
 
 use base_db::{
-    CrateGraphBuilder, CratesMap, FileSourceRootInput, FileText, Files, RootQueryDb,
+    CrateGraphBuilder, CratesMap, FileSourceRootInput, FileText, Files, Nonce, RootQueryDb,
     SourceDatabase, SourceRoot, SourceRootId, SourceRootInput, query_group,
 };
 use hir::{
@@ -83,6 +83,7 @@ pub struct RootDatabase {
     storage: ManuallyDrop<salsa::Storage<Self>>,
     files: Arc<Files>,
     crates_map: Arc<CratesMap>,
+    nonce: Nonce,
 }
 
 impl std::panic::RefUnwindSafe for RootDatabase {}
@@ -102,6 +103,7 @@ impl Clone for RootDatabase {
             storage: self.storage.clone(),
             files: self.files.clone(),
             crates_map: self.crates_map.clone(),
+            nonce: Nonce::new(),
         }
     }
 }
@@ -165,6 +167,10 @@ impl SourceDatabase for RootDatabase {
     fn crates_map(&self) -> Arc<CratesMap> {
         self.crates_map.clone()
     }
+
+    fn nonce_and_revision(&self) -> (Nonce, salsa::Revision) {
+        (self.nonce, salsa::plumbing::ZalsaDatabase::zalsa(self).current_revision())
+    }
 }
 
 impl Default for RootDatabase {
@@ -179,6 +185,7 @@ impl RootDatabase {
             storage: ManuallyDrop::new(salsa::Storage::default()),
             files: Default::default(),
             crates_map: Default::default(),
+            nonce: Nonce::new(),
         };
         // This needs to be here otherwise `CrateGraphBuilder` will panic.
         db.set_all_crates(Arc::new(Box::new([])));