Sema: Lazily build prepared overload when simplifying BindOverload constraint

Author: slavapestov

include/swift/Sema/Constraint.h

@@ -495,7 +495,6 @@ class Constraint final : public llvm::ilist_node<Constraint>,
 
   /// Construct a new overload-binding constraint, which might have a fix.
   Constraint(Type type, OverloadChoice choice,
-             PreparedOverload *preparedOverload,
              DeclContext *useDC,
              ConstraintFix *fix, ConstraintLocator *locator,
              SmallPtrSetImpl<TypeVariableType *> &typeVars);
@@ -884,6 +883,12 @@ class Constraint final : public llvm::ilist_node<Constraint>,
     return Overload.Prepared;
   }
 
+  void setPreparedOverload(PreparedOverload *preparedOverload) {
+    ASSERT(Kind == ConstraintKind::BindOverload);
+    ASSERT(!Overload.Prepared);
+    Overload.Prepared = preparedOverload;
+  }
+
   FunctionType *getAppliedFunctionType() const {
     assert(Kind == ConstraintKind::ApplicableFunction);
     return Apply.AppliedFn;

lib/Sema/CSSimplify.cpp

@@ -16648,7 +16648,7 @@ ConstraintSystem::simplifyConstraint(const Constraint &constraint) {
         constraint.getFirstType(), constraint.getSecondType(),
         constraint.getThirdType(), std::nullopt, constraint.getLocator());
 
-  case ConstraintKind::BindOverload:
+  case ConstraintKind::BindOverload: {
     if (auto *fix = constraint.getFix()) {
       // TODO(diagnostics): Impact should be associated with a fix unless
       // it's a contextual problem, then only solver can decide what the impact
@@ -16659,11 +16659,26 @@ ConstraintSystem::simplifyConstraint(const Constraint &constraint) {
         return SolutionKind::Error;
     }
 
+    // FIXME: Transitional hack.
+    bool enablePreparedOverloads = false;
+
+    auto *preparedOverload = constraint.getPreparedOverload();
+    if (!preparedOverload) {
+      if (enablePreparedOverloads &&
+          constraint.getOverloadChoice().canBePrepared()) {
+        preparedOverload = prepareOverload(constraint.getLocator(),
+                                           constraint.getOverloadChoice(),
+                                           constraint.getDeclContext());
+        const_cast<Constraint &>(constraint).setPreparedOverload(preparedOverload);
+      }
+    }
+
     resolveOverload(constraint.getLocator(), constraint.getFirstType(),
                     constraint.getOverloadChoice(),
                     constraint.getDeclContext(),
-                    constraint.getPreparedOverload());
+                    preparedOverload);
     return SolutionKind::Solved;
+  }
 
   case ConstraintKind::SubclassOf:
     return simplifySubclassOfConstraint(constraint.getFirstType(),

lib/Sema/Constraint.cpp

@@ -246,15 +246,14 @@ Constraint::Constraint(ConstraintKind kind, Type first, Type second,
 }
 
 Constraint::Constraint(Type type, OverloadChoice choice,
-                       PreparedOverload *preparedOverload,
                        DeclContext *useDC,
                        ConstraintFix *fix, ConstraintLocator *locator,
                        SmallPtrSetImpl<TypeVariableType *> &typeVars)
     : Kind(ConstraintKind::BindOverload), NumTypeVariables(typeVars.size()),
       HasFix(fix != nullptr), HasDeclContext(true), HasRestriction(false),
       IsActive(false), IsDisabled(bool(fix)), IsDisabledForPerformance(false),
       RememberChoice(false), IsFavored(false), IsIsolated(false),
-      Overload{type, preparedOverload}, Locator(locator) {
+      Overload{type, /*preparedOverload=*/nullptr}, Locator(locator) {
   std::copy(typeVars.begin(), typeVars.end(), getTypeVariablesBuffer().begin());
   if (fix)
     *getTrailingObjects<ConstraintFix *>() = fix;
@@ -899,18 +898,6 @@ Constraint *Constraint::createBindOverload(ConstraintSystem &cs, Type type,
                                            DeclContext *useDC,
                                            ConstraintFix *fix,
                                            ConstraintLocator *locator) {
-  // FIXME: Transitional hack.
-  bool enablePreparedOverloads = false;
-
-  PreparedOverload *preparedOverload = nullptr;
-
-  // Prepare the overload.
-  if (enablePreparedOverloads) {
-    if (choice.canBePrepared()) {
-      preparedOverload = cs.prepareOverload(locator, choice, useDC);
-    }
-  }
-
   // Collect type variables.
   SmallPtrSet<TypeVariableType *, 4> typeVars;
   if (type->hasTypeVariable())
@@ -926,7 +913,7 @@ Constraint *Constraint::createBindOverload(ConstraintSystem &cs, Type type,
       typeVars.size(), fix ? 1 : 0, /*hasDeclContext=*/1,
       /*hasContextualTypeInfo=*/0, /*hasOverloadChoice=*/1);
   void *mem = cs.getAllocator().Allocate(size, alignof(Constraint));
-  return new (mem) Constraint(type, choice, preparedOverload, useDC,
+  return new (mem) Constraint(type, choice, useDC,
                               fix, locator, typeVars);
 }