Improve calculation of source range for brace and do stmt

Author: ahoppen

include/swift/AST/Stmt.h

@@ -168,7 +168,8 @@ class BraceStmt final : public Stmt,
   SourceLoc getLBraceLoc() const { return LBLoc; }
   SourceLoc getRBraceLoc() const { return RBLoc; }
 
-  SourceRange getSourceRange() const { return SourceRange(LBLoc, RBLoc); }
+  SourceLoc getStartLoc() const;
+  SourceLoc getEndLoc() const;
 
   bool empty() const { return getNumElements() == 0; }
   unsigned getNumElements() const { return Bits.BraceStmt.NumElements; }
@@ -563,8 +564,8 @@ class DoStmt : public LabeledStmt {
 
   SourceLoc getDoLoc() const { return DoLoc; }
 
-  SourceLoc getStartLoc() const { return getLabelLocOrKeywordLoc(DoLoc); }
-  SourceLoc getEndLoc() const { return Body->getEndLoc(); }
+  SourceLoc getStartLoc() const;
+  SourceLoc getEndLoc() const;
 
   BraceStmt *getBody() const { return Body; }
   void setBody(BraceStmt *s) { Body = s; }

lib/AST/Stmt.cpp

@@ -156,6 +156,30 @@ BraceStmt *BraceStmt::create(ASTContext &ctx, SourceLoc lbloc,
   return ::new(Buffer) BraceStmt(lbloc, elts, rbloc, implicit);
 }
 
+SourceLoc BraceStmt::getStartLoc() const {
+  if (LBLoc) {
+    return LBLoc;
+  }
+  for (auto elt : getElements()) {
+    if (auto loc = elt.getStartLoc()) {
+      return loc;
+    }
+  }
+  return SourceLoc();
+}
+
+SourceLoc BraceStmt::getEndLoc() const {
+  if (RBLoc) {
+    return RBLoc;
+  }
+  for (auto elt : llvm::reverse(getElements())) {
+    if (auto loc = elt.getStartLoc()) {
+      return loc;
+    }
+  }
+  return SourceLoc();
+}
+
 ASTNode BraceStmt::findAsyncNode() {
   // TODO: Statements don't track their ASTContext/evaluator, so I am not making
   // this a request. It probably should be a request at some point.
@@ -508,6 +532,17 @@ DoStmt *DoStmt::createImplicit(ASTContext &C, LabeledStmtInfo labelInfo,
                         /*implicit=*/true);
 }
 
+SourceLoc DoStmt::getStartLoc() const {
+  if (DoLoc) {
+    return DoLoc;
+  }
+  return Body->getStartLoc();
+}
+
+SourceLoc DoStmt::getEndLoc() const {
+  return Body->getEndLoc();
+}
+
 namespace {
 
 template<typename CaseIterator>