[clang] 65cb4fd - [libTooling] Change `after` range-selector to operate only on source ranges

[Yitzhak Mandelbaum via cfe-commits]

Author: Yitzhak Mandelbaum
Date: 2020-10-15T20:58:30Z
New Revision: 65cb4fdd69f43b6c39a8e4ca27b509284b11d807

URL: https://github.com/llvm/llvm-project/commit/65cb4fdd69f43b6c39a8e4ca27b509284b11d807
DIFF: https://github.com/llvm/llvm-project/commit/65cb4fdd69f43b6c39a8e4ca27b509284b11d807.diff

LOG: [libTooling] Change `after` range-selector to operate only on source ranges

Currently, `after` fails when applied to locations in macro arguments.  This
change projects the subrange into a file source range and then applies `after`.

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

Added: 
    

Modified: 
    clang/lib/Tooling/Transformer/RangeSelector.cpp
    clang/unittests/Tooling/RangeSelectorTest.cpp

Removed: 
    


################################################################################
diff  --git a/clang/lib/Tooling/Transformer/RangeSelector.cpp b/clang/lib/Tooling/Transformer/RangeSelector.cpp
index 29b1a5b0372e..ce6f5fb9b444 100644
--- a/clang/lib/Tooling/Transformer/RangeSelector.cpp
+++ b/clang/lib/Tooling/Transformer/RangeSelector.cpp
@@ -116,11 +116,24 @@ RangeSelector transformer::after(RangeSelector Selector) {
     Expected<CharSourceRange> SelectedRange = Selector(Result);
     if (!SelectedRange)
       return SelectedRange.takeError();
-    if (SelectedRange->isCharRange())
-      return CharSourceRange::getCharRange(SelectedRange->getEnd());
-    return CharSourceRange::getCharRange(Lexer::getLocForEndOfToken(
-        SelectedRange->getEnd(), 0, Result.Context->getSourceManager(),
-        Result.Context->getLangOpts()));
+    SourceLocation End = SelectedRange->getEnd();
+    if (SelectedRange->isTokenRange()) {
+      // We need to find the actual (exclusive) end location from which to
+      // create a new source range. However, that's not guaranteed to be valid,
+      // even if the token location itself is valid. So, we create a token range
+      // consisting only of the last token, then map that range back to the
+      // source file. If that succeeds, we have a valid location for the end of
+      // the generated range.
+      CharSourceRange Range = Lexer::makeFileCharRange(
+          CharSourceRange::getTokenRange(SelectedRange->getEnd()),
+          *Result.SourceManager, Result.Context->getLangOpts());
+      if (Range.isInvalid())
+        return invalidArgumentError(
+            "after: can't resolve sub-range to valid source range");
+      End = Range.getEnd();
+    }
+
+    return CharSourceRange::getCharRange(End);
   };
 }
 

diff  --git a/clang/unittests/Tooling/RangeSelectorTest.cpp b/clang/unittests/Tooling/RangeSelectorTest.cpp
index 64ddee7894eb..499c9a8cc666 100644
--- a/clang/unittests/Tooling/RangeSelectorTest.cpp
+++ b/clang/unittests/Tooling/RangeSelectorTest.cpp
@@ -193,6 +193,65 @@ TEST(RangeSelectorTest, AfterOp) {
                        HasValue(EqualsCharSourceRange(ExpectedAfter)));
 }
 
+// Gets the spelling location `Length` characters after the start of AST node
+// `Id`.
+static SourceLocation getSpellingLocAfter(const MatchResult &Result,
+                                          StringRef Id, int Length) {
+  const auto *E = Result.Nodes.getNodeAs<Expr>(Id);
+  assert(E != nullptr);
+  return Result.SourceManager->getSpellingLoc(E->getBeginLoc())
+      .getLocWithOffset(Length);
+}
+
+// Test with a range that is the entire macro arg, but does not end the
+// expansion itself.
+TEST(RangeSelectorTest, AfterOpInMacroArg) {
+  StringRef Code = R"cc(
+#define ISNULL(x) x == nullptr
+    bool g() { int* y; return ISNULL(y); }
+  )cc";
+
+  TestMatch Match =
+      matchCode(Code, declRefExpr(to(namedDecl(hasName("y")))).bind("yvar"));
+  int YVarLen = 1;
+  SourceLocation After = getSpellingLocAfter(Match.Result, "yvar", YVarLen);
+  CharSourceRange Expected = CharSourceRange::getCharRange(After, After);
+  EXPECT_THAT_EXPECTED(after(node("yvar"))(Match.Result),
+                       HasValue(EqualsCharSourceRange(Expected)));
+}
+
+// Test with a range that is the entire macro arg and ends the expansion itself.
+TEST(RangeSelectorTest, AfterOpInMacroArgEndsExpansion) {
+  StringRef Code = R"cc(
+#define ISNULL(x) nullptr == x
+    bool g() { int* y; return ISNULL(y); }
+  )cc";
+
+  TestMatch Match =
+      matchCode(Code, declRefExpr(to(namedDecl(hasName("y")))).bind("yvar"));
+  int YVarLen = 1;
+  SourceLocation After = getSpellingLocAfter(Match.Result, "yvar", YVarLen);
+  CharSourceRange Expected = CharSourceRange::getCharRange(After, After);
+  EXPECT_THAT_EXPECTED(after(node("yvar"))(Match.Result),
+                       HasValue(EqualsCharSourceRange(Expected)));
+}
+
+TEST(RangeSelectorTest, AfterOpInPartOfMacroArg) {
+  StringRef Code = R"cc(
+#define ISNULL(x) x == nullptr
+    int* f(int*);
+    bool g() { int* y; return ISNULL(f(y)); }
+  )cc";
+
+  TestMatch Match =
+      matchCode(Code, declRefExpr(to(namedDecl(hasName("y")))).bind("yvar"));
+  int YVarLen = 1;
+  SourceLocation After = getSpellingLocAfter(Match.Result, "yvar", YVarLen);
+  CharSourceRange Expected = CharSourceRange::getCharRange(After, After);
+  EXPECT_THAT_EXPECTED(after(node("yvar"))(Match.Result),
+                       HasValue(EqualsCharSourceRange(Expected)));
+}
+
 TEST(RangeSelectorTest, BetweenOp) {
   StringRef Code = R"cc(
     int f(int x, int y, int z) { return 3; }