[SR-11711] [Parse] Single-expression implicit returns within #if declarations

include/swift/AST/Stmt.h

@@ -180,6 +180,7 @@ class BraceStmt final : public Stmt,
   ASTNode getLastElement() const { return getElements().back(); }
 
   void setFirstElement(ASTNode node) { getElements().front() = node; }
+  void setLastElement(ASTNode node) { getElements().back() = node; }
 
   /// The elements contained within the BraceStmt.
   MutableArrayRef<ASTNode> getElements() {

include/swift/Parse/Parser.h

@@ -700,6 +700,11 @@ 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/Decl.cpp

@@ -684,7 +684,7 @@ Expr *AbstractFunctionDecl::getSingleExpressionBody() const {
   assert(hasSingleExpressionBody() && "Not a single-expression body");
   auto braceStmt = getBody();
   assert(braceStmt != nullptr && "No body currently available.");
-  auto body = getBody()->getFirstElement();
+  auto body = getBody()->getLastElement();
   if (auto *stmt = body.dyn_cast<Stmt *>()) {
     if (auto *returnStmt = dyn_cast<ReturnStmt>(stmt)) {
       return returnStmt->getResult();
@@ -701,7 +701,7 @@ Expr *AbstractFunctionDecl::getSingleExpressionBody() const {
 
 void AbstractFunctionDecl::setSingleExpressionBody(Expr *NewBody) {
   assert(hasSingleExpressionBody() && "Not a single-expression body");
-  auto body = getBody()->getFirstElement();
+  auto body = getBody()->getLastElement();
   if (auto *stmt = body.dyn_cast<Stmt *>()) {
     if (auto *returnStmt = dyn_cast<ReturnStmt>(stmt)) {
       returnStmt->setResult(NewBody);
@@ -717,7 +717,7 @@ void AbstractFunctionDecl::setSingleExpressionBody(Expr *NewBody) {
       return;
     }
   }
-  getBody()->setFirstElement(NewBody);
+  getBody()->setLastElement(NewBody);
 }
 
 bool AbstractStorageDecl::isTransparent() const {

lib/AST/Expr.cpp

@@ -1974,10 +1974,10 @@ FORWARD_SOURCE_LOCS_TO(ClosureExpr, Body.getPointer())
 
 Expr *ClosureExpr::getSingleExpressionBody() const {
   assert(hasSingleExpressionBody() && "Not a single-expression body");
-  auto body = getBody()->getFirstElement();
+  auto body = getBody()->getLastElement();
   if (auto stmt = body.dyn_cast<Stmt *>()) {
     if (auto braceStmt = dyn_cast<BraceStmt>(stmt))
-      return braceStmt->getFirstElement().get<Expr *>();
+      return braceStmt->getLastElement().get<Expr *>();
 
     return cast<ReturnStmt>(stmt)->getResult();
   }
@@ -2003,7 +2003,7 @@ void AutoClosureExpr::setBody(Expr *E) {
 }
 
 Expr *AutoClosureExpr::getSingleExpressionBody() const {
-  return cast<ReturnStmt>(Body->getFirstElement().get<Stmt *>())->getResult();
+  return cast<ReturnStmt>(Body->getLastElement().get<Stmt *>())->getResult();
 }
 
 Expr *AutoClosureExpr::getUnwrappedCurryThunkExpr() const {

lib/IDE/CodeCompletion.cpp

@@ -6101,13 +6101,13 @@ bool CodeCompletionCallbacksImpl::trySolverCompletion(bool MaybeFuncBody) {
 // to work via TypeCheckCompletionCallback.
 static void undoSingleExpressionReturn(DeclContext *DC) {
   auto updateBody = [](BraceStmt *BS, ASTContext &Ctx) -> bool {
-    ASTNode FirstElem = BS->getFirstElement();
-    auto *RS = dyn_cast_or_null<ReturnStmt>(FirstElem.dyn_cast<Stmt*>());
+    ASTNode LastElem = BS->getLastElement();
+    auto *RS = dyn_cast_or_null<ReturnStmt>(LastElem.dyn_cast<Stmt*>());
 
     if (!RS || !RS->isImplicit())
       return false;
 
-    BS->setFirstElement(RS->getResult());
+    BS->setLastElement(RS->getResult());
     return true;
   };
 

lib/IDE/ExprContextAnalysis.cpp

@@ -1112,7 +1112,24 @@ class ExprContextAnalyzer {
   /// in order to avoid a base expression affecting the type. However, now that
   /// we've typechecked, we will take the context type into account.
   static bool isSingleExpressionBodyForCodeCompletion(BraceStmt *body) {
-    return body->getNumElements() == 1 && body->getFirstElement().is<Expr *>();
+    if (body->getNumElements() == 2) {
+      if (auto *D = body->getFirstElement().dyn_cast<Decl *>()) {
+        // Step into nested active clause.
+        while (auto *ICD = dyn_cast<IfConfigDecl>(D)) {
+          auto ACE = ICD->getActiveClauseElements();
+          if (ACE.size() == 1) {
+            return body->getLastElement().is<Expr *>();
+          } else if (ACE.size() == 2) {
+            if (auto *ND = ACE.front().dyn_cast<Decl *>()) {
+              D = ND;
+              continue;
+            }
+          }
+          break;
+        }
+      }
+    }
+    return body->getNumElements() == 1 && body->getLastElement().is<Expr *>();
   }
 
 public:

lib/Parse/ParseDecl.cpp

@@ -6656,48 +6656,45 @@ BraceStmt *Parser::parseAbstractFunctionBodyImpl(AbstractFunctionDecl *AFD) {
 
   BraceStmt *BS = Body.get();
   AFD->setBodyParsed(BS);
-
-  // If the body consists of a single expression, turn it into a return
-  // statement.
-  if (BS->getNumElements() != 1)
-    return BS;
-
-  auto Element = BS->getFirstElement();
-  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);
+
+  if (Parser::shouldReturnSingleExpressionElement(BS->getElements())) {
+    auto Element = BS->getLastElement();
+    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;
+    } 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;
+        }
       }
-    }
-    if (isa<FuncDecl>(AFD)) {
-      auto RS = new (Context) ReturnStmt(SourceLoc(), E);
-      BS->setFirstElement(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 = new (Context) ReturnStmt(E->getStartLoc(), E);
-        BS->setFirstElement(RS);
+      if (isa<FuncDecl>(AFD)) {
+        auto RS = new (Context) ReturnStmt(SourceLoc(), 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 = new (Context) ReturnStmt(E->getStartLoc(), E);
+          BS->setLastElement(RS);
+          AFD->setHasSingleExpressionBody();
+          AFD->setSingleExpressionBody(E);
+        }
       }
     }
   }

lib/Parse/ParseExpr.cpp

@@ -2839,39 +2839,30 @@ ParserResult<Expr> Parser::parseExprClosure() {
     closure->setParameterList(params);
     closure->setHasAnonymousClosureVars();
   }
-
+  
   // If the body consists of a single expression, turn it into a return
   // statement.
   bool hasSingleExpressionBody = false;
-  if (!missingRBrace && bodyElements.size() == 1) {
-    // If the closure's only body element is a single return statement,
-    // use that instead of creating a new wrapping return expression.
-    Expr *returnExpr = nullptr;
+  if (!missingRBrace &&
+      Parser::shouldReturnSingleExpressionElement(bodyElements)) {
+    auto Element = bodyElements.back();
 
-    if (bodyElements[0].is<Stmt *>()) {
-      if (auto returnStmt =
-                  dyn_cast<ReturnStmt>(bodyElements[0].get<Stmt*>())) {
-
+    if (Element.is<Stmt *>()) {
+      if (auto returnStmt = dyn_cast<ReturnStmt>(Element.get<Stmt *>())) {
+        hasSingleExpressionBody = true;
         if (!returnStmt->hasResult()) {
-
-          returnExpr = TupleExpr::createEmpty(Context,
-                                              SourceLoc(),
-                                              SourceLoc(),
-                                              /*implicit*/true);
-
+          auto returnExpr = TupleExpr::createEmpty(Context,
+                                                   SourceLoc(),
+                                                   SourceLoc(),
+                                                   /*implicit*/true);
           returnStmt->setResult(returnExpr);
         }
-
-        hasSingleExpressionBody = true;
       }
-    }
-
-    // Otherwise, create the wrapping return.
-    if (bodyElements[0].is<Expr *>()) {
+    } else if (Element.is<Expr *>()) {
+      // Create the wrapping return.
       hasSingleExpressionBody = true;
-      returnExpr = bodyElements[0].get<Expr*>();
-      bodyElements[0] = new (Context) ReturnStmt(SourceLoc(),
-                                                 returnExpr);
+      auto returnExpr = Element.get<Expr*>();
+      bodyElements.back() = new (Context) ReturnStmt(SourceLoc(), returnExpr);
     }
   }
 

lib/Parse/Parser.cpp

@@ -1185,6 +1185,31 @@ 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 {
   std::shared_ptr<SyntaxParseActions> SPActions;
   LangOptions LangOpts;

lib/Sema/CSClosure.cpp

@@ -173,7 +173,11 @@ class ClosureConstraintApplication
   }
 
   void visitDecl(Decl *decl) {
-    llvm_unreachable("Declarations not supported");
+
+    if (isa<IfConfigDecl>(decl))
+      return;
+
+    llvm_unreachable("Unimplemented case for closure body");
   }
 
   ASTNode visitBraceStmt(BraceStmt *braceStmt) {

lib/Sema/TypeCheckStmt.cpp

@@ -1926,7 +1926,7 @@ bool TypeCheckASTNodeAtLocRequest::evaluate(Evaluator &evaluator,
                 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.
-      func->getBody()->setFirstElement(func->getSingleExpressionBody());
+      func->getBody()->setLastElement(func->getSingleExpressionBody());
     }
   }
 
@@ -1991,7 +1991,7 @@ TypeCheckFunctionBodyRequest::evaluate(Evaluator &evaluator,
                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->setFirstElement(func->getSingleExpressionBody());
+      body->setLastElement(func->getSingleExpressionBody());
     }
   } else if (isa<ConstructorDecl>(AFD) &&
              (body->empty() ||
@@ -2025,12 +2025,12 @@ TypeCheckFunctionBodyRequest::evaluate(Evaluator &evaluator,
   // that we have eagerly converted something like `{ fatalError() }`
   // into `{ return fatalError() }` that has to be corrected here.
   if (isa<FuncDecl>(AFD) && cast<FuncDecl>(AFD)->hasSingleExpressionBody()) {
-    if (auto *stmt = body->getFirstElement().dyn_cast<Stmt *>()) {
+    if (auto *stmt = body->getLastElement().dyn_cast<Stmt *>()) {
       if (auto *retStmt = dyn_cast<ReturnStmt>(stmt)) {
         if (retStmt->isImplicit() && retStmt->hasResult()) {
           auto returnType = retStmt->getResult()->getType();
           if (returnType && returnType->isUninhabited())
-            body->setFirstElement(retStmt->getResult());
+            body->setLastElement(retStmt->getResult());
         }
       }
     }

test/FixCode/fixits-omit-return.swift

@@ -10,4 +10,16 @@ let cl_fixit_addreturn: () -> String = {
     "foo" // expected-warning {{string literal is unused}} expected-error {{missing return in a closure expected to return 'String'; did you mean to return the last expression?}} {{5-5=return }}
 }
 
+func ff_fixit_addreturn_ifdecl() -> String {
+    #if true
+    print("entering ff_fixit_addreturn_ifdecl()")
+    "foo" // expected-warning {{string literal is unused}} expected-error {{missing return in a function expected to return 'String'; did you mean to return the last expression?}} {{5-5=return }}
+    #endif
+}
 
+let cl_fixit_addreturn_ifdecl: () -> String = {
+    #if true
+    print("entering cl_fixit_addreturn_ifdecl()")
+    "foo" // expected-warning {{string literal is unused}} expected-error {{missing return in a closure expected to return 'String'; did you mean to return the last expression?}} {{5-5=return }}
+    #endif
+}

test/Parse/omit_return_fail.swift

@@ -7,3 +7,15 @@ func badIs<T>(_ value: Any, anInstanceOf type: T.Type) -> Bool {
 func foo() -> Int {
     return // expected-error {{non-void function should return a value}}
 }
+
+func badIs_ifdecl<T>(_ value: Any, anInstanceOf type: T.Type) -> Bool {
+    #if true
+    value is type // expected-error {{cannot find type 'type' in scope}}
+    #endif
+}
+
+func foo_ifdecl() -> Int {
+    #if true
+    return // expected-error {{non-void function should return a value}}
+    #endif
+}