Sema: Add room for a PreparedOverload to BindOverload constraint

Author: slavapestov

include/swift/Sema/Constraint.h

@@ -49,6 +49,7 @@ namespace constraints {
 class ConstraintFix;
 class ConstraintLocator;
 class ConstraintSystem;
+class PreparedOverload;
 enum class TrailingClosureMatching;
 
 /// Describes the kind of constraint placed on one or more types.
@@ -437,6 +438,9 @@ class Constraint final : public llvm::ilist_node<Constraint>,
     struct {
       /// The first type.
       Type First;
+
+      /// The prepared overload, if any.
+      PreparedOverload *Prepared;
     } Overload;
 
     struct {
@@ -490,7 +494,9 @@ class Constraint final : public llvm::ilist_node<Constraint>,
              SmallPtrSetImpl<TypeVariableType *> &typeVars);
 
   /// Construct a new overload-binding constraint, which might have a fix.
-  Constraint(Type type, OverloadChoice choice, DeclContext *useDC,
+  Constraint(Type type, OverloadChoice choice,
+             PreparedOverload *preparedOverload,
+             DeclContext *useDC,
              ConstraintFix *fix, ConstraintLocator *locator,
              SmallPtrSetImpl<TypeVariableType *> &typeVars);
 
@@ -871,6 +877,13 @@ class Constraint final : public llvm::ilist_node<Constraint>,
     return *getTrailingObjects<OverloadChoice>();
   }
 
+  /// Retrieve the prepared overload choice for an overload-binding
+  /// constraint.
+  PreparedOverload *getPreparedOverload() const {
+    ASSERT(Kind == ConstraintKind::BindOverload);
+    return Overload.Prepared;
+  }
+
   FunctionType *getAppliedFunctionType() const {
     assert(Kind == ConstraintKind::ApplicableFunction);
     return Apply.AppliedFn;

include/swift/Sema/ConstraintSystem.h

@@ -3791,9 +3791,9 @@ class ConstraintSystem {
 
   /// Add a constraint that binds an overload set to a specific choice.
   void addBindOverloadConstraint(Type boundTy, OverloadChoice choice,
-                                 ConstraintLocator *locator,
-                                 DeclContext *useDC) {
-    resolveOverload(locator, boundTy, choice, useDC);
+                                 ConstraintLocator *locator, DeclContext *useDC) {
+    resolveOverload(locator, boundTy, choice, useDC,
+                    /*preparedOverload=*/nullptr);
   }
 
   /// Add a value member constraint to the constraint system.
@@ -4907,9 +4907,9 @@ class ConstraintSystem {
                               SelectedOverload choice);
 
   /// Build and allocate a prepared overload in the solver arena.
-  const PreparedOverload *prepareOverload(ConstraintLocator *locator,
-                                          OverloadChoice choice,
-                                          DeclContext *useDC);
+  PreparedOverload *prepareOverload(ConstraintLocator *locator,
+                                    OverloadChoice choice,
+                                    DeclContext *useDC);
 
   /// Populate the prepared overload with all type variables and constraints
   /// that are to be introduced into the constraint system when this choice
@@ -4922,11 +4922,12 @@ class ConstraintSystem {
 
   void replayChanges(
       ConstraintLocator *locator,
-      const PreparedOverload *preparedOverload);
+      PreparedOverload *preparedOverload);
 
   /// Resolve the given overload set to the given choice.
   void resolveOverload(ConstraintLocator *locator, Type boundType,
-                       OverloadChoice choice, DeclContext *useDC);
+                       OverloadChoice choice, DeclContext *useDC,
+                       PreparedOverload *preparedOverload);
 
   /// Simplify a type, by replacing type variables with either their
   /// fixed types (if available) or their representatives.

include/swift/Sema/OverloadChoice.h

@@ -250,6 +250,23 @@ class OverloadChoice {
     return (OverloadChoiceKind)kind;
   }
 
+  bool canBePrepared() const {
+    switch (getKind()) {
+    case OverloadChoiceKind::Decl:
+    case OverloadChoiceKind::DeclViaBridge:
+    case OverloadChoiceKind::DeclViaDynamic:
+    case OverloadChoiceKind::DeclViaUnwrappedOptional:
+    case OverloadChoiceKind::DynamicMemberLookup:
+    case OverloadChoiceKind::KeyPathDynamicMemberLookup:
+      return true;
+    case OverloadChoiceKind::TupleIndex:
+    case OverloadChoiceKind::MaterializePack:
+    case OverloadChoiceKind::ExtractFunctionIsolation:
+    case OverloadChoiceKind::KeyPathApplication:
+      return false;
+    }
+  }
+
   /// Determine whether this choice is for a declaration.
   bool isDecl() const {
     return DeclOrKind.is<ValueDecl*>();

lib/Sema/CSGen.cpp

@@ -1660,7 +1660,7 @@ namespace {
 
       OverloadChoice choice =
           OverloadChoice(Type(), E->getDecl(), E->getFunctionRefInfo());
-      CS.resolveOverload(locator, tv, choice, CurDC);
+      CS.addBindOverloadConstraint(tv, choice, locator, CurDC);
       return tv;
     }
 

lib/Sema/CSSimplify.cpp

@@ -11834,8 +11834,8 @@ ConstraintSystem::simplifyValueWitnessConstraint(
     return fail();
 
   auto choice = OverloadChoice(resolvedBaseType, witness.getDecl(), functionRefInfo);
-  resolveOverload(getConstraintLocator(locator), memberType, choice,
-                  useDC);
+  addBindOverloadConstraint(memberType, choice, getConstraintLocator(locator),
+                            useDC);
   return SolutionKind::Solved;
 }
 
@@ -16661,7 +16661,8 @@ ConstraintSystem::simplifyConstraint(const Constraint &constraint) {
 
     resolveOverload(constraint.getLocator(), constraint.getFirstType(),
                     constraint.getOverloadChoice(),
-                    constraint.getDeclContext());
+                    constraint.getDeclContext(),
+                    constraint.getPreparedOverload());
     return SolutionKind::Solved;
 
   case ConstraintKind::SubclassOf:

lib/Sema/Constraint.cpp

@@ -245,14 +245,16 @@ Constraint::Constraint(ConstraintKind kind, Type first, Type second,
   *getTrailingObjects<DeclContext *>() = useDC;
 }
 
-Constraint::Constraint(Type type, OverloadChoice choice, DeclContext *useDC,
+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}, Locator(locator) {
+      Overload{type, preparedOverload}, Locator(locator) {
   std::copy(typeVars.begin(), typeVars.end(), getTypeVariablesBuffer().begin());
   if (fix)
     *getTrailingObjects<ConstraintFix *>() = fix;
@@ -893,10 +895,22 @@ Constraint *Constraint::createValueWitness(
 }
 
 Constraint *Constraint::createBindOverload(ConstraintSystem &cs, Type type, 
-                                           OverloadChoice choice, 
+                                           OverloadChoice choice,
                                            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())
@@ -912,7 +926,8 @@ 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, useDC, fix, locator, typeVars);
+  return new (mem) Constraint(type, choice, preparedOverload, useDC,
+                              fix, locator, typeVars);
 }
 
 Constraint *Constraint::createRestricted(ConstraintSystem &cs, 

lib/Sema/TypeOfReference.cpp

@@ -2499,7 +2499,7 @@ void ConstraintSystem::recordResolvedOverload(ConstraintLocator *locator,
 
 void ConstraintSystem::replayChanges(
     ConstraintLocator *locator,
-    const PreparedOverload *preparedOverload) {
+    PreparedOverload *preparedOverload) {
   for (auto change : preparedOverload->getChanges()) {
     switch (change.Kind) {
     case PreparedOverload::Change::AddedTypeVariable:
@@ -2581,10 +2581,9 @@ ConstraintSystem::prepareOverloadImpl(ConstraintLocator *locator,
   }
 }
 
-const PreparedOverload *
-ConstraintSystem::prepareOverload(ConstraintLocator *locator,
-                                  OverloadChoice choice,
-                                  DeclContext *useDC) {
+PreparedOverload *ConstraintSystem::prepareOverload(ConstraintLocator *locator,
+                                                    OverloadChoice choice,
+                                                    DeclContext *useDC) {
   ASSERT(!PreparingOverload);
   PreparingOverload = true;
 
@@ -2602,7 +2601,8 @@ ConstraintSystem::prepareOverload(ConstraintLocator *locator,
 
 void ConstraintSystem::resolveOverload(ConstraintLocator *locator,
                                        Type boundType, OverloadChoice choice,
-                                       DeclContext *useDC) {
+                                       DeclContext *useDC,
+                                       PreparedOverload *preparedOverload) {
   // Determine the type to which we'll bind the overload set's type.
   Type openedType;
   Type adjustedOpenedType;
@@ -2617,11 +2617,7 @@ void ConstraintSystem::resolveOverload(ConstraintLocator *locator,
   case OverloadChoiceKind::DeclViaUnwrappedOptional:
   case OverloadChoiceKind::DynamicMemberLookup:
   case OverloadChoiceKind::KeyPathDynamicMemberLookup: {
-    // FIXME: Transitional hack
-    bool enablePreparedOverloads = true;
-
-    if (enablePreparedOverloads) {
-      auto *preparedOverload = prepareOverload(locator, choice, useDC);
+    if (preparedOverload) {
       replayChanges(locator, preparedOverload);
 
       openedType = preparedOverload->getOpenedType();