Handle #if for if/switch expressions

include/swift/AST/DiagnosticsSema.def

@@ -1143,6 +1143,9 @@ ERROR(single_value_stmt_must_be_unlabeled,none,
 ERROR(if_expr_must_be_syntactically_exhaustive,none,
       "'if' must have an unconditional 'else' to be used as expression",
       ())
+ERROR(single_value_stmt_branch_empty,none,
+      "expected expression in branch of '%0' expression",
+      (StmtKind))
 ERROR(single_value_stmt_branch_must_end_in_throw,none,
       "non-expression branch of '%0' expression may only end with a 'throw'",
       (StmtKind))

include/swift/AST/Stmt.h

@@ -212,9 +212,20 @@ class BraceStmt final : public Stmt,
 
   ASTNode findAsyncNode();
 
-  /// If this brace is wrapping a single expression, returns it. Otherwise
-  /// returns \c nullptr.
-  Expr *getSingleExpressionElement() const;
+  /// If this brace contains a single ASTNode, or a \c #if that has a single active
+  /// element, returns it. This will always be the last element of the brace.
+  /// Otherwise returns \c nullptr.
+  ASTNode getSingleActiveElement() const;
+
+  /// If this brace is wrapping a single active expression, returns it. This
+  /// includes both a single expression element, or a single expression in an
+  /// active \c #if. Otherwise returns \c nullptr.
+  Expr *getSingleActiveExpression() const;
+
+  /// If this brace is wrapping a single active statement, returns it. This
+  /// includes both a single statement element, or a single statement in an
+  /// active \c #if. Otherwise returns \c nullptr.
+  Stmt *getSingleActiveStatement() const;
 
   static bool classof(const Stmt *S) { return S->getKind() == StmtKind::Brace; }
 };

include/swift/Parse/Parser.h

@@ -691,11 +691,6 @@ class Parser {
   /// \param DiagText name for the string literal in the diagnostic.
   Optional<StringRef>
   getStringLiteralIfNotInterpolated(SourceLoc Loc, StringRef DiagText);
-
-  /// Returns true when body elements are eligible as single-expression implicit returns.
-  ///
-  /// \param Body elements to search for implicit single-expression returns.
-  bool shouldReturnSingleExpressionElement(ArrayRef<ASTNode> Body); 
 
   /// Returns true to indicate that experimental concurrency syntax should be
   /// parsed if the parser is generating only a syntax tree or if the user has

lib/AST/Expr.cpp

@@ -2524,11 +2524,7 @@ SingleValueStmtExpr *SingleValueStmtExpr::createWithWrappedBranches(
     if (!BS)
       continue;
 
-    auto elts = BS->getElements();
-    if (elts.size() != 1)
-      continue;
-
-    auto *S = elts.front().dyn_cast<Stmt *>();
+    auto *S = BS->getSingleActiveStatement();
     if (!S)
       continue;
 
@@ -2615,7 +2611,7 @@ ArrayRef<Expr *> SingleValueStmtExpr::getSingleExprBranches(
     auto *BS = dyn_cast<BraceStmt>(branch);
     if (!BS)
       continue;
-    if (auto *E = BS->getSingleExpressionElement())
+    if (auto *E = BS->getSingleActiveExpression())
       scratch.push_back(E);
   }
   return scratch;

lib/AST/Stmt.cpp

@@ -299,11 +299,37 @@ ASTNode BraceStmt::findAsyncNode() {
   return asyncFinder.getAsyncNode();
 }
 
-Expr *BraceStmt::getSingleExpressionElement() const {
-  if (getElements().size() != 1)
-    return nullptr;
+static bool hasSingleActiveElement(ArrayRef<ASTNode> elts) {
+  while (true) {
+    // Single element brace.
+    if (elts.size() == 1)
+      return true;
+
+    // See if we have a #if as the first element of a 2 element brace, if so we
+    // can recuse into its active clause. If so, the second element will be the
+    // active element.
+    if (elts.size() == 2) {
+      if (auto *D = elts.front().dyn_cast<Decl *>()) {
+        if (auto *ICD = dyn_cast<IfConfigDecl>(D)) {
+          elts = ICD->getActiveClauseElements();
+          continue;
+        }
+      }
+    }
+    return false;
+  }
+}
+
+ASTNode BraceStmt::getSingleActiveElement() const {
+  return hasSingleActiveElement(getElements()) ? getLastElement() : nullptr;
+}
+
+Expr *BraceStmt::getSingleActiveExpression() const {
+  return getSingleActiveElement().dyn_cast<Expr *>();
+}
 
-  return getElements()[0].dyn_cast<Expr *>();
+Stmt *BraceStmt::getSingleActiveStatement() const {
+  return getSingleActiveElement().dyn_cast<Stmt *>();
 }
 
 IsSingleValueStmtResult Stmt::mayProduceSingleValue(Evaluator &eval) const {

lib/Parse/ParseDecl.cpp

@@ -8032,8 +8032,7 @@ Parser::parseAbstractFunctionBodyImpl(AbstractFunctionDecl *AFD) {
   AFD->setHasSingleExpressionBody(false);
   AFD->setBodyParsed(BS, fp);
 
-  if (Parser::shouldReturnSingleExpressionElement(BS->getElements())) {
-    auto Element = BS->getLastElement();
+  if (auto Element = BS->getSingleActiveElement()) {
     if (auto *stmt = Element.dyn_cast<Stmt *>()) {
       if (isa<FuncDecl>(AFD)) {
         if (auto *returnStmt = dyn_cast<ReturnStmt>(stmt)) {

lib/Parse/ParseExpr.cpp

@@ -2903,37 +2903,36 @@ ParserResult<Expr> Parser::parseExprClosure() {
     closure->setParameterList(params);
     closure->setHasAnonymousClosureVars();
   }
-
+
+  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 &&
-      Parser::shouldReturnSingleExpressionElement(bodyElements)) {
-    auto Element = bodyElements.back();
-
-    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);
+  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(new (Context) ReturnStmt(SourceLoc(), returnExpr));
       }
-    } else if (Element.is<Expr *>()) {
-      // Create the wrapping return.
-      hasSingleExpressionBody = true;
-      auto returnExpr = Element.get<Expr*>();
-      bodyElements.back() = new (Context) ReturnStmt(SourceLoc(), returnExpr);
     }
   }
 
   // Set the body of the closure.
-  closure->setBody(BraceStmt::create(Context, leftBrace, bodyElements,
-                                     rightBrace),
-                   hasSingleExpressionBody);
+  closure->setBody(BS, hasSingleExpressionBody);
 
   // If the closure includes a capture list, create an AST node for it as well.
   Expr *result = closure;

lib/Parse/Parser.cpp

@@ -1096,31 +1096,6 @@ Parser::getStringLiteralIfNotInterpolated(SourceLoc Loc,
                                                Segments.front().Length));
 }
 
-bool Parser::shouldReturnSingleExpressionElement(ArrayRef<ASTNode> Body) {
-  // If the body consists of an #if declaration with a single
-  // expression active clause, find a single expression.
-  if (Body.size() == 2) {
-    if (auto *D = Body.front().dyn_cast<Decl *>()) {
-      // Step into nested active clause.
-      while (auto *ICD = dyn_cast<IfConfigDecl>(D)) {
-        auto ACE = ICD->getActiveClauseElements();
-        if (ACE.size() == 1) {
-          assert(Body.back() == ACE.back() &&
-                 "active clause not found in body");
-          return true;
-        } else if (ACE.size() == 2) {
-          if (auto *ND = ACE.front().dyn_cast<Decl *>()) {
-            D = ND;
-            continue;
-          }
-        }
-        break;
-      }
-    }
-  }
-  return Body.size() == 1;
-}
-
 struct ParserUnit::Implementation {
   LangOptions LangOpts;
   TypeCheckerOptions TypeCheckerOpts;

lib/Sema/MiscDiagnostics.cpp

@@ -3878,6 +3878,14 @@ class SingleValueStmtUsageChecker final : public ASTWalker {
         break;
       case IsSingleValueStmtResult::Kind::UnterminatedBranches: {
         for (auto *branch : mayProduceSingleValue.getUnterminatedBranches()) {
+          if (auto *BS = dyn_cast<BraceStmt>(branch)) {
+            if (BS->empty()) {
+              Diags.diagnose(branch->getStartLoc(),
+                             diag::single_value_stmt_branch_empty,
+                             S->getKind());
+              continue;
+            }
+          }
           Diags.diagnose(branch->getEndLoc(),
                          diag::single_value_stmt_branch_must_end_in_throw,
                          S->getKind());

lib/Sema/PreCheckExpr.cpp

@@ -1354,15 +1354,13 @@ bool PreCheckExpression::walkToClosureExprPre(ClosureExpr *closure) {
   // LeaveClosureBodiesUnchecked, as the closure may become a single expression
   // closure.
   auto *body = closure->getBody();
-  if (body->getNumElements() == 1) {
-    if (auto *S = body->getLastElement().dyn_cast<Stmt *>()) {
-      if (S->mayProduceSingleValue(Ctx)) {
-        auto *SVE = SingleValueStmtExpr::createWithWrappedBranches(
-            Ctx, S, /*DC*/ closure, /*mustBeExpr*/ false);
-        auto *RS = new (Ctx) ReturnStmt(SourceLoc(), SVE);
-        body->setLastElement(RS);
-        closure->setBody(body, /*isSingleExpression*/ true);
-      }
+  if (auto *S = body->getSingleActiveStatement()) {
+    if (S->mayProduceSingleValue(Ctx)) {
+      auto *SVE = SingleValueStmtExpr::createWithWrappedBranches(
+          Ctx, S, /*DC*/ closure, /*mustBeExpr*/ false);
+      auto *RS = new (Ctx) ReturnStmt(SourceLoc(), SVE);
+      body->setLastElement(RS);
+      closure->setBody(body, /*isSingleExpression*/ true);
     }
   }
 

lib/Sema/TypeCheckStmt.cpp

@@ -2658,23 +2658,26 @@ TypeCheckFunctionBodyRequest::evaluate(Evaluator &evaluator,
 
         body->walk(ContextualizeClosuresAndMacros(AFD));
       }
-    } else 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());
-    } else if (func->getBody()->getNumElements() == 1 &&
-               !func->getResultInterfaceType()->isVoid()) {
+    } else {
+      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 (auto *S = func->getBody()->getLastElement().dyn_cast<Stmt *>()) {
-        if (S->mayProduceSingleValue(evaluator)) {
-          auto *SVE = SingleValueStmtExpr::createWithWrappedBranches(
-              ctx, S, /*DC*/ func, /*mustBeExpr*/ false);
-          auto *RS = new (ctx) ReturnStmt(SourceLoc(), SVE);
-          body->setLastElement(RS);
-          func->setHasSingleExpressionBody();
-          func->setSingleExpressionBody(SVE);
+      if (!func->getResultInterfaceType()->isVoid()) {
+        if (auto *S = body->getSingleActiveStatement()) {
+          if (S->mayProduceSingleValue(evaluator)) {
+            auto *SVE = SingleValueStmtExpr::createWithWrappedBranches(
+                ctx, S, /*DC*/ func, /*mustBeExpr*/ false);
+            auto *RS = new (ctx) ReturnStmt(SourceLoc(), SVE);
+            body->setLastElement(RS);
+            func->setHasSingleExpressionBody();
+            func->setSingleExpressionBody(SVE);
+          }
         }
       }
     }
@@ -2872,20 +2875,17 @@ areBranchesValidForSingleValueStmt(Evaluator &eval, ArrayRef<Stmt *> branches) {
     // Check to see if there are any invalid jumps.
     BS->walk(jumpFinder);
 
-    if (BS->getSingleExpressionElement()) {
+    if (BS->getSingleActiveExpression()) {
       hadSingleExpr = true;
       continue;
     }
 
     // We also allow single value statement branches, which we can wrap in
     // a SingleValueStmtExpr.
-    auto elts = BS->getElements();
-    if (elts.size() == 1) {
-      if (auto *S = elts.back().dyn_cast<Stmt *>()) {
-        if (S->mayProduceSingleValue(eval)) {
-          hadSingleExpr = true;
-          continue;
-        }
+    if (auto *S = BS->getSingleActiveStatement()) {
+      if (S->mayProduceSingleValue(eval)) {
+        hadSingleExpr = true;
+        continue;
       }
     }
     if (!doesBraceEndWithThrow(BS))