Map directly to UseId

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<Definition<'_>> {
+        fn first_binding_at_use(&self, use_id: FileUseId) -> Option<Definition<'_>> {
-            self.bindings_at_use(use_id)
+            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 =
+        let element_use_id = element.use_id();
-            element.scoped_use_id(&db, comprehension_scope_id.to_scope_id(&db, file));
 
         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 {

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()
+                self.current_use_def_map_mut().record_use(
-                    .record_use(symbol, NodeKey::from_node(name));
+                    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 {

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<ScopedNarrowingConstraintPredic
 pub(crate) enum ConstraintKey {
     NarrowingConstraint(ScopedNarrowingConstraint),
     EagerNestedScope(FileScopeId),
-    UseId(ScopedUseId),
+    UseId(FileUseId),
 }
 
 /// One of the [`Predicate`]s in a narrowing constraint, which constraints the type of the

crates/ty_python_semantic/src/semantic_index/use_def.rs

@@ -248,7 +248,6 @@ use self::place_state::{
     Bindings, Declarations, EagerSnapshot, LiveBindingsIterator, LiveDeclaration,
     LiveDeclarationsIterator, PlaceState, ScopedDefinitionId,
 };
-use crate::Db;
 use crate::node_key::NodeKey;
 use crate::place::BoundnessAnalysis;
 use crate::semantic_index::definition::{Definition, DefinitionState};
@@ -256,7 +255,7 @@ use crate::semantic_index::narrowing_constraints::{
     ConstraintKey, NarrowingConstraints, NarrowingConstraintsBuilder, NarrowingConstraintsIterator,
 };
 use crate::semantic_index::place::{
-    FileScopeId, PlaceExpr, PlaceExprWithFlags, ScopeId, ScopeKind, ScopedPlaceId,
+    FileScopeId, PlaceExpr, PlaceExprWithFlags, ScopeKind, ScopedPlaceId,
 };
 use crate::semantic_index::predicate::{
     Predicate, PredicateOrLiteral, Predicates, PredicatesBuilder, ScopedPredicateId,
@@ -265,7 +264,7 @@ use crate::semantic_index::reachability_constraints::{
     ReachabilityConstraints, ReachabilityConstraintsBuilder, ScopedReachabilityConstraintId,
 };
 use crate::semantic_index::use_def::place_state::PreviousDefinitions;
-use crate::semantic_index::{EagerSnapshotResult, SemanticIndex, use_def_map};
+use crate::semantic_index::{EagerSnapshotResult, SemanticIndex};
 use crate::types::{IntersectionBuilder, Truthiness, Type, infer_narrowing_constraint};
 
 mod place_state;
@@ -286,12 +285,8 @@ pub(crate) struct UseDefMap<'db> {
     /// Array of reachability constraints in this scope.
     reachability_constraints: ReachabilityConstraints,
 
-    /// Map from node to their use id.
+    /// [`Bindings`] reaching a [`FileUseId`].
-    /// Only contains entries for nodes implementing [`HasScopedUseId`].
+    bindings_by_use: FxHashMap<FileUseId, Bindings>,
-    uses_by_node: FxHashMap<NodeKey, ScopedUseId>,
-
-    /// [`Bindings`] reaching a [`ScopedUseId`].
-    bindings_by_use: IndexVec<ScopedUseId, Bindings>,
 
     /// Tracks whether or not a given AST node is reachable from the start of the scope.
     node_reachability: FxHashMap<NodeKey, ScopedReachabilityConstraintId>,
@@ -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 {
+    pub(crate) fn bindings_for_node(
-        self.uses_by_node[&node]
-    }
-
-    pub(crate) fn bindings_at_use(
         &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<NodeKey, ScopedUseId>,
-
     /// Live bindings at each so-far-recorded use.
-    bindings_by_use: IndexVec<ScopedUseId, Bindings>,
+    bindings_by_use: FxHashMap<FileUseId, Bindings>,
 
     /// 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(),
+            bindings_by_use: FxHashMap::default(),
-            uses_by_node: 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!(old_use, None);
-        debug_assert_eq!(self.bindings_by_use.len(), self.uses_by_node.len());
 
         // 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) {
+    pub(super) fn record_node_reachability(&mut self, node: NodeKey) {
-        self.node_reachability.insert(node_key, self.reachability);
+        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, Debug, Copy, Clone, Eq, PartialEq, Hash)]
-#[derive(get_size2::GetSize)]
+pub(crate) struct FileUseId(NodeKey);
-pub struct ScopedUseId;
 
-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 {
+impl HasFileUseId for ast::Identifier {
-    fn scoped_use_id(&self, db: &dyn Db, scope: ScopeId) -> ScopedUseId {
+    fn use_id(&self) -> FileUseId {
-        let use_def_map = use_def_map(db, scope);
+        FileUseId(NodeKey::from_node(self))
-        use_def_map.use_id(NodeKey::from_node(self))
     }
 }
 
-impl HasScopedUseId for ast::ExprName {
+impl HasFileUseId for ast::ExprName {
-    fn scoped_use_id(&self, db: &dyn Db, scope: ScopeId) -> ScopedUseId {
+    fn use_id(&self) -> FileUseId {
-        let expression_ref = ast::ExprRef::from(self);
+        FileUseId(NodeKey::from_node(self))
-        expression_ref.scoped_use_id(db, scope)
     }
 }
 
-impl HasScopedUseId for ast::ExprAttribute {
+impl HasFileUseId for ast::ExprAttribute {
-    fn scoped_use_id(&self, db: &dyn Db, scope: ScopeId) -> ScopedUseId {
+    fn use_id(&self) -> FileUseId {
-        let expression_ref = ast::ExprRef::from(self);
+        FileUseId(NodeKey::from_node(self))
-        expression_ref.scoped_use_id(db, scope)
     }
 }
 
-impl HasScopedUseId for ast::ExprSubscript {
+impl HasFileUseId for ast::ExprSubscript {
-    fn scoped_use_id(&self, db: &dyn Db, scope: ScopeId) -> ScopedUseId {
+    fn use_id(&self) -> FileUseId {
-        let expression_ref = ast::ExprRef::from(self);
+        FileUseId(NodeKey::from_node(self))
-        expression_ref.scoped_use_id(db, scope)
     }
 }
 
-impl HasScopedUseId for ast::ExprRef<'_> {
+impl HasFileUseId for ast::ExprRef<'_> {
-    fn scoped_use_id(&self, db: &dyn Db, scope: ScopeId) -> ScopedUseId {
+    fn use_id(&self) -> FileUseId {
-        let use_def_map = use_def_map(db, scope);
+        FileUseId(NodeKey::from_node(self))
-        use_def_map.use_id(NodeKey::from_node(self))
     }
 }

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;
         };

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<ScopedUseId>) {
+    ) -> (Place<'db>, Option<FileUseId>) {
         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 use_id = expr_ref.use_id();
-            let place = place_from_bindings(db, use_def.bindings_at_use(use_id));
+            let place = place_from_bindings(db, use_def.bindings_for_node(use_id));
             (place, Some(use_id))
         }
     }