[ty] Remove ScopedExpressionId (#19019)

## Summary

The motivation of `ScopedExpressionId` was that we have an expression
identifier that's local to a scope and, therefore, unlikely to change if
a user makes changes in another scope. A local identifier like this has
the advantage that query results may remain unchanged even if other
parts of the file change, which in turn allows Salsa to short-circuit
dependent queries.

However, I noticed that we aren't using `ScopedExpressionId` in a place
where it's important that the identifier is local. It's main use is
inside `infer` which we always run for the entire file. The one
exception to this is `Unpack` but unpack runs as part of `infer`.

Edit: The above isn't entirely correct. We used ScopedExpressionId in
TypeInference which is a query result. Now using ExpressionNodeKey does
mean that a change to the AST invalidates most if not all TypeInference
results of a single file. Salsa then has to run all dependent queries to
see if they're affected by this change even if the change was local to
another scope.

If this locality proves to be important I suggest that we create two
queries on top of TypeInference: one that returns the expression map
which is mainly used in the linter and type inference and a second that
returns all remaining fields. This should give us a similar optimization
at a much lower cost

I also considered remove `ScopedUseId` but I believe that one is still
useful because using `ExpressionNodeKey` for it instead would mean that
all `UseDefMap` change when a single AST node changes. Whether this is
important is something difficult to assess. I'm simply not familiar
enough with the `UseDefMap`. If the locality doesn't matter for the
`UseDefMap`, then a similar change could be made and `bindings_by_use`
could be changed to an `FxHashMap<UseId, Bindings>` where `UseId` is a
thin wrapper around `NodeKey`.

Closes astral-sh/ty#721

Author: MichaReiser

crates/ty_python_semantic/src/dunder_all.rs

@@ -6,9 +6,7 @@ use ruff_python_ast::name::Name;
 use ruff_python_ast::statement_visitor::{StatementVisitor, walk_stmt};
 use ruff_python_ast::{self as ast};
 
-use crate::semantic_index::ast_ids::HasScopedExpressionId;
-use crate::semantic_index::place::ScopeId;
-use crate::semantic_index::{SemanticIndex, global_scope, semantic_index};
+use crate::semantic_index::{SemanticIndex, semantic_index};
 use crate::types::{Truthiness, Type, infer_expression_types};
 use crate::{Db, ModuleName, resolve_module};
 
@@ -44,11 +42,6 @@ struct DunderAllNamesCollector<'db> {
     db: &'db dyn Db,
     file: File,
 
-    /// The scope in which the `__all__` names are being collected from.
-    ///
-    /// This is always going to be the global scope of the module.
-    scope: ScopeId<'db>,
-
     /// The semantic index for the module.
     index: &'db SemanticIndex<'db>,
 
@@ -68,7 +61,6 @@ impl<'db> DunderAllNamesCollector<'db> {
         Self {
             db,
             file,
-            scope: global_scope(db, file),
             index,
             origin: None,
             invalid: false,
@@ -190,8 +182,7 @@ impl<'db> DunderAllNamesCollector<'db> {
     ///
     /// This function panics if `expr` was not marked as a standalone expression during semantic indexing.
     fn standalone_expression_type(&self, expr: &ast::Expr) -> Type<'db> {
-        infer_expression_types(self.db, self.index.expression(expr))
-            .expression_type(expr.scoped_expression_id(self.db, self.scope))
+        infer_expression_types(self.db, self.index.expression(expr)).expression_type(expr)
     }
 
     /// Evaluate the given expression and return its truthiness.

crates/ty_python_semantic/src/semantic_index/ast_ids.rs

@@ -26,20 +26,11 @@ use crate::semantic_index::semantic_index;
 /// ```
 #[derive(Debug, salsa::Update, get_size2::GetSize)]
 pub(crate) struct AstIds {
-    /// Maps expressions to their expression id.
-    expressions_map: FxHashMap<ExpressionNodeKey, ScopedExpressionId>,
     /// Maps expressions which "use" a place (that is, [`ast::ExprName`], [`ast::ExprAttribute`] or [`ast::ExprSubscript`]) to a use id.
     uses_map: FxHashMap<ExpressionNodeKey, ScopedUseId>,
 }
 
 impl AstIds {
-    fn expression_id(&self, key: impl Into<ExpressionNodeKey>) -> ScopedExpressionId {
-        let key = &key.into();
-        *self.expressions_map.get(key).unwrap_or_else(|| {
-            panic!("Could not find expression ID for {key:?}");
-        })
-    }
-
     fn use_id(&self, key: impl Into<ExpressionNodeKey>) -> ScopedUseId {
         self.uses_map[&key.into()]
     }
@@ -94,90 +85,12 @@ impl HasScopedUseId for ast::ExprRef<'_> {
     }
 }
 
-/// Uniquely identifies an [`ast::Expr`] in a [`crate::semantic_index::place::FileScopeId`].
-#[newtype_index]
-#[derive(salsa::Update, get_size2::GetSize)]
-pub struct ScopedExpressionId;
-
-pub trait HasScopedExpressionId {
-    /// Returns the ID that uniquely identifies the node in `scope`.
-    fn scoped_expression_id(&self, db: &dyn Db, scope: ScopeId) -> ScopedExpressionId;
-}
-
-impl<T: HasScopedExpressionId> HasScopedExpressionId for Box<T> {
-    fn scoped_expression_id(&self, db: &dyn Db, scope: ScopeId) -> ScopedExpressionId {
-        self.as_ref().scoped_expression_id(db, scope)
-    }
-}
-
-macro_rules! impl_has_scoped_expression_id {
-    ($ty: ty) => {
-        impl HasScopedExpressionId for $ty {
-            fn scoped_expression_id(&self, db: &dyn Db, scope: ScopeId) -> ScopedExpressionId {
-                let expression_ref = ExprRef::from(self);
-                expression_ref.scoped_expression_id(db, scope)
-            }
-        }
-    };
-}
-
-impl_has_scoped_expression_id!(ast::ExprBoolOp);
-impl_has_scoped_expression_id!(ast::ExprName);
-impl_has_scoped_expression_id!(ast::ExprBinOp);
-impl_has_scoped_expression_id!(ast::ExprUnaryOp);
-impl_has_scoped_expression_id!(ast::ExprLambda);
-impl_has_scoped_expression_id!(ast::ExprIf);
-impl_has_scoped_expression_id!(ast::ExprDict);
-impl_has_scoped_expression_id!(ast::ExprSet);
-impl_has_scoped_expression_id!(ast::ExprListComp);
-impl_has_scoped_expression_id!(ast::ExprSetComp);
-impl_has_scoped_expression_id!(ast::ExprDictComp);
-impl_has_scoped_expression_id!(ast::ExprGenerator);
-impl_has_scoped_expression_id!(ast::ExprAwait);
-impl_has_scoped_expression_id!(ast::ExprYield);
-impl_has_scoped_expression_id!(ast::ExprYieldFrom);
-impl_has_scoped_expression_id!(ast::ExprCompare);
-impl_has_scoped_expression_id!(ast::ExprCall);
-impl_has_scoped_expression_id!(ast::ExprFString);
-impl_has_scoped_expression_id!(ast::ExprStringLiteral);
-impl_has_scoped_expression_id!(ast::ExprBytesLiteral);
-impl_has_scoped_expression_id!(ast::ExprNumberLiteral);
-impl_has_scoped_expression_id!(ast::ExprBooleanLiteral);
-impl_has_scoped_expression_id!(ast::ExprNoneLiteral);
-impl_has_scoped_expression_id!(ast::ExprEllipsisLiteral);
-impl_has_scoped_expression_id!(ast::ExprAttribute);
-impl_has_scoped_expression_id!(ast::ExprSubscript);
-impl_has_scoped_expression_id!(ast::ExprStarred);
-impl_has_scoped_expression_id!(ast::ExprNamed);
-impl_has_scoped_expression_id!(ast::ExprList);
-impl_has_scoped_expression_id!(ast::ExprTuple);
-impl_has_scoped_expression_id!(ast::ExprSlice);
-impl_has_scoped_expression_id!(ast::ExprIpyEscapeCommand);
-impl_has_scoped_expression_id!(ast::Expr);
-
-impl HasScopedExpressionId for ast::ExprRef<'_> {
-    fn scoped_expression_id(&self, db: &dyn Db, scope: ScopeId) -> ScopedExpressionId {
-        let ast_ids = ast_ids(db, scope);
-        ast_ids.expression_id(*self)
-    }
-}
-
 #[derive(Debug, Default)]
 pub(super) struct AstIdsBuilder {
-    expressions_map: FxHashMap<ExpressionNodeKey, ScopedExpressionId>,
     uses_map: FxHashMap<ExpressionNodeKey, ScopedUseId>,
 }
 
 impl AstIdsBuilder {
-    /// Adds `expr` to the expression ids map and returns its id.
-    pub(super) fn record_expression(&mut self, expr: &ast::Expr) -> ScopedExpressionId {
-        let expression_id = self.expressions_map.len().into();
-
-        self.expressions_map.insert(expr.into(), expression_id);
-
-        expression_id
-    }
-
     /// Adds `expr` to the use ids map and returns its id.
     pub(super) fn record_use(&mut self, expr: impl Into<ExpressionNodeKey>) -> ScopedUseId {
         let use_id = self.uses_map.len().into();
@@ -188,11 +101,9 @@ impl AstIdsBuilder {
     }
 
     pub(super) fn finish(mut self) -> AstIds {
-        self.expressions_map.shrink_to_fit();
         self.uses_map.shrink_to_fit();
 
         AstIds {
-            expressions_map: self.expressions_map,
             uses_map: self.uses_map,
         }
     }
@@ -219,6 +130,12 @@ pub(crate) mod node_key {
         }
     }
 
+    impl From<&ast::ExprCall> for ExpressionNodeKey {
+        fn from(value: &ast::ExprCall) -> Self {
+            Self(NodeKey::from_node(value))
+        }
+    }
+
     impl From<&ast::Identifier> for ExpressionNodeKey {
         fn from(value: &ast::Identifier) -> Self {
             Self(NodeKey::from_node(value))

crates/ty_python_semantic/src/semantic_index/builder.rs

@@ -1918,7 +1918,6 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> {
 
         self.scopes_by_expression
             .insert(expr.into(), self.current_scope());
-        self.current_ast_ids().record_expression(expr);
 
         let node_key = NodeKey::from_node(expr);
 

crates/ty_python_semantic/src/semantic_model.rs

@@ -7,7 +7,6 @@ use ruff_source_file::LineIndex;
 use crate::Db;
 use crate::module_name::ModuleName;
 use crate::module_resolver::{KnownModule, Module, resolve_module};
-use crate::semantic_index::ast_ids::HasScopedExpressionId;
 use crate::semantic_index::place::FileScopeId;
 use crate::semantic_index::semantic_index;
 use crate::types::ide_support::all_declarations_and_bindings;
@@ -159,8 +158,7 @@ impl HasType for ast::ExprRef<'_> {
         let file_scope = index.expression_scope_id(*self);
         let scope = file_scope.to_scope_id(model.db, model.file);
 
-        let expression_id = self.scoped_expression_id(model.db, scope);
-        infer_scope_types(model.db, scope).expression_type(expression_id)
+        infer_scope_types(model.db, scope).expression_type(*self)
     }
 }
 

crates/ty_python_semantic/src/types.rs

@@ -33,7 +33,6 @@ pub(crate) use self::subclass_of::{SubclassOfInner, SubclassOfType};
 use crate::module_name::ModuleName;
 use crate::module_resolver::{KnownModule, resolve_module};
 use crate::place::{Boundness, Place, PlaceAndQualifiers, imported_symbol};
-use crate::semantic_index::ast_ids::HasScopedExpressionId;
 use crate::semantic_index::definition::Definition;
 use crate::semantic_index::place::{ScopeId, ScopedPlaceId};
 use crate::semantic_index::{imported_modules, place_table, semantic_index};
@@ -143,18 +142,17 @@ fn definition_expression_type<'db>(
     let index = semantic_index(db, file);
     let file_scope = index.expression_scope_id(expression);
     let scope = file_scope.to_scope_id(db, file);
-    let expr_id = expression.scoped_expression_id(db, scope);
     if scope == definition.scope(db) {
         // expression is in the definition scope
         let inference = infer_definition_types(db, definition);
-        if let Some(ty) = inference.try_expression_type(expr_id) {
+        if let Some(ty) = inference.try_expression_type(expression) {
             ty
         } else {
-            infer_deferred_types(db, definition).expression_type(expr_id)
+            infer_deferred_types(db, definition).expression_type(expression)
         }
     } else {
         // expression is in a type-params sub-scope
-        infer_scope_types(db, scope).expression_type(expr_id)
+        infer_scope_types(db, scope).expression_type(expression)
     }
 }
 

crates/ty_python_semantic/src/types/class.rs

@@ -32,7 +32,6 @@ use crate::{
         known_module_symbol, place_from_bindings, place_from_declarations,
     },
     semantic_index::{
-        ast_ids::HasScopedExpressionId,
         attribute_assignments,
         definition::{DefinitionKind, TargetKind},
         place::ScopeId,
@@ -1861,10 +1860,8 @@ impl<'db> ClassLiteral<'db> {
                                 //     [.., self.name, ..] = <value>
 
                                 let unpacked = infer_unpack_types(db, unpack);
-                                let target_ast_id = assign
-                                    .target(&module)
-                                    .scoped_expression_id(db, method_scope);
-                                let inferred_ty = unpacked.expression_type(target_ast_id);
+
+                                let inferred_ty = unpacked.expression_type(assign.target(&module));
 
                                 union_of_inferred_types = union_of_inferred_types.add(inferred_ty);
                             }
@@ -1890,10 +1887,8 @@ impl<'db> ClassLiteral<'db> {
                                 //     for .., self.name, .. in <iterable>:
 
                                 let unpacked = infer_unpack_types(db, unpack);
-                                let target_ast_id = for_stmt
-                                    .target(&module)
-                                    .scoped_expression_id(db, method_scope);
-                                let inferred_ty = unpacked.expression_type(target_ast_id);
+                                let inferred_ty =
+                                    unpacked.expression_type(for_stmt.target(&module));
 
                                 union_of_inferred_types = union_of_inferred_types.add(inferred_ty);
                             }
@@ -1921,10 +1916,8 @@ impl<'db> ClassLiteral<'db> {
                                 //     with <context_manager> as .., self.name, ..:
 
                                 let unpacked = infer_unpack_types(db, unpack);
-                                let target_ast_id = with_item
-                                    .target(&module)
-                                    .scoped_expression_id(db, method_scope);
-                                let inferred_ty = unpacked.expression_type(target_ast_id);
+                                let inferred_ty =
+                                    unpacked.expression_type(with_item.target(&module));
 
                                 union_of_inferred_types = union_of_inferred_types.add(inferred_ty);
                             }
@@ -1951,10 +1944,9 @@ impl<'db> ClassLiteral<'db> {
                                 //     [... for .., self.name, .. in <iterable>]
 
                                 let unpacked = infer_unpack_types(db, unpack);
-                                let target_ast_id = comprehension
-                                    .target(&module)
-                                    .scoped_expression_id(db, unpack.target_scope(db));
-                                let inferred_ty = unpacked.expression_type(target_ast_id);
+
+                                let inferred_ty =
+                                    unpacked.expression_type(comprehension.target(&module));
 
                                 union_of_inferred_types = union_of_inferred_types.add(inferred_ty);
                             }