[ty] Clean up inherited generic contexts (#20647)

We add an `inherited_generic_context` to the constructors of a generic
class. That lets us infer specializations of the class when invoking the
constructor. The constructor might itself be generic, in which case we
have to merge the list of typevars that we are willing to infer in the
constructor call.

Before we did that by tracking the two (and their specializations)
separately, with distinct `Option` fields/parameters. This PR updates
our call binding logic such that any given function call has _one_
optional generic context that we're willing to infer a specialization
for. If needed, we use the existing `GenericContext::merge` method to
create a new combined generic context for when the class and constructor
are both generic. This simplifies the call binding code considerably,
and is no more complex in the constructor call logic.

We also have a heuristic that we will promote any literals in the
specialized types of a generic class, but we don't promote literals in
the specialized types of the function itself. To handle this, we now
track this `should_promote_literals` property within `GenericContext`.
And moreover, we track this separately for each typevar, instead of a
single property for the generic context as a whole, so that we can
correctly merge the generic context of a constructor method (where the
option should be `false`) with the inherited generic context of its
containing class (where the option should be `true`).

---------

Co-authored-by: Alex Waygood <Alex.Waygood@Gmail.com>

Author: dcreager

crates/ruff_memory_usage/src/lib.rs

@@ -1,7 +1,7 @@
 use std::sync::{LazyLock, Mutex};
 
 use get_size2::{GetSize, StandardTracker};
-use ordermap::OrderSet;
+use ordermap::{OrderMap, OrderSet};
 
 /// Returns the memory usage of the provided object, using a global tracker to avoid
 /// double-counting shared objects.
@@ -18,3 +18,11 @@ pub fn heap_size<T: GetSize>(value: &T) -> usize {
 pub fn order_set_heap_size<T: GetSize, S>(set: &OrderSet<T, S>) -> usize {
     (set.capacity() * T::get_stack_size()) + set.iter().map(heap_size).sum::<usize>()
 }
+
+/// An implementation of [`GetSize::get_heap_size`] for [`OrderMap`].
+pub fn order_map_heap_size<K: GetSize, V: GetSize, S>(map: &OrderMap<K, V, S>) -> usize {
+    (map.capacity() * (K::get_stack_size() + V::get_stack_size()))
+        + (map.iter())
+            .map(|(k, v)| heap_size(k) + heap_size(v))
+            .sum::<usize>()
+}

crates/ty_python_semantic/src/types.rs

@@ -5456,28 +5456,19 @@ impl<'db> Type<'db> {
                     }
                 }
 
-                let new_specialization = new_call_outcome
-                    .and_then(Result::ok)
-                    .as_ref()
-                    .and_then(Bindings::single_element)
-                    .into_iter()
-                    .flat_map(CallableBinding::matching_overloads)
-                    .next()
-                    .and_then(|(_, binding)| binding.inherited_specialization())
-                    .filter(|specialization| {
-                        Some(specialization.generic_context(db)) == generic_context
-                    });
-                let init_specialization = init_call_outcome
-                    .and_then(Result::ok)
-                    .as_ref()
-                    .and_then(Bindings::single_element)
-                    .into_iter()
-                    .flat_map(CallableBinding::matching_overloads)
-                    .next()
-                    .and_then(|(_, binding)| binding.inherited_specialization())
-                    .filter(|specialization| {
-                        Some(specialization.generic_context(db)) == generic_context
-                    });
+                let specialize_constructor = |outcome: Option<Bindings<'db>>| {
+                    let (_, binding) = outcome
+                        .as_ref()?
+                        .single_element()?
+                        .matching_overloads()
+                        .next()?;
+                    binding.specialization()?.restrict(db, generic_context?)
+                };
+
+                let new_specialization =
+                    specialize_constructor(new_call_outcome.and_then(Result::ok));
+                let init_specialization =
+                    specialize_constructor(init_call_outcome.and_then(Result::ok));
                 let specialization =
                     combine_specializations(db, new_specialization, init_specialization);
                 let specialized = specialization
@@ -6768,21 +6759,19 @@ impl<'db> TypeMapping<'_, 'db> {
                 db,
                 context
                     .variables(db)
-                    .iter()
-                    .filter(|var| !var.typevar(db).is_self(db))
-                    .copied(),
+                    .filter(|var| !var.typevar(db).is_self(db)),
             ),
             TypeMapping::ReplaceSelf { new_upper_bound } => GenericContext::from_typevar_instances(
                 db,
-                context.variables(db).iter().map(|typevar| {
+                context.variables(db).map(|typevar| {
                     if typevar.typevar(db).is_self(db) {
                         BoundTypeVarInstance::synthetic_self(
                             db,
                             *new_upper_bound,
                             typevar.binding_context(db),
                         )
                     } else {
-                        *typevar
+                        typevar
                     }
                 }),
             ),

crates/ty_python_semantic/src/types/call/bind.rs

@@ -32,7 +32,7 @@ use crate::types::tuple::{TupleLength, TupleType};
 use crate::types::{
     BoundMethodType, ClassLiteral, DataclassParams, FieldInstance, KnownBoundMethodType,
     KnownClass, KnownInstanceType, MemberLookupPolicy, PropertyInstanceType, SpecialFormType,
-    TrackedConstraintSet, TypeAliasType, TypeContext, TypeMapping, UnionBuilder, UnionType,
+    TrackedConstraintSet, TypeAliasType, TypeContext, UnionBuilder, UnionType,
     WrapperDescriptorKind, enums, ide_support, todo_type,
 };
 use ruff_db::diagnostic::{Annotation, Diagnostic, SubDiagnostic, SubDiagnosticSeverity};
@@ -1701,10 +1701,6 @@ impl<'db> CallableBinding<'db> {
                             parameter_type =
                                 parameter_type.apply_specialization(db, specialization);
                         }
-                        if let Some(inherited_specialization) = overload.inherited_specialization {
-                            parameter_type =
-                                parameter_type.apply_specialization(db, inherited_specialization);
-                        }
                         union_parameter_types[parameter_index.saturating_sub(skipped_parameters)]
                             .add_in_place(parameter_type);
                     }
@@ -1983,7 +1979,7 @@ impl<'db> CallableBinding<'db> {
                     for overload in overloads.iter().take(MAXIMUM_OVERLOADS) {
                         diag.info(format_args!(
                             "  {}",
-                            overload.signature(context.db(), None).display(context.db())
+                            overload.signature(context.db()).display(context.db())
                         ));
                     }
                     if overloads.len() > MAXIMUM_OVERLOADS {
@@ -2444,7 +2440,6 @@ struct ArgumentTypeChecker<'a, 'db> {
     errors: &'a mut Vec<BindingError<'db>>,
 
     specialization: Option<Specialization<'db>>,
-    inherited_specialization: Option<Specialization<'db>>,
 }
 
 impl<'a, 'db> ArgumentTypeChecker<'a, 'db> {
@@ -2466,7 +2461,6 @@ impl<'a, 'db> ArgumentTypeChecker<'a, 'db> {
             call_expression_tcx,
             errors,
             specialization: None,
-            inherited_specialization: None,
         }
     }
 
@@ -2498,9 +2492,7 @@ impl<'a, 'db> ArgumentTypeChecker<'a, 'db> {
     }
 
     fn infer_specialization(&mut self) {
-        if self.signature.generic_context.is_none()
-            && self.signature.inherited_generic_context.is_none()
-        {
+        if self.signature.generic_context.is_none() {
             return;
         }
 
@@ -2542,14 +2534,6 @@ impl<'a, 'db> ArgumentTypeChecker<'a, 'db> {
         }
 
         self.specialization = self.signature.generic_context.map(|gc| builder.build(gc));
-        self.inherited_specialization = self.signature.inherited_generic_context.map(|gc| {
-            // The inherited generic context is used when inferring the specialization of a generic
-            // class from a constructor call. In this case (only), we promote any typevars that are
-            // inferred as a literal to the corresponding instance type.
-            builder
-                .build(gc)
-                .apply_type_mapping(self.db, &TypeMapping::PromoteLiterals)
-        });
     }
 
     fn check_argument_type(
@@ -2566,11 +2550,6 @@ impl<'a, 'db> ArgumentTypeChecker<'a, 'db> {
                 argument_type = argument_type.apply_specialization(self.db, specialization);
                 expected_ty = expected_ty.apply_specialization(self.db, specialization);
             }
-            if let Some(inherited_specialization) = self.inherited_specialization {
-                argument_type =
-                    argument_type.apply_specialization(self.db, inherited_specialization);
-                expected_ty = expected_ty.apply_specialization(self.db, inherited_specialization);
-            }
             // This is one of the few places where we want to check if there's _any_ specialization
             // where assignability holds; normally we want to check that assignability holds for
             // _all_ specializations.
@@ -2742,8 +2721,8 @@ impl<'a, 'db> ArgumentTypeChecker<'a, 'db> {
         }
     }
 
-    fn finish(self) -> (Option<Specialization<'db>>, Option<Specialization<'db>>) {
-        (self.specialization, self.inherited_specialization)
+    fn finish(self) -> Option<Specialization<'db>> {
+        self.specialization
     }
 }
 
@@ -2807,10 +2786,6 @@ pub(crate) struct Binding<'db> {
     /// The specialization that was inferred from the argument types, if the callable is generic.
     specialization: Option<Specialization<'db>>,
 
-    /// The specialization that was inferred for a class method's containing generic class, if it
-    /// is being used to infer a specialization for the class.
-    inherited_specialization: Option<Specialization<'db>>,
-
     /// Information about which parameter(s) each argument was matched with, in argument source
     /// order.
     argument_matches: Box<[MatchedArgument<'db>]>,
@@ -2835,7 +2810,6 @@ impl<'db> Binding<'db> {
             signature_type,
             return_ty: Type::unknown(),
             specialization: None,
-            inherited_specialization: None,
             argument_matches: Box::from([]),
             variadic_argument_matched_to_variadic_parameter: false,
             parameter_tys: Box::from([]),
@@ -2906,15 +2880,10 @@ impl<'db> Binding<'db> {
         checker.infer_specialization();
 
         checker.check_argument_types();
-        (self.specialization, self.inherited_specialization) = checker.finish();
+        self.specialization = checker.finish();
         if let Some(specialization) = self.specialization {
             self.return_ty = self.return_ty.apply_specialization(db, specialization);
         }
-        if let Some(inherited_specialization) = self.inherited_specialization {
-            self.return_ty = self
-                .return_ty
-                .apply_specialization(db, inherited_specialization);
-        }
     }
 
     pub(crate) fn set_return_type(&mut self, return_ty: Type<'db>) {
@@ -2925,8 +2894,8 @@ impl<'db> Binding<'db> {
         self.return_ty
     }
 
-    pub(crate) fn inherited_specialization(&self) -> Option<Specialization<'db>> {
-        self.inherited_specialization
+    pub(crate) fn specialization(&self) -> Option<Specialization<'db>> {
+        self.specialization
     }
 
     /// Returns the bound types for each parameter, in parameter source order, or `None` if no
@@ -2988,7 +2957,6 @@ impl<'db> Binding<'db> {
         BindingSnapshot {
             return_ty: self.return_ty,
             specialization: self.specialization,
-            inherited_specialization: self.inherited_specialization,
             argument_matches: self.argument_matches.clone(),
             parameter_tys: self.parameter_tys.clone(),
             errors: self.errors.clone(),
@@ -2999,15 +2967,13 @@ impl<'db> Binding<'db> {
         let BindingSnapshot {
             return_ty,
             specialization,
-            inherited_specialization,
             argument_matches,
             parameter_tys,
             errors,
         } = snapshot;
 
         self.return_ty = return_ty;
         self.specialization = specialization;
-        self.inherited_specialization = inherited_specialization;
         self.argument_matches = argument_matches;
         self.parameter_tys = parameter_tys;
         self.errors = errors;
@@ -3027,7 +2993,6 @@ impl<'db> Binding<'db> {
     fn reset(&mut self) {
         self.return_ty = Type::unknown();
         self.specialization = None;
-        self.inherited_specialization = None;
         self.argument_matches = Box::from([]);
         self.parameter_tys = Box::from([]);
         self.errors.clear();
@@ -3038,7 +3003,6 @@ impl<'db> Binding<'db> {
 struct BindingSnapshot<'db> {
     return_ty: Type<'db>,
     specialization: Option<Specialization<'db>>,
-    inherited_specialization: Option<Specialization<'db>>,
     argument_matches: Box<[MatchedArgument<'db>]>,
     parameter_tys: Box<[Option<Type<'db>>]>,
     errors: Vec<BindingError<'db>>,
@@ -3078,7 +3042,6 @@ impl<'db> CallableBindingSnapshot<'db> {
                 // ... and update the snapshot with the current state of the binding.
                 snapshot.return_ty = binding.return_ty;
                 snapshot.specialization = binding.specialization;
-                snapshot.inherited_specialization = binding.inherited_specialization;
                 snapshot
                     .argument_matches
                     .clone_from(&binding.argument_matches);
@@ -3373,7 +3336,7 @@ impl<'db> BindingError<'db> {
                             }
                             diag.info(format_args!(
                                 "  {}",
-                                overload.signature(context.db(), None).display(context.db())
+                                overload.signature(context.db()).display(context.db())
                             ));
                         }
                         if overloads.len() > MAXIMUM_OVERLOADS {

crates/ty_python_semantic/src/types/class.rs

@@ -324,7 +324,6 @@ impl<'db> VarianceInferable<'db> for GenericAlias<'db> {
                 specialization
                     .generic_context(db)
                     .variables(db)
-                    .iter()
                     .zip(specialization.types(db))
                     .map(|(generic_typevar, ty)| {
                         if let Some(explicit_variance) =
@@ -346,7 +345,7 @@ impl<'db> VarianceInferable<'db> for GenericAlias<'db> {
                             let typevar_variance_in_substituted_type = ty.variance_of(db, typevar);
                             origin
                                 .with_polarity(typevar_variance_in_substituted_type)
-                                .variance_of(db, *generic_typevar)
+                                .variance_of(db, generic_typevar)
                         }
                     }),
             )
@@ -1013,8 +1012,7 @@ impl<'db> ClassType<'db> {
 
                 let synthesized_dunder = CallableType::function_like(
                     db,
-                    Signature::new(parameters, None)
-                        .with_inherited_generic_context(inherited_generic_context),
+                    Signature::new_generic(inherited_generic_context, parameters, None),
                 );
 
                 Place::bound(synthesized_dunder).into()
@@ -1454,6 +1452,16 @@ impl<'db> ClassLiteral<'db> {
         )
     }
 
+    /// Returns the generic context that should be inherited by any constructor methods of this
+    /// class.
+    ///
+    /// When inferring a specialization of the class's generic context from a constructor call, we
+    /// promote any typevars that are inferred as a literal to the corresponding instance type.
+    fn inherited_generic_context(self, db: &'db dyn Db) -> Option<GenericContext<'db>> {
+        self.generic_context(db)
+            .map(|generic_context| generic_context.promote_literals(db))
+    }
+
     fn file(self, db: &dyn Db) -> File {
         self.body_scope(db).file(db)
     }
@@ -1996,7 +2004,7 @@ impl<'db> ClassLiteral<'db> {
                     lookup_result = lookup_result.or_else(|lookup_error| {
                         lookup_error.or_fall_back_to(
                             db,
-                            class.own_class_member(db, self.generic_context(db), name),
+                            class.own_class_member(db, self.inherited_generic_context(db), name),
                         )
                     });
                 }
@@ -2246,8 +2254,14 @@ impl<'db> ClassLiteral<'db> {
             // so that the keyword-only parameters appear after positional parameters.
             parameters.sort_by_key(Parameter::is_keyword_only);
 
-            let mut signature = Signature::new(Parameters::new(parameters), return_ty);
-            signature.inherited_generic_context = self.generic_context(db);
+            let signature = match name {
+                "__new__" | "__init__" => Signature::new_generic(
+                    self.inherited_generic_context(db),
+                    Parameters::new(parameters),
+                    return_ty,
+                ),
+                _ => Signature::new(Parameters::new(parameters), return_ty),
+            };
             Some(CallableType::function_like(db, signature))
         };
 
@@ -2295,7 +2309,7 @@ impl<'db> ClassLiteral<'db> {
                 KnownClass::NamedTupleFallback
                     .to_class_literal(db)
                     .into_class_literal()?
-                    .own_class_member(db, self.generic_context(db), None, name)
+                    .own_class_member(db, self.inherited_generic_context(db), None, name)
                     .place
                     .ignore_possibly_unbound()
                     .map(|ty| {
@@ -5421,7 +5435,7 @@ enum SlotsKind {
 impl SlotsKind {
     fn from(db: &dyn Db, base: ClassLiteral) -> Self {
         let Place::Type(slots_ty, bound) = base
-            .own_class_member(db, base.generic_context(db), None, "__slots__")
+            .own_class_member(db, base.inherited_generic_context(db), None, "__slots__")
             .place
         else {
             return Self::NotSpecified;