Allow implicit last expressions for functions and closures

Gated behind the experimental feature
`ImplicitLastExprResults`.

Author: hamishknight

include/swift/AST/Stmt.h

@@ -269,8 +269,8 @@ class ReturnStmt : public Stmt {
     return createImplicit(ctx, SourceLoc(), result);
   }
 
-  /// Create an implicit implied ReturnStmt for a single expression body.
-  static ReturnStmt *forSingleExprBody(ASTContext &ctx, Expr *result) {
+  /// Create an implicit implied ReturnStmt for e.g a single expression body.
+  static ReturnStmt *createImplied(ASTContext &ctx, Expr *result) {
     assert(result && "Result must be present to be implied");
     auto *RS = createImplicit(ctx, result);
     RS->Bits.ReturnStmt.IsImplied = true;

include/swift/Basic/Features.def

@@ -231,8 +231,8 @@ EXPERIMENTAL_FEATURE(ThenStatements, false)
 /// Enable 'do' expressions.
 EXPERIMENTAL_FEATURE(DoExpressions, false)
 
-/// Enable implicitly treating the last expression in an 'if'/'switch'
-/// expression as the result
+/// Enable implicitly treating the last expression in a function, closure,
+/// and 'if'/'switch' expression as the result.
 EXPERIMENTAL_FEATURE(ImplicitLastExprResults, false)
 
 /// Enable the `@_rawLayout` attribute.

lib/Parse/ParseDecl.cpp

@@ -8862,7 +8862,7 @@ Parser::parseAbstractFunctionBodyImpl(AbstractFunctionDecl *AFD) {
       SourceLoc LBraceLoc, RBraceLoc;
       LBraceLoc = consumeToken(tok::l_brace);
       auto *CCE = new (Context) CodeCompletionExpr(Tok.getLoc());
-      auto *Return = ReturnStmt::forSingleExprBody(Context, CCE);
+      auto *Return = ReturnStmt::createImplied(Context, CCE);
       CodeCompletionCallbacks->setParsedDecl(accessor);
       CodeCompletionCallbacks->completeAccessorBeginning(CCE);
       RBraceLoc = Tok.getLoc();

lib/Sema/TypeCheckStmt.cpp

@@ -2662,15 +2662,23 @@ bool TypeCheckASTNodeAtLocRequest::evaluate(
 }
 
 /// Insert an implicit return for a single expression body function if needed.
-static void addSingleExprReturnIfNeeded(BraceStmt *body, DeclContext *dc) {
-  // Must have a single active element.
-  auto node = body->getSingleActiveElement();
-  if (!node)
+static void addImplicitReturnIfNeeded(BraceStmt *body, DeclContext *dc) {
+  if (body->empty())
     return;
 
+  // Must have a single active element (which is guarenteed to be the last
+  // element), or we must be allowing implicit last expression results.
   auto &ctx = dc->getASTContext();
+  if (!body->getSingleActiveElement() &&
+      !ctx.LangOpts.hasFeature(Feature::ImplicitLastExprResults)) {
+    return;
+  }
+  auto node = body->getLastElement();
+  if (!node)
+    return;
+
   auto makeResult = [&](Expr *E) {
-    body->setLastElement(ReturnStmt::forSingleExprBody(ctx, E));
+    body->setLastElement(ReturnStmt::createImplied(ctx, E));
   };
 
   // For a constructor, we only support nil literals as the implicit result.
@@ -2730,8 +2738,8 @@ PreCheckFunctionBodyRequest::evaluate(Evaluator &evaluator,
   assert(body && "Expected body");
   assert(!AFD->isBodyTypeChecked() && "Body already type-checked?");
 
-  // Insert an implicit return for a single expression body.
-  addSingleExprReturnIfNeeded(body, AFD);
+  // Insert an implicit return if needed.
+  addImplicitReturnIfNeeded(body, AFD);
 
   // For constructors, we make sure that the body ends with a "return"
   // stmt, which we either implicitly synthesize, or the user can write.
@@ -2765,8 +2773,8 @@ BraceStmt *PreCheckClosureBodyRequest::evaluate(Evaluator &evaluator,
       }
     }
   }
-  // Insert an implicit return for a single expression body.
-  addSingleExprReturnIfNeeded(body, closure);
+  // Insert an implicit return if needed.
+  addImplicitReturnIfNeeded(body, closure);
   return body;
 }
 

test/Constraints/if_switch_implicit_last_expr.swift

@@ -0,0 +1,64 @@
+// RUN: %target-typecheck-verify-swift -enable-experimental-feature ImplicitLastExprResults %clang-importer-sdk
+
+// REQUIRES: objc_interop
+
+// Experimental feature requires asserts
+// REQUIRES: asserts
+
+import Foundation
+
+func testReturn1(_ d: Double) -> CGFloat {
+  print("hello")
+  d
+}
+
+func testReturn2(_ d: Double) -> CGFloat {
+  ()
+  if .random() { d } else { 0 }
+}
+
+func testReturn3(_ d: CGFloat) -> Double {
+  print("hello")
+  switch Bool.random() { case true: d case false: 0 }
+}
+
+func testReturn4(_ d: Double) -> CGFloat {
+  print("hello")
+  if .random() {
+    print("hello")
+    d
+  } else {
+    if .random() {
+      print("hello")
+      d
+    } else {
+      0
+    }
+  }
+}
+
+func testClosure(_ d: CGFloat) {
+  func fn(_: () -> Double) {}
+  fn {
+    print("hello")
+    d
+  }
+  fn {
+    print("hello")
+    if .random() { 0.0 } else { d }
+  }
+  fn {
+    print("hello")
+    if .random() {
+      print("hello")
+      d
+    } else {
+      if .random() {
+        print("hello")
+        d
+      } else {
+        0
+      }
+    }
+  }
+}