diff --git a/crates/red_knot_python_semantic/src/semantic_index/ast_ids.rs b/crates/red_knot_python_semantic/src/semantic_index/ast_ids.rs index 5a69d7b507..85c197073e 100644 --- a/crates/red_knot_python_semantic/src/semantic_index/ast_ids.rs +++ b/crates/red_knot_python_semantic/src/semantic_index/ast_ids.rs @@ -8,6 +8,7 @@ use ruff_python_ast::{AnyNodeRef, ExpressionRef}; use crate::ast_node_ref::AstNodeRef; use crate::node_key::NodeKey; +use crate::semantic_index::ast_ids::node_key::ExpressionNodeKey; use crate::semantic_index::semantic_index; use crate::semantic_index::symbol::{FileScopeId, ScopeId}; use crate::Db; @@ -32,7 +33,7 @@ pub(crate) struct AstIds { expressions: IndexVec>, /// Maps expressions to their expression id. Uses `NodeKey` because it avoids cloning [`Parsed`]. - expressions_map: FxHashMap, + expressions_map: FxHashMap, statements: IndexVec>, @@ -47,11 +48,8 @@ impl AstIds { self.statements_map[&NodeKey::from_node(node.into())] } - fn expression_id<'a, N>(&self, node: N) -> ScopedExpressionId - where - N: Into>, - { - self.expressions_map[&NodeKey::from_node(node.into())] + fn expression_id(&self, key: impl Into) -> ScopedExpressionId { + self.expressions_map[&key.into()] } } @@ -294,7 +292,7 @@ impl_has_scoped_statement_id!(ast::StmtImportFrom); #[derive(Debug)] pub(super) struct AstIdsBuilder { expressions: IndexVec>, - expressions_map: FxHashMap, + expressions_map: FxHashMap, statements: IndexVec>, statements_map: FxHashMap, } @@ -341,8 +339,7 @@ impl AstIdsBuilder { ) -> ScopedExpressionId { let expression_id = self.expressions.push(AstNodeRef::new(parsed.clone(), expr)); - self.expressions_map - .insert(NodeKey::from_node(expr), expression_id); + self.expressions_map.insert(expr.into(), expression_id); expression_id } @@ -361,3 +358,25 @@ impl AstIdsBuilder { } } } + +/// Node key that can only be constructed for expressions. +mod node_key { + use ruff_python_ast as ast; + + use crate::node_key::NodeKey; + + #[derive(Copy, Clone, Eq, PartialEq, Hash, Debug)] + pub(super) struct ExpressionNodeKey(NodeKey); + + impl From> for ExpressionNodeKey { + fn from(value: ast::ExpressionRef<'_>) -> Self { + Self(NodeKey::from_node(value)) + } + } + + impl From<&ast::Expr> for ExpressionNodeKey { + fn from(value: &ast::Expr) -> Self { + Self(NodeKey::from_node(value)) + } + } +}