[ty] Simplify materialization of specialized generics (astral-sh#20121)

This is a variant of astral-sh#20076 that moves some complexity out of
`apply_type_mapping_impl` in `generics.rs`. The tradeoff is that now
every place that applies `TypeMapping::Specialization` must take care to
call `.materialize()` afterwards. (A previous version of this didn't
work because I had missed a spot where I had to call `.materialize()`.)

@carljm as asked in
astral-sh#20076 (comment) .

Author: JelleZijlstra

crates/ty_python_semantic/src/types.rs

@@ -5970,7 +5970,12 @@ impl<'db> Type<'db> {
         db: &'db dyn Db,
         specialization: Specialization<'db>,
     ) -> Type<'db> {
-        self.apply_type_mapping(db, &TypeMapping::Specialization(specialization))
+        let new_specialization =
+            self.apply_type_mapping(db, &TypeMapping::Specialization(specialization));
+        match specialization.materialization_kind(db) {
+            None => new_specialization,
+            Some(materialization_kind) => new_specialization.materialize(db, materialization_kind),
+        }
     }
 
     fn apply_type_mapping<'a>(
@@ -6713,13 +6718,6 @@ impl<'db> TypeMapping<'_, 'db> {
             }
         }
     }
-
-    fn materialization_kind(&self, db: &'db dyn Db) -> Option<MaterializationKind> {
-        match self {
-            TypeMapping::Specialization(specialization) => specialization.materialization_kind(db),
-            _ => None,
-        }
-    }
 }
 
 /// Singleton types that are heavily special-cased by ty. Despite its name,

crates/ty_python_semantic/src/types/class_base.rs

@@ -4,7 +4,8 @@ use crate::types::generics::Specialization;
 use crate::types::tuple::TupleType;
 use crate::types::{
     ApplyTypeMappingVisitor, ClassLiteral, ClassType, DynamicType, KnownClass, KnownInstanceType,
-    MroError, MroIterator, NormalizedVisitor, SpecialFormType, Type, TypeMapping, todo_type,
+    MaterializationKind, MroError, MroIterator, NormalizedVisitor, SpecialFormType, Type,
+    TypeMapping, todo_type,
 };
 
 /// Enumeration of the possible kinds of types we allow in class bases.
@@ -286,16 +287,30 @@ impl<'db> ClassBase<'db> {
         specialization: Option<Specialization<'db>>,
     ) -> Self {
         if let Some(specialization) = specialization {
-            self.apply_type_mapping_impl(
+            let new_self = self.apply_type_mapping_impl(
                 db,
                 &TypeMapping::Specialization(specialization),
                 &ApplyTypeMappingVisitor::default(),
-            )
+            );
+            match specialization.materialization_kind(db) {
+                None => new_self,
+                Some(materialization_kind) => new_self.materialize(db, materialization_kind),
+            }
         } else {
             self
         }
     }
 
+    fn materialize(self, db: &'db dyn Db, kind: MaterializationKind) -> Self {
+        match self {
+            ClassBase::Class(class) => Self::Class(class.materialize(db, kind)),
+            ClassBase::Dynamic(_)
+            | ClassBase::Generic
+            | ClassBase::Protocol
+            | ClassBase::TypedDict => self,
+        }
+    }
+
     pub(super) fn has_cyclic_mro(self, db: &'db dyn Db) -> bool {
         match self {
             ClassBase::Class(class) => {

crates/ty_python_semantic/src/types/generics.rs

@@ -581,7 +581,11 @@ impl<'db> Specialization<'db> {
     /// That lets us produce the generic alias `A[int]`, which is the corresponding entry in the
     /// MRO of `B[int]`.
     pub(crate) fn apply_specialization(self, db: &'db dyn Db, other: Specialization<'db>) -> Self {
-        self.apply_type_mapping(db, &TypeMapping::Specialization(other))
+        let new_specialization = self.apply_type_mapping(db, &TypeMapping::Specialization(other));
+        match other.materialization_kind(db) {
+            None => new_specialization,
+            Some(materialization_kind) => new_specialization.materialize(db, materialization_kind),
+        }
     }
 
     pub(crate) fn apply_type_mapping<'a>(
@@ -598,59 +602,20 @@ impl<'db> Specialization<'db> {
         type_mapping: &TypeMapping<'a, 'db>,
         visitor: &ApplyTypeMappingVisitor<'db>,
     ) -> Self {
-        // TODO it seems like this should be possible to do in a much simpler way in
-        // `Self::apply_specialization`; just apply the type mapping to create the new
-        // specialization, then materialize the new specialization appropriately, if the type
-        // mapping is a materialization. But this doesn't work; see discussion in
-        // https://github.com/astral-sh/ruff/pull/20076
-        let applied_materialization_kind = type_mapping.materialization_kind(db);
-        let mut has_dynamic_invariant_typevar = false;
         let types: Box<[_]> = self
-            .generic_context(db)
-            .variables(db)
-            .into_iter()
-            .zip(self.types(db))
-            .map(|(bound_typevar, vartype)| {
-                let ty = vartype.apply_type_mapping_impl(db, type_mapping, visitor);
-                match (applied_materialization_kind, bound_typevar.variance(db)) {
-                    (None, _) => ty,
-                    (Some(_), TypeVarVariance::Bivariant) =>
-                    // With bivariance, all specializations are subtypes of each other,
-                    // so any materialization is acceptable.
-                    {
-                        ty.materialize(db, MaterializationKind::Top)
-                    }
-                    (Some(materialization_kind), TypeVarVariance::Covariant) => {
-                        ty.materialize(db, materialization_kind)
-                    }
-                    (Some(materialization_kind), TypeVarVariance::Contravariant) => {
-                        ty.materialize(db, materialization_kind.flip())
-                    }
-                    (Some(_), TypeVarVariance::Invariant) => {
-                        let top_materialization = ty.materialize(db, MaterializationKind::Top);
-                        if !ty.is_equivalent_to(db, top_materialization) {
-                            has_dynamic_invariant_typevar = true;
-                        }
-                        ty
-                    }
-                }
-            })
+            .types(db)
+            .iter()
+            .map(|ty| ty.apply_type_mapping_impl(db, type_mapping, visitor))
             .collect();
 
         let tuple_inner = self
             .tuple_inner(db)
             .and_then(|tuple| tuple.apply_type_mapping_impl(db, type_mapping, visitor));
-        let new_materialization_kind = if has_dynamic_invariant_typevar {
-            self.materialization_kind(db)
-                .or(applied_materialization_kind)
-        } else {
-            None
-        };
         Specialization::new(
             db,
             self.generic_context(db),
             types,
-            new_materialization_kind,
+            self.materialization_kind(db),
             tuple_inner,
         )
     }