Merge pull request #83015 from slavapestov/prepared-overloads-part-3

Prepared overloads, part 3

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;

include/swift/Sema/PreparedOverload.h

@@ -71,6 +71,9 @@ struct PreparedOverloadChange {
     /// A generic requirement was opened to a constraint.
     AddedConstraint,
 
+    /// Special case of a Bind constraint.
+    AddedBindConstraint,
+
     /// A mapping of generic parameter types to type variables
     /// was recorded.
     OpenedTypes,
@@ -98,6 +101,11 @@ struct PreparedOverloadChange {
     /// For ChangeKind::AddedConstraint.
     Constraint *TheConstraint;
 
+    struct {
+      TypeBase *FirstType;
+      TypeBase * SecondType;
+    } Bind;
+
     /// For ChangeKind::OpenedTypes.
     struct {
       const OpenedType *Data;
@@ -192,6 +200,14 @@ struct PreparedOverloadBuilder {
     Changes.push_back(change);
   }
 
+  void addedBindConstraint(Type firstType, Type secondType) {
+    PreparedOverload::Change change;
+    change.Kind = PreparedOverload::Change::AddedBindConstraint;
+    change.Bind.FirstType = firstType.getPointer();
+    change.Bind.SecondType = secondType.getPointer();
+    Changes.push_back(change);
+  }
+
   void openedTypes(ArrayRef<OpenedType> replacements) {
     PreparedOverload::Change change;
     change.Kind = PreparedOverload::Change::OpenedTypes;

lib/Sema/CSSimplify.cpp

@@ -16365,6 +16365,11 @@ void ConstraintSystem::addConstraint(ConstraintKind kind, Type first,
                                      PreparedOverloadBuilder *preparedOverload) {
   if (preparedOverload) {
     ASSERT(PreparingOverload);
+    if (kind == ConstraintKind::Bind) {
+      ASSERT(!isFavored);
+      preparedOverload->addedBindConstraint(first, second);
+      return;
+    }
     auto c = Constraint::create(*this, kind, first, second,
                                 getConstraintLocator(locator));
     if (isFavored) c->setFavored();
@@ -16648,7 +16653,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 +16664,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);
 }
 

lib/Sema/ConstraintSystem.cpp

@@ -1085,6 +1085,7 @@ std::optional<Type> ConstraintSystem::isSetType(Type type) {
 
 Type ConstraintSystem::getFixedTypeRecursive(Type type, TypeMatchOptions &flags,
                                              bool wantRValue) {
+  ASSERT(!PreparingOverload);
 
   if (wantRValue)
     type = type->getRValueType();
@@ -1753,6 +1754,8 @@ Type ConstraintSystem::simplifyTypeImpl(Type type,
 }
 
 Type ConstraintSystem::simplifyType(Type type) {
+  ASSERT(!PreparingOverload);
+
   if (!type->hasTypeVariable())
     return type;
 

lib/Sema/TypeOfReference.cpp

@@ -1594,19 +1594,19 @@ DeclReferenceType ConstraintSystem::getTypeOfMemberReference(
   ASSERT(!!preparedOverload == PreparingOverload);
 
   // Figure out the instance type used for the base.
-  Type resolvedBaseTy = getFixedTypeRecursive(baseTy, /*wantRValue=*/true);
+  Type baseRValueTy = baseTy->getRValueType();
 
   // If the base is a module type, just use the type of the decl.
-  if (resolvedBaseTy->is<ModuleType>()) {
+  if (baseRValueTy->is<ModuleType>()) {
     return getTypeOfReference(value, functionRefInfo, locator, useDC,
                               preparedOverload);
   }
 
   // Check to see if the self parameter is applied, in which case we'll want to
   // strip it off later.
-  auto hasAppliedSelf = doesMemberRefApplyCurriedSelf(resolvedBaseTy, value);
+  auto hasAppliedSelf = doesMemberRefApplyCurriedSelf(baseRValueTy, value);
 
-  auto baseObjTy = resolvedBaseTy->getMetatypeInstanceType();
+  auto baseObjTy = baseRValueTy->getMetatypeInstanceType();
   FunctionType::Param baseObjParam(baseObjTy);
 
   // Indicates whether this is a valid reference to a static member on a
@@ -1617,7 +1617,7 @@ DeclReferenceType ConstraintSystem::getTypeOfMemberReference(
   bool isStaticMemberRefOnProtocol = false;
   if (baseObjTy->isExistentialType() && value->isStatic() &&
       locator->isLastElement<LocatorPathElt::UnresolvedMember>()) {
-    assert(resolvedBaseTy->is<MetatypeType>() &&
+    assert(baseRValueTy->is<MetatypeType>() &&
            "Assumed base of unresolved member access must be a metatype");
     isStaticMemberRefOnProtocol = true;
   }
@@ -1846,12 +1846,12 @@ DeclReferenceType ConstraintSystem::getTypeOfMemberReference(
   } else if (isa<AbstractFunctionDecl>(value) || isa<EnumElementDecl>(value)) {
     unsigned numApplies = getNumApplications(hasAppliedSelf, functionRefInfo);
     openedType = adjustFunctionTypeForConcurrency(
-        origOpenedType->castTo<FunctionType>(), resolvedBaseTy, value, useDC,
+        origOpenedType->castTo<FunctionType>(), baseRValueTy, value, useDC,
         numApplies, isMainDispatchQueueMember(locator), replacements, locator,
         preparedOverload);
   } else if (auto subscript = dyn_cast<SubscriptDecl>(value)) {
     openedType = adjustFunctionTypeForConcurrency(
-        origOpenedType->castTo<FunctionType>(), resolvedBaseTy, subscript, useDC,
+        origOpenedType->castTo<FunctionType>(), baseRValueTy, subscript, useDC,
         /*numApplies=*/2, /*isMainDispatchQueue=*/false, replacements, locator,
         preparedOverload);
   } else if (auto var = dyn_cast<VarDecl>(value)) {
@@ -2510,6 +2510,13 @@ void ConstraintSystem::replayChanges(
       simplifyDisjunctionChoice(change.TheConstraint);
       break;
 
+    case PreparedOverload::Change::AddedBindConstraint:
+      addConstraint(ConstraintKind::Bind,
+                    change.Bind.FirstType,
+                    change.Bind.SecondType,
+                    locator, /*isFavored=*/false);
+      break;
+
     case PreparedOverload::Change::OpenedTypes: {
       ArrayRef<OpenedType> replacements(
           change.Replacements.Data,