A couple of ReturnStmt tweaks

include/swift/AST/Stmt.h

@@ -241,12 +241,25 @@ class BraceStmt final : public Stmt,
 class ReturnStmt : public Stmt {
   SourceLoc ReturnLoc;
   Expr *Result;
-
+
+  ReturnStmt(SourceLoc returnLoc, Expr *result, bool isImplicit)
+      : Stmt(StmtKind::Return, isImplicit), ReturnLoc(returnLoc),
+        Result(result) {}
+
 public:
-  ReturnStmt(SourceLoc ReturnLoc, Expr *Result,
-             llvm::Optional<bool> implicit = llvm::None)
-      : Stmt(StmtKind::Return, getDefaultImplicitFlag(implicit, ReturnLoc)),
-        ReturnLoc(ReturnLoc), Result(Result) {}
+  static ReturnStmt *createParsed(ASTContext &ctx, SourceLoc returnLoc,
+                                  Expr *result) {
+    return new (ctx) ReturnStmt(returnLoc, result, /*isImplicit*/ false);
+  }
+
+  static ReturnStmt *createImplicit(ASTContext &ctx, SourceLoc returnLoc,
+                                    Expr *result) {
+    return new (ctx) ReturnStmt(returnLoc, result, /*isImplicit*/ true);
+  }
+
+  static ReturnStmt *createImplicit(ASTContext &ctx, Expr *result) {
+    return createImplicit(ctx, SourceLoc(), result);
+  }
 
   SourceLoc getReturnLoc() const { return ReturnLoc; }
 
@@ -1621,10 +1634,10 @@ class FailStmt : public Stmt {
       : Stmt(StmtKind::Fail, getDefaultImplicitFlag(implicit, returnLoc)),
         ReturnLoc(returnLoc), NilLoc(nilLoc) {}
 
-  SourceLoc getLoc() const { return ReturnLoc; }
-
-  SourceRange getSourceRange() const { return SourceRange(ReturnLoc, NilLoc); }
-  
+  SourceRange getSourceRange() const {
+    return SourceRange::combine(ReturnLoc, NilLoc);
+  }
+
   static bool classof(const Stmt *S) {
     return S->getKind() == StmtKind::Fail;
   }

include/swift/Basic/SourceLoc.h

@@ -121,6 +121,17 @@ class SourceRange {
   /// if (auto x = getSourceRange()) { ... }
   explicit operator bool() const { return isValid(); }
 
+  /// Combine the given source ranges into the smallest contiguous SourceRange
+  /// that includes them all, ignoring any invalid ranges present.
+  static SourceRange combine(ArrayRef<SourceRange> ranges);
+
+  /// Combine the given source ranges into the smallest contiguous SourceRange
+  /// that includes them all, ignoring any invalid ranges present.
+  template <typename ...T>
+  static SourceRange combine(T... ranges) {
+    return SourceRange::combine({ranges...});
+  }
+
   /// Extend this SourceRange to the smallest continuous SourceRange that
   /// includes both this range and the other one.
   void widen(SourceRange Other);

lib/AST/ASTBridging.cpp

@@ -1464,7 +1464,7 @@ BridgedReturnStmt BridgedReturnStmt_createParsed(BridgedASTContext cContext,
                                                  BridgedSourceLoc cLoc,
                                                  BridgedNullableExpr expr) {
   ASTContext &context = cContext.unbridged();
-  return new (context) ReturnStmt(cLoc.unbridged(), expr.unbridged());
+  return ReturnStmt::createParsed(context, cLoc.unbridged(), expr.unbridged());
 }
 
 BridgedSwitchStmt BridgedSwitchStmt_createParsed(

lib/AST/ArgumentList.cpp

@@ -24,15 +24,7 @@ using namespace swift;
 Type swift::__Expr_getType(Expr *E) { return E->getType(); }
 
 SourceRange Argument::getSourceRange() const {
-  auto labelLoc = getLabelLoc();
-  if (labelLoc.isInvalid())
-    return getExpr()->getSourceRange();
-
-  auto exprEndLoc = getExpr()->getEndLoc();
-  if (exprEndLoc.isInvalid())
-    return labelLoc;
-
-  return SourceRange(labelLoc, exprEndLoc);
+  return SourceRange::combine(getLabelLoc(), getExpr()->getSourceRange());
 }
 
 Argument Argument::implicitInOut(ASTContext &ctx, Expr *expr) {

lib/AST/Expr.cpp

@@ -2072,7 +2072,7 @@ FORWARD_SOURCE_LOCS_TO(AutoClosureExpr, Body)
 
 void AutoClosureExpr::setBody(Expr *E) {
   auto &Context = getASTContext();
-  auto *RS = new (Context) ReturnStmt(SourceLoc(), E);
+  auto *RS = ReturnStmt::createImplicit(Context, E);
   Body = BraceStmt::create(Context, E->getStartLoc(), { RS }, E->getEndLoc());
 }
 
@@ -2715,19 +2715,8 @@ SourceRange TapExpr::getSourceRange() const {
   if (!SubExpr)
     return Body->getSourceRange();
 
-  SourceLoc start = SubExpr->getStartLoc();
-  if (!start.isValid())
-    start = Body->getStartLoc();
-  if (!start.isValid())
-    return SourceRange();
-
-  SourceLoc end = Body->getEndLoc();
-  if (!end.isValid())
-    end = SubExpr->getEndLoc();
-  if (!end.isValid())
-    return SourceRange();
-
-  return SourceRange(start, end);
+  return SourceRange::combine(SubExpr->getSourceRange(),
+                              Body->getSourceRange());
 }
 
 RegexLiteralExpr *

lib/AST/Stmt.cpp

@@ -371,14 +371,7 @@ ThenStmt *ThenStmt::createImplicit(ASTContext &ctx, Expr *result) {
 }
 
 SourceRange ThenStmt::getSourceRange() const {
-  auto range = getResult()->getSourceRange();
-  if (!range)
-    return ThenLoc;
-
-  if (ThenLoc)
-    range.widen(ThenLoc);
-
-  return range;
+  return SourceRange::combine(ThenLoc, getResult()->getSourceRange());
 }
 
 SourceLoc ThrowStmt::getEndLoc() const { return SubExpr->getEndLoc(); }

lib/Basic/SourceLoc.cpp

@@ -533,6 +533,23 @@ bool SourceManager::isOwning(SourceLoc Loc) const {
   return findBufferContainingLocInternal(Loc).has_value();
 }
 
+SourceRange SourceRange::combine(ArrayRef<SourceRange> ranges) {
+  if (ranges.empty())
+    return SourceRange();
+
+  SourceRange result = ranges.front();
+  for (auto other : ranges.drop_front()) {
+    if (!other)
+      continue;
+    if (!result) {
+      result = other;
+      continue;
+    }
+    result.widen(other);
+  }
+  return result;
+}
+
 void SourceRange::widen(SourceRange Other) {
   if (Other.Start.Value.getPointer() < Start.Value.getPointer())
     Start = Other.Start;

lib/ClangImporter/ClangImporter.cpp

@@ -5155,8 +5155,7 @@ synthesizeBaseClassMethodBody(AbstractFunctionDecl *afd, void *context) {
   baseMemberCallExpr->setType(baseMember->getResultInterfaceType());
   baseMemberCallExpr->setThrows(nullptr);
 
-  auto returnStmt = new (ctx) ReturnStmt(SourceLoc(), baseMemberCallExpr,
-                                         /*implicit=*/true);
+  auto *returnStmt = ReturnStmt::createImplicit(ctx, baseMemberCallExpr);
 
   auto body = BraceStmt::create(ctx, SourceLoc(), {returnStmt}, SourceLoc(),
                                 /*implicit=*/true);
@@ -5450,8 +5449,7 @@ synthesizeBaseClassFieldGetterOrAddressGetterBody(AbstractFunctionDecl *afd,
         ctx, baseGetterMethod->getResultInterfaceType(), resultType,
         returnExpr);
 
-  auto returnStmt = new (ctx) ReturnStmt(SourceLoc(), returnExpr,
-                                         /*implicit=*/true);
+  auto *returnStmt = ReturnStmt::createImplicit(ctx, returnExpr);
 
   auto body = BraceStmt::create(ctx, SourceLoc(), {returnStmt}, SourceLoc(),
                                 /*implicit=*/true);
@@ -6763,8 +6761,7 @@ synthesizeDependentTypeThunkParamForwarding(AbstractFunctionDecl *afd, void *con
         ctx, specializedFuncCallExpr, thunkDecl->getResultInterfaceType());
   }
 
-  auto returnStmt = new (ctx)
-      ReturnStmt(SourceLoc(), resultExpr, /*implicit=*/true);
+  auto *returnStmt = ReturnStmt::createImplicit(ctx, resultExpr);
   auto body = BraceStmt::create(ctx, SourceLoc(), {returnStmt}, SourceLoc(),
                                 /*implicit=*/true);
   return {body, /*isTypeChecked=*/true};
@@ -6886,8 +6883,7 @@ synthesizeForwardingThunkBody(AbstractFunctionDecl *afd, void *context) {
   specializedFuncCallExpr->setType(thunkDecl->getResultInterfaceType());
   specializedFuncCallExpr->setThrows(nullptr);
 
-  auto returnStmt = new (ctx) ReturnStmt(SourceLoc(), specializedFuncCallExpr,
-                                         /*implicit=*/true);
+  auto *returnStmt = ReturnStmt::createImplicit(ctx, specializedFuncCallExpr);
 
   auto body = BraceStmt::create(ctx, SourceLoc(), {returnStmt}, SourceLoc(),
                                 /*implicit=*/true);

lib/ClangImporter/ImportDecl.cpp

@@ -594,7 +594,7 @@ synthesizeErrorDomainGetterBody(AbstractFunctionDecl *afd, void *context) {
   domainDeclRef->setType(
     getterDecl->mapTypeIntoContext(swiftValueDecl->getInterfaceType()));
 
-  auto ret = new (ctx) ReturnStmt(SourceLoc(), domainDeclRef);
+  auto *ret = ReturnStmt::createImplicit(ctx, domainDeclRef);
   return { BraceStmt::create(ctx, SourceLoc(), {ret}, SourceLoc(), isImplicit),
            /*isTypeChecked=*/true };
 }

lib/ClangImporter/SwiftDeclSynthesizer.cpp

@@ -373,7 +373,7 @@ synthesizeConstantGetterBody(AbstractFunctionDecl *afd, void *voidContext) {
   }
 
   // Create the return statement.
-  auto ret = new (ctx) ReturnStmt(SourceLoc(), expr);
+  auto ret = ReturnStmt::createImplicit(ctx, expr);
 
   return {BraceStmt::create(ctx, SourceLoc(), ASTNode(ret), SourceLoc()),
           /*isTypeChecked=*/true};
@@ -471,7 +471,7 @@ synthesizeStructDefaultConstructorBody(AbstractFunctionDecl *afd,
   auto assign = new (ctx) AssignExpr(lhs, SourceLoc(), call, /*implicit*/ true);
   assign->setType(emptyTuple);
 
-  auto ret = new (ctx) ReturnStmt(SourceLoc(), nullptr, /*Implicit=*/true);
+  auto *ret = ReturnStmt::createImplicit(ctx, /*expr*/ nullptr);
 
   // Create the function body.
   auto body = BraceStmt::create(ctx, SourceLoc(), {assign, ret}, SourceLoc());
@@ -565,8 +565,7 @@ synthesizeValueConstructorBody(AbstractFunctionDecl *afd, void *context) {
     }
   }
 
-  auto ret = new (ctx) ReturnStmt(SourceLoc(), nullptr, /*Implicit=*/true);
-  stmts.push_back(ret);
+  stmts.push_back(ReturnStmt::createImplicit(ctx, /*expr*/ nullptr));
 
   // Create the function body.
   auto body = BraceStmt::create(ctx, SourceLoc(), stmts, SourceLoc());
@@ -689,7 +688,7 @@ synthesizeRawValueBridgingConstructorBody(AbstractFunctionDecl *afd,
                                      /*Implicit=*/true);
   assign->setType(TupleType::getEmpty(ctx));
 
-  auto ret = new (ctx) ReturnStmt(SourceLoc(), nullptr, /*Implicit=*/true);
+  auto *ret = ReturnStmt::createImplicit(ctx, /*expr*/ nullptr);
 
   auto body = BraceStmt::create(ctx, SourceLoc(), {assign, ret}, SourceLoc());
   return {body, /*isTypeChecked=*/true};
@@ -852,7 +851,7 @@ synthesizeUnionFieldGetterBody(AbstractFunctionDecl *afd, void *context) {
       CallExpr::createImplicit(ctx, reinterpretCastRefExpr, argList);
   reinterpreted->setType(importedFieldDecl->getInterfaceType());
   reinterpreted->setThrows(nullptr);
-  auto ret = new (ctx) ReturnStmt(SourceLoc(), reinterpreted);
+  auto *ret = ReturnStmt::createImplicit(ctx, reinterpreted);
   auto body = BraceStmt::create(ctx, SourceLoc(), ASTNode(ret), SourceLoc(),
                                 /*implicit*/ true);
   return {body, /*isTypeChecked*/ true};
@@ -1118,7 +1117,7 @@ synthesizeIndirectFieldGetterBody(AbstractFunctionDecl *afd, void *context) {
                                  DeclNameLoc(), /*implicit*/ true);
   expr->setType(anonymousInnerFieldDecl->getInterfaceType());
 
-  auto ret = new (ctx) ReturnStmt(SourceLoc(), expr);
+  auto *ret = ReturnStmt::createImplicit(ctx, expr);
   auto body = BraceStmt::create(ctx, SourceLoc(), ASTNode(ret), SourceLoc(),
                                 /*implicit*/ true);
   return {body, /*isTypeChecked=*/true};
@@ -1242,7 +1241,7 @@ synthesizeEnumRawValueConstructorBody(AbstractFunctionDecl *afd,
                                      /*implicit*/ true);
   assign->setType(TupleType::getEmpty(ctx));
 
-  auto ret = new (ctx) ReturnStmt(SourceLoc(), nullptr, /*Implicit=*/true);
+  auto *ret = ReturnStmt::createImplicit(ctx, /*expr*/ nullptr);
 
   auto body = BraceStmt::create(ctx, SourceLoc(), {assign, ret}, SourceLoc(),
                                 /*implicit*/ true);
@@ -1310,7 +1309,7 @@ synthesizeEnumRawValueGetterBody(AbstractFunctionDecl *afd, void *context) {
   reinterpreted->setType(rawTy);
   reinterpreted->setThrows(nullptr);
 
-  auto ret = new (ctx) ReturnStmt(SourceLoc(), reinterpreted);
+  auto *ret = ReturnStmt::createImplicit(ctx, reinterpreted);
   auto body = BraceStmt::create(ctx, SourceLoc(), ASTNode(ret), SourceLoc(),
                                 /*implicit*/ true);
   return {body, /*isTypeChecked=*/true};
@@ -1380,7 +1379,7 @@ synthesizeStructRawValueGetterBody(AbstractFunctionDecl *afd, void *context) {
     result = CoerceExpr::createImplicit(ctx, bridge, computedType);
   }
 
-  auto ret = new (ctx) ReturnStmt(SourceLoc(), result);
+  auto ret = ReturnStmt::createImplicit(ctx, result);
   auto body = BraceStmt::create(ctx, SourceLoc(), ASTNode(ret), SourceLoc(),
                                 /*implicit*/ true);
   return {body, /*isTypeChecked=*/true};
@@ -1565,8 +1564,7 @@ synthesizeUnwrappingGetterOrAddressGetterBody(AbstractFunctionDecl *afd,
     propertyExpr = SwiftDeclSynthesizer::synthesizeReturnReinterpretCast(
         ctx, getterImpl->getResultInterfaceType(), elementTy, propertyExpr);
 
-  auto returnStmt = new (ctx) ReturnStmt(SourceLoc(), propertyExpr,
-                                         /*implicit*/ true);
+  auto *returnStmt = ReturnStmt::createImplicit(ctx, propertyExpr);
 
   auto body = BraceStmt::create(ctx, SourceLoc(), {returnStmt}, SourceLoc(),
                                 /*implicit*/ true);
@@ -1645,8 +1643,7 @@ synthesizeUnwrappingAddressSetterBody(AbstractFunctionDecl *afd,
   auto *setterImplCallExpr =
       createAccessorImplCallExpr(setterImpl, selfArg, nullptr);
 
-  auto returnStmt = new (ctx) ReturnStmt(SourceLoc(), setterImplCallExpr,
-                                         /*implicit*/ true);
+  auto *returnStmt = ReturnStmt::createImplicit(ctx, setterImplCallExpr);
 
   auto body = BraceStmt::create(ctx, SourceLoc(), {returnStmt}, SourceLoc(),
                                 /*implicit*/ true);
@@ -1884,8 +1881,7 @@ synthesizeSuccessorFuncBody(AbstractFunctionDecl *afd, void *context) {
                                                  /*implicit*/ true);
   copyRefRValueExpr->setType(copyDecl->getInterfaceType());
 
-  auto returnStmt = new (ctx) ReturnStmt(SourceLoc(), copyRefRValueExpr,
-                                         /*implicit*/ true);
+  auto *returnStmt = ReturnStmt::createImplicit(ctx, copyRefRValueExpr);
 
   auto body = BraceStmt::create(ctx, SourceLoc(),
                                 {
@@ -1974,8 +1970,7 @@ synthesizeOperatorMethodBody(AbstractFunctionDecl *afd, void *context) {
   callExpr->setType(funcDecl->getResultInterfaceType());
   callExpr->setThrows(nullptr);
 
-  auto returnStmt = new (ctx) ReturnStmt(SourceLoc(), callExpr,
-                                         /*implicit*/ true);
+  auto *returnStmt = ReturnStmt::createImplicit(ctx, callExpr);
 
   auto body = BraceStmt::create(ctx, SourceLoc(), {returnStmt}, SourceLoc(),
                                 /*implicit*/ true);
@@ -2061,10 +2056,9 @@ synthesizeComputedGetterFromCXXMethod(AbstractFunctionDecl *afd,
   auto selfArg = createSelfArg(accessor);
 
   auto *getterImplCallExpr = createAccessorImplCallExpr(method, selfArg);
-  auto returnStmt =
-      new (method->getASTContext()) ReturnStmt(SourceLoc(), getterImplCallExpr);
-  auto body = BraceStmt::create(method->getASTContext(), SourceLoc(),
-                                {returnStmt}, SourceLoc());
+  auto &ctx = method->getASTContext();
+  auto *returnStmt = ReturnStmt::createImplicit(ctx, getterImplCallExpr);
+  auto *body = BraceStmt::create(ctx, SourceLoc(), {returnStmt}, SourceLoc());
 
   return {body, /*isTypeChecked*/ true};
 }
@@ -2218,8 +2212,7 @@ synthesizeDefaultArgumentBody(AbstractFunctionDecl *afd, void *context) {
   initCall->setThrows(nullptr);
 
   // Synthesize `return __cxx__defaultArg_XYZ()`.
-  auto returnStmt = new (ctx) ReturnStmt(SourceLoc(), initCall,
-                                         /*implicit=*/true);
+  auto *returnStmt = ReturnStmt::createImplicit(ctx, initCall);
 
   auto body = BraceStmt::create(ctx, SourceLoc(), {returnStmt}, SourceLoc(),
                                 /*implicit=*/true);

lib/Parse/ParseDecl.cpp

@@ -8670,8 +8670,7 @@ Parser::parseAbstractFunctionBodyImpl(AbstractFunctionDecl *AFD) {
       SourceLoc LBraceLoc, RBraceLoc;
       LBraceLoc = consumeToken(tok::l_brace);
       auto *CCE = new (Context) CodeCompletionExpr(Tok.getLoc());
-      auto *Return =
-          new (Context) ReturnStmt(SourceLoc(), CCE, /*implicit=*/true);
+      auto *Return = ReturnStmt::createImplicit(Context, CCE);
       CodeCompletionCallbacks->setParsedDecl(accessor);
       CodeCompletionCallbacks->completeAccessorBeginning(CCE);
       RBraceLoc = Tok.getLoc();
@@ -8725,15 +8724,15 @@ Parser::parseAbstractFunctionBodyImpl(AbstractFunctionDecl *AFD) {
         }
       }
       if (isa<FuncDecl>(AFD)) {
-        auto RS = new (Context) ReturnStmt(SourceLoc(), E);
+        auto RS = ReturnStmt::createImplicit(Context, 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);
+          auto RS = ReturnStmt::createImplicit(Context, E);
           BS->setLastElement(RS);
           AFD->setHasSingleExpressionBody();
           AFD->setSingleExpressionBody(E);

lib/Parse/ParseExpr.cpp

@@ -3072,7 +3072,7 @@ ParserResult<Expr> Parser::parseExprClosure() {
         // Create the wrapping return.
         hasSingleExpressionBody = true;
         auto returnExpr = Element.get<Expr*>();
-        BS->setLastElement(new (Context) ReturnStmt(SourceLoc(), returnExpr));
+        BS->setLastElement(ReturnStmt::createImplicit(Context, returnExpr));
       }
     }
   }