build from constraint set

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

@@ -2921,9 +2921,7 @@ impl<'a, 'db> ArgumentTypeChecker<'a, 'db> {
         };
 
         // Build the specialization first without inferring the complete type context.
-        let isolated_specialization = builder
-            .mapped(generic_context, maybe_promote)
-            .build(generic_context);
+        let isolated_specialization = builder.build_mapped(generic_context, maybe_promote);
         let isolated_return_ty = self
             .return_ty
             .apply_specialization(self.db, isolated_specialization);
@@ -2948,9 +2946,7 @@ impl<'a, 'db> ArgumentTypeChecker<'a, 'db> {
             builder.infer(return_ty, call_expression_tcx).ok()?;
 
             // Otherwise, build the specialization again after inferring the complete type context.
-            let specialization = builder
-                .mapped(generic_context, maybe_promote)
-                .build(generic_context);
+            let specialization = builder.build_mapped(generic_context, maybe_promote);
             let return_ty = return_ty.apply_specialization(self.db, specialization);
 
             Some((Some(specialization), return_ty))

crates/ty_python_semantic/src/types/constraints.rs

@@ -3059,6 +3059,15 @@ impl<'db> GenericContext<'db> {
         self,
         db: &'db dyn Db,
         constraints: ConstraintSet<'db>,
+    ) -> Result<Specialization<'db>, ()> {
+        self.specialize_constrained_mapped(db, constraints, |_, _, ty| ty)
+    }
+
+    pub(crate) fn specialize_constrained_mapped(
+        self,
+        db: &'db dyn Db,
+        constraints: ConstraintSet<'db>,
+        f: impl Fn(BoundTypeVarIdentity<'db>, BoundTypeVarInstance<'db>, Type<'db>) -> Type<'db>,
     ) -> Result<Specialization<'db>, ()> {
         // If the constraint set is cyclic, don't even try to construct a specialization.
         if constraints.is_cyclic(db) {
@@ -3091,17 +3100,14 @@ impl<'db> GenericContext<'db> {
             let mut satisfied = false;
             let mut greatest_lower_bound = Type::Never;
             let mut least_upper_bound = Type::object();
-            abstracted.find_representative_types(
-                db,
-                bound_typevar.identity(db),
-                |lower_bound, upper_bound| {
-                    satisfied = true;
-                    greatest_lower_bound =
-                        UnionType::from_elements(db, [greatest_lower_bound, lower_bound]);
-                    least_upper_bound =
-                        IntersectionType::from_elements(db, [least_upper_bound, upper_bound]);
-                },
-            );
+            let identity = bound_typevar.identity(db);
+            abstracted.find_representative_types(db, identity, |lower_bound, upper_bound| {
+                satisfied = true;
+                greatest_lower_bound =
+                    UnionType::from_elements(db, [greatest_lower_bound, lower_bound]);
+                least_upper_bound =
+                    IntersectionType::from_elements(db, [least_upper_bound, upper_bound]);
+            });
 
             // If there are no satisfiable paths in the BDD, then there is no valid specialization
             // for this constraint set.
@@ -3119,7 +3125,7 @@ impl<'db> GenericContext<'db> {
 
             // Of all of the types that satisfy all of the paths in the BDD, we choose the
             // "largest" one (i.e., "closest to `object`") as the specialization.
-            types[i] = least_upper_bound;
+            types[i] = f(identity, bound_typevar, least_upper_bound);
         }
 
         Ok(self.specialize_recursive(db, types.into_boxed_slice()))

crates/ty_python_semantic/src/types/generics.rs

@@ -1330,6 +1330,7 @@ pub(crate) struct SpecializationBuilder<'db> {
     db: &'db dyn Db,
     inferable: InferableTypeVars<'db, 'db>,
     types: FxHashMap<BoundTypeVarIdentity<'db>, Type<'db>>,
+    constraints: ConstraintSet<'db>,
 }
 
 /// An assignment from a bound type variable to a given type, along with the variance of the outermost
@@ -1342,6 +1343,7 @@ impl<'db> SpecializationBuilder<'db> {
             db,
             inferable,
             types: FxHashMap::default(),
+            constraints: ConstraintSet::from(true),
         }
     }
 
@@ -1350,34 +1352,25 @@ impl<'db> SpecializationBuilder<'db> {
         &self.types
     }
 
-    /// Map the types that have been assigned in this specialization.
-    pub(crate) fn mapped(
-        &self,
-        generic_context: GenericContext<'db>,
-        f: impl Fn(BoundTypeVarIdentity<'db>, BoundTypeVarInstance<'db>, Type<'db>) -> Type<'db>,
-    ) -> Self {
-        let mut types = self.types.clone();
-        for (identity, variable) in generic_context.variables_inner(self.db) {
-            if let Some(ty) = types.get_mut(identity) {
-                *ty = f(*identity, *variable, *ty);
-            }
-        }
-
-        Self {
-            db: self.db,
-            inferable: self.inferable,
-            types,
-        }
+    pub(crate) fn build(&mut self, generic_context: GenericContext<'db>) -> Specialization<'db> {
+        self.build_mapped(generic_context, |_, _, ty| ty)
     }
 
-    pub(crate) fn build(&mut self, generic_context: GenericContext<'db>) -> Specialization<'db> {
-        let types = generic_context
-            .variables_inner(self.db)
-            .iter()
-            .map(|(identity, _)| self.types.get(identity).copied());
-
+    pub(crate) fn build_mapped(
+        &mut self,
+        generic_context: GenericContext<'db>,
+        f: impl Fn(BoundTypeVarIdentity<'db>, BoundTypeVarInstance<'db>, Type<'db>) -> Type<'db>,
+    ) -> Specialization<'db> {
         // TODO Infer the tuple spec for a tuple type
-        generic_context.specialize_partial(self.db, types)
+        let Ok(specialization) =
+            generic_context.specialize_constrained_mapped(self.db, self.constraints, f)
+        else {
+            panic!(
+                "should not be able to create unrealizable specialization from {}",
+                self.constraints.display(self.db),
+            );
+        };
+        specialization
     }
 
     fn add_type_mapping(
@@ -1392,6 +1385,32 @@ impl<'db> SpecializationBuilder<'db> {
             return;
         };
 
+        let constraint = match variance {
+            TypeVarVariance::Covariant => ConstraintSet::constrain_typevar(
+                self.db,
+                bound_typevar,
+                Type::Never,
+                ty,
+                TypeRelation::Assignability,
+            ),
+            TypeVarVariance::Contravariant => ConstraintSet::constrain_typevar(
+                self.db,
+                bound_typevar,
+                ty,
+                Type::object(),
+                TypeRelation::Assignability,
+            ),
+            TypeVarVariance::Invariant => ConstraintSet::constrain_typevar(
+                self.db,
+                bound_typevar,
+                ty,
+                ty,
+                TypeRelation::Assignability,
+            ),
+            TypeVarVariance::Bivariant => ConstraintSet::from(true),
+        };
+        self.constraints.intersect(self.db, constraint);
+
         match self.types.entry(identity) {
             Entry::Occupied(mut entry) => {
                 *entry.get_mut() = UnionType::from_elements(self.db, [*entry.get(), ty]);