improve performance of `find_type_var_from`

crates/ty_python_semantic/resources/mdtest/literal_promotion.md

@@ -376,10 +376,13 @@ reveal_type(x6)  # revealed: Sub1[Literal[1]]
 x7: Sup1[Literal[1]] | None = sub1(1)
 reveal_type(x7)  # revealed: Sub1[Literal[1]]
 
-class Sup2[T, U]:
+class Sup2A[T, U]:
     value: tuple[T, U]
 
-class Sub2[T, U](Sup2[T, Any], Sup2[Any, U]): ...
+class Sup2B[T, U]:
+    value: tuple[T, U]
+
+class Sub2[T, U](Sup2A[T, Any], Sup2B[Any, U]): ...
 
 def sub2[T, U](x: T, y: U) -> Sub2[T, U]:
     return Sub2()
@@ -387,6 +390,9 @@ def sub2[T, U](x: T, y: U) -> Sub2[T, U]:
 x8 = sub2(1, 2)
 reveal_type(x8)  # revealed: Sub2[int, int]
 
-x9: Sup2[Literal[1], Literal[2]] = sub2(1, 2)
+x9: Sup2A[Literal[1], Literal[2]] = sub2(1, 2)
-reveal_type(x9)  # revealed: Sub2[Literal[1], Literal[2]]
+reveal_type(x9)  # revealed: Sub2[Literal[1], int]
+
+x10: Sup2B[Literal[1], Literal[2]] = sub2(1, 2)
+reveal_type(x10)  # revealed: Sub2[int, Literal[2]]
 ```

crates/ty_python_semantic/src/types.rs

@@ -2,7 +2,6 @@ use compact_str::{CompactString, ToCompactString};
 use infer::nearest_enclosing_class;
 use itertools::{Either, Itertools};
 use ruff_diagnostics::{Edit, Fix};
-use rustc_hash::FxHashSet;
 
 use std::borrow::Cow;
 use std::time::Duration;
@@ -1041,55 +1040,40 @@ impl<'db> Type<'db> {
 
     /// Given a type variable `T` from the generic context of a class `C`:
     /// - If `self` is a specialized instance of `C`, returns the type assigned to `T` on `self`.
-    /// - If `self` is a specialized instance of some class `A`, and `C` is a subclass of `A`
+    /// - If `self` is a specialized instance of some class `A[T]`, and `C[T]` is a subclass of
-    ///   such that the type variable `U` on `A` is specialized to `T`, returns the type
+    ///   `A[T]`, returns the type assigned to `T` on `self`.
-    ///   assigned to `U` on `self`.
     pub(crate) fn find_type_var_from(
         self,
         db: &'db dyn Db,
         bound_typevar: BoundTypeVarInstance<'db>,
         class: ClassLiteral<'db>,
-    ) -> Option<Type<'db>> {
-        self.find_type_var_from_impl(db, bound_typevar, class, &mut FxHashSet::default())
-    }
-
-    pub(crate) fn find_type_var_from_impl(
-        self,
-        db: &'db dyn Db,
-        bound_typevar: BoundTypeVarInstance<'db>,
-        class: ClassLiteral<'db>,
-        visited: &mut FxHashSet<(ClassLiteral<'db>, BoundTypeVarIdentity<'db>)>,
     ) -> Option<Type<'db>> {
         if let Some(specialization) = self.specialization_of(db, class) {
             return specialization.get(db, bound_typevar);
         }
 
         // TODO: We should use the constraint solver here to determine the type mappings for more
-        // complex subtyping relationships, e.g., `type[C[T]]` to `Callable[..., T]`, or unions
+        // complex subtyping relationships, e.g., callables, protocols, or unions containing multiple
-        // containing multiple generic elements.
+        // generic elements.
-        for base in class.iter_mro(db, None) {
+        for base in class.iter_mro(db, None).skip(1) {
-            let Some((origin, Some(specialization))) =
+            let Some((base, Some(base_specialization))) =
                 base.into_class().map(|class| class.class_literal(db))
             else {
                 continue;
             };
 
-            for (base_typevar, base_ty) in specialization
+            if let Some(specialization) = self.specialization_of(db, base) {
+                for (base_typevar, base_ty) in base_specialization
                     .generic_context(db)
                     .variables(db)
-                .zip(specialization.types(db))
+                    .zip(base_specialization.types(db))
                 {
                     if *base_ty == Type::TypeVar(bound_typevar) {
-                    if !visited.insert((origin, base_typevar.identity(db))) {
+                        return specialization.get(db, base_typevar);
-                        return None;
+                    }
                 }
 
-                    if let Some(ty) =
+                return None;
-                        self.find_type_var_from_impl(db, base_typevar, origin, visited)
-                    {
-                        return Some(ty);
-                    }
-                }
             }
         }
 

crates/ty_python_semantic/src/types/infer/builder.rs

@@ -8029,7 +8029,6 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
         call_expression: &ast::ExprCall,
         tcx: TypeContext<'db>,
     ) -> Type<'db> {
-        // TODO: Use the type context for more precise inference.
         let callable_type =
             self.infer_maybe_standalone_expression(&call_expression.func, TypeContext::default());