Sema: Record opened type variables in PreparedOverload

slavapestov · slavapestov · commit 543793564eae · 2025-07-09T15:00:02.000-04:00
diff --git a/include/swift/Sema/ConstraintSystem.h b/include/swift/Sema/ConstraintSystem.h
@@ -2924,15 +2924,15 @@ class ConstraintSystem {
   SolverTrail *getTrail() const {
     return solverState ? &solverState->Trail : nullptr;
   }
-
-  /// Add a new type variable that was already created.
-  void addTypeVariable(TypeVariableType *typeVar);
   
   /// Add a constraint from the subscript base to the root of the key
   /// path literal to the constraint system.
   void addKeyPathApplicationRootConstraint(Type root, ConstraintLocatorBuilder locator);
 
 public:
+  /// Add a new type variable that was already created.
+  void addTypeVariable(TypeVariableType *typeVar);
+
   /// Lookup for a member with the given name which is in the given base type.
   ///
   /// This routine caches the results of member lookups in the top constraint
@@ -2953,7 +2953,9 @@ class ConstraintSystem {
 
   /// Create a new type variable.
   TypeVariableType *createTypeVariable(ConstraintLocator *locator,
-                                       unsigned options);
+                                       unsigned options,
+                                       PreparedOverload *preparedOverload
+                                           = nullptr);
 
   /// Retrieve the set of active type variables.
   ArrayRef<TypeVariableType *> getTypeVariables() const {
@@ -4333,7 +4335,8 @@ class ConstraintSystem {
   ///
   /// \returns The opened type, or \c type if there are no archetypes in it.
   Type openType(Type type, ArrayRef<OpenedType> replacements,
-                ConstraintLocatorBuilder locator);
+                ConstraintLocatorBuilder locator,
+                PreparedOverload *preparedOverload);
 
   /// "Open" an opaque archetype type, similar to \c openType.
   Type openOpaqueType(OpaqueTypeArchetypeType *type,
@@ -4349,7 +4352,8 @@ class ConstraintSystem {
   /// aforementioned variable via special constraints.
   Type openPackExpansionType(PackExpansionType *expansion,
                              ArrayRef<OpenedType> replacements,
-                             ConstraintLocatorBuilder locator);
+                             ConstraintLocatorBuilder locator,
+                             PreparedOverload *preparedOverload);
 
   /// Update OpenedPackExpansionTypes and record a change in the trail.
   void recordOpenedPackExpansionType(PackExpansionType *expansion,
@@ -4378,26 +4382,30 @@ class ConstraintSystem {
   FunctionType *openFunctionType(AnyFunctionType *funcType,
                                  ConstraintLocatorBuilder locator,
                                  SmallVectorImpl<OpenedType> &replacements,
-                                 DeclContext *outerDC);
+                                 DeclContext *outerDC,
+                                 PreparedOverload *preparedOverload);
 
   /// Open the generic parameter list and its requirements,
   /// creating type variables for each of the type parameters.
   void openGeneric(DeclContext *outerDC,
                    GenericSignature signature,
                    ConstraintLocatorBuilder locator,
-                   SmallVectorImpl<OpenedType> &replacements);
+                   SmallVectorImpl<OpenedType> &replacements,
+                   PreparedOverload *preparedOverload);
 
   /// Open the generic parameter list creating type variables for each of the
   /// type parameters.
   void openGenericParameters(DeclContext *outerDC,
                              GenericSignature signature,
                              SmallVectorImpl<OpenedType> &replacements,
-                             ConstraintLocatorBuilder locator);
+                             ConstraintLocatorBuilder locator,
+                             PreparedOverload *preparedOverload);
 
   /// Open a generic parameter into a type variable and record
   /// it in \c replacements.
   TypeVariableType *openGenericParameter(GenericTypeParamType *parameter,
-                                         ConstraintLocatorBuilder locator);
+                                         ConstraintLocatorBuilder locator,
+                                         PreparedOverload *preparedOverload);
 
   /// Given generic signature open its generic requirements,
   /// using substitution function, and record them in the
@@ -4436,7 +4444,8 @@ class ConstraintSystem {
   FunctionType *adjustFunctionTypeForConcurrency(
       FunctionType *fnType, Type baseType, ValueDecl *decl, DeclContext *dc,
       unsigned numApplies, bool isMainDispatchQueue,
-      ArrayRef<OpenedType> replacements, ConstraintLocatorBuilder locator);
+      ArrayRef<OpenedType> replacements, ConstraintLocatorBuilder locator,
+      PreparedOverload *preparedOverload);
 
   /// Retrieve the type of a reference to the given value declaration.
   ///
diff --git a/include/swift/Sema/PreparedOverload.h b/include/swift/Sema/PreparedOverload.h
@@ -12,17 +12,21 @@
 #ifndef SWIFT_SEMA_PREPAREDOVERLOAD_H
 #define SWIFT_SEMA_PREPAREDOVERLOAD_H
 
-#include "swift/AST/AvailabilityRange.h"
-#include "swift/AST/FunctionRefInfo.h"
-#include "swift/AST/Types.h"
 #include "llvm/ADT/PointerIntPair.h"
-#include "llvm/Support/ErrorHandling.h"
+#include "llvm/ADT/SmallVector.h"
 
 namespace swift {
+
+class TypeVariableType;
+
 namespace constraints {
 
+class ConstraintSystem;
+
 struct PreparedOverload {
+  SmallVector<TypeVariableType *, 2> TypeVariables;
 
+  void discharge(ConstraintSystem &cs) const;
 };
 
 }
diff --git a/lib/Sema/CSRanking.cpp b/lib/Sema/CSRanking.cpp
@@ -423,7 +423,8 @@ static bool isProtocolExtensionAsSpecializedAs(DeclContext *dc1,
   // the second protocol extension.
   ConstraintSystem cs(dc1, std::nullopt);
   SmallVector<OpenedType, 4> replacements;
-  cs.openGeneric(dc2, sig2, ConstraintLocatorBuilder(nullptr), replacements);
+  cs.openGeneric(dc2, sig2, ConstraintLocatorBuilder(nullptr), replacements,
+                 /*preparedOverload=*/nullptr);
 
   // Bind the 'Self' type from the first extension to the type parameter from
   // opening 'Self' of the second extension.
@@ -581,13 +582,15 @@ bool CompareDeclSpecializationRequest::evaluate(
                       SmallVectorImpl<OpenedType> &replacements,
                       ConstraintLocator *locator) -> Type {
     if (auto *funcType = type->getAs<AnyFunctionType>()) {
-      return cs.openFunctionType(funcType, locator, replacements, outerDC);
+      return cs.openFunctionType(funcType, locator, replacements, outerDC,
+                                 /*preparedOverload=*/nullptr);
     }
 
     cs.openGeneric(outerDC, innerDC->getGenericSignatureOfContext(), locator,
-                   replacements);
+                   replacements, /*preparedOverload=*/nullptr);
 
-    return cs.openType(type, replacements, locator);
+    return cs.openType(type, replacements, locator,
+                       /*preparedOverload=*/nullptr);
   };
 
   bool knownNonSubtype = false;
diff --git a/lib/Sema/CSSimplify.cpp b/lib/Sema/CSSimplify.cpp
@@ -12046,8 +12046,10 @@ static Type getOpenedResultBuilderTypeFor(ConstraintSystem &cs,
     // Find the opened type for this callee and substitute in the type
     // parameters.
     auto substitutions = cs.getOpenedTypes(calleeLocator);
-    if (!substitutions.empty())
-      builderType = cs.openType(builderType, substitutions, locator);
+    if (!substitutions.empty()) {
+      builderType = cs.openType(builderType, substitutions, locator,
+                                /*preparedOverload=*/nullptr);
+    }
 
     assert(!builderType->hasTypeParameter());
   }
diff --git a/lib/Sema/CSSolver.cpp b/lib/Sema/CSSolver.cpp
@@ -22,6 +22,7 @@
 #include "swift/Basic/Defer.h"
 #include "swift/Sema/ConstraintGraph.h"
 #include "swift/Sema/ConstraintSystem.h"
+#include "swift/Sema/PreparedOverload.h"
 #include "swift/Sema/SolutionResult.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SetVector.h"
@@ -58,11 +59,15 @@ STATISTIC(LargestSolutionAttemptNumber, "# of the largest solution attempt");
 
 TypeVariableType *ConstraintSystem::createTypeVariable(
                                      ConstraintLocator *locator,
-                                     unsigned options) {
+                                     unsigned options,
+                                     PreparedOverload *preparedOverload) {
   ++TotalNumTypeVariables;
   auto tv = TypeVariableType::getNew(getASTContext(), assignTypeVariableID(),
                                      locator, options);
-  addTypeVariable(tv);
+  if (preparedOverload)
+    preparedOverload->TypeVariables.push_back(tv);
+  else
+    addTypeVariable(tv);
   return tv;
 }
 
diff --git a/lib/Sema/TypeCheckConstraints.cpp b/lib/Sema/TypeCheckConstraints.cpp
@@ -627,7 +627,7 @@ Type TypeChecker::typeCheckParameterDefault(Expr *&defaultValue,
       if (auto *typeVar = findParam(GP))
         return typeVar;
 
-      auto *typeVar = cs.openGenericParameter(GP, locator);
+      auto *typeVar = cs.openGenericParameter(GP, locator, nullptr);
       genericParameters.emplace_back(GP, typeVar);
 
       return typeVar;
@@ -720,8 +720,8 @@ Type TypeChecker::typeCheckParameterDefault(Expr *&defaultValue,
       cs.openGenericRequirement(DC->getParent(), signature, index, requirement,
                                 /*skipSelfProtocolConstraint=*/false, locator,
                                 [&](Type type) -> Type {
-                                  return cs.openType(type, genericParameters,
-                                                     locator);
+                                  return cs.openType(type, genericParameters, locator,
+                                                     /*preparedOverload=*/nullptr);
                                 });
     };
 
diff --git a/lib/Sema/TypeCheckProtocol.cpp b/lib/Sema/TypeCheckProtocol.cpp
@@ -1251,11 +1251,11 @@ swift::matchWitness(WitnessChecker::RequirementEnvironmentCache &reqEnvCache,
 
       reqThrownError = getThrownErrorType(reqASD);
       reqThrownError = cs->openType(reqThrownError, reqReplacements,
-                                    reqLocator);
+                                    reqLocator, /*preparedOverload=*/nullptr);
 
       witnessThrownError = getThrownErrorType(witnessASD);
       witnessThrownError = cs->openType(witnessThrownError, witnessReplacements,
-                                        witnessLocator);
+                                        witnessLocator, /*preparedOverload=*/nullptr);
     }
 
     return std::make_tuple(std::nullopt, reqType, openWitnessType,
diff --git a/lib/Sema/TypeOfReference.cpp b/lib/Sema/TypeOfReference.cpp
