1 files changed, 108 insertions, 2 deletions

diff --git a/clang/unittests/Tooling/SourceCodeTest.cpp b/clang/unittests/Tooling/SourceCodeTest.cpp
index 258947a1a70..e3da9bf14b6 100644
--- a/clang/unittests/Tooling/SourceCodeTest.cpp
+++ b/clang/unittests/Tooling/SourceCodeTest.cpp
@@ -6,17 +6,32 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "clang/Tooling/Refactoring/SourceCode.h"
 #include "TestVisitor.h"
 #include "clang/Basic/Diagnostic.h"
-#include "clang/Tooling/Refactoring/SourceCode.h"
+#include "llvm/Testing/Support/Annotations.h"
+#include "llvm/Testing/Support/SupportHelpers.h"
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
 
 using namespace clang;
 
-using tooling::getText;
+using llvm::ValueIs;
 using tooling::getExtendedText;
+using tooling::getRangeForEdit;
+using tooling::getText;
 
 namespace {
 
+struct IntLitVisitor : TestVisitor<IntLitVisitor> {
+  bool VisitIntegerLiteral(IntegerLiteral *Expr) {
+    OnIntLit(Expr, Context);
+    return true;
+  }
+
+  std::function<void(IntegerLiteral *, ASTContext *Context)> OnIntLit;
+};
+
 struct CallsVisitor : TestVisitor<CallsVisitor> {
   bool VisitCallExpr(CallExpr *Expr) {
     OnCall(Expr, Context);
@@ -26,6 +41,19 @@ struct CallsVisitor : TestVisitor<CallsVisitor> {
   std::function<void(CallExpr *, ASTContext *Context)> OnCall;
 };
 
+// Equality matcher for `clang::CharSourceRange`, which lacks `operator==`.
+MATCHER_P(EqualsRange, R, "") {
+  return arg.isTokenRange() == R.isTokenRange() &&
+         arg.getBegin() == R.getBegin() && arg.getEnd() == R.getEnd();
+}
+
+static ::testing::Matcher<CharSourceRange> AsRange(const SourceManager &SM,
+                                                   llvm::Annotations::Range R) {
+  return EqualsRange(CharSourceRange::getCharRange(
+      SM.getLocForStartOfFile(SM.getMainFileID()).getLocWithOffset(R.Begin),
+      SM.getLocForStartOfFile(SM.getMainFileID()).getLocWithOffset(R.End)));
+}
+
 TEST(SourceCodeTest, getText) {
   CallsVisitor Visitor;
 
@@ -94,4 +122,82 @@ TEST(SourceCodeTest, getExtendedText) {
   Visitor.runOver("int foo() { return foo() + 3; }");
 }
 
+TEST(SourceCodeTest, EditRangeWithMacroExpansionsShouldSucceed) {
+  // The call expression, whose range we are extracting, includes two macro
+  // expansions.
+  llvm::Annotations Code(R"cpp(
+#define M(a) a * 13
+int foo(int x, int y);
+int a = $r[[foo(M(1), M(2))]];
+)cpp");
+
+  CallsVisitor Visitor;
+
+  Visitor.OnCall = [&Code](CallExpr *CE, ASTContext *Context) {
+    auto Range = CharSourceRange::getTokenRange(CE->getSourceRange());
+    EXPECT_THAT(getRangeForEdit(Range, *Context),
+                ValueIs(AsRange(Context->getSourceManager(), Code.range("r"))));
+  };
+  Visitor.runOver(Code.code());
+}
+
+TEST(SourceCodeTest, EditWholeMacroExpansionShouldSucceed) {
+  llvm::Annotations Code(R"cpp(
+#define FOO 10
+int a = $r[[FOO]];
+)cpp");
+
+  IntLitVisitor Visitor;
+  Visitor.OnIntLit = [&Code](IntegerLiteral *Expr, ASTContext *Context) {
+    auto Range = CharSourceRange::getTokenRange(Expr->getSourceRange());
+    EXPECT_THAT(getRangeForEdit(Range, *Context),
+                ValueIs(AsRange(Context->getSourceManager(), Code.range("r"))));
+  };
+  Visitor.runOver(Code.code());
+}
+
+TEST(SourceCodeTest, EditPartialMacroExpansionShouldFail) {
+  std::string Code = R"cpp(
+#define BAR 10+
+int c = BAR 3.0;
+)cpp";
+
+  IntLitVisitor Visitor;
+  Visitor.OnIntLit = [](IntegerLiteral *Expr, ASTContext *Context) {
+    auto Range = CharSourceRange::getTokenRange(Expr->getSourceRange());
+    EXPECT_FALSE(getRangeForEdit(Range, *Context).hasValue());
+  };
+  Visitor.runOver(Code);
+}
+
+TEST(SourceCodeTest, EditWholeMacroArgShouldSucceed) {
+  llvm::Annotations Code(R"cpp(
+#define FOO(a) a + 7.0;
+int a = FOO($r[[10]]);
+)cpp");
+
+  IntLitVisitor Visitor;
+  Visitor.OnIntLit = [&Code](IntegerLiteral *Expr, ASTContext *Context) {
+    auto Range = CharSourceRange::getTokenRange(Expr->getSourceRange());
+    EXPECT_THAT(getRangeForEdit(Range, *Context),
+                ValueIs(AsRange(Context->getSourceManager(), Code.range("r"))));
+  };
+  Visitor.runOver(Code.code());
+}
+
+TEST(SourceCodeTest, EditPartialMacroArgShouldSucceed) {
+  llvm::Annotations Code(R"cpp(
+#define FOO(a) a + 7.0;
+int a = FOO($r[[10]] + 10.0);
+)cpp");
+
+  IntLitVisitor Visitor;
+  Visitor.OnIntLit = [&Code](IntegerLiteral *Expr, ASTContext *Context) {
+    auto Range = CharSourceRange::getTokenRange(Expr->getSourceRange());
+    EXPECT_THAT(getRangeForEdit(Range, *Context),
+                ValueIs(AsRange(Context->getSourceManager(), Code.range("r"))));
+  };
+  Visitor.runOver(Code.code());
+}
+
 } // end anonymous namespace