Make infer_local_place_load a query

crates/ty_python_semantic/src/semantic_index/ast_ids.rs

@@ -42,7 +42,7 @@ fn ast_ids<'db>(db: &'db dyn Db, scope: ScopeId) -> &'db AstIds {
 
 /// Uniquely identifies a use of a name in a [`crate::semantic_index::FileScopeId`].
 #[newtype_index]
-#[derive(get_size2::GetSize)]
+#[derive(get_size2::GetSize, salsa::Update)]
 pub struct ScopedUseId;
 
 pub trait HasScopedUseId {

crates/ty_python_semantic/src/types.rs

@@ -3337,7 +3337,7 @@ impl<'db> Type<'db> {
         name: Name,
         policy: MemberLookupPolicy,
     ) -> PlaceAndQualifiers<'db> {
-        tracing::trace!("member_lookup_with_policy: {}.{}", self.display(db), name);
+        let _span = tracing::trace_span!("member_lookup_with_policy: {}", ?name).entered();
         if name == "__class__" {
             return Place::bound(self.dunder_class(db)).into();
         }

crates/ty_python_semantic/src/types/class.rs

@@ -2816,6 +2816,8 @@ impl<'db> ClassLiteral<'db> {
         name: String,
         target_method_decorator: MethodDecorator,
     ) -> PlaceAndQualifiers<'db> {
+        let _span =
+            tracing::trace_span!("implicit_attribute_inner", ?class_body_scope, ?name).entered();
         // If we do not see any declarations of an attribute, neither in the class body nor in
         // any method, we build a union of `Unknown` with the inferred types of all bindings of
         // that attribute. We include `Unknown` in that union to account for the fact that the

crates/ty_python_semantic/src/types/infer.rs

@@ -83,7 +83,7 @@ use crate::semantic_index::definition::{
 };
 use crate::semantic_index::expression::{Expression, ExpressionKind};
 use crate::semantic_index::narrowing_constraints::ConstraintKey;
-use crate::semantic_index::place::{PlaceExpr, PlaceExprRef};
+use crate::semantic_index::place::{PlaceExpr, PlaceExprRef, ScopedPlaceId};
 use crate::semantic_index::scope::{
     FileScopeId, NodeWithScopeKind, NodeWithScopeRef, ScopeId, ScopeKind,
 };
@@ -137,7 +137,7 @@ use crate::types::{
 use crate::unpack::{EvaluationMode, Unpack, UnpackPosition};
 use crate::util::diagnostics::format_enumeration;
 use crate::util::subscript::{PyIndex, PySlice};
-use crate::{Db, FxOrderSet, Program};
+use crate::{Db, FxOrderSet, Program, db};
 
 /// Infer all types for a [`ScopeId`], including all definitions and expressions in that scope.
 /// Use when checking a scope, or needing to provide a type for an arbitrary expression in the
@@ -171,7 +171,6 @@ fn scope_cycle_initial<'db>(_db: &'db dyn Db, scope: ScopeId<'db>) -> ScopeInfer
 
 /// Infer all types for a [`Definition`] (including sub-expressions).
 /// Use when resolving a place use or public type of a place.
-#[salsa::tracked(returns(ref), cycle_fn=definition_cycle_recover, cycle_initial=definition_cycle_initial, heap_size=ruff_memory_usage::heap_size)]
 pub(crate) fn infer_definition_types<'db>(
     db: &'db dyn Db,
     definition: Definition<'db>,
@@ -252,7 +251,6 @@ fn deferred_cycle_initial<'db>(
 /// Use rarely; only for cases where we'd otherwise risk double-inferring an expression: RHS of an
 /// assignment, which might be unpacking/multi-target and thus part of multiple definitions, or a
 /// type narrowing guard expression (e.g. if statement test node).
-#[salsa::tracked(returns(ref), cycle_fn=expression_cycle_recover, cycle_initial=expression_cycle_initial, heap_size=ruff_memory_usage::heap_size)]
 pub(crate) fn infer_expression_types<'db>(
     db: &'db dyn Db,
     expression: Expression<'db>,
@@ -2676,7 +2674,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
     fn infer_definition(&mut self, node: impl Into<DefinitionNodeKey> + std::fmt::Debug + Copy) {
         let definition = self.index.expect_single_definition(node);
         let result = infer_definition_types(self.db(), definition);
-        self.extend_definition(result);
+        self.extend_definition(&result);
     }
 
     fn infer_function_definition_statement(&mut self, function: &ast::StmtFunctionDef) {
@@ -5142,7 +5140,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
                         self.check_deprecated(alias, ty.inner);
                     }
                 }
-                self.extend_definition(inferred);
+                self.extend_definition(&inferred);
             }
         }
     }
@@ -5636,7 +5634,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
         standalone_expression: Expression<'db>,
     ) -> Type<'db> {
         let types = infer_expression_types(self.db(), standalone_expression);
-        self.extend_expression(types);
+        self.extend_expression(&types);
 
         // Instead of calling `self.expression_type(expr)` after extending here, we get
         // the result from `types` directly because we might be in cycle recovery where
@@ -6102,7 +6100,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
             if comprehension.is_first() && target.is_name_expr() {
                 result.expression_type(iterable)
             } else {
-                self.extend_expression_unchecked(result);
+                self.extend_expression_unchecked(&result);
                 result.expression_type(iterable)
             }
         };
@@ -6141,7 +6139,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
         if named.target.is_name_expr() {
             let definition = self.index.expect_single_definition(named);
             let result = infer_definition_types(self.db(), definition);
-            self.extend_definition(result);
+            self.extend_definition(&result);
             result.binding_type(definition)
         } else {
             // For syntactically invalid targets, we still need to run type inference:
@@ -6797,37 +6795,27 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
         expr: PlaceExprRef,
         expr_ref: ast::ExprRef,
     ) -> (Place<'db>, Option<ScopedUseId>) {
+        if expr_ref
+            .as_name_expr()
+            .is_some_and(|name| name.is_invalid())
+        {
+            return (Place::Unbound, None);
+        }
+
         let db = self.db();
         let scope = self.scope();
-        let file_scope_id = scope.file_scope_id(db);
-        let place_table = self.index.place_table(file_scope_id);
-        let use_def = self.index.use_def_map(file_scope_id);
-
-        // If we're inferring types of deferred expressions, look them up from end-of-scope.
-        if self.is_deferred() {
-            let place = if let Some(place_id) = place_table.place_id(expr) {
-                place_from_bindings(db, use_def.all_reachable_bindings(place_id))
-            } else {
-                assert!(
-                    self.deferred_state.in_string_annotation(),
-                    "Expected the place table to create a place for every valid PlaceExpr node"
-                );
-                Place::Unbound
-            };
-            (place, None)
+        let is_deferred = self.is_deferred();
+        let deferred_state = self.deferred_state;
+        let use_id = if is_deferred {
+            None
         } else {
-            if expr_ref
-                .as_name_expr()
-                .is_some_and(|name| name.is_invalid())
-            {
-                return (Place::Unbound, None);
-            }
+            Some(expr_ref.scoped_use_id(db, scope))
+        };
 
-            let use_id = expr_ref.scoped_use_id(db, scope);
-            let place = place_from_bindings(db, use_def.bindings_at_use(use_id));
+        let place_table = self.index.place_table(scope.file_scope_id(db));
+        let place_id = place_table.place_id(expr);
 
-            (place, Some(use_id))
-        }
+        infer_local_place_load(db, scope, deferred_state, place_id, use_id)
     }
 
     /// Infer the type of a place expression from definitions, assuming a load context.
@@ -7341,6 +7329,8 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
     }
 
     fn infer_attribute_expression(&mut self, attribute: &ast::ExprAttribute) -> Type<'db> {
+        let _span = tracing::debug_span!("infer_attribute_expression", ?attribute).entered();
+
         let ast::ExprAttribute {
             value,
             attr: _,
@@ -11296,6 +11286,60 @@ impl<'db> TypeInferenceBuilder<'db, '_> {
     }
 }
 
+fn infer_local_place_load_cycle_recover<'db>(
+    _db: &'db dyn Db,
+    _value: &(Place<'db>, Option<ScopedUseId>),
+    _count: u32,
+    _scope: ScopeId<'db>,
+    _deferred_state: DeferredExpressionState,
+    _place_id: Option<ScopedPlaceId>,
+    _use_id: Option<ScopedUseId>,
+) -> salsa::CycleRecoveryAction<(Place<'db>, Option<ScopedUseId>)> {
+    salsa::CycleRecoveryAction::Iterate
+}
+
+fn infer_local_place_load_cycle_initial<'db>(
+    _db: &'db dyn Db,
+    _scope: ScopeId<'db>,
+    _deferred_state: DeferredExpressionState,
+    _place_id: Option<ScopedPlaceId>,
+    _use_id: Option<ScopedUseId>,
+) -> (Place<'db>, Option<ScopedUseId>) {
+    (Place::Unbound, None)
+}
+
+#[salsa::tracked(cycle_fn=infer_local_place_load_cycle_recover, cycle_initial=infer_local_place_load_cycle_initial, heap_size=ruff_memory_usage::heap_size)]
+fn infer_local_place_load<'db>(
+    db: &'db dyn Db,
+    scope: ScopeId<'db>,
+    deferred_state: DeferredExpressionState,
+    place_id: Option<ScopedPlaceId>,
+    use_id: Option<ScopedUseId>,
+) -> (Place<'db>, Option<ScopedUseId>) {
+    let file_scope_id = scope.file_scope_id(db);
+    let index = semantic_index(db, scope.file(db));
+    let use_def = index.use_def_map(file_scope_id);
+
+    if let Some(use_id) = use_id {
+        // Non-deferred load
+        let place = place_from_bindings(db, use_def.bindings_at_use(use_id));
+
+        (place, Some(use_id))
+    } else {
+        // If we're inferring types of deferred expressions, look them up from end-of-scope.
+        let place = if let Some(place_id) = place_id {
+            place_from_bindings(db, use_def.all_reachable_bindings(place_id))
+        } else {
+            assert!(
+                deferred_state.in_string_annotation(),
+                "Expected the place table to create a place for every valid PlaceExpr node"
+            );
+            Place::Unbound
+        };
+        (place, None)
+    }
+}
+
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 enum GenericContextError {
     /// It's invalid to subscript `Generic` or `Protocol` with this type
@@ -11320,7 +11364,7 @@ impl GenericContextError {
 }
 
 /// The deferred state of a specific expression in an inference region.
-#[derive(Default, Debug, Clone, Copy)]
+#[derive(Default, Debug, Clone, Copy, Hash, PartialEq, Eq)]
 enum DeferredExpressionState {
     /// The expression is not deferred.
     #[default]