[ty] Move node_reachability from UseDefMap to SemanticIndex

crates/ty_python_semantic/src/semantic_index.rs

@@ -23,6 +23,7 @@ use crate::semantic_index::place::{
     FileScopeId, NodeWithScopeKey, NodeWithScopeRef, PlaceExpr, PlaceTable, Scope, ScopeId,
     ScopeKind, ScopedPlaceId,
 };
+use crate::semantic_index::reachability_constraints::ScopedReachabilityConstraintId;
 use crate::semantic_index::use_def::{EagerSnapshotKey, ScopedEagerSnapshotId, UseDefMap};
 use crate::util::get_size::untracked_arc_size;
 
@@ -211,6 +212,9 @@ pub(crate) struct SemanticIndex<'db> {
     /// Map from a standalone expression to its [`Expression`] ingredient.
     expressions_by_node: FxHashMap<ExpressionNodeKey, Expression<'db>>,
 
+    /// Tracks whether or not a given AST node is reachable from the start of the scope.
+    node_reachability: FxHashMap<NodeKey, ScopedReachabilityConstraintId>,
+
     /// Map from nodes that create a scope to the scope they create.
     scopes_by_node: FxHashMap<NodeWithScopeKey, FileScopeId>,
 
@@ -364,8 +368,15 @@ impl<'db> SemanticIndex<'db> {
         scope_id: FileScopeId,
         node_key: NodeKey,
     ) -> bool {
-        self.is_scope_reachable(db, scope_id)
-            && self.use_def_map(scope_id).is_node_reachable(db, node_key)
+        if !self.is_scope_reachable(db, scope_id) {
+            return false;
+        }
+
+        let constraint = *self.node_reachability.get(&node_key).expect(
+            "`is_node_reachable` should only be called on AST nodes with recorded reachability",
+        );
+
+        self.use_def_map(scope_id).is_node_reachable(db, constraint)
     }
 
     /// Returns an iterator over the descendent scopes of `scope`.

crates/ty_python_semantic/src/semantic_index/builder.rs

@@ -20,8 +20,8 @@ use crate::ast_node_ref::AstNodeRef;
 use crate::module_name::ModuleName;
 use crate::module_resolver::resolve_module;
 use crate::node_key::NodeKey;
-use crate::semantic_index::ast_ids::AstIdsBuilder;
 use crate::semantic_index::ast_ids::node_key::ExpressionNodeKey;
+use crate::semantic_index::ast_ids::{AstIdsBuilder, ScopedUseId};
 use crate::semantic_index::definition::{
     AnnotatedAssignmentDefinitionNodeRef, AssignmentDefinitionNodeRef,
     ComprehensionDefinitionNodeRef, Definition, DefinitionCategory, DefinitionNodeKey,
@@ -105,6 +105,8 @@ pub(super) struct SemanticIndexBuilder<'db, 'ast> {
     scopes_by_expression: FxHashMap<ExpressionNodeKey, FileScopeId>,
     definitions_by_node: FxHashMap<DefinitionNodeKey, Definitions<'db>>,
     expressions_by_node: FxHashMap<ExpressionNodeKey, Expression<'db>>,
+    /// Tracks whether or not a given AST node is reachable from the start of the scope.
+    node_reachability: FxHashMap<NodeKey, ScopedReachabilityConstraintId>,
     imported_modules: FxHashSet<ModuleName>,
     /// Hashset of all [`FileScopeId`]s that correspond to [generator functions].
     ///
@@ -140,6 +142,7 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
             scopes_by_node: FxHashMap::default(),
             definitions_by_node: FxHashMap::default(),
             expressions_by_node: FxHashMap::default(),
+            node_reachability: FxHashMap::default(),
 
             imported_modules: FxHashSet::default(),
             generator_functions: FxHashSet::default(),
@@ -664,6 +667,19 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
             .record_reachability_constraint(negated_constraint);
     }
 
+    fn record_node_reachability(&mut self, node: NodeKey) {
+        self.node_reachability
+            .insert(node, self.current_use_def_map().reachability);
+    }
+
+    fn record_use(&mut self, place: ScopedPlaceId, use_id: ScopedUseId, node_key: NodeKey) {
+        self.current_use_def_map_mut().record_use(place, use_id);
+
+        // Track reachability of all uses of places to silence `unresolved-reference`
+        // diagnostics in unreachable code.
+        self.record_node_reachability(node_key);
+    }
+
     fn push_assignment(&mut self, assignment: CurrentAssignment<'ast, 'db>) {
         self.current_assignments.push(assignment);
     }
@@ -1040,6 +1056,7 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
         ast_ids.shrink_to_fit();
         self.scopes_by_expression.shrink_to_fit();
         self.definitions_by_node.shrink_to_fit();
+        self.node_reachability.shrink_to_fit();
 
         self.scope_ids_by_scope.shrink_to_fit();
         self.scopes_by_node.shrink_to_fit();
@@ -1055,6 +1072,7 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
             ast_ids,
             scopes_by_expression: self.scopes_by_expression,
             scopes_by_node: self.scopes_by_node,
+            node_reachability: self.node_reachability,
             use_def_maps,
             imported_modules: Arc::new(self.imported_modules),
             has_future_annotations: self.has_future_annotations,
@@ -1145,8 +1163,7 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> {
                 // AST uses.
                 self.mark_place_used(symbol);
                 let use_id = self.current_ast_ids().record_use(name);
-                self.current_use_def_map_mut()
-                    .record_use(symbol, use_id, NodeKey::from_node(name));
+                self.record_use(symbol, use_id, NodeKey::from_node(name));
 
                 self.add_definition(symbol, function_def);
             }
@@ -1196,8 +1213,7 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> {
                 );
             }
             ast::Stmt::Import(node) => {
-                self.current_use_def_map_mut()
-                    .record_node_reachability(NodeKey::from_node(node));
+                self.record_node_reachability(NodeKey::from_node(node));
 
                 for (alias_index, alias) in node.names.iter().enumerate() {
                     // Mark the imported module, and all of its parents, as being imported in this
@@ -1224,8 +1240,7 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> {
                 }
             }
             ast::Stmt::ImportFrom(node) => {
-                self.current_use_def_map_mut()
-                    .record_node_reachability(NodeKey::from_node(node));
+                self.record_node_reachability(NodeKey::from_node(node));
 
                 let mut found_star = false;
                 for (alias_index, alias) in node.names.iter().enumerate() {
@@ -2055,8 +2070,7 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> {
                     if is_use {
                         self.mark_place_used(place_id);
                         let use_id = self.current_ast_ids().record_use(expr);
-                        self.current_use_def_map_mut()
-                            .record_use(place_id, use_id, node_key);
+                        self.record_use(place_id, use_id, node_key);
                     }
 
                     if is_definition {
@@ -2149,8 +2163,7 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> {
                 // Track reachability of attribute expressions to silence `unresolved-attribute`
                 // diagnostics in unreachable code.
                 if expr.is_attribute_expr() {
-                    self.current_use_def_map_mut()
-                        .record_node_reachability(node_key);
+                    self.record_node_reachability(node_key);
                 }
 
                 walk_expr(self, expr);
@@ -2311,8 +2324,7 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> {
             ast::Expr::StringLiteral(_) => {
                 // Track reachability of string literals, as they could be a stringified annotation
                 // with child expressions whose reachability we are interested in.
-                self.current_use_def_map_mut()
-                    .record_node_reachability(node_key);
+                self.record_node_reachability(node_key);
 
                 walk_expr(self, expr);
             }

crates/ty_python_semantic/src/semantic_index/use_def.rs

@@ -247,7 +247,6 @@ use self::place_state::{
     Bindings, Declarations, EagerSnapshot, LiveBindingsIterator, LiveDeclaration,
     LiveDeclarationsIterator, PlaceState, ScopedDefinitionId,
 };
-use crate::node_key::NodeKey;
 use crate::place::BoundnessAnalysis;
 use crate::semantic_index::ast_ids::ScopedUseId;
 use crate::semantic_index::definition::{Definition, DefinitionState};
@@ -288,9 +287,6 @@ pub(crate) struct UseDefMap<'db> {
     /// [`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>,
-
     /// If the definition is a binding (only) -- `x = 1` for example -- then we need
     /// [`Declarations`] to know whether this binding is permitted by the live declarations.
     ///
@@ -402,17 +398,13 @@ impl<'db> UseDefMap<'db> {
     /// be unreachable. Use [`super::SemanticIndex::is_node_reachable`] for the global
     /// analysis.
     #[track_caller]
-    pub(super) fn is_node_reachable(&self, db: &dyn crate::Db, node_key: NodeKey) -> bool {
-        self
-            .reachability_constraints
-            .evaluate(
-                db,
-                &self.predicates,
-                *self
-                    .node_reachability
-                    .get(&node_key)
-                    .expect("`is_node_reachable` should only be called on AST nodes with recorded reachability"),
-            )
+    pub(super) fn is_node_reachable(
+        &self,
+        db: &dyn crate::Db,
+        constraint: ScopedReachabilityConstraintId,
+    ) -> bool {
+        self.reachability_constraints
+            .evaluate(db, &self.predicates, constraint)
             .may_be_true()
     }
 
@@ -741,9 +733,6 @@ pub(super) struct UseDefMapBuilder<'db> {
     /// start of the scope.
     pub(super) reachability: ScopedReachabilityConstraintId,
 
-    /// Tracks whether or not a given AST node is reachable from the start of the scope.
-    node_reachability: FxHashMap<NodeKey, ScopedReachabilityConstraintId>,
-
     /// Live declarations for each so-far-recorded binding.
     declarations_by_binding: FxHashMap<Definition<'db>, Declarations>,
 
@@ -773,7 +762,6 @@ impl<'db> UseDefMapBuilder<'db> {
             reachability_constraints: ReachabilityConstraintsBuilder::default(),
             bindings_by_use: IndexVec::new(),
             reachability: ScopedReachabilityConstraintId::ALWAYS_TRUE,
-            node_reachability: FxHashMap::default(),
             declarations_by_binding: FxHashMap::default(),
             bindings_by_definition: FxHashMap::default(),
             place_states: IndexVec::new(),
@@ -1000,26 +988,13 @@ impl<'db> UseDefMapBuilder<'db> {
         );
     }
 
-    pub(super) fn record_use(
-        &mut self,
-        place: ScopedPlaceId,
-        use_id: ScopedUseId,
-        node_key: NodeKey,
-    ) {
+    pub(super) fn record_use(&mut self, place: ScopedPlaceId, use_id: ScopedUseId) {
         // 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
             .bindings_by_use
             .push(self.place_states[place].bindings().clone());
         debug_assert_eq!(use_id, new_use);
-
-        // 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 snapshot_eager_state(
@@ -1123,7 +1098,6 @@ impl<'db> UseDefMapBuilder<'db> {
         self.place_states.shrink_to_fit();
         self.reachable_definitions.shrink_to_fit();
         self.bindings_by_use.shrink_to_fit();
-        self.node_reachability.shrink_to_fit();
         self.declarations_by_binding.shrink_to_fit();
         self.bindings_by_definition.shrink_to_fit();
         self.eager_snapshots.shrink_to_fit();
@@ -1134,7 +1108,6 @@ impl<'db> UseDefMapBuilder<'db> {
             narrowing_constraints: self.narrowing_constraints.build(),
             reachability_constraints: self.reachability_constraints.build(),
             bindings_by_use: self.bindings_by_use,
-            node_reachability: self.node_reachability,
             end_of_scope_places: self.place_states,
             reachable_definitions: self.reachable_definitions,
             declarations_by_binding: self.declarations_by_binding,