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

[ymandel]

Summary

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

[tdl-g]

Just one comment about the tests.

clang/unittests/Tooling/RangeSelectorTest.cpp

196

If this helper function took an "expected" parameter I might consider it self-explanatory, but as it is, it's extremely non-obvious what it does without reading the body in detail--which means the tests that use it are not particularly readable either. (Each test reads like "here's the input data, make sure something unspecified is true about it.") Specifically, this function searches for a reference to the variable "y" and ensures that after() finds the character after it. So at a minimum, document that.

I'm also trying to decide if this helper is too similar to the implementation--tests should not just restate what the production code does, they should find some other way to validate the behavior. Do you think this is sufficiently different from the prod implementation to be meaningful? If so, that's fine. If not, maybe the helper should just take an expected byte offset, be renamed to "afterYHasByteOffset", and then each test would be a bit more readable?

Up to you.

[ymandel]

cleaned up test code

[ymandel]

clang/unittests/Tooling/RangeSelectorTest.cpp

196

I reduced the helper to just getting the spelling location at the given offset. the rest of the code I've inlined into the tests. That helper is still substantially less complicated than what happens in the production code, and I don't see any better way to find the location, unless we want to do an offset from, say, the braces of the function. that avoids macro-related calculations, but also disconnects from the var. So, I think it's just a tradeoff rather than an improvement.

clang/lib/Tooling/Transformer/RangeSelector.cpp

RangeSelector transformer::before(RangeSelector Selector) {
   };
 }
 
 RangeSelector transformer::after(RangeSelector Selector) {
   return [Selector](const MatchResult &Result) -> Expected<CharSourceRange> {
     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);
   };
 }
 
 RangeSelector transformer::node(std::string ID) {
   return [ID](const MatchResult &Result) -> Expected<CharSourceRange> {
     Expected<DynTypedNode> Node = getNode(Result.Nodes, ID);
     if (!Node)
       return Node.takeError();

clang/unittests/Tooling/RangeSelectorTest.cpp

EXPECT_THAT_EXPECTED(after(charRange(CharRange))(Match.Result),
                        HasValue(EqualsCharSourceRange(ExpectedAfter)));
 
   // Test with a token range.
   auto TokenRange = CharSourceRange::getTokenRange(Range);
   EXPECT_THAT_EXPECTED(after(charRange(TokenRange))(Match.Result),
                        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; }
     int g() { return f(3, /* comment */ 7 /* comment */, 9); }
   )cc";
   auto Matcher = callExpr(hasArgument(0, expr().bind("a0")),
                           hasArgument(1, expr().bind("a1")));
   RangeSelector R = between(node("a0"), node("a1"));