[red-knot] Upgrade to the *new* *new* salsa (#12406)

crates/red_knot_python_semantic/src/db.rs

@@ -1,55 +1,23 @@
-use salsa::DbWithJar;
-
 use red_knot_module_resolver::Db as ResolverDb;
-use ruff_db::{Db as SourceDb, Upcast};
-
-use crate::builtins::builtins_scope;
-use crate::semantic_index::definition::Definition;
-use crate::semantic_index::expression::Expression;
-use crate::semantic_index::symbol::ScopeId;
-use crate::semantic_index::{global_scope, semantic_index, symbol_table, use_def_map};
-use crate::types::{
-    infer_definition_types, infer_expression_types, infer_scope_types, ClassType, FunctionType,
-    IntersectionType, UnionType,
-};
-
-#[salsa::jar(db=Db)]
-pub struct Jar(
-    ScopeId<'_>,
-    Definition<'_>,
-    Expression<'_>,
-    FunctionType<'_>,
-    ClassType<'_>,
-    UnionType<'_>,
-    IntersectionType<'_>,
-    symbol_table,
-    use_def_map,
-    global_scope,
-    semantic_index,
-    infer_definition_types,
-    infer_expression_types,
-    infer_scope_types,
-    builtins_scope,
-);
+use ruff_db::Upcast;
 
 /// Database giving access to semantic information about a Python program.
-pub trait Db: SourceDb + ResolverDb + DbWithJar<Jar> + Upcast<dyn ResolverDb> {}
+#[salsa::db]
+pub trait Db: ResolverDb + Upcast<dyn ResolverDb> {}
 
 #[cfg(test)]
 pub(crate) mod tests {
     use std::sync::Arc;
 
-    use salsa::DebugWithDb;
-
-    use red_knot_module_resolver::{vendored_typeshed_stubs, Db as ResolverDb, Jar as ResolverJar};
+    use red_knot_module_resolver::{vendored_typeshed_stubs, Db as ResolverDb};
     use ruff_db::files::Files;
     use ruff_db::system::{DbWithTestSystem, System, TestSystem};
     use ruff_db::vendored::VendoredFileSystem;
-    use ruff_db::{Db as SourceDb, Jar as SourceJar, Upcast};
+    use ruff_db::{Db as SourceDb, Upcast};
 
-    use super::{Db, Jar};
+    use super::Db;
 
-    #[salsa::db(Jar, ResolverJar, SourceJar)]
+    #[salsa::db]
     pub(crate) struct TestDb {
         storage: salsa::Storage<Self>,
         files: Files,
@@ -63,7 +31,7 @@ pub(crate) mod tests {
             Self {
                 storage: salsa::Storage::default(),
                 system: TestSystem::default(),
-                vendored: vendored_typeshed_stubs().snapshot(),
+                vendored: vendored_typeshed_stubs().clone(),
                 events: std::sync::Arc::default(),
                 files: Files::default(),
             }
@@ -99,6 +67,7 @@ pub(crate) mod tests {
         }
     }
 
+    #[salsa::db]
     impl SourceDb for TestDb {
         fn vendored(&self) -> &VendoredFileSystem {
             &self.vendored
@@ -131,26 +100,20 @@ pub(crate) mod tests {
         }
     }
 
+    #[salsa::db]
     impl red_knot_module_resolver::Db for TestDb {}
+
+    #[salsa::db]
     impl Db for TestDb {}
 
+    #[salsa::db]
     impl salsa::Database for TestDb {
         fn salsa_event(&self, event: salsa::Event) {
-            tracing::trace!("event: {:?}", event.debug(self));
-            let mut events = self.events.lock().unwrap();
-            events.push(event);
-        }
-    }
-
-    impl salsa::ParallelDatabase for TestDb {
-        fn snapshot(&self) -> salsa::Snapshot<Self> {
-            salsa::Snapshot::new(Self {
-                storage: self.storage.snapshot(),
-                files: self.files.snapshot(),
-                system: self.system.snapshot(),
-                vendored: self.vendored.snapshot(),
-                events: self.events.clone(),
-            })
+            self.attach(|_| {
+                tracing::trace!("event: {event:?}");
+                let mut events = self.events.lock().unwrap();
+                events.push(event);
+            });
         }
     }
 }

crates/red_knot_python_semantic/src/lib.rs

@@ -2,7 +2,7 @@ use std::hash::BuildHasherDefault;
 
 use rustc_hash::FxHasher;
 
-pub use db::{Db, Jar};
+pub use db::Db;
 pub use semantic_model::{HasTy, SemanticModel};
 
 pub mod ast_node_ref;

crates/red_knot_python_semantic/src/semantic_index.rs

@@ -2,6 +2,7 @@ use std::iter::FusedIterator;
 use std::sync::Arc;
 
 use rustc_hash::FxHashMap;
+use salsa::plumbing::AsId;
 
 use ruff_db::files::File;
 use ruff_db::parsed::parsed_module;
@@ -17,6 +18,8 @@ use crate::semantic_index::symbol::{
 };
 use crate::Db;
 
+pub(crate) use self::use_def::UseDefMap;
+
 pub mod ast_ids;
 mod builder;
 pub mod definition;
@@ -24,8 +27,6 @@ pub mod expression;
 pub mod symbol;
 mod use_def;
 
-pub(crate) use self::use_def::UseDefMap;
-
 type SymbolMap = hashbrown::HashMap<ScopedSymbolId, (), ()>;
 
 /// Returns the semantic index for `file`.
@@ -33,7 +34,7 @@ type SymbolMap = hashbrown::HashMap<ScopedSymbolId, (), ()>;
 /// Prefer using [`symbol_table`] when working with symbols from a single scope.
 #[salsa::tracked(return_ref, no_eq)]
 pub(crate) fn semantic_index(db: &dyn Db, file: File) -> SemanticIndex<'_> {
-    let _span = tracing::trace_span!("semantic_index", ?file).entered();
+    let _span = tracing::trace_span!("semantic_index", file=?file.path(db)).entered();
 
     let parsed = parsed_module(db.upcast(), file);
 
@@ -47,8 +48,10 @@ pub(crate) fn semantic_index(db: &dyn Db, file: File) -> SemanticIndex<'_> {
 /// is unchanged.
 #[salsa::tracked]
 pub(crate) fn symbol_table<'db>(db: &'db dyn Db, scope: ScopeId<'db>) -> Arc<SymbolTable> {
-    let _span = tracing::trace_span!("symbol_table", ?scope).entered();
-    let index = semantic_index(db, scope.file(db));
+    let file = scope.file(db);
+    let _span =
+        tracing::trace_span!("symbol_table", scope=?scope.as_id(), file=?file.path(db)).entered();
+    let index = semantic_index(db, file);
 
     index.symbol_table(scope.file_scope_id(db))
 }
@@ -60,8 +63,10 @@ pub(crate) fn symbol_table<'db>(db: &'db dyn Db, scope: ScopeId<'db>) -> Arc<Sym
 /// is unchanged.
 #[salsa::tracked]
 pub(crate) fn use_def_map<'db>(db: &'db dyn Db, scope: ScopeId<'db>) -> Arc<UseDefMap<'db>> {
-    let _span = tracing::trace_span!("use_def_map", ?scope).entered();
-    let index = semantic_index(db, scope.file(db));
+    let file = scope.file(db);
+    let _span =
+        tracing::trace_span!("use_def_map", scope=?scope.as_id(), file=?file.path(db)).entered();
+    let index = semantic_index(db, file);
 
     index.use_def_map(scope.file_scope_id(db))
 }
@@ -69,7 +74,7 @@ pub(crate) fn use_def_map<'db>(db: &'db dyn Db, scope: ScopeId<'db>) -> Arc<UseD
 /// Returns the module global scope of `file`.
 #[salsa::tracked]
 pub(crate) fn global_scope(db: &dyn Db, file: File) -> ScopeId<'_> {
-    let _span = tracing::trace_span!("global_scope", ?file).entered();
+    let _span = tracing::trace_span!("global_scope", file=?file.path(db)).entered();
 
     FileScopeId::global().to_scope_id(db, file)
 }

crates/red_knot_python_semantic/src/types.rs

@@ -10,7 +10,7 @@ use crate::{Db, FxOrderSet};
 mod display;
 mod infer;
 
-pub(crate) use self::infer::{infer_definition_types, infer_expression_types, infer_scope_types};
+pub(crate) use self::infer::{infer_definition_types, infer_scope_types};
 
 /// Infer the public type of a symbol (its type as seen from outside its scope).
 pub(crate) fn symbol_ty<'db>(

crates/red_knot_python_semantic/src/types/display.rs

@@ -26,7 +26,7 @@ impl Display for DisplayType<'_> {
             Type::Unbound => f.write_str("Unbound"),
             Type::None => f.write_str("None"),
             Type::Module(file) => {
-                write!(f, "<module '{:?}'>", file.path(self.db.upcast()))
+                write!(f, "<module '{:?}'>", file.path(self.db))
             }
             // TODO functions and classes should display using a fully qualified name
             Type::Class(class) => write!(f, "Literal[{}]", class.name(self.db)),

crates/red_knot_python_semantic/src/types/infer.rs

@@ -22,6 +22,7 @@
 //! holds types for every [`Definition`] and expression within the inferred region.
 use rustc_hash::FxHashMap;
 use salsa;
+use salsa::plumbing::AsId;
 
 use red_knot_module_resolver::{resolve_module, ModuleName};
 use ruff_db::files::File;
@@ -48,7 +49,9 @@ use crate::Db;
 #[salsa::tracked(return_ref)]
 pub(crate) fn infer_scope_types<'db>(db: &'db dyn Db, scope: ScopeId<'db>) -> TypeInference<'db> {
     let file = scope.file(db);
-    let _span = tracing::trace_span!("infer_scope_types", ?scope, ?file).entered();
+    let _span =
+        tracing::trace_span!("infer_scope_types", scope=?scope.as_id(), file=?file.path(db))
+            .entered();
 
     // Using the index here is fine because the code below depends on the AST anyway.
     // The isolation of the query is by the return inferred types.
@@ -77,7 +80,12 @@ pub(crate) fn infer_definition_types<'db>(
     definition: Definition<'db>,
 ) -> TypeInference<'db> {
     let file = definition.file(db);
-    let _span = tracing::trace_span!("infer_definition_types", ?definition, ?file,).entered();
+    let _span = tracing::trace_span!(
+        "infer_definition_types",
+        definition = ?definition.as_id(),
+        file = ?file.path(db)
+    )
+    .entered();
 
     let index = semantic_index(db, file);
 
@@ -95,7 +103,9 @@ pub(crate) fn infer_expression_types<'db>(
     expression: Expression<'db>,
 ) -> TypeInference<'db> {
     let file = expression.file(db);
-    let _span = tracing::trace_span!("infer_expression_types", ?expression, ?file).entered();
+    let _span =
+        tracing::trace_span!("infer_expression_types", expression=?expression.as_id(), file=?file.path(db))
+            .entered();
 
     let index = semantic_index(db, file);
 
@@ -1491,12 +1501,9 @@ mod tests {
     use crate::builtins::builtins_scope;
     use crate::db::tests::TestDb;
     use crate::semantic_index::definition::Definition;
-    use crate::semantic_index::semantic_index;
     use crate::semantic_index::symbol::FileScopeId;
-    use crate::types::{
-        global_scope, global_symbol_ty_by_name, infer_definition_types, symbol_table,
-        symbol_ty_by_name, use_def_map, Type,
-    };
+    use crate::semantic_index::{global_scope, semantic_index, symbol_table, use_def_map};
+    use crate::types::{global_symbol_ty_by_name, infer_definition_types, symbol_ty_by_name, Type};
     use crate::{HasTy, SemanticModel};
 
     fn setup_db() -> TestDb {
@@ -2231,6 +2238,14 @@ mod tests {
         Ok(())
     }
 
+    fn first_public_def<'db>(db: &'db TestDb, file: File, name: &str) -> Definition<'db> {
+        let scope = global_scope(db, file);
+        *use_def_map(db, scope)
+            .public_definitions(symbol_table(db, scope).symbol_id_by_name(name).unwrap())
+            .first()
+            .unwrap()
+    }
+
     #[test]
     fn big_int() -> anyhow::Result<()> {
         let mut db = setup_db();
@@ -2315,14 +2330,6 @@ mod tests {
         Ok(())
     }
 
-    fn first_public_def<'db>(db: &'db TestDb, file: File, name: &str) -> Definition<'db> {
-        let scope = global_scope(db, file);
-        *use_def_map(db, scope)
-            .public_definitions(symbol_table(db, scope).symbol_id_by_name(name).unwrap())
-            .first()
-            .unwrap()
-    }
-
     #[test]
     fn dependency_public_symbol_type_change() -> anyhow::Result<()> {
         let mut db = setup_db();
@@ -2375,10 +2382,10 @@ mod tests {
 
         let events = db.take_salsa_events();
 
-        assert_function_query_was_not_run::<infer_definition_types, _, _>(
+        assert_function_query_was_not_run(
             &db,
-            |ty| &ty.function,
-            &first_public_def(&db, a, "x"),
+            infer_definition_types,
+            first_public_def(&db, a, "x"),
             &events,
         );
 
@@ -2411,10 +2418,10 @@ mod tests {
 
         let events = db.take_salsa_events();
 
-        assert_function_query_was_not_run::<infer_definition_types, _, _>(
+        assert_function_query_was_not_run(
             &db,
-            |ty| &ty.function,
-            &first_public_def(&db, a, "x"),
+            infer_definition_types,
+            first_public_def(&db, a, "x"),
             &events,
         );
         Ok(())