[Constraint system] Adopt rewriteTarget for function builder transform.

When applying a function builder to a closure to produce the final,
type-checked closure, use the new rewriteTarget() so it’s performed on
a per-target basis. Use this to eliminate some duplicating in the handling
of return types.

Author: DougGregor

lib/Sema/BuilderTransform.cpp

@@ -627,8 +627,9 @@ class BuilderClosureRewriter
   const Solution &solution;
   DeclContext *dc;
   AppliedBuilderTransform builderTransform;
-  std::function<Expr *(Expr *)> rewriteExpr;
-  std::function<Expr *(Expr *, Type, ConstraintLocator *)> coerceToType;
+  std::function<
+      Optional<SolutionApplicationTarget> (SolutionApplicationTarget)>
+        rewriteTarget;
 
   /// Retrieve the temporary variable that will be used to capture the
   /// value of the given expression.
@@ -648,6 +649,17 @@ class BuilderClosureRewriter
     return recorded;
   }
 
+  /// Rewrite an expression without any particularly special context.
+  Expr *rewriteExpr(Expr *expr) {
+    auto result = rewriteTarget(
+      SolutionApplicationTarget(expr, dc, CTP_Unused, Type(),
+                                /*isDiscarded=*/false));
+    if (result)
+      return result->getAsExpr();
+
+    return nullptr;
+  }
+
 public:
   /// Retrieve information about a captured statement.
   std::pair<VarDecl *, llvm::TinyPtrVector<Expr *>>
@@ -673,19 +685,21 @@ class BuilderClosureRewriter
   ASTNode initializeTarget(FunctionBuilderTarget target) {
     assert(target.captured.second.size() == 1);
     auto capturedExpr = target.captured.second.front();
-    auto finalCapturedExpr = rewriteExpr(capturedExpr);
     SourceLoc implicitLoc = capturedExpr->getEndLoc();
     switch (target.kind) {
     case FunctionBuilderTarget::ReturnValue: {
       // Return the expression.
-      ConstraintSystem &cs = solution.getConstraintSystem();
       Type bodyResultType =
           solution.simplifyType(builderTransform.bodyResultType);
-      finalCapturedExpr = coerceToType(
-          finalCapturedExpr,
-          bodyResultType,
-          cs.getConstraintLocator(capturedExpr));
-      return new (ctx) ReturnStmt(implicitLoc, finalCapturedExpr);
+
+      SolutionApplicationTarget returnTarget(
+          capturedExpr, dc, CTP_ReturnStmt, bodyResultType,
+          /*isDiscarded=*/false);
+      Expr *resultExpr = nullptr;
+      if (auto resultTarget = rewriteTarget(returnTarget))
+        resultExpr = resultTarget->getAsExpr();
+
+      return new (ctx) ReturnStmt(implicitLoc, resultExpr);
     }
 
     case FunctionBuilderTarget::TemporaryVar: {
@@ -696,6 +710,7 @@ class BuilderClosureRewriter
       declRef->setType(LValueType::get(temporaryVar->getType()));
 
       // Load the right-hand side if needed.
+      auto finalCapturedExpr = rewriteExpr(capturedExpr);
       if (finalCapturedExpr->getType()->hasLValueType()) {
         finalCapturedExpr =
             TypeChecker::addImplicitLoadExpr(ctx, finalCapturedExpr);
@@ -734,12 +749,12 @@ class BuilderClosureRewriter
       const Solution &solution,
       DeclContext *dc,
       const AppliedBuilderTransform &builderTransform,
-      std::function<Expr *(Expr *)> rewriteExpr,
-      std::function<Expr *(Expr *, Type, ConstraintLocator *)> coerceToType
+      std::function<
+          Optional<SolutionApplicationTarget> (SolutionApplicationTarget)>
+            rewriteTarget
     ) : ctx(solution.getConstraintSystem().getASTContext()),
         solution(solution), dc(dc), builderTransform(builderTransform),
-        rewriteExpr(rewriteExpr),
-        coerceToType(coerceToType){ }
+        rewriteTarget(rewriteTarget) { }
 
   Stmt *visitBraceStmt(BraceStmt *braceStmt, FunctionBuilderTarget target,
                        Optional<FunctionBuilderTarget> innerTarget = None) {
@@ -954,9 +969,10 @@ BraceStmt *swift::applyFunctionBuilderTransform(
     AppliedBuilderTransform applied,
     BraceStmt *body,
     DeclContext *dc,
-    std::function<Expr *(Expr *)> rewriteExpr,
-    std::function<Expr *(Expr *, Type, ConstraintLocator *)> coerceToType) {
-  BuilderClosureRewriter rewriter(solution, dc, applied, rewriteExpr, coerceToType);
+    std::function<
+        Optional<SolutionApplicationTarget> (SolutionApplicationTarget)>
+          rewriteTarget) {
+  BuilderClosureRewriter rewriter(solution, dc, applied, rewriteTarget);
   auto captured = rewriter.takeCapturedStmt(body);
   return cast<BraceStmt>(
     rewriter.visitBraceStmt(

lib/Sema/CSApply.cpp

@@ -7063,14 +7063,14 @@ namespace {
           llvm::SaveAndRestore<DeclContext *> savedDC(Rewriter.dc, closure);
           auto newBody = applyFunctionBuilderTransform(
               Rewriter.solution, *transform, closure->getBody(), closure,
-              [&](Expr *expr) {
-                Expr *result = expr->walk(*this);
-                if (result)
-                  Rewriter.solution.setExprTypes(result);
-                return result;
-              },
-              [&](Expr *expr, Type toType, ConstraintLocator *locator) {
-                return Rewriter.coerceToType(expr, toType, locator);
+              [&](SolutionApplicationTarget target) {
+                auto resultTarget = rewriteTarget(target);
+                if (resultTarget) {
+                  if (auto expr = resultTarget->getAsExpr())
+                    Rewriter.solution.setExprTypes(expr);
+                }
+
+                return resultTarget;
               });
           closure->setBody(newBody, /*isSingleExpression=*/false);
           closure->setAppliedFunctionBuilder();
@@ -7340,14 +7340,14 @@ ExprWalker::rewriteTarget(SolutionApplicationTarget target) {
 
     auto newBody = applyFunctionBuilderTransform(
         solution, *transform, fn.getBody(), fn.getAsDeclContext(),
-        [&](Expr *expr) {
-          Expr *result = expr->walk(*this);
-          if (result)
-            solution.setExprTypes(result);
-          return result;
-        },
-        [&](Expr *expr, Type toType, ConstraintLocator *locator) {
-          return Rewriter.coerceToType(expr, toType, locator);
+        [&](SolutionApplicationTarget target) {
+          auto resultTarget = rewriteTarget(target);
+          if (resultTarget) {
+            if (auto expr = resultTarget->getAsExpr())
+              Rewriter.solution.setExprTypes(expr);
+          }
+
+          return resultTarget;
         });
 
     if (!newBody)

lib/Sema/ConstraintSystem.h

@@ -5067,20 +5067,19 @@ bool isSIMDOperator(ValueDecl *value);
 /// \param applied The applied builder transform.
 /// \param body The body to transform
 /// \param dc The context in which the transform occurs.
-/// \param rewriteExpr Rewrites expressions that show up in the transform
-/// to their final, type-checked versions.
-/// \param coerceToType Coerce the given expression to the specified type,
-/// which may introduce implicit conversions.
+/// \param rewriteTarget Rewrites a solution application target to its final,
+/// type-checked version.
 ///
 /// \returns the transformed body
 BraceStmt *applyFunctionBuilderTransform(
     const constraints::Solution &solution,
     constraints::AppliedBuilderTransform applied,
     BraceStmt *body,
     DeclContext *dc,
-    std::function<Expr *(Expr *)> rewriteExpr,
-    std::function<Expr *(Expr *, Type, constraints::ConstraintLocator *)>
-      coerceToType);
+    std::function<
+        Optional<constraints::SolutionApplicationTarget> (
+          constraints::SolutionApplicationTarget)>
+            rewriteTarget);
 
 } // end namespace swift