[Constraint solver] Fix function builders with single-expression closures

Fix a few related issues involving the interaction with
single-expression closures:

* A single-expression closure can have a "return" in it; in such
cases, disable the function-builder transform.
* Have the function builder constraint generator look through the
"return" statement in a single-expression closure the same way as
solution application does

Fixes rdar://problem/59045763, where we rejected some well-formed code
involving a single-expression closure with a "return" keyword.

Author: DougGregor

lib/Sema/BuilderTransform.cpp

@@ -244,7 +244,15 @@ class BuilderClosureVisitor
       expressions.push_back(buildVarRef(childVar, childVar->getLoc()));
     };
 
-    for (const auto &node : braceStmt->getElements()) {
+    for (auto node : braceStmt->getElements()) {
+      // Implicit returns in single-expression function bodies are treated
+      // as the expression.
+      if (auto returnStmt =
+              dyn_cast_or_null<ReturnStmt>(node.dyn_cast<Stmt *>())) {
+        assert(returnStmt->isImplicit());
+        node = returnStmt->getResult();
+      }
+
       if (auto stmt = node.dyn_cast<Stmt *>()) {
         addChild(visit(stmt));
         continue;
@@ -291,14 +299,9 @@ class BuilderClosureVisitor
   }
 
   VarDecl *visitReturnStmt(ReturnStmt *stmt) {
-    // Allow implicit returns due to 'return' elision.
-    if (!stmt->isImplicit() || !stmt->hasResult()) {
-      if (!unhandledNode)
-        unhandledNode = stmt;
-      return nullptr;
-    }
-
-    return captureExpr(stmt->getResult(), /*oneWay=*/true);
+    if (!unhandledNode)
+      unhandledNode = stmt;
+    return nullptr;
   }
 
   VarDecl *visitDoStmt(DoStmt *doStmt) {
@@ -1117,7 +1120,8 @@ Optional<BraceStmt *> TypeChecker::applyFunctionBuilderBodyTransform(
   return nullptr;
 }
 
-ConstraintSystem::TypeMatchResult ConstraintSystem::matchFunctionBuilder(
+Optional<ConstraintSystem::TypeMatchResult>
+ConstraintSystem::matchFunctionBuilder(
     AnyFunctionRef fn, Type builderType, Type bodyResultType,
     ConstraintKind bodyResultConstraintKind,
     ConstraintLocator *calleeLocator, ConstraintLocatorBuilder locator) {
@@ -1141,7 +1145,7 @@ ConstraintSystem::TypeMatchResult ConstraintSystem::matchFunctionBuilder(
   case FunctionBuilderBodyPreCheck::HasReturnStmt:
     // If the body has a return statement, suppress the transform but
     // continue solving the constraint system.
-    return getTypeMatchSuccess();
+    return None;
   }
 
   // Check the form of this body to see if we can apply the
@@ -1287,12 +1291,6 @@ class PreCheckFunctionBuilderApplication : public ASTWalker {
 llvm::Expected<FunctionBuilderBodyPreCheck>
 PreCheckFunctionBuilderRequest::evaluate(Evaluator &eval,
                                          AnyFunctionRef fn) const {
-  // Single-expression closures should already have been pre-checked.
-  if (auto closure = fn.getAbstractClosureExpr()) {
-    if (closure->hasSingleExpressionBody())
-      return FunctionBuilderBodyPreCheck::Okay;
-  }
-
   return PreCheckFunctionBuilderApplication(fn, false).run();
 }
 

lib/Sema/CSSimplify.cpp

@@ -6635,10 +6635,11 @@ bool ConstraintSystem::resolveClosure(TypeVariableType *typeVar,
       auto *calleeLocator = getCalleeLocator(getConstraintLocator(locator));
       if (auto functionBuilderType = getFunctionBuilderTypeFor(
               *this, argToParam->getParamIdx(), calleeLocator)) {
-        auto result = matchFunctionBuilder(
-            closure, functionBuilderType, closureType->getResult(),
-            ConstraintKind::Conversion, calleeLocator, locator);
-        return result.isSuccess();
+        if (auto result = matchFunctionBuilder(
+                closure, functionBuilderType, closureType->getResult(),
+                ConstraintKind::Conversion, calleeLocator, locator)) {
+          return result->isSuccess();
+        }
       }
     }
   }

lib/Sema/ConstraintSystem.h

@@ -3726,7 +3726,10 @@ class ConstraintSystem {
   void simplifyDisjunctionChoice(Constraint *choice);
 
   /// Apply the given function builder to the closure expression.
-  TypeMatchResult matchFunctionBuilder(
+  ///
+  /// \returns \c None when the function builder cannot be applied at all,
+  /// otherwise the result of applying the function builder.
+  Optional<TypeMatchResult> matchFunctionBuilder(
       AnyFunctionRef fn, Type builderType, Type bodyResultType,
       ConstraintKind bodyResultConstraintKind,
       ConstraintLocator *calleeLocator, ConstraintLocatorBuilder locator);

test/Constraints/function_builder_diags.swift

@@ -284,3 +284,18 @@ func checkConditions(cond: Bool) {
     }
   }
 }
+
+// Check that a closure with a single "return" works with function builders.
+func checkSingleReturn(cond: Bool) {
+  tuplify(cond) { value in
+    return (value, 17)
+  }
+
+  tuplify(cond) { value in
+    (value, 17)
+  }
+
+  tuplify(cond) {
+    ($0, 17)
+  }
+}