[Constraint system] Sink initialization logic into SolutionApplicationTarget

lib/Sema/CSGen.cpp

@@ -3782,6 +3782,54 @@ static Expr *generateConstraintsFor(ConstraintSystem &cs, Expr *expr,
   return result;
 }
 
+bool ConstraintSystem::generateConstraints(
+    SolutionApplicationTarget &target,
+    FreeTypeVariableBinding allowFreeTypeVariables) {
+  if (Expr *expr = target.getAsExpr()) {
+    // Try to shrink the system by reducing disjunction domains. This
+    // goes through every sub-expression and generate its own sub-system, to
+    // try to reduce the domains of those subexpressions.
+    shrink(expr);
+    target.setExpr(expr);
+
+    // Generate constraints for the main system.
+    expr = generateConstraints(expr, target.getDeclContext());
+    if (!expr)
+      return true;
+    target.setExpr(expr);
+
+    // If there is a type that we're expected to convert to, add the conversion
+    // constraint.
+    if (Type convertType = target.getExprConversionType()) {
+      // Determine whether we know more about the contextual type.
+      ContextualTypePurpose ctp = target.getExprContextualTypePurpose();
+      bool isOpaqueReturnType = target.infersOpaqueReturnType();
+
+      // Substitute type variables in for unresolved types.
+      if (allowFreeTypeVariables == FreeTypeVariableBinding::UnresolvedType) {
+        bool isForSingleExprFunction = (ctp == CTP_ReturnSingleExpr);
+        auto *convertTypeLocator = getConstraintLocator(
+            expr, LocatorPathElt::ContextualType(isForSingleExprFunction));
+
+        convertType = convertType.transform([&](Type type) -> Type {
+          if (type->is<UnresolvedType>()) {
+            return createTypeVariable(
+                convertTypeLocator, TVO_CanBindToNoEscape);
+          }
+          return type;
+        });
+      }
+
+      addContextualConversionConstraint(expr, convertType, ctp,
+                                        isOpaqueReturnType);
+    }
+
+    return false;
+  }
+
+  llvm_unreachable("BOOM");
+}
+
 Expr *ConstraintSystem::generateConstraints(ClosureExpr *closure) {
   assert(closure->hasSingleExpressionBody());
   return generateConstraintsFor(*this, closure->getSingleExpressionBody(),

lib/Sema/CSSolver.cpp

@@ -1239,46 +1239,14 @@ ConstraintSystem::solveImpl(SolutionApplicationTarget &target,
   assert(!solverState && "cannot be used directly");
 
   // Set up the expression type checker timer.
-  Expr *expr = target.getAsExpr();
-  Timer.emplace(expr, *this);
-
-  // Try to shrink the system by reducing disjunction domains. This
-  // goes through every sub-expression and generate its own sub-system, to
-  // try to reduce the domains of those subexpressions.
-  shrink(expr);
-
-  // Generate constraints for the main system.
-  if (auto generatedExpr = generateConstraints(expr, DC))
-    expr = generatedExpr;
-  else {
-    return SolutionResult::forError();
-  }
+  if (Expr *expr = target.getAsExpr())
+    Timer.emplace(expr, *this);
 
-  // If there is a type that we're expected to convert to, add the conversion
-  // constraint.
-  if (Type convertType = target.getExprConversionType()) {
-    // Determine whether we know more about the contextual type.
-    ContextualTypePurpose ctp = target.getExprContextualTypePurpose();
-    bool isOpaqueReturnType = target.infersOpaqueReturnType();
-
-    // Substitute type variables in for unresolved types.
-    if (allowFreeTypeVariables == FreeTypeVariableBinding::UnresolvedType) {
-      bool isForSingleExprFunction = (ctp == CTP_ReturnSingleExpr);
-      auto *convertTypeLocator = getConstraintLocator(
-          expr, LocatorPathElt::ContextualType(isForSingleExprFunction));
-
-      convertType = convertType.transform([&](Type type) -> Type {
-        if (type->is<UnresolvedType>())
-          return createTypeVariable(convertTypeLocator, TVO_CanBindToNoEscape);
-        return type;
-      });
-    }
-
-    addContextualConversionConstraint(expr, convertType, ctp,
-                                      isOpaqueReturnType);
-  }
+  if (generateConstraints(target, allowFreeTypeVariables))
+    return SolutionResult::forError();;
 
   // Notify the listener that we've built the constraint system.
+  Expr *expr = target.getAsExpr();
   if (listener && listener->builtConstraints(*this, expr)) {
     return SolutionResult::forError();
   }

lib/Sema/ConstraintSystem.cpp

@@ -3994,9 +3994,62 @@ SolutionApplicationTarget::SolutionApplicationTarget(
   expression.dc = dc;
   expression.contextualPurpose = contextualPurpose;
   expression.convertType = convertType;
+  expression.pattern = nullptr;
+  expression.wrappedVar = nullptr;
   expression.isDiscarded = isDiscarded;
 }
 
+void SolutionApplicationTarget::maybeApplyPropertyWrapper() {
+  assert(kind == Kind::expression);
+  assert(expression.contextualPurpose == CTP_Initialization);
+  auto singleVar = expression.pattern->getSingleVar();
+  if (!singleVar)
+    return;
+
+  auto wrapperAttrs = singleVar->getAttachedPropertyWrappers();
+  if (wrapperAttrs.empty())
+    return;
+
+  // If the outermost property wrapper is directly initialized, form the
+  // call.
+  auto &ctx = singleVar->getASTContext();
+  auto outermostWrapperAttr = wrapperAttrs.front();
+  Expr *backingInitializer;
+  if (Expr *initializer = expression.expression) {
+    // Form init(wrappedValue:) call(s).
+    Expr *wrappedInitializer =
+        buildPropertyWrapperInitialValueCall(
+            singleVar, Type(), initializer, /*ignoreAttributeArgs=*/false);
+    if (!wrappedInitializer)
+      return;
+
+    backingInitializer = wrappedInitializer;
+  } else if (auto outermostArg = outermostWrapperAttr->getArg()) {
+    Type outermostWrapperType =
+        singleVar->getAttachedPropertyWrapperType(0);
+    if (!outermostWrapperType)
+      return;
+
+    auto typeExpr = TypeExpr::createImplicitHack(
+        outermostWrapperAttr->getTypeLoc().getLoc(),
+        outermostWrapperType, ctx);
+    backingInitializer = CallExpr::create(
+        ctx, typeExpr, outermostArg,
+        outermostWrapperAttr->getArgumentLabels(),
+        outermostWrapperAttr->getArgumentLabelLocs(),
+        /*hasTrailingClosure=*/false,
+        /*implicit=*/false);
+  } else {
+    llvm_unreachable("No initializer anywhere?");
+  }
+  wrapperAttrs[0]->setSemanticInit(backingInitializer);
+
+  // Note that we have applied to property wrapper, so we can adjust
+  // the initializer type later.
+  expression.wrappedVar = singleVar;
+  expression.expression = backingInitializer;
+}
+
 SolutionApplicationTarget SolutionApplicationTarget::forInitialization(
     Expr *initializer, DeclContext *dc, Type patternType, Pattern *pattern) {
   // Determine the contextual type for the initialization.
@@ -4020,6 +4073,7 @@ SolutionApplicationTarget SolutionApplicationTarget::forInitialization(
       initializer, dc, CTP_Initialization, contextualType,
       /*isDiscarded=*/false);
   target.expression.pattern = pattern;
+  target.maybeApplyPropertyWrapper();
   return target;
 }
 

lib/Sema/ConstraintSystem.h

@@ -774,8 +774,8 @@ struct Score {
 
 /// An AST node that can gain type information while solving.
 using TypedNode =
-    llvm::PointerUnion3<const Expr *, const TypeLoc *,
-                        const VarDecl *>;
+    llvm::PointerUnion4<const Expr *, const TypeLoc *,
+                        const VarDecl *, const Pattern *>;
 
 /// Display a score.
 llvm::raw_ostream &operator<<(llvm::raw_ostream &out, const Score &score);
@@ -1161,7 +1161,11 @@ class SolutionApplicationTarget {
 
       /// When initializing a pattern from the expression, this is the
       /// pattern.
-      Pattern *pattern = nullptr;
+      Pattern *pattern;
+
+      /// The variable to which property wrappers have been applied, if
+      /// this is an initialization involving a property wrapper.
+      VarDecl *wrappedVar;
 
       /// Whether the expression result will be discarded at the end.
       bool isDiscarded;
@@ -1173,6 +1177,11 @@ class SolutionApplicationTarget {
     } function;
   };
 
+  // If the pattern contains a single variable that has an attached
+  // property wrapper, set up the initializer expression to initialize
+  // the backing storage.
+  void maybeApplyPropertyWrapper();
+
 public:
   SolutionApplicationTarget(Expr *expr, DeclContext *dc,
                             ContextualTypePurpose contextualPurpose,
@@ -1281,6 +1290,14 @@ class SolutionApplicationTarget {
         isa<OptionalSomePattern>(expression.pattern);
   }
 
+  /// Retrieve the wrapped variable when initializing a pattern with a
+  /// property wrapper.
+  VarDecl *getInitializationWrappedVar() const {
+    assert(kind == Kind::expression);
+    assert(expression.contextualPurpose == CTP_Initialization);
+    return expression.wrappedVar;
+  }
+
   /// Whether this context infers an opaque return type.
   bool infersOpaqueReturnType() const;
 
@@ -3231,6 +3248,12 @@ class ConstraintSystem {
     return allocateCopy(vec.begin(), vec.end());
   }
 
+  /// Generate constraints for the given solution target.
+  ///
+  /// \returns true if an error occurred, false otherwise.
+  bool generateConstraints(SolutionApplicationTarget &target,
+                           FreeTypeVariableBinding allowFreeTypeVariables);
+
   /// Generate constraints for the body of the given single-statement closure.
   ///
   /// \returns a possibly-sanitized expression, or null if an error occurred.

lib/Sema/TypeCheckConstraints.cpp

@@ -2474,8 +2474,8 @@ bool TypeChecker::typeCheckBinding(Pattern *&pattern, Expr *&initializer,
   /// Type checking listener for pattern binding initializers.
   class BindingListener : public ExprTypeCheckListener {
     ASTContext &context;
-    Pattern *&pattern;
-    Expr *&initializer;
+
+    SolutionApplicationTarget target;
 
     /// The locator we're using.
     ConstraintLocator *Locator;
@@ -2487,11 +2487,9 @@ bool TypeChecker::typeCheckBinding(Pattern *&pattern, Expr *&initializer,
     VarDecl *wrappedVar = nullptr;
 
   public:
-    explicit BindingListener(ASTContext &ctx, Pattern *&pattern,
-                             Expr *&initializer)
-        : context(ctx), pattern(pattern), initializer(initializer),
-          Locator(nullptr) {
-      maybeApplyPropertyWrapper();
+    explicit BindingListener(ASTContext &ctx, SolutionApplicationTarget target)
+        : context(ctx), target(target), Locator(nullptr) {
+      wrappedVar = target.getInitializationWrappedVar();
     }
 
     /// Retrieve the type to which the pattern should be coerced.
@@ -2505,7 +2503,7 @@ bool TypeChecker::typeCheckBinding(Pattern *&pattern, Expr *&initializer,
       if (cs) {
         Type valueType = LValueType::get(initType);
         auto dc = wrappedVar->getInnermostDeclContext();
-        auto *loc = cs->getConstraintLocator(initializer);
+        auto *loc = cs->getConstraintLocator(target.getAsExpr());
 
         for (unsigned i : indices(wrappedVar->getAttachedPropertyWrappers())) {
           auto wrapperInfo = wrappedVar->getAttachedPropertyWrapperTypeInfo(i);
@@ -2538,7 +2536,8 @@ bool TypeChecker::typeCheckBinding(Pattern *&pattern, Expr *&initializer,
       Locator = cs.getConstraintLocator(expr, LocatorPathElt::ContextualType());
 
       // Collect constraints from the pattern.
-      Type patternType = cs.generateConstraints(pattern, Locator);
+      Type patternType =
+          cs.generateConstraints(target.getInitializationPattern(), Locator);
       if (!patternType)
         return true;
 
@@ -2563,7 +2562,6 @@ bool TypeChecker::typeCheckBinding(Pattern *&pattern, Expr *&initializer,
       }
 
       // The expression has been pre-checked; save it in case we fail later.
-      initializer = expr;
       return false;
     }
 
@@ -2593,70 +2591,20 @@ bool TypeChecker::typeCheckBinding(Pattern *&pattern, Expr *&initializer,
             ->setSemanticInit(expr);
       }
 
-      initializer = expr;
       return expr;
     }
-
-  private:
-    // If the pattern contains a single variable that has an attached
-    // property wrapper, set up the initializer expression to initialize
-    // the backing storage.
-    void maybeApplyPropertyWrapper() {
-      auto singleVar = pattern->getSingleVar();
-      if (!singleVar)
-        return;
-
-      auto wrapperAttrs = singleVar->getAttachedPropertyWrappers();
-      if (wrapperAttrs.empty())
-        return;
-
-      // If the outermost property wrapper is directly initialized, form the
-      // call.
-      auto &ctx = singleVar->getASTContext();
-      auto outermostWrapperAttr = wrapperAttrs.front();
-      if (initializer) {
-        // Form init(wrappedValue:) call(s).
-        Expr *wrappedInitializer =
-            buildPropertyWrapperInitialValueCall(
-                singleVar, Type(), initializer, /*ignoreAttributeArgs=*/false);
-        if (!wrappedInitializer)
-          return;
-
-        initializer = wrappedInitializer;
-      } else if (auto outermostArg = outermostWrapperAttr->getArg()) {
-        Type outermostWrapperType =
-            singleVar->getAttachedPropertyWrapperType(0);
-        if (!outermostWrapperType)
-          return;
-
-        auto typeExpr = TypeExpr::createImplicitHack(
-            outermostWrapperAttr->getTypeLoc().getLoc(),
-            outermostWrapperType, ctx);
-        initializer = CallExpr::create(
-            ctx, typeExpr, outermostArg,
-            outermostWrapperAttr->getArgumentLabels(),
-            outermostWrapperAttr->getArgumentLabelLocs(),
-            /*hasTrailingClosure=*/false,
-            /*implicit=*/false);
-      } else {
-        llvm_unreachable("No initializer anywhere?");
-      }
-      wrapperAttrs[0]->setSemanticInit(initializer);
-
-      // Note that we have applied to property wrapper, so we can adjust
-      // the initializer type later.
-      wrappedVar = singleVar;
-    }
   };
 
   auto &Context = DC->getASTContext();
-  BindingListener listener(Context, pattern, initializer);
+  auto target = SolutionApplicationTarget::forInitialization(
+      initializer, DC, patternType, pattern);
+  initializer = target.getAsExpr();
+
+  BindingListener listener(Context, target);
   if (!initializer)
     return true;
 
   // Type-check the initializer.
-  auto target = SolutionApplicationTarget::forInitialization(
-      initializer, DC, patternType, pattern);
   bool unresolvedTypeExprs = false;
   auto resultTarget = typeCheckExpression(target, unresolvedTypeExprs,
                                           None, &listener);