[SR-11711] Single-expression returns for #if declarations within functions

Author: maustinstar

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()->getElements().back();
   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()->getElements().back();
   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()->getElements().back() = 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()->getElements().back();
   if (auto stmt = body.dyn_cast<Stmt *>()) {
     if (auto braceStmt = dyn_cast<BraceStmt>(stmt))
-      return braceStmt->getFirstElement().get<Expr *>();
+      return braceStmt->getElements().back().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->getElements().back().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->getElements().back() = RS->getResult();
     return true;
   };
 

lib/IDE/ExprContextAnalysis.cpp

@@ -1112,7 +1112,15 @@ 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 *>()) {
+        if (auto *ICD = dyn_cast<IfConfigDecl>(D)) {
+          auto ACE = ICD->getActiveClauseElements();
+          return ACE.size() == 1;
+        }
+      }
+    }
+    return body->getNumElements() == 1 && body->getElements().back().is<Expr *>();
   }
 
 public:

lib/Parse/ParseDecl.cpp

@@ -6656,13 +6656,41 @@ BraceStmt *Parser::parseAbstractFunctionBodyImpl(AbstractFunctionDecl *AFD) {
 
   BraceStmt *BS = Body.get();
   AFD->setBodyParsed(BS);
+  
+  if (BS->getNumElements() == 0)
+    return BS;
+  
+  auto Element = BS->getFirstElement();
+  
+  // If the body consists of an #if declaration with a single
+  // expression active clause, turn it into a return statement.
+  bool ReturnsLastElement = false;
+  if (BS->getNumElements() == 2) {
+    if (auto *D = Element.dyn_cast<Decl *>()) {
+      // Step into nested active clause.
+      while (auto *ICD = dyn_cast<IfConfigDecl>(D)) {
+        auto ACE = ICD->getActiveClauseElements();
+        if (ACE.size() == 1) {
+          Element = BS->getLastElement();
+          ReturnsLastElement = true;
+          break;
+        } else if (ACE.size() == 2) {
+          if (auto *ND = ACE.front().dyn_cast<Decl *>()) {
+            D = ND;
+            continue;
+          }
+        }
+        break;
+      }
+    }
+  }
 
   // If the body consists of a single expression, turn it into a return
   // statement.
-  if (BS->getNumElements() != 1)
+  if (BS->getNumElements() != 1 &&
+      !ReturnsLastElement)
     return BS;
 
-  auto Element = BS->getFirstElement();
   if (auto *stmt = Element.dyn_cast<Stmt *>()) {
     if (isa<FuncDecl>(AFD)) {
       if (auto *returnStmt = dyn_cast<ReturnStmt>(stmt)) {
@@ -6687,15 +6715,15 @@ BraceStmt *Parser::parseAbstractFunctionBodyImpl(AbstractFunctionDecl *AFD) {
     }
     if (isa<FuncDecl>(AFD)) {
       auto RS = new (Context) ReturnStmt(SourceLoc(), E);
-      BS->setFirstElement(RS);
+      BS->getElements().back() = 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);
+        BS->getElements().back() = RS;
         AFD->setHasSingleExpressionBody();
         AFD->setSingleExpressionBody(E);
       }

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()->getElements().back() = 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->getElements().back() = 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->getElements().back().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->getElements().back() = 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()")
+    "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()")
+    "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
+}