Transition to salsa coarse-grained tracked structs (#15763)

## Summary

Transition to using coarse-grained tracked structs (depends on
https://github.com/salsa-rs/salsa/pull/657). For now, this PR doesn't
add any `#[tracked]` fields, meaning that any changes cause the entire
struct to be invalidated. It also changes `AstNodeRef` to be
compared/hashed by pointer address, instead of performing a deep AST
comparison.

## Test Plan

This yields a 10-15% improvement on my machine (though weirdly some runs
were 5-10% without being flagged as inconsistent by criterion, is there
some non-determinism involved?). It's possible that some of this is
unrelated, I'll try applying the patch to the current salsa version to
make sure.

---------

Co-authored-by: Micha Reiser <micha@reiser.io>

Cargo.lock

@@ -29,6 +29,12 @@ dependencies = [
  "memchr",
 ]
 
+[[package]]
+name = "allocator-api2"
+version = "0.2.21"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923"
+
 [[package]]
 name = "android-tzdata"
 version = "0.1.1"
@@ -1096,6 +1102,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1"
 dependencies = [
  "ahash",
+ "allocator-api2",
 ]
 
 [[package]]
@@ -2870,6 +2877,7 @@ name = "ruff_index"
 version = "0.0.0"
 dependencies = [
  "ruff_macros",
+ "salsa",
  "static_assertions",
 ]
 
@@ -2980,6 +2988,7 @@ dependencies = [
  "ruff_source_file",
  "ruff_text_size",
  "rustc-hash 2.1.1",
+ "salsa",
  "schemars",
  "serde",
 ]
@@ -3304,12 +3313,14 @@ checksum = "6ea1a2d0a644769cc99faa24c3ad26b379b786fe7c36fd3c546254801650e6dd"
 [[package]]
 name = "salsa"
 version = "0.18.0"
-source = "git+https://github.com/salsa-rs/salsa.git?rev=88a1d7774d78f048fbd77d40abca9ebd729fd1f0#88a1d7774d78f048fbd77d40abca9ebd729fd1f0"
+source = "git+https://github.com/salsa-rs/salsa.git?rev=351d9cf0037be949d17800d0c7b4838e533c2ed6#351d9cf0037be949d17800d0c7b4838e533c2ed6"
 dependencies = [
  "append-only-vec",
  "arc-swap",
+ "compact_str",
  "crossbeam",
  "dashmap 6.1.0",
+ "hashbrown 0.14.5",
  "hashlink",
  "indexmap",
  "parking_lot",
@@ -3324,12 +3335,12 @@ dependencies = [
 [[package]]
 name = "salsa-macro-rules"
 version = "0.1.0"
-source = "git+https://github.com/salsa-rs/salsa.git?rev=88a1d7774d78f048fbd77d40abca9ebd729fd1f0#88a1d7774d78f048fbd77d40abca9ebd729fd1f0"
+source = "git+https://github.com/salsa-rs/salsa.git?rev=351d9cf0037be949d17800d0c7b4838e533c2ed6#351d9cf0037be949d17800d0c7b4838e533c2ed6"
 
 [[package]]
 name = "salsa-macros"
 version = "0.18.0"
-source = "git+https://github.com/salsa-rs/salsa.git?rev=88a1d7774d78f048fbd77d40abca9ebd729fd1f0#88a1d7774d78f048fbd77d40abca9ebd729fd1f0"
+source = "git+https://github.com/salsa-rs/salsa.git?rev=351d9cf0037be949d17800d0c7b4838e533c2ed6#351d9cf0037be949d17800d0c7b4838e533c2ed6"
 dependencies = [
  "heck",
  "proc-macro2",

Cargo.toml

@@ -123,7 +123,7 @@ rayon = { version = "1.10.0" }
 regex = { version = "1.10.2" }
 rustc-hash = { version = "2.0.0" }
 # When updating salsa, make sure to also update the revision in `fuzz/Cargo.toml`
-salsa = { git = "https://github.com/salsa-rs/salsa.git", rev = "88a1d7774d78f048fbd77d40abca9ebd729fd1f0" }
+salsa = { git = "https://github.com/salsa-rs/salsa.git", rev = "351d9cf0037be949d17800d0c7b4838e533c2ed6" }
 schemars = { version = "0.8.16" }
 seahash = { version = "4.1.0" }
 serde = { version = "1.0.197", features = ["derive"] }

crates/red_knot_python_semantic/Cargo.toml

@@ -12,9 +12,9 @@ license = { workspace = true }
 
 [dependencies]
 ruff_db = { workspace = true }
-ruff_index = { workspace = true }
+ruff_index = { workspace = true, features = ["salsa"] }
 ruff_macros = { workspace = true }
-ruff_python_ast = { workspace = true }
+ruff_python_ast = { workspace = true, features = ["salsa"] }
 ruff_python_parser = { workspace = true }
 ruff_python_stdlib = { workspace = true }
 ruff_source_file = { workspace = true }
@@ -31,7 +31,7 @@ drop_bomb = { workspace = true }
 indexmap = { workspace = true }
 itertools = { workspace = true }
 ordermap = { workspace = true }
-salsa = { workspace = true }
+salsa = { workspace = true, features = ["compact_str"] }
 thiserror = { workspace = true }
 tracing = { workspace = true }
 rustc-hash = { workspace = true }

crates/red_knot_python_semantic/src/ast_node_ref.rs

@@ -12,7 +12,29 @@ use ruff_db::parsed::ParsedModule;
 /// Holding on to any [`AstNodeRef`] prevents the [`ParsedModule`] from being released.
 ///
 /// ## Equality
-/// Two `AstNodeRef` are considered equal if their wrapped nodes are equal.
+/// Two `AstNodeRef` are considered equal if their pointer addresses are equal.
+///
+/// ## Usage in salsa tracked structs
+/// It's important that [`AstNodeRef`] fields in salsa tracked structs are tracked fields
+/// (attributed with `#[tracked`]). It prevents that the tracked struct gets a new ID
+/// everytime the AST changes, which in turn, invalidates the result of any query
+/// that takes said tracked struct as a query argument or returns the tracked struct as part of its result.
+///
+/// For example, marking the [`AstNodeRef`] as tracked on `Expression`
+/// has the effect that salsa will consider the expression as "unchanged" for as long as it:
+///
+/// * belongs to the same file
+/// * belongs to the same scope
+/// * has the same kind
+/// * was created in the same order
+///
+/// This means that changes to expressions in other scopes don't invalidate the expression's id, giving
+/// us some form of scope-stable identity for expressions. Only queries accessing the node field
+/// run on every AST change. All other queries only run when the expression's identity changes.
+///
+/// The one exception to this is if it is known that all queries tacking the tracked struct
+/// as argument or returning it as part of their result are known to access the node field.
+/// Marking the field tracked is then unnecessary.
 #[derive(Clone)]
 pub struct AstNodeRef<T> {
     /// Owned reference to the node's [`ParsedModule`].
@@ -67,23 +89,17 @@ where
     }
 }
 
-impl<T> PartialEq for AstNodeRef<T>
+impl<T> PartialEq for AstNodeRef<T> {
-where
-    T: PartialEq,
-{
     fn eq(&self, other: &Self) -> bool {
-        self.node().eq(other.node())
+        self.node.eq(&other.node)
     }
 }
 
-impl<T> Eq for AstNodeRef<T> where T: Eq {}
+impl<T> Eq for AstNodeRef<T> {}
 
-impl<T> Hash for AstNodeRef<T>
+impl<T> Hash for AstNodeRef<T> {
-where
-    T: Hash,
-{
     fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
-        self.node().hash(state);
+        self.node.hash(state);
     }
 }
 
@@ -117,7 +133,7 @@ mod tests {
         let stmt_cloned = &cloned.syntax().body[0];
         let cloned_node = unsafe { AstNodeRef::new(cloned.clone(), stmt_cloned) };
 
-        assert_eq!(node1, cloned_node);
+        assert_ne!(node1, cloned_node);
 
         let other_raw = parse_unchecked_source("2 + 2", PySourceType::Python);
         let other = ParsedModule::new(other_raw);

crates/red_knot_python_semantic/src/module_resolver/resolver.rs

@@ -133,7 +133,7 @@ pub(crate) fn search_paths(db: &dyn Db) -> SearchPathIterator {
 }
 
 #[derive(Debug, PartialEq, Eq)]
-pub(crate) struct SearchPaths {
+pub struct SearchPaths {
     /// Search paths that have been statically determined purely from reading Ruff's configuration settings.
     /// These shouldn't ever change unless the config settings themselves change.
     static_paths: Vec<SearchPath>,

crates/red_knot_python_semantic/src/semantic_index.rs

@@ -1,13 +1,14 @@
 use std::iter::FusedIterator;
 use std::sync::Arc;
 
-use rustc_hash::{FxBuildHasher, FxHashMap, FxHashSet};
-use salsa::plumbing::AsId;
-
 use ruff_db::files::File;
 use ruff_db::parsed::parsed_module;
 use ruff_index::{IndexSlice, IndexVec};
 
+use rustc_hash::{FxBuildHasher, FxHashMap, FxHashSet};
+use salsa::plumbing::AsId;
+use salsa::Update;
+
 use crate::module_name::ModuleName;
 use crate::semantic_index::ast_ids::node_key::ExpressionNodeKey;
 use crate::semantic_index::ast_ids::AstIds;
@@ -123,7 +124,7 @@ pub(crate) fn global_scope(db: &dyn Db, file: File) -> ScopeId<'_> {
 }
 
 /// The symbol tables and use-def maps for all scopes in a file.
-#[derive(Debug)]
+#[derive(Debug, Update)]
 pub(crate) struct SemanticIndex<'db> {
     /// List of all symbol tables in this file, indexed by scope.
     symbol_tables: IndexVec<FileScopeId, Arc<SymbolTable>>,

crates/red_knot_python_semantic/src/semantic_index/ast_ids.rs

@@ -24,7 +24,7 @@ use crate::Db;
 ///
 /// x = foo()
 /// ```
-#[derive(Debug)]
+#[derive(Debug, salsa::Update)]
 pub(crate) struct AstIds {
     /// Maps expressions to their expression id.
     expressions_map: FxHashMap<ExpressionNodeKey, ScopedExpressionId>,
@@ -74,6 +74,7 @@ impl HasScopedUseId for ast::ExprRef<'_> {
 
 /// Uniquely identifies an [`ast::Expr`] in a [`crate::semantic_index::symbol::FileScopeId`].
 #[newtype_index]
+#[derive(salsa::Update)]
 pub struct ScopedExpressionId;
 
 pub trait HasScopedExpressionId {
@@ -181,7 +182,7 @@ pub(crate) mod node_key {
 
     use crate::node_key::NodeKey;
 
-    #[derive(Copy, Clone, Eq, PartialEq, Hash, Debug)]
+    #[derive(Copy, Clone, Eq, PartialEq, Hash, Debug, salsa::Update)]
     pub(crate) struct ExpressionNodeKey(NodeKey);
 
     impl From<ast::ExprRef<'_>> for ExpressionNodeKey {

crates/red_knot_python_semantic/src/semantic_index/attribute_assignment.rs

@@ -9,7 +9,7 @@ use rustc_hash::FxHashMap;
 
 /// Describes an (annotated) attribute assignment that we discovered in a method
 /// body, typically of the form `self.x: int`, `self.x: int = …` or `self.x = …`.
-#[derive(Debug, Clone, PartialEq, Eq)]
+#[derive(Debug, Clone, PartialEq, Eq, salsa::Update)]
 pub(crate) enum AttributeAssignment<'db> {
     /// An attribute assignment with an explicit type annotation, either
     /// `self.x: <annotation>` or `self.x: <annotation> = …`.

crates/red_knot_python_semantic/src/semantic_index/constraint.rs

@@ -5,20 +5,20 @@ use crate::db::Db;
 use crate::semantic_index::expression::Expression;
 use crate::semantic_index::symbol::{FileScopeId, ScopeId};
 
-#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)]
+#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq, salsa::Update)]
 pub(crate) struct Constraint<'db> {
     pub(crate) node: ConstraintNode<'db>,
     pub(crate) is_positive: bool,
 }
 
-#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)]
+#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq, salsa::Update)]
 pub(crate) enum ConstraintNode<'db> {
     Expression(Expression<'db>),
     Pattern(PatternConstraint<'db>),
 }
 
 /// Pattern kinds for which we support type narrowing and/or static visibility analysis.
-#[derive(Debug, Clone, Hash, PartialEq)]
+#[derive(Debug, Clone, Hash, PartialEq, salsa::Update)]
 pub(crate) enum PatternConstraintKind<'db> {
     Singleton(Singleton, Option<Expression<'db>>),
     Value(Expression<'db>, Option<Expression<'db>>),
@@ -28,21 +28,15 @@ pub(crate) enum PatternConstraintKind<'db> {
 
 #[salsa::tracked]
 pub(crate) struct PatternConstraint<'db> {
-    #[id]
     pub(crate) file: File,
 
-    #[id]
     pub(crate) file_scope: FileScopeId,
 
-    #[no_eq]
-    #[return_ref]
     pub(crate) subject: Expression<'db>,
 
-    #[no_eq]
     #[return_ref]
     pub(crate) kind: PatternConstraintKind<'db>,
 
-    #[no_eq]
     count: countme::Count<PatternConstraint<'static>>,
 }
 

crates/red_knot_python_semantic/src/semantic_index/definition.rs

@@ -25,22 +25,19 @@ use crate::Db;
 #[salsa::tracked]
 pub struct Definition<'db> {
     /// The file in which the definition occurs.
-    #[id]
     pub(crate) file: File,
 
     /// The scope in which the definition occurs.
-    #[id]
     pub(crate) file_scope: FileScopeId,
 
     /// The symbol defined.
-    #[id]
     pub(crate) symbol: ScopedSymbolId,
 
     #[no_eq]
     #[return_ref]
+    #[tracked]
     pub(crate) kind: DefinitionKind<'db>,
 
-    #[no_eq]
     count: countme::Count<Definition<'static>>,
 }
 
@@ -435,7 +432,14 @@ impl DefinitionCategory {
     }
 }
 
-#[derive(Clone, Debug)]
+/// The kind of a definition.
+///
+/// ## Usage in salsa tracked structs
+///
+/// [`DefinitionKind`] fields in salsa tracked structs should be tracked (attributed with `#[tracked]`)
+/// because the kind is a thin wrapper around [`AstNodeRef`]. See the [`AstNodeRef`] documentation
+/// for an in-depth explanation of why this is necessary.
+#[derive(Clone, Debug, Hash)]
 pub enum DefinitionKind<'db> {
     Import(AstNodeRef<ast::Alias>),
     ImportFrom(ImportFromDefinitionKind),
@@ -540,7 +544,7 @@ impl DefinitionKind<'_> {
     }
 }
 
-#[derive(Copy, Clone, Debug, PartialEq)]
+#[derive(Copy, Clone, Debug, PartialEq, Hash)]
 pub(crate) enum TargetKind<'db> {
     Sequence(Unpack<'db>),
     Name,
@@ -555,7 +559,7 @@ impl<'db> From<Option<Unpack<'db>>> for TargetKind<'db> {
     }
 }
 
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, Hash)]
 #[allow(dead_code)]
 pub struct MatchPatternDefinitionKind {
     pattern: AstNodeRef<ast::Pattern>,
@@ -573,7 +577,7 @@ impl MatchPatternDefinitionKind {
     }
 }
 
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, Hash)]
 pub struct ComprehensionDefinitionKind {
     iterable: AstNodeRef<ast::Expr>,
     target: AstNodeRef<ast::ExprName>,
@@ -599,7 +603,7 @@ impl ComprehensionDefinitionKind {
     }
 }
 
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, Hash)]
 pub struct ImportFromDefinitionKind {
     node: AstNodeRef<ast::StmtImportFrom>,
     alias_index: usize,
@@ -615,7 +619,7 @@ impl ImportFromDefinitionKind {
     }
 }
 
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, Hash)]
 pub struct AssignmentDefinitionKind<'db> {
     target: TargetKind<'db>,
     value: AstNodeRef<ast::Expr>,
@@ -641,7 +645,7 @@ impl<'db> AssignmentDefinitionKind<'db> {
     }
 }
 
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, Hash)]
 pub struct WithItemDefinitionKind {
     node: AstNodeRef<ast::WithItem>,
     target: AstNodeRef<ast::ExprName>,
@@ -662,7 +666,7 @@ impl WithItemDefinitionKind {
     }
 }
 
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, Hash)]
 pub struct ForStmtDefinitionKind<'db> {
     target: TargetKind<'db>,
     iterable: AstNodeRef<ast::Expr>,
@@ -693,7 +697,7 @@ impl<'db> ForStmtDefinitionKind<'db> {
     }
 }
 
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, Hash)]
 pub struct ExceptHandlerDefinitionKind {
     handler: AstNodeRef<ast::ExceptHandlerExceptHandler>,
     is_star: bool,
@@ -713,7 +717,7 @@ impl ExceptHandlerDefinitionKind {
     }
 }
 
-#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug)]
+#[derive(Copy, Clone, Eq, PartialEq, Hash, Debug, salsa::Update)]
 pub(crate) struct DefinitionNodeKey(NodeKey);
 
 impl From<&ast::Alias> for DefinitionNodeKey {

crates/red_knot_python_semantic/src/semantic_index/expression.rs

@@ -33,23 +33,20 @@ pub(crate) enum ExpressionKind {
 #[salsa::tracked]
 pub(crate) struct Expression<'db> {
     /// The file in which the expression occurs.
-    #[id]
     pub(crate) file: File,
 
     /// The scope in which the expression occurs.
-    #[id]
     pub(crate) file_scope: FileScopeId,
 
     /// The expression node.
     #[no_eq]
+    #[tracked]
     #[return_ref]
     pub(crate) node_ref: AstNodeRef<ast::Expr>,
 
     /// Should this expression be inferred as a normal expression or a type expression?
-    #[id]
     pub(crate) kind: ExpressionKind,
 
-    #[no_eq]
     count: countme::Count<Expression<'static>>,
 }
 

crates/red_knot_python_semantic/src/semantic_index/symbol.rs

@@ -96,18 +96,16 @@ impl From<FileSymbolId> for ScopedSymbolId {
 
 /// Symbol ID that uniquely identifies a symbol inside a [`Scope`].
 #[newtype_index]
+#[derive(salsa::Update)]
 pub struct ScopedSymbolId;
 
 /// A cross-module identifier of a scope that can be used as a salsa query parameter.
 #[salsa::tracked]
 pub struct ScopeId<'db> {
-    #[id]
     pub file: File,
 
-    #[id]
     pub file_scope_id: FileScopeId,
 
-    #[no_eq]
     count: countme::Count<ScopeId<'static>>,
 }
 
@@ -159,6 +157,7 @@ impl<'db> ScopeId<'db> {
 
 /// ID that uniquely identifies a scope inside of a module.
 #[newtype_index]
+#[derive(salsa::Update)]
 pub struct FileScopeId;
 
 impl FileScopeId {
@@ -177,7 +176,7 @@ impl FileScopeId {
     }
 }
 
-#[derive(Debug)]
+#[derive(Debug, salsa::Update)]
 pub struct Scope {
     pub(super) parent: Option<FileScopeId>,
     pub(super) node: NodeWithScopeKind,
@@ -216,7 +215,7 @@ impl ScopeKind {
 }
 
 /// Symbol table for a specific [`Scope`].
-#[derive(Debug, Default)]
+#[derive(Debug, Default, salsa::Update)]
 pub struct SymbolTable {
     /// The symbols in this scope.
     symbols: IndexVec<ScopedSymbolId, Symbol>,
@@ -424,7 +423,7 @@ impl NodeWithScopeRef<'_> {
 }
 
 /// Node that introduces a new scope.
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, salsa::Update)]
 pub enum NodeWithScopeKind {
     Module,
     Class(AstNodeRef<ast::StmtClassDef>),

crates/red_knot_python_semantic/src/semantic_index/use_def.rs

@@ -278,7 +278,7 @@ mod symbol_state;
 type AllConstraints<'db> = IndexVec<ScopedConstraintId, Constraint<'db>>;
 
 /// Applicable definitions and constraints for every use of a name.
-#[derive(Debug, PartialEq, Eq)]
+#[derive(Debug, PartialEq, Eq, salsa::Update)]
 pub(crate) struct UseDefMap<'db> {
     /// Array of [`Definition`] in this scope. Only the first entry should be `None`;
     /// this represents the implicit "unbound"/"undeclared" definition of every symbol.
@@ -384,7 +384,7 @@ impl<'db> UseDefMap<'db> {
 }
 
 /// Either live bindings or live declarations for a symbol.
-#[derive(Debug, PartialEq, Eq)]
+#[derive(Debug, PartialEq, Eq, salsa::Update)]
 enum SymbolDefinitions {
     Bindings(SymbolBindings),
     Declarations(SymbolDeclarations),

crates/red_knot_python_semantic/src/symbol.rs

@@ -26,7 +26,7 @@ pub(crate) enum Boundness {
 /// possibly_unbound:  Symbol::Type(Type::IntLiteral(2), Boundness::PossiblyUnbound),
 /// non_existent:      Symbol::Unbound,
 /// ```
-#[derive(Debug, Clone, PartialEq, Eq)]
+#[derive(Debug, Clone, PartialEq, Eq, salsa::Update)]
 pub(crate) enum Symbol<'db> {
     Type(Type<'db>, Boundness),
     Unbound,

crates/red_knot_python_semantic/src/types.rs

@@ -56,7 +56,6 @@ mod mro;
 mod narrow;
 mod signatures;
 mod slots;
-mod statistics;
 mod string_annotation;
 mod subclass_of;
 mod type_ordering;
@@ -2589,7 +2588,7 @@ bitflags! {
 ///
 /// Example: `Annotated[ClassVar[tuple[int]], "metadata"]` would have type `tuple[int]` and the
 /// qualifier `ClassVar`.
-#[derive(Clone, Debug, Copy, Eq, PartialEq)]
+#[derive(Clone, Debug, Copy, Eq, PartialEq, salsa::Update)]
 pub(crate) struct TypeAndQualifiers<'db> {
     inner: Type<'db>,
     qualifiers: TypeQualifiers,
@@ -4400,7 +4399,7 @@ impl<'db> TypeAliasType<'db> {
 }
 
 /// Either the explicit `metaclass=` keyword of the class, or the inferred metaclass of one of its base classes.
-#[derive(Debug, Clone, PartialEq, Eq)]
+#[derive(Debug, Clone, PartialEq, Eq, salsa::Update)]
 pub(super) struct MetaclassCandidate<'db> {
     metaclass: Class<'db>,
     explicit_metaclass_of: Class<'db>,
@@ -4443,7 +4442,7 @@ impl<'db> From<InstanceType<'db>> for Type<'db> {
     }
 }
 
-#[derive(Debug, Clone, PartialEq, Eq)]
+#[derive(Debug, Clone, PartialEq, Eq, salsa::Update)]
 pub(super) struct MetaclassError<'db> {
     kind: MetaclassErrorKind<'db>,
 }
@@ -4455,7 +4454,7 @@ impl<'db> MetaclassError<'db> {
     }
 }
 
-#[derive(Debug, Clone, PartialEq, Eq)]
+#[derive(Debug, Clone, PartialEq, Eq, salsa::Update)]
 pub(super) enum MetaclassErrorKind<'db> {
     /// The class has incompatible metaclasses in its inheritance hierarchy.
     ///

crates/red_knot_python_semantic/src/types/infer.rs

@@ -61,7 +61,6 @@ use crate::types::diagnostic::{
     UNDEFINED_REVEAL, UNRESOLVED_ATTRIBUTE, UNRESOLVED_IMPORT, UNSUPPORTED_OPERATOR,
 };
 use crate::types::mro::MroErrorKind;
-use crate::types::statistics::TypeStatistics;
 use crate::types::unpacker::{UnpackResult, Unpacker};
 use crate::types::{
     builtins_symbol, global_symbol, symbol, symbol_from_bindings, symbol_from_declarations,
@@ -237,7 +236,7 @@ impl<'db> InferenceRegion<'db> {
 }
 
 /// The inferred types for a single region.
-#[derive(Debug, Eq, PartialEq)]
+#[derive(Debug, Eq, PartialEq, salsa::Update)]
 pub(crate) struct TypeInference<'db> {
     /// The types of every expression in this region.
     expressions: FxHashMap<ScopedExpressionId, Type<'db>>,
@@ -300,14 +299,6 @@ impl<'db> TypeInference<'db> {
         self.diagnostics.shrink_to_fit();
         self.deferred.shrink_to_fit();
     }
-
-    pub(super) fn statistics(&self) -> TypeStatistics {
-        let mut statistics = TypeStatistics::default();
-        for ty in self.expressions.values() {
-            statistics.increment(*ty);
-        }
-        statistics
-    }
 }
 
 impl WithDiagnostics for TypeInference<'_> {
@@ -6433,6 +6424,7 @@ mod tests {
             assert_eq!(attr_ty.display(&db).to_string(), "Unknown | str | None");
             db.take_salsa_events()
         };
+
         assert_function_query_was_not_run(
             &db,
             infer_expression_types,

crates/red_knot_python_semantic/src/types/mro.rs

@@ -10,7 +10,7 @@ use crate::Db;
 /// The inferred method resolution order of a given class.
 ///
 /// See [`Class::iter_mro`] for more details.
-#[derive(PartialEq, Eq, Clone, Debug)]
+#[derive(PartialEq, Eq, Clone, Debug, salsa::Update)]
 pub(super) struct Mro<'db>(Box<[ClassBase<'db>]>);
 
 impl<'db> Mro<'db> {
@@ -236,7 +236,7 @@ impl<'db> Iterator for MroIterator<'db> {
 
 impl std::iter::FusedIterator for MroIterator<'_> {}
 
-#[derive(Debug, PartialEq, Eq)]
+#[derive(Debug, PartialEq, Eq, salsa::Update)]
 pub(super) struct MroError<'db> {
     kind: MroErrorKind<'db>,
     fallback_mro: Mro<'db>,
@@ -256,7 +256,7 @@ impl<'db> MroError<'db> {
 }
 
 /// Possible ways in which attempting to resolve the MRO of a class might fail.
-#[derive(Debug, PartialEq, Eq)]
+#[derive(Debug, PartialEq, Eq, salsa::Update)]
 pub(super) enum MroErrorKind<'db> {
     /// The class inherits from one or more invalid bases.
     ///

crates/red_knot_python_semantic/src/types/narrow.rs

@@ -231,10 +231,10 @@ impl<'db> NarrowingConstraintsBuilder<'db> {
 
         match pattern.kind(self.db) {
             PatternConstraintKind::Singleton(singleton, _guard) => {
-                self.evaluate_match_pattern_singleton(*subject, *singleton)
+                self.evaluate_match_pattern_singleton(subject, *singleton)
             }
             PatternConstraintKind::Class(cls, _guard) => {
-                self.evaluate_match_pattern_class(*subject, *cls)
+                self.evaluate_match_pattern_class(subject, *cls)
             }
             // TODO: support more pattern kinds
             PatternConstraintKind::Value(..) | PatternConstraintKind::Unsupported => None,

crates/red_knot_python_semantic/src/types/signatures.rs

@@ -4,7 +4,7 @@ use crate::{semantic_index::definition::Definition, types::todo_type};
 use ruff_python_ast::{self as ast, name::Name};
 
 /// A typed callable signature.
-#[derive(Clone, Debug, PartialEq, Eq)]
+#[derive(Clone, Debug, PartialEq, Eq, salsa::Update)]
 pub(crate) struct Signature<'db> {
     /// Parameters, in source order.
     ///
@@ -60,7 +60,7 @@ impl<'db> Signature<'db> {
 }
 
 // TODO: use SmallVec here once invariance bug is fixed
-#[derive(Clone, Debug, PartialEq, Eq)]
+#[derive(Clone, Debug, PartialEq, Eq, salsa::Update)]
 pub(crate) struct Parameters<'db>(Vec<Parameter<'db>>);
 
 impl<'db> Parameters<'db> {
@@ -218,7 +218,7 @@ impl<'db> std::ops::Index<usize> for Parameters<'db> {
     }
 }
 
-#[derive(Clone, Debug, PartialEq, Eq)]
+#[derive(Clone, Debug, PartialEq, Eq, salsa::Update)]
 pub(crate) struct Parameter<'db> {
     /// Parameter name.
     ///
@@ -304,7 +304,7 @@ impl<'db> Parameter<'db> {
     }
 }
 
-#[derive(Clone, Debug, PartialEq, Eq)]
+#[derive(Clone, Debug, PartialEq, Eq, salsa::Update)]
 pub(crate) enum ParameterKind<'db> {
     /// Positional-only parameter, e.g. `def f(x, /): ...`
     PositionalOnly { default_ty: Option<Type<'db>> },

crates/red_knot_python_semantic/src/types/statistics.rs

@@ -1,121 +0,0 @@
-use crate::types::{infer_scope_types, semantic_index, Type};
-use crate::Db;
-use ruff_db::files::File;
-use rustc_hash::FxHashMap;
-
-/// Get type-coverage statistics for a file.
-#[salsa::tracked(return_ref)]
-pub fn type_statistics<'db>(db: &'db dyn Db, file: File) -> TypeStatistics<'db> {
-    let _span = tracing::trace_span!("type_statistics", file=?file.path(db)).entered();
-
-    tracing::debug!(
-        "Gathering statistics for file '{path}'",
-        path = file.path(db)
-    );
-
-    let index = semantic_index(db, file);
-    let mut statistics = TypeStatistics::default();
-
-    for scope_id in index.scope_ids() {
-        let result = infer_scope_types(db, scope_id);
-        statistics.extend(&result.statistics());
-    }
-
-    statistics
-}
-
-/// Map each type to count of expressions with that type.
-#[derive(Debug, Default, Eq, PartialEq)]
-pub(super) struct TypeStatistics<'db>(FxHashMap<Type<'db>, u32>);
-
-impl<'db> TypeStatistics<'db> {
-    fn extend(&mut self, other: &TypeStatistics<'db>) {
-        for (ty, count) in &other.0 {
-            self.0
-                .entry(*ty)
-                .and_modify(|my_count| *my_count += count)
-                .or_insert(*count);
-        }
-    }
-
-    pub(super) fn increment(&mut self, ty: Type<'db>) {
-        self.0
-            .entry(ty)
-            .and_modify(|count| *count += 1)
-            .or_insert(1);
-    }
-
-    #[allow(unused)]
-    fn expression_count(&self) -> u32 {
-        self.0.values().sum()
-    }
-
-    #[allow(unused)]
-    fn todo_count(&self) -> u32 {
-        self.0
-            .iter()
-            .filter(|(key, _)| key.is_todo())
-            .map(|(_, count)| count)
-            .sum()
-    }
-}
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-    use crate::db::tests::{setup_db, TestDb};
-    use ruff_db::files::system_path_to_file;
-    use ruff_db::system::DbWithTestSystem;
-
-    fn get_stats<'db>(
-        db: &'db mut TestDb,
-        filename: &str,
-        source: &str,
-    ) -> &'db TypeStatistics<'db> {
-        db.write_dedented(filename, source).unwrap();
-
-        type_statistics(db, system_path_to_file(db, filename).unwrap())
-    }
-
-    #[test]
-    fn all_static() {
-        let mut db = setup_db();
-
-        let stats = get_stats(&mut db, "src/foo.py", "1");
-
-        assert_eq!(stats.0, FxHashMap::from_iter([(Type::IntLiteral(1), 1)]));
-    }
-
-    #[test]
-    fn todo_and_expression_count() {
-        let mut db = setup_db();
-
-        let stats = get_stats(
-            &mut db,
-            "src/foo.py",
-            r#"
-                x = [x for x in [1]]
-            "#,
-        );
-
-        assert_eq!(stats.todo_count(), 4);
-        assert_eq!(stats.expression_count(), 6);
-    }
-
-    #[test]
-    fn sum() {
-        let mut db = setup_db();
-
-        let stats = get_stats(
-            &mut db,
-            "src/foo.py",
-            r#"
-                1
-                def f():
-                    1
-            "#,
-        );
-
-        assert_eq!(stats.0[&Type::IntLiteral(1)], 2);
-    }
-}

crates/red_knot_python_semantic/src/types/unpacker.rs

@@ -261,7 +261,7 @@ impl<'db> Unpacker<'db> {
     }
 }
 
-#[derive(Debug, Default, PartialEq, Eq)]
+#[derive(Debug, Default, PartialEq, Eq, salsa::Update)]
 pub(crate) struct UnpackResult<'db> {
     targets: FxHashMap<ScopedExpressionId, Type<'db>>,
     diagnostics: TypeCheckDiagnostics,

crates/red_knot_python_semantic/src/unpack.rs

@@ -28,10 +28,8 @@ use crate::Db;
 /// * an argument of a cross-module query
 #[salsa::tracked]
 pub(crate) struct Unpack<'db> {
-    #[id]
     pub(crate) file: File,
 
-    #[id]
     pub(crate) file_scope: FileScopeId,
 
     /// The target expression that is being unpacked. For example, in `(a, b) = (1, 2)`, the target
@@ -45,7 +43,6 @@ pub(crate) struct Unpack<'db> {
     #[no_eq]
     pub(crate) value: UnpackValue<'db>,
 
-    #[no_eq]
     count: countme::Count<Unpack<'static>>,
 }
 
@@ -62,7 +59,7 @@ impl<'db> Unpack<'db> {
 }
 
 /// The expression that is being unpacked.
-#[derive(Clone, Copy, Debug)]
+#[derive(Clone, Copy, Debug, Hash)]
 pub(crate) enum UnpackValue<'db> {
     /// An iterable expression like the one in a `for` loop or a comprehension.
     Iterable(Expression<'db>),

crates/red_knot_python_semantic/src/visibility_constraints.rs

@@ -338,7 +338,7 @@ const SMALLEST_TERMINAL: ScopedVisibilityConstraintId = ALWAYS_FALSE;
 
 /// A collection of visibility constraints. This is currently stored in `UseDefMap`, which means we
 /// maintain a separate set of visibility constraints for each scope in file.
-#[derive(Debug, PartialEq, Eq)]
+#[derive(Debug, PartialEq, Eq, salsa::Update)]
 pub(crate) struct VisibilityConstraints<'db> {
     constraints: IndexVec<Atom, Constraint<'db>>,
     interiors: IndexVec<ScopedVisibilityConstraintId, InteriorNode>,
@@ -627,7 +627,7 @@ impl<'db> VisibilityConstraints<'db> {
             ConstraintNode::Pattern(inner) => match inner.kind(db) {
                 PatternConstraintKind::Value(value, guard) => {
                     let subject_expression = inner.subject(db);
-                    let inference = infer_expression_types(db, *subject_expression);
+                    let inference = infer_expression_types(db, subject_expression);
                     let scope = subject_expression.scope(db);
                     let subject_ty = inference.expression_type(
                         subject_expression

crates/ruff_db/src/parsed.rs

@@ -73,6 +73,14 @@ impl std::fmt::Debug for ParsedModule {
     }
 }
 
+impl PartialEq for ParsedModule {
+    fn eq(&self, other: &Self) -> bool {
+        Arc::ptr_eq(&self.inner, &other.inner)
+    }
+}
+
+impl Eq for ParsedModule {}
+
 #[cfg(test)]
 mod tests {
     use crate::files::{system_path_to_file, vendored_path_to_file};

crates/ruff_db/src/testing.rs

@@ -18,7 +18,7 @@ pub fn assert_function_query_was_not_run<Db, Q, QDb, I, R>(
 
     db.attach(|_| {
         if let Some(will_execute_event) = will_execute_event {
-            panic!("Expected query {query_name}({id}) not to have run but it did: {will_execute_event:?}");
+            panic!("Expected query {query_name}({id}) not to have run but it did: {will_execute_event:?}\n\n{events:#?}");
         }
     });
 }
@@ -46,7 +46,7 @@ pub fn assert_const_function_query_was_not_run<Db, Q, QDb, R>(
     db.attach(|_| {
         if let Some(will_execute_event) = event {
             panic!(
-                "Expected query {query_name}() not to have run but it did: {will_execute_event:?}"
+                "Expected query {query_name}() not to have run but it did: {will_execute_event:?}\n\n{events:#?}"
             );
         }
     });

crates/ruff_index/Cargo.toml

@@ -15,6 +15,7 @@ doctest = false
 
 [dependencies]
 ruff_macros = { workspace = true }
+salsa = { workspace = true, optional = true }
 
 [dev-dependencies]
 static_assertions = { workspace = true }

crates/ruff_index/src/vec.rs

@@ -181,3 +181,16 @@ impl<I: Idx, T, const N: usize> From<[T; N]> for IndexVec<I, T> {
 // not the phantom data.
 #[allow(unsafe_code)]
 unsafe impl<I: Idx, T> Send for IndexVec<I, T> where T: Send {}
+
+#[allow(unsafe_code)]
+#[cfg(feature = "salsa")]
+unsafe impl<I, T> salsa::Update for IndexVec<I, T>
+where
+    T: salsa::Update,
+{
+    #[allow(unsafe_code)]
+    unsafe fn maybe_update(old_pointer: *mut Self, new_value: Self) -> bool {
+        let old_vec: &mut IndexVec<I, T> = unsafe { &mut *old_pointer };
+        salsa::Update::maybe_update(&mut old_vec.raw, new_value.raw)
+    }
+}

crates/ruff_python_ast/Cargo.toml

@@ -26,6 +26,7 @@ is-macro = { workspace = true }
 itertools = { workspace = true }
 memchr = { workspace = true }
 rustc-hash = { workspace = true }
+salsa = { workspace = true, optional = true }
 schemars = { workspace = true, optional = true }
 serde = { workspace = true, optional = true }
 

crates/ruff_python_ast/src/name.rs

@@ -8,6 +8,7 @@ use crate::{nodes, Expr};
 #[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Default)]
 #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
 #[cfg_attr(feature = "cache", derive(ruff_macros::CacheKey))]
+#[cfg_attr(feature = "salsa", derive(salsa::Update))]
 pub struct Name(compact_str::CompactString);
 
 impl Name {

fuzz/Cargo.toml

@@ -29,7 +29,7 @@ ruff_python_formatter = { path = "../crates/ruff_python_formatter" }
 ruff_text_size = { path = "../crates/ruff_text_size" }
 
 libfuzzer-sys = { git = "https://github.com/rust-fuzz/libfuzzer", default-features = false }
-salsa = { git = "https://github.com/salsa-rs/salsa.git", rev = "88a1d7774d78f048fbd77d40abca9ebd729fd1f0" }
+salsa = { git = "https://github.com/salsa-rs/salsa.git", rev = "351d9cf0037be949d17800d0c7b4838e533c2ed6" }
 similar = { version = "2.5.0" }
 tracing = { version = "0.1.40" }