[Rewrite] Extend to further accept CharSourceRange

Some Rewrite functions are already overloaded to accept
CharSourceRange, and this extends others in the same manner.  I'm
calling these in code that's not ready to upstream, but I figure they
might be useful to others in the meantime.

Reviewed By: jdoerfert

Differential Revision: https://reviews.llvm.org/D61467

git-svn-id: https://llvm.org/svn/llvm-project/cfe/trunk@365258 91177308-0d34-0410-b5e6-96231b3b80d8

4 files changed, 105 insertions, 4 deletions

diff --git a/include/clang/Rewrite/Core/Rewriter.h b/include/clang/Rewrite/Core/Rewriter.h
index 5a3ff6c798..84c5ac3d72 100644
--- a/include/clang/Rewrite/Core/Rewriter.h
+++ b/include/clang/Rewrite/Core/Rewriter.h
@@ -84,7 +84,16 @@ public:
   /// in different buffers, this returns an empty string.
   ///
   /// Note that this method is not particularly efficient.
-  std::string getRewrittenText(SourceRange Range) const;
+  std::string getRewrittenText(CharSourceRange Range) const;
+
+  /// getRewrittenText - Return the rewritten form of the text in the specified
+  /// range.  If the start or end of the range was unrewritable or if they are
+  /// in different buffers, this returns an empty string.
+  ///
+  /// Note that this method is not particularly efficient.
+  std::string getRewrittenText(SourceRange Range) const {
+    return getRewrittenText(CharSourceRange::getTokenRange(Range));
+  }
 
   /// InsertText - Insert the specified string at the specified location in the
   /// original buffer.  This method returns true (and does nothing) if the input
@@ -141,6 +150,13 @@ public:
   /// ReplaceText - This method replaces a range of characters in the input
   /// buffer with a new string.  This is effectively a combined "remove/insert"
   /// operation.
+  bool ReplaceText(CharSourceRange range, StringRef NewStr) {
+    return ReplaceText(range.getBegin(), getRangeSize(range), NewStr);
+  }
+
+  /// ReplaceText - This method replaces a range of characters in the input
+  /// buffer with a new string.  This is effectively a combined "remove/insert"
+  /// operation.
   bool ReplaceText(SourceRange range, StringRef NewStr) {
     return ReplaceText(range.getBegin(), getRangeSize(range), NewStr);
   }
diff --git a/lib/Rewrite/Rewriter.cpp b/lib/Rewrite/Rewriter.cpp
index 281cf6f1be..881399e98e 100644
--- a/lib/Rewrite/Rewriter.cpp
+++ b/lib/Rewrite/Rewriter.cpp
@@ -170,7 +170,7 @@ int Rewriter::getRangeSize(SourceRange Range, RewriteOptions opts) const {
 /// in different buffers, this returns an empty string.
 ///
 /// Note that this method is not particularly efficient.
-std::string Rewriter::getRewrittenText(SourceRange Range) const {
+std::string Rewriter::getRewrittenText(CharSourceRange Range) const {
   if (!isRewritable(Range.getBegin()) ||
       !isRewritable(Range.getEnd()))
     return {};
@@ -193,7 +193,9 @@ std::string Rewriter::getRewrittenText(SourceRange Range) const {
 
     // Adjust the end offset to the end of the last token, instead of being the
     // start of the last token.
-    EndOff += Lexer::MeasureTokenLength(Range.getEnd(), *SourceMgr, *LangOpts);
+    if (Range.isTokenRange())
+      EndOff +=
+          Lexer::MeasureTokenLength(Range.getEnd(), *SourceMgr, *LangOpts);
     return std::string(Ptr, Ptr+EndOff-StartOff);
   }
 
@@ -203,7 +205,8 @@ std::string Rewriter::getRewrittenText(SourceRange Range) const {
 
   // Adjust the end offset to the end of the last token, instead of being the
   // start of the last token.
-  EndOff += Lexer::MeasureTokenLength(Range.getEnd(), *SourceMgr, *LangOpts);
+  if (Range.isTokenRange())
+    EndOff += Lexer::MeasureTokenLength(Range.getEnd(), *SourceMgr, *LangOpts);
 
   // Advance the iterators to the right spot, yay for linear time algorithms.
   RewriteBuffer::iterator Start = RB.begin();
diff --git a/unittests/Rewrite/CMakeLists.txt b/unittests/Rewrite/CMakeLists.txt
index a4c6e37245..613d716cb9 100644
--- a/unittests/Rewrite/CMakeLists.txt
+++ b/unittests/Rewrite/CMakeLists.txt
@@ -4,8 +4,10 @@ set(LLVM_LINK_COMPONENTS
 
 add_clang_unittest(RewriteTests
   RewriteBufferTest.cpp
+  RewriterTest.cpp
   )
 clang_target_link_libraries(RewriteTests
   PRIVATE
   clangRewrite
+  clangTooling
   )
diff --git a/unittests/Rewrite/RewriterTest.cpp b/unittests/Rewrite/RewriterTest.cpp
new file mode 100644
index 0000000000..ca72dde46f
--- /dev/null
+++ b/unittests/Rewrite/RewriterTest.cpp
@@ -0,0 +1,80 @@
+//===- unittests/Rewrite/RewriterTest.cpp - Rewriter tests ----------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "clang/Rewrite/Core/Rewriter.h"
+#include "clang/Tooling/Tooling.h"
+#include "gtest/gtest.h"
+
+using namespace clang;
+
+namespace {
+
+struct RangeTypeTest {
+  std::unique_ptr<ASTUnit> AST;
+  Rewriter Rewrite;
+  SourceLocation FileStart;
+  CharSourceRange CRange; // covers exact char range
+  CharSourceRange TRange; // extends CRange to whole tokens
+  SourceRange SRange;     // different type but behaves like TRange
+  SourceLocation makeLoc(int Off) { return FileStart.getLocWithOffset(Off); }
+  CharSourceRange makeCharRange(int StartOff, int EndOff) {
+    return CharSourceRange::getCharRange(makeLoc(StartOff), makeLoc(EndOff));
+  }
+  RangeTypeTest(StringRef Code, int StartOff, int EndOff) {
+    AST = tooling::buildASTFromCode(Code);
+    ASTContext &C = AST->getASTContext();
+    Rewrite = Rewriter(C.getSourceManager(), C.getLangOpts());
+    FileStart = AST->getStartOfMainFileID();
+    CRange = makeCharRange(StartOff, EndOff);
+    SRange = CRange.getAsRange();
+    TRange = CharSourceRange::getTokenRange(SRange);
+  }
+};
+
+TEST(Rewriter, GetRewrittenTextRangeTypes) {
+  // Check that correct text is retrieved for each range type.  Check again
+  // after a modification.  Ranges remain in terms of the original text but
+  // include the new text.
+  StringRef Code = "int main() { return 0; }";
+  //              get char range ^~~    = "ret"
+  //             get token range ^~~+++ = "return"
+  //            get source range ^~~+++ = "return"
+  //                  insert "x" ^
+  //              get char range ^~~    = "xret"
+  //             get token range ^~~+++ = "xreturn"
+  //            get source range ^~~+++ = "xreturn"
+  RangeTypeTest T(Code, 13, 16);
+  EXPECT_EQ(T.Rewrite.getRewrittenText(T.CRange), "ret");
+  EXPECT_EQ(T.Rewrite.getRewrittenText(T.TRange), "return");
+  EXPECT_EQ(T.Rewrite.getRewrittenText(T.SRange), "return");
+  T.Rewrite.InsertText(T.makeLoc(13), "x");
+  EXPECT_EQ(T.Rewrite.getRewrittenText(T.CRange), "xret");
+  EXPECT_EQ(T.Rewrite.getRewrittenText(T.TRange), "xreturn");
+  EXPECT_EQ(T.Rewrite.getRewrittenText(T.SRange), "xreturn");
+}
+
+TEST(Rewriter, ReplaceTextRangeTypes) {
+  // Check that correct text is replaced for each range type.  Ranges remain in
+  // terms of the original text but include the new text.
+  StringRef Code = "int main(int argc, char *argv[]) { return argc; }";
+  //                            replace char range with "foo" ^~
+  //                                                      get ^~~~~ = "foogc;"
+  //                           replace token range with "bar" ^~++
+  //                                                      get ^~~~~ = "bar;"
+  //                            replace source range with "0" ^~++
+  //                                                      get ^~~~~ = "0;"
+  RangeTypeTest T(Code, 42, 44);
+  T.Rewrite.ReplaceText(T.CRange, "foo");
+  EXPECT_EQ(T.Rewrite.getRewrittenText(T.makeCharRange(42, 47)), "foogc;");
+  T.Rewrite.ReplaceText(T.TRange, "bar");
+  EXPECT_EQ(T.Rewrite.getRewrittenText(T.makeCharRange(42, 47)), "bar;");
+  T.Rewrite.ReplaceText(T.SRange, "0");
+  EXPECT_EQ(T.Rewrite.getRewrittenText(T.makeCharRange(42, 47)), "0;");
+}
+
+} // anonymous namespace