Sema: Collect constraints in PreparedOverload

Author: slavapestov

include/swift/Sema/ConstraintSystem.h

@@ -3695,12 +3695,14 @@ class ConstraintSystem {
   /// Add a constraint to the constraint system.
   void addConstraint(ConstraintKind kind, Type first, Type second,
                      ConstraintLocatorBuilder locator,
-                     bool isFavored = false);
+                     bool isFavored = false,
+                     PreparedOverload *preparedOverload = nullptr);
 
   /// Add a requirement as a constraint to the constraint system.
   void addConstraint(Requirement req, ConstraintLocatorBuilder locator,
                      bool isFavored,
-                     bool prohibitNonisolatedConformance);
+                     bool prohibitNonisolatedConformance,
+                     PreparedOverload *preparedOverload = nullptr);
 
   void addApplicationConstraint(
       FunctionType *appliedFn, Type calleeType,
@@ -4414,7 +4416,8 @@ class ConstraintSystem {
                                GenericSignature signature,
                                bool skipProtocolSelfConstraint,
                                ConstraintLocatorBuilder locator,
-                               llvm::function_ref<Type(Type)> subst);
+                               llvm::function_ref<Type(Type)> subst,
+                               PreparedOverload *preparedOverload);
 
   // Record the given requirement in the constraint system.
   void openGenericRequirement(DeclContext *outerDC,
@@ -4423,7 +4426,8 @@ class ConstraintSystem {
                               const Requirement &requirement,
                               bool skipProtocolSelfConstraint,
                               ConstraintLocatorBuilder locator,
-                              llvm::function_ref<Type(Type)> subst);
+                              llvm::function_ref<Type(Type)> subst,
+                              PreparedOverload *preparedOverload);
 
   /// Update OpenedTypes and record a change in the trail.
   void recordOpenedType(
@@ -4432,8 +4436,9 @@ class ConstraintSystem {
   /// Record the set of opened types for the given locator.
   void recordOpenedTypes(
          ConstraintLocatorBuilder locator,
-         SmallVectorImpl<OpenedType> &replacements,
-         bool fixmeAllowDuplicates=false);
+         const SmallVectorImpl<OpenedType> &replacements,
+         PreparedOverload *preparedOverload = nullptr,
+         bool fixmeAllowDuplicates = false);
 
   /// Check whether the given type conforms to the given protocol and if
   /// so return a valid conformance reference.

include/swift/Sema/PreparedOverload.h

@@ -17,16 +17,24 @@
 
 namespace swift {
 
+class GenericTypeParamType;
 class TypeVariableType;
 
 namespace constraints {
 
+class ConstraintLocatorBuilder;
 class ConstraintSystem;
 
+/// Describes a dependent type that has been opened to a particular type
+/// variable.
+using OpenedType = std::pair<GenericTypeParamType *, TypeVariableType *>;
+
 struct PreparedOverload {
   SmallVector<TypeVariableType *, 2> TypeVariables;
+  SmallVector<Constraint *, 2> Constraints;
+  SmallVector<OpenedType, 2> Replacements;
 
-  void discharge(ConstraintSystem &cs) const;
+  void discharge(ConstraintSystem &cs, ConstraintLocatorBuilder locator) const;
 };
 
 }

lib/Sema/CSSimplify.cpp

@@ -40,6 +40,7 @@
 #include "swift/Sema/CSFix.h"
 #include "swift/Sema/ConstraintSystem.h"
 #include "swift/Sema/IDETypeChecking.h"
+#include "swift/Sema/PreparedOverload.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/Support/Compiler.h"
 
@@ -16269,15 +16270,17 @@ ConstraintSystem::addKeyPathConstraint(
 void ConstraintSystem::addConstraint(Requirement req,
                                      ConstraintLocatorBuilder locator,
                                      bool isFavored,
-                                     bool prohibitNonisolatedConformance) {
+                                     bool prohibitNonisolatedConformance,
+                                     PreparedOverload *preparedOverload) {
   bool conformsToAnyObject = false;
   std::optional<ConstraintKind> kind;
   switch (req.getKind()) {
   case RequirementKind::SameShape: {
     auto type1 = req.getFirstType();
     auto type2 = req.getSecondType();
 
-    addConstraint(ConstraintKind::SameShape, type1, type2, locator);
+    addConstraint(ConstraintKind::SameShape, type1, type2, locator,
+                  /*isFavored=*/false, preparedOverload);
     return;
   }
 
@@ -16318,19 +16321,32 @@ void ConstraintSystem::addConstraint(Requirement req,
   auto firstType = req.getFirstType();
   if (kind) {
     addConstraint(*kind, req.getFirstType(), req.getSecondType(), locator,
-                  isFavored);
+                  isFavored, preparedOverload);
   }
 
   if (conformsToAnyObject) {
     auto anyObject = getASTContext().getAnyObjectConstraint();
-    addConstraint(ConstraintKind::ConformsTo, firstType, anyObject, locator);
+    addConstraint(ConstraintKind::ConformsTo, firstType, anyObject, locator,
+                  /*isFavored=*/false, preparedOverload);
   }
 }
 
 void ConstraintSystem::addConstraint(ConstraintKind kind, Type first,
                                      Type second,
                                      ConstraintLocatorBuilder locator,
-                                     bool isFavored) {
+                                     bool isFavored,
+                                     PreparedOverload *preparedOverload) {
+  if (preparedOverload) {
+    ASSERT(PreparingOverload);
+    auto c = Constraint::create(*this, kind, first, second,
+                                getConstraintLocator(locator));
+    if (isFavored) c->setFavored();
+    preparedOverload->Constraints.push_back(c);
+    return;
+  }
+
+  ASSERT(!PreparingOverload);
+
   switch (addConstraintImpl(kind, first, second, locator, isFavored)) {
   case SolutionKind::Error:
     // Add a failing constraint, if needed.

lib/Sema/CSSolver.cpp

@@ -64,10 +64,12 @@ TypeVariableType *ConstraintSystem::createTypeVariable(
   ++TotalNumTypeVariables;
   auto tv = TypeVariableType::getNew(getASTContext(), assignTypeVariableID(),
                                      locator, options);
-  if (preparedOverload)
+  if (preparedOverload) {
+    ASSERT(PreparingOverload);
     preparedOverload->TypeVariables.push_back(tv);
-  else
+  } else {
     addTypeVariable(tv);
+  }
   return tv;
 }
 

lib/Sema/TypeCheckConstraints.cpp

@@ -722,7 +722,7 @@ Type TypeChecker::typeCheckParameterDefault(Expr *&defaultValue,
                                 [&](Type type) -> Type {
                                   return cs.openType(type, genericParameters, locator,
                                                      /*preparedOverload=*/nullptr);
-                                });
+                                }, /*preparedOverload=*/nullptr);
     };
 
     auto diagnoseInvalidRequirement = [&](Requirement requirement) {