[AST] Improve source range info for TapExpr

Previously we would only base the start loc on the
`SubExpr`, but that isn't set until CSApply. Change
it to take both `SubExpr` and `Body`'s source range
into account.

Also tighten up the invariant that a TapExpr must
be created with a non-null BraceStmt.

Author: hamishknight

include/swift/AST/Expr.h

@@ -880,12 +880,8 @@ class TapExpr : public Expr {
   void setBody(BraceStmt * b) { Body = b; }
 
   SourceLoc getLoc() const { return SubExpr ? SubExpr->getLoc() : SourceLoc(); }
-  
-  SourceLoc getStartLoc() const {
-    return SubExpr ? SubExpr->getStartLoc() : SourceLoc();
-  }
 
-  SourceLoc getEndLoc() const;
+  SourceRange getSourceRange() const;
 
   static bool classof(const Expr *E) {
     return E->getKind() == ExprKind::Tap;

lib/AST/Expr.cpp

@@ -2611,29 +2611,33 @@ void InterpolatedStringLiteralExpr::forEachSegment(ASTContext &Ctx,
 TapExpr::TapExpr(Expr * SubExpr, BraceStmt *Body)
     : Expr(ExprKind::Tap, /*Implicit=*/true),
       SubExpr(SubExpr), Body(Body) {
-  if (Body) {
-    assert(!Body->empty() &&
-         Body->getFirstElement().isDecl(DeclKind::Var) &&
-         "First element of Body should be a variable to init with the subExpr");
-  }
+  assert(Body);
+  assert(!Body->empty() &&
+       Body->getFirstElement().isDecl(DeclKind::Var) &&
+       "First element of Body should be a variable to init with the subExpr");
 }
 
 VarDecl * TapExpr::getVar() const {
   return dyn_cast<VarDecl>(Body->getFirstElement().dyn_cast<Decl *>());
 }
 
-SourceLoc TapExpr::getEndLoc() const {
-  // Include the body in the range, assuming the body follows the SubExpr.
-  // Also, be (perhaps overly) defensive about null pointers & invalid
-  // locations.
-  if (auto *const b = getBody()) {
-    const auto be = b->getEndLoc();
-    if (be.isValid())
-      return be;
-  }
-  if (auto *const se = getSubExpr())
-    return se->getEndLoc();
-  return SourceLoc();
+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);
 }
 
 RegexLiteralExpr *