WIP: Allow return to be omitted from single expression functions.

nate-chandler · nate-chandler · commit 74b462b3ba1e · 2019-04-24T09:59:53.000-07:00
diff --git a/include/swift/AST/Decl.h b/include/swift/AST/Decl.h
@@ -401,8 +401,7 @@ class alignas(1 << DeclAlignInBits) Decl {
   SWIFT_INLINE_BITFIELD(SubscriptDecl, VarDecl, 2,
     StaticSpelling : 2
   );
-
-  SWIFT_INLINE_BITFIELD(AbstractFunctionDecl, ValueDecl, 3+8+1+1+1+1+1+1+1,
+  SWIFT_INLINE_BITFIELD(AbstractFunctionDecl, ValueDecl, 3+8+1+1+1+1+1+1+1+1,
     /// \see AbstractFunctionDecl::BodyKind
     BodyKind : 3,
 
@@ -426,7 +425,10 @@ class alignas(1 << DeclAlignInBits) Decl {
 
     /// Whether this member was synthesized as part of a derived
     /// protocol conformance.
-    Synthesized : 1
+    Synthesized : 1,
+
+    /// Whether this member's body consists of a single expression.
+    HasSingleExpressionBody : 1
   );
 
   SWIFT_INLINE_BITFIELD(FuncDecl, AbstractFunctionDecl, 1+2+1+1+2,
@@ -5433,13 +5435,25 @@ class AbstractFunctionDecl : public GenericContext, public ValueDecl {
     Bits.AbstractFunctionDecl.NeedsNewVTableEntry = false;
     Bits.AbstractFunctionDecl.HasComputedNeedsNewVTableEntry = false;
     Bits.AbstractFunctionDecl.Synthesized = false;
+    Bits.AbstractFunctionDecl.HasSingleExpressionBody = false;
   }
 
   void setBodyKind(BodyKind K) {
     Bits.AbstractFunctionDecl.BodyKind = unsigned(K);
   }
 
 public:
+  void setHasSingleExpressionBody(bool Has = true) { 
+    Bits.AbstractFunctionDecl.HasSingleExpressionBody = Has;
+  }
+
+  bool hasSingleExpressionBody() const {
+    return Bits.AbstractFunctionDecl.HasSingleExpressionBody;
+  }
+
+  Expr *getSingleExpressionBody() const;
+  void setSingleExpressionBody(Expr *NewBody);
+
   /// Returns the string for the base name, or "_" if this is unnamed.
   StringRef getNameStr() const {
     assert(!getFullName().isSpecial() && "Cannot get string for special names");
diff --git a/lib/AST/ASTDumper.cpp b/lib/AST/ASTDumper.cpp
@@ -1070,7 +1070,10 @@ namespace {
           Indent -= 2;
         }
       }
-      if (auto Body = D->getBody(/*canSynthesize=*/false)) {
+      if (D->hasSingleExpressionBody()) {
+        OS << '\n';
+        printRec(D->getSingleExpressionBody());
+      } else if (auto Body = D->getBody(/*canSynthesize=*/false)) {
         OS << '\n';
         printRec(Body, D->getASTContext());
       }
diff --git a/lib/AST/ASTWalker.cpp b/lib/AST/ASTWalker.cpp
@@ -339,6 +339,15 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
       }
     }
 
+    // FIXME: FIXME_NOW: DETERMINE: Is this or something like it needed?  Right
+    //        now it's causing assertion failures during AST verification.
+    //     
+    //     if (AFD->hasSingleExpressionBody()) {
+    //       if (Expr *body = doIt(AFD->getSingleExpressionBody()))
+    //         AFD->setSingleExpressionBody(body);
+    //       else
+    //         return true;
+    //     }
     if (AFD->getBody(/*canSynthesize=*/false)) {
       AbstractFunctionDecl::BodyKind PreservedKind = AFD->getBodyKind();
       if (BraceStmt *S = cast_or_null<BraceStmt>(doIt(AFD->getBody())))
diff --git a/lib/AST/Decl.cpp b/lib/AST/Decl.cpp
@@ -471,6 +471,26 @@ case DeclKind::ID: return cast<ID##Decl>(this)->getLoc();
   llvm_unreachable("Unknown decl kind");
 }
 
+Expr *AbstractFunctionDecl::getSingleExpressionBody() const {
+  assert(hasSingleExpressionBody() && "Not a single-expression body");
+  auto braceStmt = getBody();
+  assert(braceStmt != nullptr && "No body currently available.");
+  auto body = getBody()->getElement(0);
+  if (body.is<Stmt *>())
+    return cast<ReturnStmt>(body.get<Stmt *>())->getResult();
+  return body.get<Expr *>();
+}
+
+void AbstractFunctionDecl::setSingleExpressionBody(Expr *NewBody) {
+  assert(hasSingleExpressionBody() && "Not a single-expression body");
+  auto body = getBody()->getElement(0);
+  if (body.is<Stmt *>()) {
+    cast<ReturnStmt>(body.get<Stmt *>())->setResult(NewBody);
+    return;
+  }
+  getBody()->setElement(0, NewBody);
+}
+
 bool AbstractStorageDecl::isTransparent() const {
   return getAttrs().hasAttribute<TransparentAttr>();
 }
diff --git a/lib/Parse/ParseDecl.cpp b/lib/Parse/ParseDecl.cpp
@@ -5655,8 +5655,56 @@ void Parser::parseAbstractFunctionBody(AbstractFunctionDecl *AFD) {
     Context.Stats->getFrontendCounters().NumFunctionsParsed++;
 
   ParserResult<BraceStmt> Body = parseBraceItemList(diag::invalid_diagnostic);
-  if (!Body.isNull())
-    AFD->setBody(Body.get());
+  if (!Body.isNull()) {
+    // If the body consists of a single expression, turn it into a return 
+    // statement.
+    BraceStmt * BS = Body.get();
+    AFD->setBody(BS);
+    // FIXME: FIXME_NOW: DETERMINE: Do we need to handle the code completion 
+    //        case here?  It is being handled for closure 
+    //        (Parser::parseExprClosure).
+    if (BS->getNumElements() == 1) {
+      auto Element = BS->getElement(0);
+      // FIXME: FIXME_NOW: DETERMINE: Do we need to do this?  It's being done
+      //        for closures (Parser::parseExprClosure) but it may not be 
+      //        relevant here.
+      //
+      //           if (Element.is<Stmt *>()) {
+      //             if (auto returnStmt =
+      //                         dyn_cast_or_null<ReturnStmt>(Element.get<Stmt*>())) {
+      //               
+      //               if (!returnStmt->hasResult()) {
+      //                 Expr *returnExpr = TupleExpr::createEmpty(Context,
+      //                                                           SourceLoc(),
+      //                                                           SourceLoc(),
+      //                                                           /*implicit*/true);
+      //                 
+      //                 returnStmt->setResult(returnExpr);
+      //                 AFD->setHasSingleExpressionBody();
+      //                 AFD->setSingleExpressionBody(returnExpr);
+      //               }
+      //             }
+      //           } else 
+      if (Element.is<Expr *>()) {
+        auto E = Element.get<Expr *>();
+        if (auto F = dyn_cast_or_null<FuncDecl>(AFD)) {
+          auto RS = new (Context) ReturnStmt(SourceLoc(), E);
+          BS->setElement(0, RS); 
+          AFD->setHasSingleExpressionBody();
+          AFD->setSingleExpressionBody(E);
+        } else if (auto F = dyn_cast_or_null<ConstructorDecl>(AFD)) {
+          if (F->getFailability() != OTK_None && 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->setElement(0, RS); 
+            AFD->setHasSingleExpressionBody();
+            AFD->setSingleExpressionBody(E);
+          }
+        }
+      }
+    }
+  }
 }
 
 bool Parser::parseAbstractFunctionBodyDelayed(AbstractFunctionDecl *AFD) {
diff --git a/lib/Sema/TypeCheckStmt.cpp b/lib/Sema/TypeCheckStmt.cpp
@@ -1905,6 +1905,48 @@ bool TypeChecker::typeCheckFunctionBodyUntil(FuncDecl *FD,
   BraceStmt *BS = FD->getBody();
   assert(BS && "Should have a body");
 
+  // FIXME: FIXME_NOW: DETERMINE: What is the appropriate place for this code to
+  //        live?
+  if (FD->hasSingleExpressionBody()) {
+    auto fnType = FD->getBodyResultTypeLoc().getType();
+    auto returnStmt = cast<ReturnStmt>(BS->getElement(0).get<Stmt *>());
+    auto E = returnStmt->getResult();
+
+    // FIXME: FIXME_NOW: DETERMINE: Does this actually need to be here?  It's 
+    //        being done for closures, but there may be a reason that applies
+    //        there but not here.
+    //
+    //           if (E != FD->getSingleExpressionBody()) {
+    //             FD->setSingleExpressionBody(E);
+    //           }
+
+    if (FD->getBodyResultTypeLoc().isNull() || fnType->isVoid()) {
+      auto voidExpr = TupleExpr::createEmpty(Context,
+                                             E->getStartLoc(),
+                                             E->getEndLoc(),
+                                             /*implicit*/true);
+      returnStmt->setResult(voidExpr);
+      BS = BraceStmt::create(Context,
+                             BS->getStartLoc(),
+                             { E, returnStmt },
+                             BS->getEndLoc(),
+                             /*implicit=*/true);
+    } else {
+      // FIXME: FIXME_NOW: DETERMINE: Is this the appropriate mechanism to use
+      //        to divine the type of E?
+      auto exprTy = typeCheckExpression(E, FD, TypeLoc(),
+                                        CTP_ReturnStmt);
+      if (!exprTy.isNull() && exprTy->isUninhabited() &&
+          exprTy->getCanonicalType() != fnType->getCanonicalType()) {
+        BS = BraceStmt::create(Context,
+                               BS->getStartLoc(),
+                               { E },
+                               BS->getEndLoc(),
+                               /*implicit=*/true);
+      }
+    }
+  }
+
   StmtChecker SC(*this, static_cast<AbstractFunctionDecl *>(FD));
   SC.EndTypeCheckLoc = EndTypeCheckLoc;
   bool HadError = SC.typeCheckBody(BS);
diff --git a/test/Parse/omit_return.swift b/test/Parse/omit_return.swift

Original file line number	Diff line number	Diff line change
`@@ -1070,7 +1070,10 @@ namespace {`
`1070`	`1070`	`Indent -= 2;`
`1071`	`1071`	`}`
`1072`	`1072`	`}`
`1073`		`- if (auto Body = D->getBody(/canSynthesize=/false)) {`
	`1073`	`+ if (D->hasSingleExpressionBody()) {`
	`1074`	`+ OS << '\n';`
	`1075`	`+ printRec(D->getSingleExpressionBody());`
	`1076`	`+ } else if (auto Body = D->getBody(/canSynthesize=/false)) {`
`1074`	`1077`	`OS << '\n';`
`1075`	`1078`	`printRec(Body, D->getASTContext());`
`1076`	`1079`	`}`