[AST] Improve FailStmt source location computation

Avoid forming invalid source ranges when
`ReturnLoc` is invalid. Also introduce a utility
to make this kind of range computation easier,
and use it in a couple of other cases.

Author: hamishknight

include/swift/AST/Stmt.h

@@ -1634,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/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

@@ -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;