[CSClosure] Start using recorded targets for ReturnStmts

Multi-statement closures and result builder transformed entities
record a target for `return` statement during constraint generation.

Single-statement closure do not but they'd be handled in a special
way.

Author: xedin

lib/Sema/CSClosure.cpp

@@ -1584,6 +1584,8 @@ class SyntacticElementSolutionApplication
   }
 
   ASTNode visitReturnStmt(ReturnStmt *returnStmt) {
+    auto &cs = solution.getConstraintSystem();
+
     if (!returnStmt->hasResult()) {
       // If contextual is not optional, there is nothing to do here.
       if (resultType->isVoid())
@@ -1598,8 +1600,6 @@ class SyntacticElementSolutionApplication
       assert(resultType->getOptionalObjectType() &&
              resultType->lookThroughAllOptionalTypes()->isVoid());
 
-      auto &cs = solution.getConstraintSystem();
-
       auto target = *cs.getSolutionApplicationTarget(returnStmt);
       returnStmt->setResult(target.getAsExpr());
     }
@@ -1629,13 +1629,24 @@ class SyntacticElementSolutionApplication
       mode = convertToResult;
     }
 
-    SolutionApplicationTarget resultTarget(
-        resultExpr, context.getAsDeclContext(),
-        mode == convertToResult ? CTP_ReturnStmt : CTP_Unused,
-        mode == convertToResult ? resultType : Type(),
-        /*isDiscarded=*/false);
-    if (auto newResultTarget = rewriteTarget(resultTarget))
+    Optional<SolutionApplicationTarget> resultTarget;
+    if (auto target = cs.getSolutionApplicationTarget(returnStmt)) {
+      resultTarget = *target;
+    } else {
+      // Single-expression closures have to handle returns in a special
+      // way so the target has to be created for them during solution
+      // application based on the resolved type.
+      assert(isSingleExpression);
+      resultTarget = SolutionApplicationTarget(
+          resultExpr, context.getAsDeclContext(),
+          mode == convertToResult ? CTP_ReturnStmt : CTP_Unused,
+          mode == convertToResult ? resultType : Type(),
+          /*isDiscarded=*/false);
+    }
+
+    if (auto newResultTarget = rewriteTarget(*resultTarget)) {
       resultExpr = newResultTarget->getAsExpr();
+    }
 
     switch (mode) {
     case convertToResult: