[AST] Remove the "single expression body" bit

Remove this bit from function decls and closures.
Instead, for closures, infer it from the presence
of a single return or single expression AST node
in the body, which ought to be equivalent, and
automatically takes result builders into
consideration. We can also completely drop this
query from AbstractFunctionDecl, replacing it
instead with a bit on ReturnStmt.

Author: hamishknight

include/swift/AST/AnyFunctionRef.h

@@ -94,12 +94,6 @@ class AnyFunctionRef {
     return !TheFunction.get<AbstractClosureExpr *>()->getType().isNull();
   }
 
-  bool hasSingleExpressionBody() const {
-    if (auto *AFD = TheFunction.dyn_cast<AbstractFunctionDecl *>())
-      return AFD->hasSingleExpressionBody();
-    return TheFunction.get<AbstractClosureExpr *>()->hasSingleExpressionBody();
-  }
-
   ParameterList *getParameters() const {
     if (auto *AFD = TheFunction.dyn_cast<AbstractFunctionDecl *>())
       return AFD->getParameters();
@@ -146,33 +140,31 @@ class AnyFunctionRef {
     return cast<AutoClosureExpr>(ACE)->getBody();
   }
 
-  void setParsedBody(BraceStmt *stmt, bool isSingleExpression) {
+  void setParsedBody(BraceStmt *stmt) {
     if (auto *AFD = TheFunction.dyn_cast<AbstractFunctionDecl *>()) {
       AFD->setBody(stmt, AbstractFunctionDecl::BodyKind::Parsed);
-      AFD->setHasSingleExpressionBody(isSingleExpression);
       return;
     }
 
     auto *ACE = TheFunction.get<AbstractClosureExpr *>();
     if (auto *CE = dyn_cast<ClosureExpr>(ACE)) {
-      CE->setBody(stmt, isSingleExpression);
+      CE->setBody(stmt);
       CE->setBodyState(ClosureExpr::BodyState::ReadyForTypeChecking);
       return;
     }
 
     llvm_unreachable("autoclosures don't have statement bodies");
   }
 
-  void setTypecheckedBody(BraceStmt *stmt, bool isSingleExpression) {
+  void setTypecheckedBody(BraceStmt *stmt) {
     if (auto *AFD = TheFunction.dyn_cast<AbstractFunctionDecl *>()) {
       AFD->setBody(stmt, AbstractFunctionDecl::BodyKind::TypeChecked);
-      AFD->setHasSingleExpressionBody(isSingleExpression);
       return;
     }
 
     auto *ACE = TheFunction.get<AbstractClosureExpr *>();
     if (auto *CE = dyn_cast<ClosureExpr>(ACE)) {
-      CE->setBody(stmt, isSingleExpression);
+      CE->setBody(stmt);
       CE->setBodyState(ClosureExpr::BodyState::TypeCheckedWithSignature);
       return;
     }

include/swift/AST/Decl.h

@@ -464,7 +464,7 @@ class alignas(1 << DeclAlignInBits) Decl : public ASTAllocated<Decl> {
   SWIFT_INLINE_BITFIELD(SubscriptDecl, VarDecl, 2,
     StaticSpelling : 2
   );
-  SWIFT_INLINE_BITFIELD(AbstractFunctionDecl, ValueDecl, 3+2+2+2+8+1+1+1+1+1+1+1,
+  SWIFT_INLINE_BITFIELD(AbstractFunctionDecl, ValueDecl, 3+2+2+2+8+1+1+1+1+1+1,
     /// \see AbstractFunctionDecl::BodyKind
     BodyKind : 3,
 
@@ -492,9 +492,6 @@ class alignas(1 << DeclAlignInBits) Decl : public ASTAllocated<Decl> {
     /// Whether the function body throws.
     Throws : 1,
 
-    /// Whether this member's body consists of a single expression.
-    HasSingleExpressionBody : 1,
-
     /// Whether peeking into this function detected nested type declarations.
     /// This is set when skipping over the decl at parsing.
     HasNestedTypeDeclarations : 1,
@@ -7177,7 +7174,6 @@ class AbstractFunctionDecl : public GenericContext, public ValueDecl {
     Bits.AbstractFunctionDecl.Overridden = false;
     Bits.AbstractFunctionDecl.Async = Async;
     Bits.AbstractFunctionDecl.Throws = Throws;
-    Bits.AbstractFunctionDecl.HasSingleExpressionBody = false;
     Bits.AbstractFunctionDecl.HasNestedTypeDeclarations = false;
     Bits.AbstractFunctionDecl.DistributedThunk = false;
   }
@@ -7211,14 +7207,6 @@ class AbstractFunctionDecl : public GenericContext, public ValueDecl {
   }
 
 public:
-  void setHasSingleExpressionBody(bool Has = true) { 
-    Bits.AbstractFunctionDecl.HasSingleExpressionBody = Has;
-  }
-
-  bool hasSingleExpressionBody() const {
-    return Bits.AbstractFunctionDecl.HasSingleExpressionBody;
-  }
-
   /// Returns the string for the base name, or "_" if this is unnamed.
   StringRef getNameStr() const {
     assert(!getName().isSpecial() && "Cannot get string for special names");

include/swift/AST/Expr.h

@@ -3927,9 +3927,7 @@ class AbstractClosureExpr : public DeclContext, public Expr {
   bool hasSingleExpressionBody() const;
 
   /// Retrieve the body for closure that has a single expression for
-  /// its body.
-  ///
-  /// Only valid when \c hasSingleExpressionBody() is true.
+  /// its body, or \c nullptr if there is no single expression body.
   Expr *getSingleExpressionBody() const;
 
   /// Returns the body of closures that have a body
@@ -4055,9 +4053,9 @@ class ClosureExpr : public AbstractClosureExpr {
   /// The explicitly-specified result type.
   llvm::PointerIntPair<TypeExpr *, 2, BodyState> ExplicitResultTypeAndBodyState;
 
-  /// The body of the closure, along with a bit indicating whether it
-  /// was originally just a single expression.
-  llvm::PointerIntPair<BraceStmt *, 1, bool> Body;
+  /// The body of the closure.
+  BraceStmt *Body;
+
 public:
   ClosureExpr(const DeclAttributes &attributes,
               SourceRange bracketRange, VarDecl *capturedSelfDecl,
@@ -4083,11 +4081,8 @@ class ClosureExpr : public AbstractClosureExpr {
   SourceLoc getEndLoc() const;
   SourceLoc getLoc() const;
 
-  BraceStmt *getBody() const { return Body.getPointer(); }
-  void setBody(BraceStmt *S, bool isSingleExpression) {
-    Body.setPointer(S);
-    Body.setInt(isSingleExpression);
-  }
+  BraceStmt *getBody() const { return Body; }
+  void setBody(BraceStmt *S) { Body = S; }
 
   DeclAttributes &getAttrs() { return Attributes; }
   const DeclAttributes &getAttrs() const { return Attributes; }
@@ -4202,13 +4197,11 @@ class ClosureExpr : public AbstractClosureExpr {
   /// ... even if the closure has been coerced to return Void by the type
   /// checker. This function does not return true for empty closures.
   bool hasSingleExpressionBody() const {
-    return Body.getInt();
+    return getSingleExpressionBody();
   }
 
   /// Retrieve the body for closure that has a single expression for
-  /// its body.
-  ///
-  /// Only valid when \c hasSingleExpressionBody() is true.
+  /// its body, or \c nullptr if there is no single expression body.
   Expr *getSingleExpressionBody() const;
 
   /// Is this a completely empty closure?

include/swift/AST/Stmt.h

@@ -99,6 +99,12 @@ class alignas(8) Stmt : public ASTAllocated<Stmt> {
     : NumPadBits,
     CaseCount : 32
   );
+    
+  SWIFT_INLINE_BITFIELD(ReturnStmt, Stmt, 1,
+    /// Whether the result is an implied return, e.g for an implicit single
+    /// expression return.
+    IsImplied : 1
+  );
 
   SWIFT_INLINE_BITFIELD_FULL(YieldStmt, Stmt, 32,
     : NumPadBits,
@@ -244,7 +250,9 @@ class ReturnStmt : public Stmt {
 
   ReturnStmt(SourceLoc returnLoc, Expr *result, bool isImplicit)
       : Stmt(StmtKind::Return, isImplicit), ReturnLoc(returnLoc),
-        Result(result) {}
+        Result(result) {
+    Bits.ReturnStmt.IsImplied = false;
+  }
 
 public:
   static ReturnStmt *createParsed(ASTContext &ctx, SourceLoc returnLoc,
@@ -261,11 +269,25 @@ 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) {
+    assert(result && "Result must be present to be implied");
+    auto *RS = createImplicit(ctx, result);
+    RS->Bits.ReturnStmt.IsImplied = true;
+    return RS;
+  }
+
   SourceLoc getReturnLoc() const { return ReturnLoc; }
 
   SourceLoc getStartLoc() const;
   SourceLoc getEndLoc() const;
 
+  /// Whether the result is an implied return, e.g for an implicit single
+  /// expression return.
+  bool isImplied() const {
+    return Bits.ReturnStmt.IsImplied;
+  }
+
   bool hasResult() const { return Result != 0; }
   Expr *getResult() const {
     assert(Result && "ReturnStmt doesn't have a result");

include/swift/Sema/ConstraintLocator.h

@@ -47,16 +47,17 @@ class ProtocolConformance;
 /// specified error messages when the conversion fails.
 ///
 enum ContextualTypePurpose : uint8_t {
-  CTP_Unused,           ///< No contextual type is specified.
-  CTP_Initialization,   ///< Pattern binding initialization.
-  CTP_ReturnStmt,       ///< Value specified to a 'return' statement.
-  CTP_ReturnSingleExpr, ///< Value implicitly returned from a function.
-  CTP_YieldByValue,     ///< By-value yield operand.
-  CTP_YieldByReference, ///< By-reference yield operand.
-  CTP_ThrowStmt,        ///< Value specified to a 'throw' statement.
-  CTP_DiscardStmt,      ///< Value specified to a 'discard' statement.
-  CTP_EnumCaseRawValue, ///< Raw value specified for "case X = 42" in enum.
-  CTP_DefaultParameter, ///< Default value in parameter 'foo(a : Int = 42)'.
+  CTP_Unused,            ///< No contextual type is specified.
+  CTP_Initialization,    ///< Pattern binding initialization.
+  CTP_ReturnStmt,        ///< Value specified to a 'return' statement.
+  CTP_ImpliedReturnStmt, ///< Value from an implied 'return', e.g a single expr
+                         ///  function body.
+  CTP_YieldByValue,      ///< By-value yield operand.
+  CTP_YieldByReference,  ///< By-reference yield operand.
+  CTP_ThrowStmt,         ///< Value specified to a 'throw' statement.
+  CTP_DiscardStmt,       ///< Value specified to a 'discard' statement.
+  CTP_EnumCaseRawValue,  ///< Raw value specified for "case X = 42" in enum.
+  CTP_DefaultParameter,  ///< Default value in parameter 'foo(a : Int = 42)'.
 
   /// Default value in @autoclosure parameter
   /// 'foo(a : @autoclosure () -> Int = 42)'.
@@ -819,11 +820,13 @@ class LocatorPathElt::TypeParameterRequirement final
 
 class LocatorPathElt::ClosureBody final : public StoredIntegerElement<1> {
   public:
-  ClosureBody(bool hasExplicitReturn = false)
-    : StoredIntegerElement(ConstraintLocator::ClosureBody, hasExplicitReturn) {}
+  ClosureBody(bool hasImpliedReturn)
+    : StoredIntegerElement(ConstraintLocator::ClosureBody, hasImpliedReturn) {}
 
-  /// Indicates whether body of the closure has any `return` statements.
-  bool hasExplicitReturn() const { return bool(getValue()); }
+  /// Indicates whether body of the closure has an implied `return` statement,
+  /// this is the case for single expression bodies where the `return` was not
+  /// written explicitly.
+  bool hasImpliedReturn() const { return bool(getValue()); }
 
   static bool classof(const LocatorPathElt *elt) {
     return elt->getKind() == ConstraintLocator::ClosureBody;

lib/AST/ASTBridging.cpp

@@ -1019,7 +1019,7 @@ BridgedClosureExpr_createParsed(BridgedASTContext cContext,
   auto *out = new (context) ClosureExpr(
       attributes, bracketRange, nullptr, nullptr, asyncLoc, throwsLoc,
       /*FIXME:thrownType=*/nullptr, arrowLoc, inLoc, nullptr, declContext);
-  out->setBody(body.unbridged(), /*isSingleExpression*/ false);
+  out->setBody(body.unbridged());
   out->setParameterList(params);
   return out;
 }

lib/AST/ASTWalker.cpp

@@ -1019,7 +1019,7 @@ class Traversal : public ASTVisitor<Traversal, Expr*, Stmt*,
 
     // Handle other closures.
     if (BraceStmt *body = cast_or_null<BraceStmt>(doIt(expr->getBody()))) {
-      expr->setBody(body, expr->hasSingleExpressionBody());
+      expr->setBody(body);
       return expr;
     }
     return nullptr;

lib/AST/Expr.cpp

@@ -2034,18 +2034,28 @@ swift::__AbstractClosureExpr_getActorIsolation(AbstractClosureExpr *CE) {
     return (NODE)->getStartLoc();                 \
   }
 
-FORWARD_SOURCE_LOCS_TO(ClosureExpr, Body.getPointer())
+FORWARD_SOURCE_LOCS_TO(ClosureExpr, Body)
 
 Expr *ClosureExpr::getSingleExpressionBody() const {
-  assert(hasSingleExpressionBody() && "Not a single-expression body");
-  auto body = getBody()->getLastElement();
-  if (auto stmt = body.dyn_cast<Stmt *>()) {
-    if (auto braceStmt = dyn_cast<BraceStmt>(stmt))
-      return braceStmt->getLastElement().get<Expr *>();
+  auto *body = getBody();
+  if (!body)
+    return nullptr;
 
-    return cast<ReturnStmt>(stmt)->getResult();
-  }
-  return body.get<Expr *>();
+  // If we have a lone expression, use it.
+  if (auto *E = body->getSingleActiveExpression())
+    return E;
+
+  // Otherwise we must have a return of an expression.
+  auto elt = body->getSingleActiveStatement();
+  if (!elt)
+    return nullptr;
+
+  // Any lone ReturnStmt, even explicit, is treated as a single expression body.
+  auto *RS = dyn_cast<ReturnStmt>(elt);
+  if (!RS || !RS->hasResult())
+    return nullptr;
+
+  return RS->getResult();
 }
 
 bool ClosureExpr::hasEmptyBody() const {

lib/IDE/TypeCheckCompletionCallback.cpp

@@ -159,7 +159,7 @@ bool swift::ide::isImplicitSingleExpressionReturn(ConstraintSystem &CS,
                                                   Expr *CompletionExpr) {
   Expr *ParentExpr = CS.getParentExpr(CompletionExpr);
   if (!ParentExpr)
-    return CS.getContextualTypePurpose(CompletionExpr) == CTP_ReturnSingleExpr;
+    return CS.getContextualTypePurpose(CompletionExpr) == CTP_ImpliedReturnStmt;
 
   if (auto *ParentCE = dyn_cast<ClosureExpr>(ParentExpr)) {
     if (ParentCE->hasSingleExpressionBody() &&

lib/Parse/ParseDecl.cpp

@@ -8814,15 +8814,14 @@ Parser::parseAbstractFunctionBodyImpl(AbstractFunctionDecl *AFD) {
       SourceLoc LBraceLoc, RBraceLoc;
       LBraceLoc = consumeToken(tok::l_brace);
       auto *CCE = new (Context) CodeCompletionExpr(Tok.getLoc());
-      auto *Return = ReturnStmt::createImplicit(Context, CCE);
+      auto *Return = ReturnStmt::forSingleExprBody(Context, CCE);
       CodeCompletionCallbacks->setParsedDecl(accessor);
       CodeCompletionCallbacks->completeAccessorBeginning(CCE);
       RBraceLoc = Tok.getLoc();
       consumeToken(tok::code_complete);
       auto *BS =
           BraceStmt::create(Context, LBraceLoc, ASTNode(Return), RBraceLoc,
                             /*implicit*/ true);
-      AFD->setHasSingleExpressionBody();
       AFD->setBodyParsed(BS);
       return {BS, Fingerprint::ZERO()};
     }