[red-knot] Separate 'infer_expression_type' query

crates/red_knot_python_semantic/src/types.rs

@@ -15,7 +15,8 @@ pub(crate) use self::diagnostic::register_lints;
 pub use self::diagnostic::{TypeCheckDiagnostic, TypeCheckDiagnostics};
 pub(crate) use self::display::TypeArrayDisplay;
 pub(crate) use self::infer::{
-    infer_deferred_types, infer_definition_types, infer_expression_types, infer_scope_types,
+    infer_deferred_types, infer_definition_types, infer_expression_type, infer_expression_types,
+    infer_scope_types,
 };
 pub use self::narrow::KnownConstraintFunction;
 pub(crate) use self::signatures::Signature;
@@ -25,7 +26,6 @@ use crate::module_resolver::{file_to_module, resolve_module, KnownModule};
 use crate::semantic_index::ast_ids::HasScopedExpressionId;
 use crate::semantic_index::attribute_assignment::AttributeAssignment;
 use crate::semantic_index::definition::Definition;
-use crate::semantic_index::expression::Expression;
 use crate::semantic_index::symbol::{self as symbol, ScopeId, ScopedSymbolId};
 use crate::semantic_index::{
     attribute_assignments, global_scope, imported_modules, semantic_index, symbol_table,
@@ -4147,16 +4147,6 @@ impl<'db> Class<'db> {
         name: &str,
         inferred_type_from_class_body: Option<Type<'db>>,
     ) -> Symbol<'db> {
-        // We use a separate salsa query here to prevent unrelated changes in the AST of an external
-        // file from triggering re-evaluations of downstream queries.
-        // See the `dependency_implicit_instance_attribute` test for more information.
-        #[salsa::tracked]
-        fn infer_expression_type<'db>(db: &'db dyn Db, expression: Expression<'db>) -> Type<'db> {
-            let inference = infer_expression_types(db, expression);
-            let expr_scope = expression.scope(db);
-            inference.expression_type(expression.node_ref(db).scoped_expression_id(db, expr_scope))
-        }
-
         // 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/red_knot_python_semantic/src/types/infer.rs

@@ -193,6 +193,20 @@ pub(crate) fn infer_expression_types<'db>(
     TypeInferenceBuilder::new(db, InferenceRegion::Expression(expression), index).finish()
 }
 
+// Similar to `infer_expression_types` (with the same restrictions). Directly returns the
+// type of the overall expression. This is a salsa query because it accesses `node_ref`,
+// which is sensitive to changes in the AST. Making it a query allows downstream queries
+// to short-circuit if the result type has not changed.
+#[salsa::tracked]
+pub(crate) fn infer_expression_type<'db>(
+    db: &'db dyn Db,
+    expression: Expression<'db>,
+) -> Type<'db> {
+    let inference = infer_expression_types(db, expression);
+    let expr_scope = expression.scope(db);
+    inference.expression_type(expression.node_ref(db).scoped_expression_id(db, expr_scope))
+}
+
 /// Infer the types for an [`Unpack`] operation.
 ///
 /// This infers the expression type and performs structural match against the target expression

crates/red_knot_python_semantic/src/types/unpacker.rs

@@ -7,7 +7,7 @@ use ruff_python_ast::{self as ast, AnyNodeRef};
 
 use crate::semantic_index::ast_ids::{HasScopedExpressionId, ScopedExpressionId};
 use crate::semantic_index::symbol::ScopeId;
-use crate::types::{infer_expression_types, todo_type, Type, TypeCheckDiagnostics};
+use crate::types::{infer_expression_type, todo_type, Type, TypeCheckDiagnostics};
 use crate::unpack::UnpackValue;
 use crate::Db;
 
@@ -42,8 +42,7 @@ impl<'db> Unpacker<'db> {
             "Unpacking target must be a list or tuple expression"
         );
 
-        let mut value_ty = infer_expression_types(self.db(), value.expression())
-            .expression_type(value.scoped_expression_id(self.db(), self.scope));
+        let mut value_ty = infer_expression_type(self.db(), value.expression());
 
         if value.is_assign()
             && self.context.in_stub()

crates/red_knot_python_semantic/src/unpack.rs

@@ -3,7 +3,6 @@ use ruff_python_ast::{self as ast, AnyNodeRef};
 use ruff_text_size::{Ranged, TextRange};
 
 use crate::ast_node_ref::AstNodeRef;
-use crate::semantic_index::ast_ids::{HasScopedExpressionId, ScopedExpressionId};
 use crate::semantic_index::expression::Expression;
 use crate::semantic_index::symbol::{FileScopeId, ScopeId};
 use crate::Db;
@@ -88,17 +87,6 @@ impl<'db> UnpackValue<'db> {
         }
     }
 
-    /// Returns the [`ScopedExpressionId`] of the underlying expression.
-    pub(crate) fn scoped_expression_id(
-        self,
-        db: &'db dyn Db,
-        scope: ScopeId<'db>,
-    ) -> ScopedExpressionId {
-        self.expression()
-            .node_ref(db)
-            .scoped_expression_id(db, scope)
-    }
-
     /// Returns the expression as an [`AnyNodeRef`].
     pub(crate) fn as_any_node_ref(self, db: &'db dyn Db) -> AnyNodeRef<'db> {
         self.expression().node_ref(db).node().into()

crates/red_knot_python_semantic/src/visibility_constraints.rs

@@ -152,11 +152,8 @@
 
 use ruff_index::{newtype_index, IndexVec};
 
-use crate::semantic_index::{
-    ast_ids::HasScopedExpressionId,
-    constraint::{Constraint, ConstraintNode, PatternConstraintKind},
-};
-use crate::types::{infer_expression_types, Truthiness};
+use crate::semantic_index::constraint::{Constraint, ConstraintNode, PatternConstraintKind};
+use crate::types::{infer_expression_type, Truthiness};
 use crate::Db;
 
 /// The maximum depth of recursion when evaluating visibility constraints.
@@ -394,28 +391,15 @@ impl<'db> VisibilityConstraints<'db> {
     fn analyze_single(db: &dyn Db, constraint: &Constraint) -> Truthiness {
         match constraint.node {
             ConstraintNode::Expression(test_expr) => {
-                let inference = infer_expression_types(db, test_expr);
-                let scope = test_expr.scope(db);
-                let ty = inference
-                    .expression_type(test_expr.node_ref(db).scoped_expression_id(db, scope));
+                let ty = infer_expression_type(db, test_expr);
 
                 ty.bool(db).negate_if(!constraint.is_positive)
             }
             ConstraintNode::Pattern(inner) => match inner.kind(db) {
                 PatternConstraintKind::Value(value, guard) => {
                     let subject_expression = inner.subject(db);
-                    let inference = infer_expression_types(db, *subject_expression);
-                    let scope = subject_expression.scope(db);
-                    let subject_ty = inference.expression_type(
-                        subject_expression
-                            .node_ref(db)
-                            .scoped_expression_id(db, scope),
-                    );
-
-                    let inference = infer_expression_types(db, *value);
-                    let scope = value.scope(db);
-                    let value_ty = inference
-                        .expression_type(value.node_ref(db).scoped_expression_id(db, scope));
+                    let subject_ty = infer_expression_type(db, *subject_expression);
+                    let value_ty = infer_expression_type(db, *value);
 
                     if subject_ty.is_single_valued(db) {
                         let truthiness =