Move single-expr return transform from Parse to Sema

Rather than doing the transform in the parser, and
then potentially undoing it in Sema, move the
entire transform into Sema. This also lets us
unify the logic between function decls and
closures, and allows ASTGen to benefit from it.

Author: hamishknight

include/swift/AST/TypeCheckRequests.h

@@ -4128,9 +4128,9 @@ class PreCheckReturnStmtRequest
   bool isCached() const { return true; }
 };
 
-/// Performs some pre-checking of a function body, including inserting and
-/// removing implict returns where needed. This request is currently
-/// side-effectful, as it will mutate the body in-place.
+/// Performs some pre-checking of a function body, including inserting implicit
+/// returns where needed. This request is currently side-effectful, as it will
+/// mutate the body in-place.
 ///
 /// Note this request is currently uncached as:
 /// - The result will ultimately be cached by TypeCheckFunctionBodyRequest.
@@ -4152,6 +4152,23 @@ class PreCheckFunctionBodyRequest
   BraceStmt *evaluate(Evaluator &evaluator, AbstractFunctionDecl *AFD) const;
 };
 
+/// Performs some pre-checking of a closure's body, including inserting implicit
+/// returns where needed. This request is currently side-effectful, as it will
+/// mutate the body in-place. Like PreCheckFunctionBodyRequest, it is currently
+/// uncached.
+class PreCheckClosureBodyRequest
+    : public SimpleRequest<PreCheckClosureBodyRequest,
+                           BraceStmt *(ClosureExpr *),
+                           RequestFlags::Uncached> {
+public:
+  using SimpleRequest::SimpleRequest;
+
+private:
+  friend SimpleRequest;
+
+  BraceStmt *evaluate(Evaluator &evaluator, ClosureExpr *closure) const;
+};
+
 /// The result of the query for whether a statement can produce a single value.
 class IsSingleValueStmtResult {
 public:

include/swift/AST/TypeCheckerTypeIDZone.def

@@ -350,6 +350,9 @@ SWIFT_REQUEST(TypeChecker, TangentStoredPropertyRequest,
 SWIFT_REQUEST(TypeChecker, PreCheckFunctionBodyRequest,
               BraceStmt *(AbstractFunctionDecl *),
               Uncached, NoLocationInfo)
+SWIFT_REQUEST(TypeChecker, PreCheckClosureBodyRequest,
+              BraceStmt *(ClosureExpr *),
+              Uncached, NoLocationInfo)
 SWIFT_REQUEST(TypeChecker, TypeCheckFunctionBodyRequest,
               BraceStmt *(AbstractFunctionDecl *), SeparatelyCached,
               NoLocationInfo)

lib/AST/ASTBridging.cpp

@@ -1019,7 +1019,7 @@ BridgedClosureExpr_createParsed(BridgedASTContext cContext,
   auto *out = new (context) ClosureExpr(
       attributes, bracketRange, nullptr, nullptr, asyncLoc, throwsLoc,
       /*FIXME:thrownType=*/nullptr, arrowLoc, inLoc, nullptr, declContext);
-  out->setBody(body.unbridged(), true);
+  out->setBody(body.unbridged(), /*isSingleExpression*/ false);
   out->setParameterList(params);
   return out;
 }

lib/Parse/ParseDecl.cpp

@@ -8840,51 +8840,7 @@ Parser::parseAbstractFunctionBodyImpl(AbstractFunctionDecl *AFD) {
   Fingerprint fp(std::move(currentHash));
 
   BraceStmt *BS = Body.get();
-  // Reset the single expression body status.
-  AFD->setHasSingleExpressionBody(false);
   AFD->setBodyParsed(BS, fp);
-
-  if (auto Element = BS->getSingleActiveElement()) {
-    if (auto *stmt = Element.dyn_cast<Stmt *>()) {
-      if (isa<FuncDecl>(AFD)) {
-        if (auto *returnStmt = dyn_cast<ReturnStmt>(stmt)) {
-          if (!returnStmt->hasResult()) {
-            auto returnExpr = TupleExpr::createEmpty(Context,
-                                                     SourceLoc(),
-                                                     SourceLoc(),
-                                                     /*implicit*/true);
-            returnStmt->setResult(returnExpr);
-            AFD->setHasSingleExpressionBody();
-            AFD->setSingleExpressionBody(returnExpr);
-          }
-        }
-      }
-    } else if (auto *E = Element.dyn_cast<Expr *>()) {
-      if (auto SE = dyn_cast<SequenceExpr>(E->getSemanticsProvidingExpr())) {
-        if (SE->getNumElements() > 1 && isa<AssignExpr>(SE->getElement(1))) {
-          // This is an assignment.  We don't want to implicitly return
-          // it.
-          return {BS, fp};
-        }
-      }
-      if (isa<FuncDecl>(AFD)) {
-        auto RS = ReturnStmt::createImplicit(Context, E);
-        BS->setLastElement(RS);
-        AFD->setHasSingleExpressionBody();
-        AFD->setSingleExpressionBody(E);
-      } else if (auto *F = dyn_cast<ConstructorDecl>(AFD)) {
-        if (F->isFailable() && isa<NilLiteralExpr>(E)) {
-          // If it's a nil literal, just insert return.  This is the only
-          // legal thing to return.
-          auto RS = ReturnStmt::createImplicit(Context, E);
-          BS->setLastElement(RS);
-          AFD->setHasSingleExpressionBody();
-          AFD->setSingleExpressionBody(E);
-        }
-      }
-    }
-  }
-
   return {BS, fp};
 }
 

lib/Parse/ParseExpr.cpp

@@ -3050,35 +3050,10 @@ ParserResult<Expr> Parser::parseExprClosure() {
     closure->setHasAnonymousClosureVars();
   }
 
+  // Set the body of the closure. The computation of the single expression body
+  // is left up to the type-checker.
   auto *BS = BraceStmt::create(Context, leftBrace, bodyElements, rightBrace);
-
-  // If the body consists of a single expression, turn it into a return
-  // statement.
-  bool hasSingleExpressionBody = false;
-  if (!missingRBrace) {
-    if (auto Element = BS->getSingleActiveElement()) {
-      if (Element.is<Stmt *>()) {
-        if (auto returnStmt = dyn_cast<ReturnStmt>(Element.get<Stmt *>())) {
-          hasSingleExpressionBody = true;
-          if (!returnStmt->hasResult()) {
-            auto returnExpr = TupleExpr::createEmpty(Context,
-                                                     SourceLoc(),
-                                                     SourceLoc(),
-                                                     /*implicit*/true);
-            returnStmt->setResult(returnExpr);
-          }
-        }
-      } else if (Element.is<Expr *>()) {
-        // Create the wrapping return.
-        hasSingleExpressionBody = true;
-        auto returnExpr = Element.get<Expr*>();
-        BS->setLastElement(ReturnStmt::createImplicit(Context, returnExpr));
-      }
-    }
-  }
-
-  // Set the body of the closure.
-  closure->setBody(BS, hasSingleExpressionBody);
+  closure->setBody(BS, /*hasSingleExpressionBody*/ false);
 
   // If the closure includes a capture list, create an AST node for it as well.
   Expr *result = closure;

lib/Sema/PreCheckExpr.cpp

@@ -1462,17 +1462,9 @@ namespace {
 /// Perform prechecking of a ClosureExpr before we dive into it.  This returns
 /// true when we want the body to be considered part of this larger expression.
 bool PreCheckExpression::walkToClosureExprPre(ClosureExpr *closure) {
-  // If we have a single statement that can become an expression, turn it
-  // into an expression now.
-  auto *body = closure->getBody();
-  if (auto *S = body->getSingleActiveStatement()) {
-    if (S->mayProduceSingleValue(Ctx)) {
-      auto *SVE = SingleValueStmtExpr::createWithWrappedBranches(
-          Ctx, S, /*DC*/ closure, /*mustBeExpr*/ false);
-      body->setLastElement(ReturnStmt::createImplicit(Ctx, SVE));
-      closure->setBody(body, /*isSingleExpression*/ true);
-    }
-  }
+  // Pre-check the closure body.
+  (void)evaluateOrDefault(Ctx.evaluator, PreCheckClosureBodyRequest{closure},
+                          nullptr);
 
   // Update the current DeclContext to be the closure we're about to
   // recurse into.

lib/Sema/TypeCheckStmt.cpp

@@ -1065,30 +1065,6 @@ class StmtChecker : public StmtVisitor<StmtChecker, Stmt*> {
       return RS;
     }
 
-    // If the body consisted of a single return without a result
-    // 
-    //   func foo() -> Int {
-    //     return
-    //   }
-    // 
-    // in parseAbstractFunctionBody the return is given an empty, implicit tuple
-    // as its result
-    //
-    //   func foo() -> Int {
-    //     return ()
-    //   }
-    //
-    // Look for that case and diagnose it as missing return expression.
-    if (!ResultTy->isVoid() && TheFunc->hasSingleExpressionBody()) {
-      auto expr = TheFunc->getSingleExpressionBody();
-      if (expr->isImplicit() && isa<TupleExpr>(expr) &&
-          cast<TupleExpr>(expr)->getNumElements() == 0) {
-        getASTContext().Diags.diagnose(RS->getReturnLoc(),
-                                       diag::return_expr_missing);
-        return RS;
-      }
-    }
-
     Expr *E = RS->getResult();
     TypeCheckExprOptions options = {};
 
@@ -2677,6 +2653,83 @@ bool TypeCheckASTNodeAtLocRequest::evaluate(
   return false;
 }
 
+/// Insert an implicit return for a single expression body function if needed,
+/// returning \c true if the implicit return was added, \c false otherwise.
+static bool addSingleExprReturnIfNeeded(BraceStmt *body, DeclContext *dc) {
+  // Must have a single active element.
+  auto node = body->getSingleActiveElement();
+  if (!node)
+    return false;
+
+  auto &ctx = dc->getASTContext();
+  auto makeResult = [&](Expr *E) -> bool {
+    body->setLastElement(ReturnStmt::createImplicit(ctx, E));
+    return true;
+  };
+
+  // For a constructor, we only support nil literals as the implicit result.
+  if (auto *ctor = dyn_cast<ConstructorDecl>(dc)) {
+    if (auto *E = node.dyn_cast<Expr *>()) {
+      if (ctor->isFailable() && isa<NilLiteralExpr>(E))
+        return makeResult(E);
+    }
+    return false;
+  }
+
+  // Otherwise, we only support implicit results for FuncDecls and ClosureExprs.
+  if (!isa<FuncDecl>(dc) && !isa<ClosureExpr>(dc))
+    return false;
+
+  if (auto *fd = dyn_cast<FuncDecl>(dc)) {
+    // Don't apply if we have a result builder, or a Void return type.
+    if (getResultBuilderType(fd) || fd->getResultInterfaceType()->isVoid())
+      return false;
+  }
+
+  // A statement can potentially become an expression.
+  if (auto *S = node.dyn_cast<Stmt *>()) {
+    if (S->mayProduceSingleValue(ctx)) {
+      auto *SVE = SingleValueStmtExpr::createWithWrappedBranches(
+          ctx, S, dc, /*mustBeExpr*/ false);
+      return makeResult(SVE);
+    }
+    // If we have a single statement 'return' in a closure, treat it as a
+    // single expression body. If we don't have a return value, synthesize
+    // 'return ()'. Don't do this in a FuncDecl since we know we have a
+    // non-Void return.
+    if (isa<ClosureExpr>(dc)) {
+      if (auto *returnStmt = dyn_cast<ReturnStmt>(S)) {
+        // Treat the existing result as a single expression return.
+        if (returnStmt->hasResult())
+          return true;
+
+        auto *returnExpr = TupleExpr::createEmpty(ctx, SourceLoc(), SourceLoc(),
+                                                  /*implicit*/ true);
+        returnStmt->setResult(returnExpr);
+        return true;
+      }
+    }
+  }
+
+  if (auto *E = node.dyn_cast<Expr *>()) {
+    // Take any expression, except for assignments. This helps improves
+    // diagnostics.
+    // TODO: We probably ought to apply this to closures too, but that currently
+    // regresses a couple of diagnostics.
+    if (!isa<ClosureExpr>(dc)) {
+      auto *SemanticExpr = E->getSemanticsProvidingExpr();
+      if (auto *SE = dyn_cast<SequenceExpr>(SemanticExpr)) {
+        if (SE->getNumElements() > 1 && isa<AssignExpr>(SE->getElement(1)))
+          return false;
+      }
+      if (isa<AssignExpr>(SemanticExpr))
+        return false;
+    }
+    return makeResult(E);
+  }
+  return false;
+}
+
 BraceStmt *
 PreCheckFunctionBodyRequest::evaluate(Evaluator &evaluator,
                                       AbstractFunctionDecl *AFD) const {
@@ -2687,38 +2740,15 @@ PreCheckFunctionBodyRequest::evaluate(Evaluator &evaluator,
   assert(body && "Expected body");
   assert(!AFD->isBodyTypeChecked() && "Body already type-checked?");
 
-  if (auto *func = dyn_cast<FuncDecl>(AFD)) {
-    // Don't apply this pre-checking to functions with result builders.
-    if (getResultBuilderType(func))
-      return body;
-
-    if (func->hasSingleExpressionBody() &&
-        func->getResultInterfaceType()->isVoid()) {
-      // The function returns void.  We don't need an explicit return, no
-      // matter what the type of the expression is. Take the inserted return
-      // back out.
-      body->setLastElement(func->getSingleExpressionBody());
-    }
-    // If there is a single statement in the body that can be turned into a
-    // single expression return, do so now.
-    if (!func->getResultInterfaceType()->isVoid()) {
-      if (auto *S = body->getSingleActiveStatement()) {
-        if (S->mayProduceSingleValue(ctx)) {
-          auto *SVE = SingleValueStmtExpr::createWithWrappedBranches(
-              ctx, S, /*DC*/ func, /*mustBeExpr*/ false);
-          body->setLastElement(ReturnStmt::createImplicit(ctx, SVE));
-          func->setHasSingleExpressionBody();
-          func->setSingleExpressionBody(SVE);
-        }
-      }
-    }
-  }
+  // Insert an implicit return for a single expression body.
+  if (addSingleExprReturnIfNeeded(body, AFD))
+    AFD->setHasSingleExpressionBody();
 
+  // For constructors, we make sure that the body ends with a "return"
+  // stmt, which we either implicitly synthesize, or the user can write.
+  // This simplifies SILGen.
   if (auto *ctor = dyn_cast<ConstructorDecl>(AFD)) {
     if (body->empty() || !isKnownEndOfConstructor(body->getLastElement())) {
-      // For constructors, we make sure that the body ends with a "return"
-      // stmt, which we either implicitly synthesize, or the user can write.
-      // This simplifies SILGen.
       SmallVector<ASTNode, 8> Elts(body->getElements().begin(),
                                    body->getElements().end());
       Elts.push_back(ReturnStmt::createImplicit(ctx, body->getRBraceLoc(),
@@ -2730,6 +2760,17 @@ PreCheckFunctionBodyRequest::evaluate(Evaluator &evaluator,
   return body;
 }
 
+BraceStmt *PreCheckClosureBodyRequest::evaluate(Evaluator &evaluator,
+                                                ClosureExpr *closure) const {
+  auto *body = closure->getBody();
+
+  // Insert an implicit return for a single expression body.
+  if (addSingleExprReturnIfNeeded(body, closure))
+    closure->setBody(body, /*isSingleExpression*/ true);
+
+  return body;
+}
+
 /// Determine whether the given declaration requires a definition.
 ///
 /// Only valid for declarations that can have definitions, i.e.,

test/SourceKit/CodeExpand/code-expand.swift

@@ -146,6 +146,19 @@ braced2(x: {<#T##() -> Void#>}, y: Int)
 // CHECK-NEXT:  <#code#>
 // CHECK-NEXT:  }, y: Int)
 
+braced3({
+  #if true
+  <#T##() -> Int#>
+  #endif
+})
+// CHECK:      braced3({
+// CHECK-NEXT:   #if true
+// CHECK-NEXT:    {
+// CHECK-NEXT:      <#code#>
+// CHECK-NEXT:    }
+// CHECK-NEXT:    #endif
+// CHECK-NEXT:  })
+
 func returnTrailing() -> Int {
   return withtrail(<#T##() -> ()#>)
 // CHECK: return withtrail {