3 files changed, 97 insertions, 15 deletions

diff --git a/clang/include/clang/Tooling/FixIt.h b/clang/include/clang/Tooling/FixIt.h
index 8c762367f30..ae8c45b0d57 100644
--- a/clang/include/clang/Tooling/FixIt.h
+++ b/clang/include/clang/Tooling/FixIt.h
@@ -20,36 +20,81 @@
 #define LLVM_CLANG_TOOLING_FIXIT_H
 
 #include "clang/AST/ASTContext.h"
+#include "clang/Basic/TokenKinds.h"
 
 namespace clang {
 namespace tooling {
 namespace fixit {
 
 namespace internal {
-StringRef getText(SourceRange Range, const ASTContext &Context);
+StringRef getText(CharSourceRange Range, const ASTContext &Context);
 
-/// Returns the SourceRange of a SourceRange. This identity function is
-///        used by the following template abstractions.
-inline SourceRange getSourceRange(const SourceRange &Range) { return Range; }
+/// Returns the token CharSourceRange corresponding to \p Range.
+inline CharSourceRange getSourceRange(const SourceRange &Range) {
+  return CharSourceRange::getTokenRange(Range);
+}
 
-/// Returns the SourceRange of the token at Location \p Loc.
-inline SourceRange getSourceRange(const SourceLocation &Loc) {
-  return SourceRange(Loc);
+/// Returns the CharSourceRange of the token at Location \p Loc.
+inline CharSourceRange getSourceRange(const SourceLocation &Loc) {
+  return CharSourceRange::getTokenRange(Loc, Loc);
 }
 
-/// Returns the SourceRange of an given Node. \p Node is typically a
+/// Returns the CharSourceRange of an given Node. \p Node is typically a
 ///        'Stmt', 'Expr' or a 'Decl'.
-template <typename T> SourceRange getSourceRange(const T &Node) {
-  return Node.getSourceRange();
+template <typename T> CharSourceRange getSourceRange(const T &Node) {
+  return CharSourceRange::getTokenRange(Node.getSourceRange());
 }
+
+/// Extends \p Range to include the token \p Next, if it immediately follows the
+/// end of the range. Otherwise, returns \p Range unchanged.
+CharSourceRange maybeExtendRange(CharSourceRange Range, tok::TokenKind Next,
+                                 ASTContext &Context);
 } // end namespace internal
 
-// Returns a textual representation of \p Node.
+/// Returns a textual representation of \p Node.
 template <typename T>
 StringRef getText(const T &Node, const ASTContext &Context) {
   return internal::getText(internal::getSourceRange(Node), Context);
 }
 
+/// Returns the source range spanning the node, extended to include \p Next, if
+/// it immediately follows \p Node. Otherwise, returns the normal range of \p
+/// Node.  See comments on `getExtendedText()` for examples.
+template <typename T>
+CharSourceRange getExtendedRange(const T &Node, tok::TokenKind Next,
+                                 ASTContext &Context) {
+  return internal::maybeExtendRange(internal::getSourceRange(Node), Next,
+                                    Context);
+}
+
+/// Returns the source text of the node, extended to include \p Next, if it
+/// immediately follows the node. Otherwise, returns the text of just \p Node.
+///
+/// For example, given statements S1 and S2 below:
+/// \code
+///   {
+///     // S1:
+///     if (!x) return foo();
+///     // S2:
+///     if (!x) { return 3; }
+//    }
+/// \endcode
+/// then
+/// \code
+///   getText(S1, Context) = "if (!x) return foo()"
+///   getExtendedText(S1, tok::TokenKind::semi, Context)
+///     = "if (!x) return foo();"
+///   getExtendedText(*S1.getThen(), tok::TokenKind::semi, Context)
+///     = "return foo();"
+///   getExtendedText(*S2.getThen(), tok::TokenKind::semi, Context)
+///     = getText(S2, Context) = "{ return 3; }"
+/// \endcode
+template <typename T>
+StringRef getExtendedText(const T &Node, tok::TokenKind Next,
+                          ASTContext &Context) {
+  return internal::getText(getExtendedRange(Node, Next, Context), Context);
+}
+
 // Returns a FixItHint to remove \p Node.
 // TODO: Add support for related syntactical elements (i.e. comments, ...).
 template <typename T> FixItHint createRemoval(const T &Node) {
diff --git a/clang/lib/Tooling/FixIt.cpp b/clang/lib/Tooling/FixIt.cpp
index c828c18871f..ca12a5642a1 100644
--- a/clang/lib/Tooling/FixIt.cpp
+++ b/clang/lib/Tooling/FixIt.cpp
@@ -18,12 +18,20 @@ namespace tooling {
 namespace fixit {
 
 namespace internal {
-StringRef getText(SourceRange Range, const ASTContext &Context) {
-  return Lexer::getSourceText(CharSourceRange::getTokenRange(Range),
-                              Context.getSourceManager(),
+StringRef getText(CharSourceRange Range, const ASTContext &Context) {
+  return Lexer::getSourceText(Range, Context.getSourceManager(),
                               Context.getLangOpts());
 }
-} // end namespace internal
+
+CharSourceRange maybeExtendRange(CharSourceRange Range, tok::TokenKind Next,
+                                 ASTContext &Context) {
+  Optional<Token> Tok = Lexer::findNextToken(
+      Range.getEnd(), Context.getSourceManager(), Context.getLangOpts());
+  if (!Tok || !Tok->is(Next))
+    return Range;
+  return CharSourceRange::getTokenRange(Range.getBegin(), Tok->getLocation());
+}
+} // namespace internal
 
 } // end namespace fixit
 } // end namespace tooling
diff --git a/clang/unittests/Tooling/FixItTest.cpp b/clang/unittests/Tooling/FixItTest.cpp
index ec9801d3455..bb3711e30b2 100644
--- a/clang/unittests/Tooling/FixItTest.cpp
+++ b/clang/unittests/Tooling/FixItTest.cpp
@@ -13,6 +13,7 @@
 using namespace clang;
 
 using tooling::fixit::getText;
+using tooling::fixit::getExtendedText;
 using tooling::fixit::createRemoval;
 using tooling::fixit::createReplacement;
 
@@ -77,6 +78,34 @@ TEST(FixItTest, getTextWithMacro) {
                   "void foo(int x, int y) { FOO(x,y) }");
 }
 
+TEST(FixItTest, getExtendedText) {
+  CallsVisitor Visitor;
+
+  Visitor.OnCall = [](CallExpr *CE, ASTContext *Context) {
+    EXPECT_EQ("foo(x, y);",
+              getExtendedText(*CE, tok::TokenKind::semi, *Context));
+
+    Expr *P0 = CE->getArg(0);
+    Expr *P1 = CE->getArg(1);
+    EXPECT_EQ("x", getExtendedText(*P0, tok::TokenKind::semi, *Context));
+    EXPECT_EQ("x,", getExtendedText(*P0, tok::TokenKind::comma, *Context));
+    EXPECT_EQ("y", getExtendedText(*P1, tok::TokenKind::semi, *Context));
+  };
+  Visitor.runOver("void foo(int x, int y) { foo(x, y); }");
+  Visitor.runOver("void foo(int x, int y) { if (true) foo(x, y); }");
+  Visitor.runOver("int foo(int x, int y) { if (true) return 3 + foo(x, y); }");
+  Visitor.runOver("void foo(int x, int y) { for (foo(x, y);;) ++x; }");
+  Visitor.runOver(
+      "bool foo(int x, int y) { for (;foo(x, y);) x = 1; return true; }");
+
+  Visitor.OnCall = [](CallExpr *CE, ASTContext *Context) {
+    EXPECT_EQ("foo()", getExtendedText(*CE, tok::TokenKind::semi, *Context));
+  };
+  Visitor.runOver("bool foo() { if (foo()) return true; return false; }");
+  Visitor.runOver("void foo() { int x; for (;; foo()) ++x; }");
+  Visitor.runOver("int foo() { return foo() + 3; }");
+}
+
 TEST(FixItTest, createRemoval) {
   CallsVisitor Visitor;