From 56cb92f486a6680584f51cda5ac37708676b592d Mon Sep 17 00:00:00 2001 From: Micha Reiser Date: Sat, 12 Jul 2025 12:58:23 +0200 Subject: [PATCH] Map directly to UseId --- .../ty_python_semantic/src/semantic_index.rs | 13 ++- .../src/semantic_index/builder.rs | 18 +++- .../semantic_index/narrowing_constraints.rs | 4 +- .../src/semantic_index/use_def.rs | 98 ++++++++----------- .../ty_python_semantic/src/types/function.rs | 6 +- crates/ty_python_semantic/src/types/infer.rs | 8 +- 6 files changed, 69 insertions(+), 78 deletions(-) diff --git a/crates/ty_python_semantic/src/semantic_index.rs b/crates/ty_python_semantic/src/semantic_index.rs index 01c04fca7a..858bf4b4e5 100644 --- a/crates/ty_python_semantic/src/semantic_index.rs +++ b/crates/ty_python_semantic/src/semantic_index.rs @@ -22,7 +22,7 @@ use crate::semantic_index::place::{ ScopeKind, ScopedPlaceId, }; use crate::semantic_index::use_def::{EagerSnapshotKey, ScopedEagerSnapshotId, UseDefMap}; -pub(crate) use crate::semantic_index::use_def::{HasScopedUseId, ScopedUseId}; +pub(crate) use crate::semantic_index::use_def::{FileUseId, HasFileUseId}; use crate::util::get_size::untracked_arc_size; mod builder; @@ -614,7 +614,7 @@ mod tests { use crate::semantic_index::place::{FileScopeId, PlaceTable, Scope, ScopeKind, ScopedPlaceId}; use crate::semantic_index::use_def::UseDefMap; use crate::semantic_index::{ - HasScopedUseId, ScopedUseId, global_scope, place_table, semantic_index, use_def_map, + FileUseId, HasFileUseId, global_scope, place_table, semantic_index, use_def_map, }; impl UseDefMap<'_> { @@ -623,8 +623,8 @@ mod tests { .find_map(|constrained_binding| constrained_binding.binding.definition()) } - fn first_binding_at_use(&self, use_id: ScopedUseId) -> Option> { - self.bindings_at_use(use_id) + fn first_binding_at_use(&self, use_id: FileUseId) -> Option> { + self.bindings_for_node(use_id) .find_map(|constrained_binding| constrained_binding.binding.definition()) } } @@ -1047,8 +1047,7 @@ def f(a: str, /, b: str, c: int = 1, *args, d: int = 2, **kwargs): .elt .as_name_expr() .unwrap(); - let element_use_id = - element.scoped_use_id(&db, comprehension_scope_id.to_scope_id(&db, file)); + let element_use_id = element.use_id(); let binding = use_def.first_binding_at_use(element_use_id).unwrap(); let DefinitionKind::Comprehension(comprehension) = binding.kind(&db) else { @@ -1322,7 +1321,7 @@ class C[T]: let ast::Expr::Name(x_use_expr_name) = x_use_expr.as_ref() else { panic!("expected a Name"); }; - let x_use_id = x_use_expr_name.scoped_use_id(&db, scope); + let x_use_id = x_use_expr_name.use_id(); let use_def = use_def_map(&db, scope); let binding = use_def.first_binding_at_use(x_use_id).unwrap(); let DefinitionKind::Assignment(assignment) = binding.kind(&db) else { diff --git a/crates/ty_python_semantic/src/semantic_index/builder.rs b/crates/ty_python_semantic/src/semantic_index/builder.rs index 4b5e424e74..b735db5427 100644 --- a/crates/ty_python_semantic/src/semantic_index/builder.rs +++ b/crates/ty_python_semantic/src/semantic_index/builder.rs @@ -44,7 +44,7 @@ use crate::semantic_index::reachability_constraints::{ use crate::semantic_index::use_def::{ EagerSnapshotKey, FlowSnapshot, ScopedEagerSnapshotId, UseDefMapBuilder, }; -use crate::semantic_index::{ArcUseDefMap, SemanticIndex}; +use crate::semantic_index::{ArcUseDefMap, HasFileUseId, SemanticIndex}; use crate::unpack::{Unpack, UnpackKind, UnpackPosition, UnpackValue}; use crate::{Db, Program}; @@ -1125,8 +1125,11 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> { // done on the `Identifier` node as opposed to `ExprName` because that's what the // AST uses. self.mark_place_used(symbol); - self.current_use_def_map_mut() - .record_use(symbol, NodeKey::from_node(name)); + self.current_use_def_map_mut().record_use( + symbol, + name.use_id(), + NodeKey::from_node(name), + ); self.add_definition(symbol, function_def); } @@ -2033,9 +2036,16 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> { let place_id = self.add_place(place_expr); if is_use { + let use_id = match expr { + ast::Expr::Name(name) => name.use_id(), + ast::Expr::Attribute(attribute) => attribute.use_id(), + ast::Expr::Subscript(subscript) => subscript.use_id(), + _ => unreachable!(), + }; + self.mark_place_used(place_id); self.current_use_def_map_mut() - .record_use(place_id, node_key); + .record_use(place_id, use_id, node_key); } if is_definition { diff --git a/crates/ty_python_semantic/src/semantic_index/narrowing_constraints.rs b/crates/ty_python_semantic/src/semantic_index/narrowing_constraints.rs index 59547b0f5c..6ea10c0d72 100644 --- a/crates/ty_python_semantic/src/semantic_index/narrowing_constraints.rs +++ b/crates/ty_python_semantic/src/semantic_index/narrowing_constraints.rs @@ -31,7 +31,7 @@ use crate::list::{List, ListBuilder, ListSetReverseIterator, ListStorage}; use crate::semantic_index::place::FileScopeId; use crate::semantic_index::predicate::ScopedPredicateId; -use crate::semantic_index::use_def::ScopedUseId; +use crate::semantic_index::use_def::FileUseId; /// A narrowing constraint associated with a live binding. /// @@ -44,7 +44,7 @@ pub(crate) type ScopedNarrowingConstraint = List { /// Array of reachability constraints in this scope. reachability_constraints: ReachabilityConstraints, - /// Map from node to their use id. - /// Only contains entries for nodes implementing [`HasScopedUseId`]. - uses_by_node: FxHashMap, - - /// [`Bindings`] reaching a [`ScopedUseId`]. - bindings_by_use: IndexVec, + /// [`Bindings`] reaching a [`FileUseId`]. + bindings_by_use: FxHashMap, /// Tracks whether or not a given AST node is reachable from the start of the scope. node_reachability: FxHashMap, @@ -353,16 +348,12 @@ pub(crate) enum ApplicableConstraints<'map, 'db> { impl<'db> UseDefMap<'db> { #[track_caller] - pub(crate) fn use_id(&self, node: NodeKey) -> ScopedUseId { - self.uses_by_node[&node] - } - - pub(crate) fn bindings_at_use( + pub(crate) fn bindings_for_node( &self, - use_id: ScopedUseId, + use_id: FileUseId, ) -> BindingWithConstraintsIterator<'_, 'db> { self.bindings_iterator( - &self.bindings_by_use[use_id], + &self.bindings_by_use[&use_id], BoundnessAnalysis::BasedOnUnboundVisibility, ) } @@ -392,7 +383,7 @@ impl<'db> UseDefMap<'db> { ApplicableConstraints::ConstrainedBindings(bindings) } ConstraintKey::UseId(use_id) => { - ApplicableConstraints::ConstrainedBindings(self.bindings_at_use(use_id)) + ApplicableConstraints::ConstrainedBindings(self.bindings_for_node(use_id)) } } } @@ -744,12 +735,8 @@ pub(super) struct UseDefMapBuilder<'db> { /// Builder of reachability constraints. pub(super) reachability_constraints: ReachabilityConstraintsBuilder, - /// Map from node to their use id. - /// Only contains entries for nodes implementing [`HasScopedUseId`]. - uses_by_node: FxHashMap, - /// Live bindings at each so-far-recorded use. - bindings_by_use: IndexVec, + bindings_by_use: FxHashMap, /// Tracks whether or not the current point in control flow is reachable from the /// start of the scope. @@ -785,8 +772,7 @@ impl<'db> UseDefMapBuilder<'db> { predicates: PredicatesBuilder::default(), narrowing_constraints: NarrowingConstraintsBuilder::default(), reachability_constraints: ReachabilityConstraintsBuilder::default(), - bindings_by_use: IndexVec::new(), - uses_by_node: FxHashMap::default(), + bindings_by_use: FxHashMap::default(), reachability: ScopedReachabilityConstraintId::ALWAYS_TRUE, node_reachability: FxHashMap::default(), declarations_by_binding: FxHashMap::default(), @@ -1015,23 +1001,27 @@ impl<'db> UseDefMapBuilder<'db> { ); } - pub(super) fn record_use(&mut self, place: ScopedPlaceId, node_key: NodeKey) { + pub(super) fn record_use( + &mut self, + place: ScopedPlaceId, + use_id: FileUseId, + node_key: NodeKey, + ) { // We have a use of a place; clone the current bindings for that place, and record them // as the live bindings for this use. - let new_use = self + let old_use = self .bindings_by_use - .push(self.place_states[place].bindings().clone()); + .insert(use_id, self.place_states[place].bindings().clone()); - self.uses_by_node.insert(node_key, new_use); - debug_assert_eq!(self.bindings_by_use.len(), self.uses_by_node.len()); + debug_assert_eq!(old_use, None); // Track reachability of all uses of places to silence `unresolved-reference` // diagnostics in unreachable code. self.record_node_reachability(node_key); } - pub(super) fn record_node_reachability(&mut self, node_key: NodeKey) { - self.node_reachability.insert(node_key, self.reachability); + pub(super) fn record_node_reachability(&mut self, node: NodeKey) { + self.node_reachability.insert(node, self.reachability); } pub(super) fn snapshot_eager_state( @@ -1134,7 +1124,6 @@ impl<'db> UseDefMapBuilder<'db> { self.all_definitions.shrink_to_fit(); self.place_states.shrink_to_fit(); self.reachable_definitions.shrink_to_fit(); - self.uses_by_node.shrink_to_fit(); self.bindings_by_use.shrink_to_fit(); self.node_reachability.shrink_to_fit(); self.declarations_by_binding.shrink_to_fit(); @@ -1146,7 +1135,6 @@ impl<'db> UseDefMapBuilder<'db> { predicates: self.predicates.build(), narrowing_constraints: self.narrowing_constraints.build(), reachability_constraints: self.reachability_constraints.build(), - uses_by_node: self.uses_by_node, bindings_by_use: self.bindings_by_use, node_reachability: self.node_reachability, end_of_scope_places: self.place_states, @@ -1160,46 +1148,40 @@ impl<'db> UseDefMapBuilder<'db> { } /// Uniquely identifies a use of a name in a [`crate::semantic_index::place::FileScopeId`]. -#[newtype_index] -#[derive(get_size2::GetSize)] -pub struct ScopedUseId; +#[derive(get_size2::GetSize, Debug, Copy, Clone, Eq, PartialEq, Hash)] +pub(crate) struct FileUseId(NodeKey); -pub(crate) trait HasScopedUseId { +pub(crate) trait HasFileUseId { /// Returns the ID that uniquely identifies the use in `scope`. - fn scoped_use_id(&self, db: &dyn Db, scope: ScopeId) -> ScopedUseId; + fn use_id(&self) -> FileUseId; } -impl HasScopedUseId for ast::Identifier { - fn scoped_use_id(&self, db: &dyn Db, scope: ScopeId) -> ScopedUseId { - let use_def_map = use_def_map(db, scope); - use_def_map.use_id(NodeKey::from_node(self)) +impl HasFileUseId for ast::Identifier { + fn use_id(&self) -> FileUseId { + FileUseId(NodeKey::from_node(self)) } } -impl HasScopedUseId for ast::ExprName { - fn scoped_use_id(&self, db: &dyn Db, scope: ScopeId) -> ScopedUseId { - let expression_ref = ast::ExprRef::from(self); - expression_ref.scoped_use_id(db, scope) +impl HasFileUseId for ast::ExprName { + fn use_id(&self) -> FileUseId { + FileUseId(NodeKey::from_node(self)) } } -impl HasScopedUseId for ast::ExprAttribute { - fn scoped_use_id(&self, db: &dyn Db, scope: ScopeId) -> ScopedUseId { - let expression_ref = ast::ExprRef::from(self); - expression_ref.scoped_use_id(db, scope) +impl HasFileUseId for ast::ExprAttribute { + fn use_id(&self) -> FileUseId { + FileUseId(NodeKey::from_node(self)) } } -impl HasScopedUseId for ast::ExprSubscript { - fn scoped_use_id(&self, db: &dyn Db, scope: ScopeId) -> ScopedUseId { - let expression_ref = ast::ExprRef::from(self); - expression_ref.scoped_use_id(db, scope) +impl HasFileUseId for ast::ExprSubscript { + fn use_id(&self) -> FileUseId { + FileUseId(NodeKey::from_node(self)) } } -impl HasScopedUseId for ast::ExprRef<'_> { - fn scoped_use_id(&self, db: &dyn Db, scope: ScopeId) -> ScopedUseId { - let use_def_map = use_def_map(db, scope); - use_def_map.use_id(NodeKey::from_node(self)) +impl HasFileUseId for ast::ExprRef<'_> { + fn use_id(&self) -> FileUseId { + FileUseId(NodeKey::from_node(self)) } } diff --git a/crates/ty_python_semantic/src/types/function.rs b/crates/ty_python_semantic/src/types/function.rs index 97531281d5..7f8e3e75ab 100644 --- a/crates/ty_python_semantic/src/types/function.rs +++ b/crates/ty_python_semantic/src/types/function.rs @@ -62,7 +62,7 @@ use crate::module_resolver::{KnownModule, file_to_module}; use crate::place::{Boundness, Place, place_from_bindings}; use crate::semantic_index::definition::Definition; use crate::semantic_index::place::ScopeId; -use crate::semantic_index::{HasScopedUseId, semantic_index}; +use crate::semantic_index::{HasFileUseId, semantic_index}; use crate::types::context::InferContext; use crate::types::diagnostic::{ REDUNDANT_CAST, STATIC_ASSERT_ERROR, TYPE_ASSERTION_FAILURE, @@ -292,10 +292,10 @@ impl<'db> OverloadLiteral<'db> { .node(db) .expect_function(&module) .name - .scoped_use_id(db, scope); + .use_id(); let Place::Type(Type::FunctionLiteral(previous_type), Boundness::Bound) = - place_from_bindings(db, use_def.bindings_at_use(use_id)) + place_from_bindings(db, use_def.bindings_for_node(use_id)) else { return None; }; diff --git a/crates/ty_python_semantic/src/types/infer.rs b/crates/ty_python_semantic/src/types/infer.rs index 6d520956b0..00a170bf2c 100644 --- a/crates/ty_python_semantic/src/types/infer.rs +++ b/crates/ty_python_semantic/src/types/infer.rs @@ -82,7 +82,7 @@ use crate::semantic_index::place::{ FileScopeId, NodeWithScopeKind, NodeWithScopeRef, PlaceExpr, ScopeId, ScopeKind, ScopedPlaceId, }; use crate::semantic_index::{ - ApplicableConstraints, EagerSnapshotResult, HasScopedUseId, ScopedUseId, SemanticIndex, + ApplicableConstraints, EagerSnapshotResult, FileUseId, HasFileUseId, SemanticIndex, place_table, semantic_index, }; use crate::types::call::{Binding, Bindings, CallArgumentTypes, CallArguments, CallError}; @@ -5823,7 +5823,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { &self, expr: &PlaceExpr, expr_ref: ast::ExprRef, - ) -> (Place<'db>, Option) { + ) -> (Place<'db>, Option) { let db = self.db(); let scope = self.scope(); let file_scope_id = scope.file_scope_id(db); @@ -5850,8 +5850,8 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { return (Place::Unbound, None); } - let use_id = expr_ref.scoped_use_id(db, scope); - let place = place_from_bindings(db, use_def.bindings_at_use(use_id)); + let use_id = expr_ref.use_id(); + let place = place_from_bindings(db, use_def.bindings_for_node(use_id)); (place, Some(use_id)) } }