Replace ScopedExpressionId with ExpressionNodeKey

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

@@ -1902,7 +1902,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

@@ -32,7 +32,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};
@@ -141,18 +140,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,
@@ -1867,10 +1866,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);
                             }
@@ -1896,10 +1893,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);
                             }
@@ -1927,10 +1922,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);
                             }
@@ -1957,10 +1950,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);
                             }