[ty] Conservative narrowing places optimization (#22734)

## Summary

Add helpers to compute a conservative upper bound on the set of places a
predicate might actually narrow, and use this to save work when
computing narrowing constraints.

## Test Plan

Existing tests pass

---------

Co-authored-by: David Peter <mail@david-peter.de>

Author: ericmarkmartin

crates/ty_python_semantic/src/semantic_index/builder.rs

@@ -51,6 +51,7 @@ use crate::semantic_index::use_def::{
 };
 use crate::semantic_index::{ExpressionsScopeMap, SemanticIndex, VisibleAncestorsIter};
 use crate::semantic_model::HasTrackedScope;
+use crate::types::PossiblyNarrowedPlaces;
 use crate::unpack::{EvaluationMode, Unpack, UnpackKind, UnpackPosition, UnpackValue};
 use crate::{Db, Program};
 
@@ -857,7 +858,7 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
 
     /// Adds a new predicate to the list of all predicates, but does not record it. Returns the
     /// predicate ID for later recording using
-    /// [`SemanticIndexBuilder::record_narrowing_constraint_id`].
+    /// [`SemanticIndexBuilder::record_narrowing_constraint_id_for_places`].
     fn add_predicate(&mut self, predicate: PredicateOrLiteral<'db>) -> ScopedPredicateId {
         self.current_use_def_map_mut().add_predicate(predicate)
     }
@@ -868,27 +869,70 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
             .add_predicate(predicate.negated())
     }
 
-    /// Records a previously added narrowing constraint by adding it to all live bindings.
-    fn record_narrowing_constraint_id(&mut self, predicate: ScopedPredicateId) {
+    /// Records a previously added narrowing constraint by adding it to the live bindings
+    /// of the specified places.
+    fn record_narrowing_constraint_id_for_places(
+        &mut self,
+        predicate: ScopedPredicateId,
+        places: &PossiblyNarrowedPlaces,
+    ) {
         self.current_use_def_map_mut()
-            .record_narrowing_constraint(predicate);
+            .record_narrowing_constraint_for_places(predicate, places);
     }
 
-    /// Adds and records a narrowing constraint, i.e. adds it to all live bindings.
+    /// Adds and records a narrowing constraint for only the places that could possibly be narrowed.
     fn record_narrowing_constraint(&mut self, predicate: PredicateOrLiteral<'db>) {
+        let possibly_narrowed = self.compute_possibly_narrowed_places(&predicate);
         let use_def = self.current_use_def_map_mut();
         let predicate_id = use_def.add_predicate(predicate);
-        use_def.record_narrowing_constraint(predicate_id);
+        use_def.record_narrowing_constraint_for_places(predicate_id, &possibly_narrowed);
     }
 
-    /// Negates the given predicate and then adds it as a narrowing constraint to all live
-    /// bindings.
+    /// Computes the conservative set of places that could possibly be narrowed by a predicate.
+    ///
+    /// This uses the closure-based approach to avoid calling Salsa queries that depend on
+    /// the semantic index (which is still being built).
+    fn compute_possibly_narrowed_places(
+        &self,
+        predicate: &PredicateOrLiteral<'db>,
+    ) -> PossiblyNarrowedPlaces {
+        use crate::types::PossiblyNarrowedPlacesBuilder;
+
+        match predicate {
+            PredicateOrLiteral::Literal(_) => PossiblyNarrowedPlaces::default(),
+            PredicateOrLiteral::Predicate(pred) => {
+                let place_table = self.current_place_table();
+
+                match pred.node {
+                    PredicateNode::Expression(expression) => {
+                        let module = self.module;
+                        let expression_node = expression.node_ref(self.db, module);
+                        PossiblyNarrowedPlacesBuilder::new(self.db, place_table)
+                            .expression(expression_node)
+                    }
+                    PredicateNode::Pattern(pattern) => {
+                        let module = self.module;
+                        PossiblyNarrowedPlacesBuilder::new(self.db, place_table)
+                            .pattern(pattern, module)
+                    }
+                    PredicateNode::ReturnsNever(_) | PredicateNode::StarImportPlaceholder(_) => {
+                        // These predicates don't narrow any places
+                        PossiblyNarrowedPlaces::default()
+                    }
+                }
+            }
+        }
+    }
+
+    /// Negates the given predicate and then adds it as a narrowing constraint to the places
+    /// that could possibly be narrowed.
     fn record_negated_narrowing_constraint(
         &mut self,
         predicate: PredicateOrLiteral<'db>,
     ) -> ScopedPredicateId {
+        let possibly_narrowed = self.compute_possibly_narrowed_places(&predicate);
         let id = self.add_negated_predicate(predicate);
-        self.record_narrowing_constraint_id(id);
+        self.record_narrowing_constraint_id_for_places(id, &possibly_narrowed);
         id
     }
 
@@ -2767,6 +2811,7 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> {
                     // anymore.
                     if index < values.len() - 1 {
                         let predicate = self.build_predicate(value);
+                        let possibly_narrowed = self.compute_possibly_narrowed_places(&predicate);
                         let predicate_id = match op {
                             ast::BoolOp::And => self.add_predicate(predicate),
                             ast::BoolOp::Or => self.add_negated_predicate(predicate),
@@ -2789,7 +2834,10 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> {
                         // the application of the reachability constraint until after the expression
                         // has been evaluated, so we only push it onto the stack here.
                         self.flow_restore(after_expr);
-                        self.record_narrowing_constraint_id(predicate_id);
+                        self.record_narrowing_constraint_id_for_places(
+                            predicate_id,
+                            &possibly_narrowed,
+                        );
                         reachability_constraints.push(reachability_constraint);
                     }
                 }

crates/ty_python_semantic/src/semantic_index/place.rs

@@ -261,7 +261,7 @@ pub(crate) struct PlaceTableBuilder {
 
 impl PlaceTableBuilder {
     /// Looks up a place ID by its expression.
-    pub(super) fn place_id(&self, expression: PlaceExprRef) -> Option<ScopedPlaceId> {
+    pub(crate) fn place_id(&self, expression: PlaceExprRef) -> Option<ScopedPlaceId> {
         match expression {
             PlaceExprRef::Symbol(symbol) => self.symbols.symbol_id(symbol.name()).map(Into::into),
             PlaceExprRef::Member(member) => {

crates/ty_python_semantic/src/semantic_index/use_def.rs

@@ -266,7 +266,9 @@ use crate::semantic_index::use_def::place_state::{
     LiveDeclarationsIterator, PlaceState, PreviousDefinitions, ScopedDefinitionId,
 };
 use crate::semantic_index::{EnclosingSnapshotResult, SemanticIndex};
-use crate::types::{NarrowingConstraint, Truthiness, Type, infer_narrowing_constraint};
+use crate::types::{
+    NarrowingConstraint, PossiblyNarrowedPlaces, Truthiness, Type, infer_narrowing_constraint,
+};
 
 mod place_state;
 
@@ -1002,7 +1004,12 @@ impl<'db> UseDefMapBuilder<'db> {
         }
     }
 
-    pub(super) fn record_narrowing_constraint(&mut self, predicate: ScopedPredicateId) {
+    /// Records a narrowing constraint for only the specified places.
+    pub(super) fn record_narrowing_constraint_for_places(
+        &mut self,
+        predicate: ScopedPredicateId,
+        places: &PossiblyNarrowedPlaces,
+    ) {
         if predicate == ScopedPredicateId::ALWAYS_TRUE
             || predicate == ScopedPredicateId::ALWAYS_FALSE
         {
@@ -1011,14 +1018,25 @@ impl<'db> UseDefMapBuilder<'db> {
         }
 
         let narrowing_constraint = predicate.into();
-        for state in &mut self.symbol_states {
-            state
-                .record_narrowing_constraint(&mut self.narrowing_constraints, narrowing_constraint);
-        }
-
-        for state in &mut self.member_states {
-            state
-                .record_narrowing_constraint(&mut self.narrowing_constraints, narrowing_constraint);
+        for place in places {
+            match place {
+                ScopedPlaceId::Symbol(symbol_id) => {
+                    if let Some(state) = self.symbol_states.get_mut(*symbol_id) {
+                        state.record_narrowing_constraint(
+                            &mut self.narrowing_constraints,
+                            narrowing_constraint,
+                        );
+                    }
+                }
+                ScopedPlaceId::Member(member_id) => {
+                    if let Some(state) = self.member_states.get_mut(*member_id) {
+                        state.record_narrowing_constraint(
+                            &mut self.narrowing_constraints,
+                            narrowing_constraint,
+                        );
+                    }
+                }
+            }
         }
     }
 

crates/ty_python_semantic/src/types.rs

@@ -68,7 +68,10 @@ use crate::types::generics::{
     typing_self, walk_generic_context,
 };
 use crate::types::mro::{Mro, MroIterator, StaticMroError};
-pub(crate) use crate::types::narrow::{NarrowingConstraint, infer_narrowing_constraint};
+pub(crate) use crate::types::narrow::{
+    NarrowingConstraint, PossiblyNarrowedPlaces, PossiblyNarrowedPlacesBuilder,
+    infer_narrowing_constraint,
+};
 use crate::types::newtype::NewType;
 pub(crate) use crate::types::signatures::{Parameter, Parameters};
 use crate::types::signatures::{ParameterForm, walk_signature};