[ty] Apply type mappings to functions eagerly (#20596)

`TypeMapping` is no longer cow-shaped.

Before, `TypeMapping` defined a `to_owned` method, which would make an
owned copy of the type mapping. This let us apply type mappings to
function literals lazily. The primary part of a function that you have
to apply the type mapping to is its signature. The hypothesis was that
doing this lazily would prevent us from constructing the signature of a
function just to apply a type mapping; if you never ended up needed the
updated function signature, that would be extraneous work.

But looking at the CI for this PR, it looks like that hypothesis is
wrong! And this definitely cleans up the code quite a bit. It also means
that over time we can consider replacing all of these `TypeMapping` enum
variants with separate `TypeTransformer` impls.

---------

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

Author: dcreager

crates/ty_python_semantic/src/types.rs

@@ -53,7 +53,6 @@ use crate::types::function::{
 };
 use crate::types::generics::{
     GenericContext, PartialSpecialization, Specialization, bind_typevar, walk_generic_context,
-    walk_partial_specialization, walk_specialization,
 };
 pub use crate::types::ide_support::{
     CallSignatureDetails, Member, MemberWithDefinition, all_members, call_signature_details,
@@ -6109,7 +6108,7 @@ impl<'db> Type<'db> {
             }
 
             Type::FunctionLiteral(function) => {
-                let function = Type::FunctionLiteral(function.with_type_mapping(db, type_mapping));
+                let function = Type::FunctionLiteral(function.apply_type_mapping_impl(db, type_mapping, visitor));
 
                 match type_mapping {
                     TypeMapping::PromoteLiterals => function.literal_promotion_type(db)
@@ -6120,8 +6119,8 @@ impl<'db> Type<'db> {
 
             Type::BoundMethod(method) => Type::BoundMethod(BoundMethodType::new(
                 db,
-                method.function(db).with_type_mapping(db, type_mapping),
-                method.self_instance(db).apply_type_mapping(db, type_mapping),
+                method.function(db).apply_type_mapping_impl(db, type_mapping, visitor),
+                method.self_instance(db).apply_type_mapping_impl(db, type_mapping, visitor),
             )),
 
             Type::NominalInstance(instance) =>
@@ -6140,13 +6139,13 @@ impl<'db> Type<'db> {
 
             Type::KnownBoundMethod(KnownBoundMethodType::FunctionTypeDunderGet(function)) => {
                 Type::KnownBoundMethod(KnownBoundMethodType::FunctionTypeDunderGet(
-                    function.with_type_mapping(db, type_mapping),
+                    function.apply_type_mapping_impl(db, type_mapping, visitor),
                 ))
             }
 
             Type::KnownBoundMethod(KnownBoundMethodType::FunctionTypeDunderCall(function)) => {
                 Type::KnownBoundMethod(KnownBoundMethodType::FunctionTypeDunderCall(
-                    function.with_type_mapping(db, type_mapping),
+                    function.apply_type_mapping_impl(db, type_mapping, visitor),
                 ))
             }
 
@@ -6782,84 +6781,7 @@ pub enum TypeMapping<'a, 'db> {
     Materialize(MaterializationKind),
 }
 
-fn walk_type_mapping<'db, V: visitor::TypeVisitor<'db> + ?Sized>(
-    db: &'db dyn Db,
-    mapping: &TypeMapping<'_, 'db>,
-    visitor: &V,
-) {
-    match mapping {
-        TypeMapping::Specialization(specialization) => {
-            walk_specialization(db, *specialization, visitor);
-        }
-        TypeMapping::PartialSpecialization(specialization) => {
-            walk_partial_specialization(db, specialization, visitor);
-        }
-        TypeMapping::BindSelf(self_type) => {
-            visitor.visit_type(db, *self_type);
-        }
-        TypeMapping::ReplaceSelf { new_upper_bound } => {
-            visitor.visit_type(db, *new_upper_bound);
-        }
-        TypeMapping::PromoteLiterals
-        | TypeMapping::BindLegacyTypevars(_)
-        | TypeMapping::MarkTypeVarsInferable(_)
-        | TypeMapping::Materialize(_) => {}
-    }
-}
-
 impl<'db> TypeMapping<'_, 'db> {
-    fn to_owned(&self) -> TypeMapping<'db, 'db> {
-        match self {
-            TypeMapping::Specialization(specialization) => {
-                TypeMapping::Specialization(*specialization)
-            }
-            TypeMapping::PartialSpecialization(partial) => {
-                TypeMapping::PartialSpecialization(partial.to_owned())
-            }
-            TypeMapping::PromoteLiterals => TypeMapping::PromoteLiterals,
-            TypeMapping::BindLegacyTypevars(binding_context) => {
-                TypeMapping::BindLegacyTypevars(*binding_context)
-            }
-            TypeMapping::BindSelf(self_type) => TypeMapping::BindSelf(*self_type),
-            TypeMapping::ReplaceSelf { new_upper_bound } => TypeMapping::ReplaceSelf {
-                new_upper_bound: *new_upper_bound,
-            },
-            TypeMapping::MarkTypeVarsInferable(binding_context) => {
-                TypeMapping::MarkTypeVarsInferable(*binding_context)
-            }
-            TypeMapping::Materialize(materialization_kind) => {
-                TypeMapping::Materialize(*materialization_kind)
-            }
-        }
-    }
-
-    fn normalized_impl(&self, db: &'db dyn Db, visitor: &NormalizedVisitor<'db>) -> Self {
-        match self {
-            TypeMapping::Specialization(specialization) => {
-                TypeMapping::Specialization(specialization.normalized_impl(db, visitor))
-            }
-            TypeMapping::PartialSpecialization(partial) => {
-                TypeMapping::PartialSpecialization(partial.normalized_impl(db, visitor))
-            }
-            TypeMapping::PromoteLiterals => TypeMapping::PromoteLiterals,
-            TypeMapping::BindLegacyTypevars(binding_context) => {
-                TypeMapping::BindLegacyTypevars(*binding_context)
-            }
-            TypeMapping::BindSelf(self_type) => {
-                TypeMapping::BindSelf(self_type.normalized_impl(db, visitor))
-            }
-            TypeMapping::ReplaceSelf { new_upper_bound } => TypeMapping::ReplaceSelf {
-                new_upper_bound: new_upper_bound.normalized_impl(db, visitor),
-            },
-            TypeMapping::MarkTypeVarsInferable(binding_context) => {
-                TypeMapping::MarkTypeVarsInferable(*binding_context)
-            }
-            TypeMapping::Materialize(materialization_kind) => {
-                TypeMapping::Materialize(*materialization_kind)
-            }
-        }
-    }
-
     /// Update the generic context of a [`Signature`] according to the current type mapping
     pub(crate) fn update_signature_generic_context(
         &self,

crates/ty_python_semantic/src/types/function.rs

@@ -77,11 +77,11 @@ use crate::types::narrow::ClassInfoConstraintFunction;
 use crate::types::signatures::{CallableSignature, Signature};
 use crate::types::visitor::any_over_type;
 use crate::types::{
-    BoundMethodType, BoundTypeVarInstance, CallableType, ClassBase, ClassLiteral, ClassType,
-    DeprecatedInstance, DynamicType, FindLegacyTypeVarsVisitor, HasRelationToVisitor,
-    IsEquivalentVisitor, KnownClass, KnownInstanceType, NormalizedVisitor, SpecialFormType,
-    TrackedConstraintSet, Truthiness, Type, TypeMapping, TypeRelation, UnionBuilder, all_members,
-    binding_type, todo_type, walk_type_mapping,
+    ApplyTypeMappingVisitor, BoundMethodType, BoundTypeVarInstance, CallableType, ClassBase,
+    ClassLiteral, ClassType, DeprecatedInstance, DynamicType, FindLegacyTypeVarsVisitor,
+    HasRelationToVisitor, IsEquivalentVisitor, KnownClass, KnownInstanceType, NormalizedVisitor,
+    SpecialFormType, TrackedConstraintSet, Truthiness, Type, TypeMapping, TypeRelation,
+    UnionBuilder, all_members, binding_type, todo_type, walk_signature,
 };
 use crate::{Db, FxOrderSet, ModuleName, resolve_module};
 
@@ -623,33 +623,24 @@ impl<'db> FunctionLiteral<'db> {
     /// calling query is not in the same file as this function is defined in, then this will create
     /// a cross-module dependency directly on the full AST which will lead to cache
     /// over-invalidation.
-    fn signature<'a>(
-        self,
-        db: &'db dyn Db,
-        type_mappings: &'a [TypeMapping<'a, 'db>],
-    ) -> CallableSignature<'db>
-    where
-        'db: 'a,
-    {
+    fn signature(self, db: &'db dyn Db) -> CallableSignature<'db> {
         // We only include an implementation (i.e. a definition not decorated with `@overload`) if
         // it's the only definition.
         let inherited_generic_context = self.inherited_generic_context(db);
         let (overloads, implementation) = self.overloads_and_implementation(db);
         if let Some(implementation) = implementation {
             if overloads.is_empty() {
-                return CallableSignature::single(type_mappings.iter().fold(
+                return CallableSignature::single(
                     implementation.signature(db, inherited_generic_context),
-                    |sig, mapping| sig.apply_type_mapping(db, mapping),
-                ));
+                );
             }
         }
 
-        CallableSignature::from_overloads(overloads.iter().map(|overload| {
-            type_mappings.iter().fold(
-                overload.signature(db, inherited_generic_context),
-                |sig, mapping| sig.apply_type_mapping(db, mapping),
-            )
-        }))
+        CallableSignature::from_overloads(
+            overloads
+                .iter()
+                .map(|overload| overload.signature(db, inherited_generic_context)),
+        )
     }
 
     /// Typed externally-visible signature of the last overload or implementation of this function.
@@ -660,20 +651,10 @@ impl<'db> FunctionLiteral<'db> {
     /// calling query is not in the same file as this function is defined in, then this will create
     /// a cross-module dependency directly on the full AST which will lead to cache
     /// over-invalidation.
-    fn last_definition_signature<'a>(
-        self,
-        db: &'db dyn Db,
-        type_mappings: &'a [TypeMapping<'a, 'db>],
-    ) -> Signature<'db>
-    where
-        'db: 'a,
-    {
+    fn last_definition_signature(self, db: &'db dyn Db) -> Signature<'db> {
         let inherited_generic_context = self.inherited_generic_context(db);
-        type_mappings.iter().fold(
-            self.last_definition(db)
-                .signature(db, inherited_generic_context),
-            |sig, mapping| sig.apply_type_mapping(db, mapping),
-        )
+        self.last_definition(db)
+            .signature(db, inherited_generic_context)
     }
 
     fn normalized_impl(self, db: &'db dyn Db, visitor: &NormalizedVisitor<'db>) -> Self {
@@ -691,11 +672,19 @@ impl<'db> FunctionLiteral<'db> {
 pub struct FunctionType<'db> {
     pub(crate) literal: FunctionLiteral<'db>,
 
-    /// Type mappings that should be applied to the function's parameter and return types. This
-    /// might include specializations of enclosing generic contexts (e.g. for non-generic methods
-    /// of a specialized generic class).
-    #[returns(deref)]
-    type_mappings: Box<[TypeMapping<'db, 'db>]>,
+    /// Contains a potentially modified signature for this function literal, in case certain operations
+    /// (like type mappings) have been applied to it.
+    ///
+    /// See also: [`FunctionLiteral::updated_signature`].
+    #[returns(as_ref)]
+    updated_signature: Option<CallableSignature<'db>>,
+
+    /// Contains a potentially modified signature for the last overload or the implementation of this
+    /// function literal, in case certain operations (like type mappings) have been applied to it.
+    ///
+    /// See also: [`FunctionLiteral::last_definition_signature`].
+    #[returns(as_ref)]
+    updated_last_definition_signature: Option<Signature<'db>>,
 }
 
 // The Salsa heap is tracked separately.
@@ -707,8 +696,13 @@ pub(super) fn walk_function_type<'db, V: super::visitor::TypeVisitor<'db> + ?Siz
     visitor: &V,
 ) {
     walk_function_literal(db, function.literal(db), visitor);
-    for mapping in function.type_mappings(db) {
-        walk_type_mapping(db, mapping, visitor);
+    if let Some(callable_signature) = function.updated_signature(db) {
+        for signature in &callable_signature.overloads {
+            walk_signature(db, signature, visitor);
+        }
+    }
+    if let Some(signature) = function.updated_last_definition_signature(db) {
+        walk_signature(db, signature, visitor);
     }
 }
 
@@ -722,21 +716,41 @@ impl<'db> FunctionType<'db> {
         let literal = self
             .literal(db)
             .with_inherited_generic_context(db, inherited_generic_context);
-        Self::new(db, literal, self.type_mappings(db))
+        let updated_signature = self.updated_signature(db).map(|signature| {
+            signature.with_inherited_generic_context(Some(inherited_generic_context))
+        });
+        let updated_last_definition_signature =
+            self.updated_last_definition_signature(db).map(|signature| {
+                signature
+                    .clone()
+                    .with_inherited_generic_context(Some(inherited_generic_context))
+            });
+        Self::new(
+            db,
+            literal,
+            updated_signature,
+            updated_last_definition_signature,
+        )
     }
 
-    pub(crate) fn with_type_mapping<'a>(
+    pub(crate) fn apply_type_mapping_impl<'a>(
         self,
         db: &'db dyn Db,
         type_mapping: &TypeMapping<'a, 'db>,
+        visitor: &ApplyTypeMappingVisitor<'db>,
     ) -> Self {
-        let type_mappings: Box<[_]> = self
-            .type_mappings(db)
-            .iter()
-            .cloned()
-            .chain(std::iter::once(type_mapping.to_owned()))
-            .collect();
-        Self::new(db, self.literal(db), type_mappings)
+        let updated_signature =
+            self.signature(db)
+                .apply_type_mapping_impl(db, type_mapping, visitor);
+        let updated_last_definition_signature = self
+            .last_definition_signature(db)
+            .apply_type_mapping_impl(db, type_mapping, visitor);
+        Self::new(
+            db,
+            self.literal(db),
+            Some(updated_signature),
+            Some(updated_last_definition_signature),
+        )
     }
 
     pub(crate) fn with_dataclass_transformer_params(
@@ -752,7 +766,7 @@ impl<'db> FunctionType<'db> {
             .with_dataclass_transformer_params(db, params);
         let literal =
             FunctionLiteral::new(db, last_definition, literal.inherited_generic_context(db));
-        Self::new(db, literal, self.type_mappings(db))
+        Self::new(db, literal, None, None)
     }
 
     /// Returns the [`File`] in which this function is defined.
@@ -907,7 +921,9 @@ impl<'db> FunctionType<'db> {
     /// would depend on the function's AST and rerun for every change in that file.
     #[salsa::tracked(returns(ref), cycle_fn=signature_cycle_recover, cycle_initial=signature_cycle_initial, heap_size=ruff_memory_usage::heap_size)]
     pub(crate) fn signature(self, db: &'db dyn Db) -> CallableSignature<'db> {
-        self.literal(db).signature(db, self.type_mappings(db))
+        self.updated_signature(db)
+            .cloned()
+            .unwrap_or_else(|| self.literal(db).signature(db))
     }
 
     /// Typed externally-visible signature of the last overload or implementation of this function.
@@ -926,8 +942,9 @@ impl<'db> FunctionType<'db> {
         heap_size=ruff_memory_usage::heap_size,
     )]
     pub(crate) fn last_definition_signature(self, db: &'db dyn Db) -> Signature<'db> {
-        self.literal(db)
-            .last_definition_signature(db, self.type_mappings(db))
+        self.updated_last_definition_signature(db)
+            .cloned()
+            .unwrap_or_else(|| self.literal(db).last_definition_signature(db))
     }
 
     /// Convert the `FunctionType` into a [`CallableType`].
@@ -1017,12 +1034,19 @@ impl<'db> FunctionType<'db> {
     }
 
     pub(crate) fn normalized_impl(self, db: &'db dyn Db, visitor: &NormalizedVisitor<'db>) -> Self {
-        let mappings: Box<_> = self
-            .type_mappings(db)
-            .iter()
-            .map(|mapping| mapping.normalized_impl(db, visitor))
-            .collect();
-        Self::new(db, self.literal(db).normalized_impl(db, visitor), mappings)
+        let literal = self.literal(db).normalized_impl(db, visitor);
+        let updated_signature = self
+            .updated_signature(db)
+            .map(|signature| signature.normalized_impl(db, visitor));
+        let updated_last_definition_signature = self
+            .updated_last_definition_signature(db)
+            .map(|signature| signature.normalized_impl(db, visitor));
+        Self::new(
+            db,
+            literal,
+            updated_signature,
+            updated_last_definition_signature,
+        )
     }
 }