Sema: Collect property wrappers in PreparedOverload

Author: slavapestov

include/swift/Sema/ConstraintSystem.h

@@ -3621,7 +3621,8 @@ class ConstraintSystem {
 
   /// Log and record the application of the fix. Return true iff any
   /// subsequent solution would be worse than the best known solution.
-  bool recordFix(ConstraintFix *fix, unsigned impact = 1);
+  bool recordFix(ConstraintFix *fix, unsigned impact = 1,
+                 PreparedOverload *preparedOverload = nullptr);
 
   void recordPotentialHole(TypeVariableType *typeVar);
   void recordAnyTypeVarAsPotentialHole(Type type);
@@ -5333,13 +5334,20 @@ class ConstraintSystem {
   /// Matches a wrapped or projected value parameter type to its backing
   /// property wrapper type by applying the property wrapper.
   TypeMatchResult applyPropertyWrapperToParameter(
-      Type wrapperType, Type paramType, ParamDecl *param, Identifier argLabel,
-      ConstraintKind matchKind, ConstraintLocator *locator,
-      ConstraintLocator *calleeLocator);
+      Type wrapperType,
+      Type paramType,
+      ParamDecl *param,
+      Identifier argLabel,
+      ConstraintKind matchKind,
+      ConstraintLocator *locator,
+      ConstraintLocator *calleeLocator,
+      PreparedOverload *preparedOverload = nullptr);
 
   /// Used by applyPropertyWrapperToParameter() to update appliedPropertyWrappers
   /// and record a change in the trail.
-  void applyPropertyWrapper(Expr *anchor, AppliedPropertyWrapper applied);
+  void applyPropertyWrapper(Expr *anchor,
+                            AppliedPropertyWrapper applied,
+                            PreparedOverload *preparedOverload = nullptr);
 
   /// Undo the above change.
   void removePropertyWrapper(Expr *anchor);

include/swift/Sema/PreparedOverload.h

@@ -12,6 +12,7 @@
 #ifndef SWIFT_SEMA_PREPAREDOVERLOAD_H
 #define SWIFT_SEMA_PREPAREDOVERLOAD_H
 
+#include "swift/AST/PropertyWrappers.h"
 #include "llvm/ADT/PointerIntPair.h"
 #include "llvm/ADT/SmallVector.h"
 
@@ -37,6 +38,8 @@ struct PreparedOverload {
   ExistentialArchetypeType *OpenedExistential = nullptr;
   SmallVector<std::pair<PackExpansionType *, TypeVariableType *>>
     OpenedPackExpansionTypes;
+  SmallVector<AppliedPropertyWrapper, 2> PropertyWrappers;
+  SmallVector<std::pair<ConstraintFix *, unsigned>, 2> Fixes;
 
   void discharge(ConstraintSystem &cs, ConstraintLocatorBuilder locator) const;
 };

lib/Sema/CSGen.cpp

@@ -32,6 +32,7 @@
 #include "swift/Sema/ConstraintGraph.h"
 #include "swift/Sema/ConstraintSystem.h"
 #include "swift/Sema/IDETypeChecking.h"
+#include "swift/Sema/PreparedOverload.h"
 #include "swift/Subsystems.h"
 #include "llvm/ADT/APInt.h"
 #include "llvm/ADT/SetVector.h"
@@ -5117,7 +5118,16 @@ bool ConstraintSystem::generateConstraints(StmtCondition condition,
 }
 
 void ConstraintSystem::applyPropertyWrapper(
-    Expr *anchor, AppliedPropertyWrapper applied) {
+    Expr *anchor, AppliedPropertyWrapper applied,
+    PreparedOverload *preparedOverload) {
+  if (preparedOverload) {
+    ASSERT(PreparingOverload);
+    preparedOverload->PropertyWrappers.push_back(applied);
+    return;
+  }
+
+  ASSERT(!PreparingOverload);
+
   appliedPropertyWrappers[anchor].push_back(applied);
 
   if (solverState)
@@ -5138,9 +5148,14 @@ void ConstraintSystem::removePropertyWrapper(Expr *anchor) {
 
 ConstraintSystem::TypeMatchResult
 ConstraintSystem::applyPropertyWrapperToParameter(
-    Type wrapperType, Type paramType, ParamDecl *param, Identifier argLabel,
-    ConstraintKind matchKind, ConstraintLocator *locator,
-    ConstraintLocator *calleeLocator) {
+    Type wrapperType,
+    Type paramType,
+    ParamDecl *param,
+    Identifier argLabel,
+    ConstraintKind matchKind,
+    ConstraintLocator *locator,
+    ConstraintLocator *calleeLocator,
+    PreparedOverload *preparedOverload) {
   Expr *anchor = getAsExpr(calleeLocator->getAnchor());
 
   auto recordPropertyWrapperFix = [&](ConstraintFix *fix) -> TypeMatchResult {
@@ -5149,7 +5164,7 @@ ConstraintSystem::applyPropertyWrapperToParameter(
 
     recordAnyTypeVarAsPotentialHole(paramType);
 
-    if (recordFix(fix))
+    if (recordFix(fix, /*impact=*/1, preparedOverload))
       return getTypeMatchFailure(locator);
 
     return getTypeMatchSuccess();
@@ -5176,21 +5191,33 @@ ConstraintSystem::applyPropertyWrapperToParameter(
 
   if (argLabel.hasDollarPrefix()) {
     Type projectionType = computeProjectedValueType(param, wrapperType);
-    addConstraint(matchKind, paramType, projectionType, locator);
+    addConstraint(matchKind, paramType, projectionType, locator,
+                  /*isFavored=*/false,
+                  preparedOverload);
     if (param->hasImplicitPropertyWrapper()) {
       auto wrappedValueType = getType(param->getPropertyWrapperWrappedValueVar());
-      addConstraint(ConstraintKind::PropertyWrapper, projectionType, wrappedValueType,
-                    getConstraintLocator(param));
+      addConstraint(ConstraintKind::PropertyWrapper,
+                    projectionType, wrappedValueType,
+                    getConstraintLocator(param),
+                    /*isFavored=*/false,
+                    preparedOverload);
       setType(param->getPropertyWrapperProjectionVar(), projectionType);
     }
 
-    applyPropertyWrapper(anchor, { wrapperType, PropertyWrapperInitKind::ProjectedValue });
+    applyPropertyWrapper(anchor,
+                         { wrapperType, PropertyWrapperInitKind::ProjectedValue },
+                         preparedOverload);
   } else if (param->hasExternalPropertyWrapper()) {
     Type wrappedValueType = computeWrappedValueType(param, wrapperType);
-    addConstraint(matchKind, paramType, wrappedValueType, locator);
+    addConstraint(matchKind, paramType, wrappedValueType,
+                  locator,
+                  /*isFavored=*/false,
+                  preparedOverload);
     setType(param->getPropertyWrapperWrappedValueVar(), wrappedValueType);
 
-    applyPropertyWrapper(anchor, { wrapperType, PropertyWrapperInitKind::WrappedValue });
+    applyPropertyWrapper(anchor,
+                         { wrapperType, PropertyWrapperInitKind::WrappedValue },
+                         preparedOverload);
   } else {
     return getTypeMatchFailure(locator);
   }

lib/Sema/CSSimplify.cpp

@@ -15403,7 +15403,16 @@ static bool isAugmentingFix(ConstraintFix *fix) {
   }
 }
 
-bool ConstraintSystem::recordFix(ConstraintFix *fix, unsigned impact) {
+bool ConstraintSystem::recordFix(ConstraintFix *fix, unsigned impact,
+                                 PreparedOverload *preparedOverload) {
+  if (preparedOverload) {
+    ASSERT(PreparingOverload);
+    preparedOverload->Fixes.push_back({fix, impact});
+    return true;
+  }
+
+  ASSERT(!PreparingOverload);
+
   if (isDebugMode()) {
     auto &log = llvm::errs();
     log.indent(solverState ? solverState->getCurrentIndent() : 0)

lib/Sema/ConstraintSystem.cpp

@@ -307,6 +307,8 @@ void ConstraintSystem::removeConversionRestriction(
 }
 
 void ConstraintSystem::addFix(ConstraintFix *fix) {
+  ASSERT(!PreparingOverload);
+
   bool inserted = Fixes.insert(fix);
   ASSERT(inserted);
 

lib/Sema/TypeOfReference.cpp

@@ -804,9 +804,12 @@ unsigned constraints::getNumApplications(bool hasAppliedSelf,
 /// Replaces property wrapper types in the parameter list of the given function type
 /// with the wrapped-value or projected-value types (depending on argument label).
 static FunctionType *
-unwrapPropertyWrapperParameterTypes(ConstraintSystem &cs, AbstractFunctionDecl *funcDecl,
-                                    FunctionRefInfo functionRefInfo, FunctionType *functionType,
-                                    ConstraintLocatorBuilder locator) {
+unwrapPropertyWrapperParameterTypes(ConstraintSystem &cs,
+                                    AbstractFunctionDecl *funcDecl,
+                                    FunctionRefInfo functionRefInfo,
+                                    FunctionType *functionType,
+                                    ConstraintLocatorBuilder locator,
+                                    PreparedOverload *preparedOverload) {
   // Only apply property wrappers to unapplied references to functions.
   if (!functionRefInfo.isUnapplied())
     return functionType;
@@ -842,14 +845,15 @@ unwrapPropertyWrapperParameterTypes(ConstraintSystem &cs, AbstractFunctionDecl *
     }
 
     auto *loc = cs.getConstraintLocator(locator);
-    auto *wrappedType = cs.createTypeVariable(loc, 0);
+    auto *wrappedType = cs.createTypeVariable(loc, 0, preparedOverload);
     auto paramType = paramTypes[i].getParameterType();
     auto paramLabel = paramTypes[i].getLabel();
     auto paramInternalLabel = paramTypes[i].getInternalLabel();
     adjustedParamTypes.push_back(AnyFunctionType::Param(
         wrappedType, paramLabel, ParameterTypeFlags(), paramInternalLabel));
-    cs.applyPropertyWrapperToParameter(paramType, wrappedType, paramDecl, argLabel,
-                                       ConstraintKind::Equal, loc, loc);
+    cs.applyPropertyWrapperToParameter(
+        paramType, wrappedType, paramDecl, argLabel, ConstraintKind::Equal,
+        loc, loc, preparedOverload);
   }
 
   return FunctionType::get(adjustedParamTypes, functionType->getResult(),
@@ -1049,7 +1053,7 @@ ConstraintSystem::getTypeOfReference(ValueDecl *value,
                           ->removeArgumentLabels(numLabelsToRemove);
     openedType = unwrapPropertyWrapperParameterTypes(
         *this, funcDecl, functionRefInfo, openedType->castTo<FunctionType>(),
-        locator);
+        locator, preparedOverload);
 
     auto origOpenedType = openedType;
     if (!isRequirementOrWitness(locator)) {
@@ -1821,8 +1825,9 @@ DeclReferenceType ConstraintSystem::getTypeOfMemberReference(
 
     // Strip off the 'self' parameter
     auto *functionType = fullFunctionType->getResult()->getAs<FunctionType>();
-    functionType = unwrapPropertyWrapperParameterTypes(*this, funcDecl, functionRefInfo,
-                                                       functionType, locator);
+    functionType = unwrapPropertyWrapperParameterTypes(
+        *this, funcDecl, functionRefInfo, functionType,
+        locator, preparedOverload);
     // FIXME: Verify ExtInfo state is correct, not working by accident.
     FunctionType::ExtInfo info;
 
@@ -2494,18 +2499,33 @@ void PreparedOverload::discharge(ConstraintSystem &cs,
   for (auto *tv : TypeVariables) {
     cs.addTypeVariable(tv);
   }
+
   for (auto *c : Constraints) {
     cs.addUnsolvedConstraint(c);
     cs.activateConstraint(c);
   }
+
   cs.recordOpenedTypes(locator, Replacements);
+
   if (OpenedExistential) {
     cs.recordOpenedExistentialType(cs.getConstraintLocator(locator),
                                    OpenedExistential);
   }
+
   for (auto pair : OpenedPackExpansionTypes) {
     cs.recordOpenedPackExpansionType(pair.first, pair.second);
   }
+
+  if (!PropertyWrappers.empty()) {
+    Expr *anchor = getAsExpr(cs.getConstraintLocator(locator)->getAnchor());
+    for (auto applied : PropertyWrappers) {
+      cs.applyPropertyWrapper(anchor, applied);
+    }
+  }
+
+  for (auto pair : Fixes) {
+    cs.recordFix(pair.first, pair.second);
+  }
 }
 
 void ConstraintSystem::resolveOverload(ConstraintLocator *locator,