[libTooling] Fix `maybeExtendRange` to support `CharRange`s.

[ymandel]

Summary

Currently, maybeExtendRange takes a CharSourceRange, but only works
correctly for the TokenRange case. This change adds proper support for the
CharRange case.

clang/lib/Tooling/Transformer/SourceCode.cpp

StringRef clang::tooling::getText(CharSourceRange Range,
                                   const ASTContext &Context) {
   return Lexer::getSourceText(Range, Context.getSourceManager(),
                               Context.getLangOpts());
 }
 
 CharSourceRange clang::tooling::maybeExtendRange(CharSourceRange Range,
                                                  tok::TokenKind Next,
                                                  ASTContext &Context) {
-  Optional<Token> Tok = Lexer::findNextToken(
-      Range.getEnd(), Context.getSourceManager(), Context.getLangOpts());
-  if (!Tok || !Tok->is(Next))
+  CharSourceRange R = Lexer::getAsCharRange(Range, Context.getSourceManager(),
+                                            Context.getLangOpts());
+  if (R.isInvalid())
+    return Range;
+  Token Tok;
+  bool Err =
+      Lexer::getRawToken(R.getEnd(), Tok, Context.getSourceManager(),
+                         Context.getLangOpts(), /*IgnoreWhiteSpace=*/true);
+  if (Err || !Tok.is(Next))
     return Range;
-  return CharSourceRange::getTokenRange(Range.getBegin(), Tok->getLocation());
+  return CharSourceRange::getTokenRange(Range.getBegin(), Tok.getLocation());
 }
 
 llvm::Error clang::tooling::validateEditRange(const CharSourceRange &Range,
                                               const SourceManager &SM) {
   if (Range.isInvalid())
     return llvm::make_error<StringError>(errc::invalid_argument,
                                          "Invalid range");
 

clang/unittests/Tooling/SourceCodeTest.cpp

 //===- unittest/Tooling/SourceCodeTest.cpp --------------------------------===//
 //
 // 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/Tooling/Transformer/SourceCode.h"
 #include "TestVisitor.h"
 #include "clang/Basic/Diagnostic.h"
+#include "clang/Basic/SourceLocation.h"
+#include "clang/Lex/Lexer.h"
 #include "llvm/Testing/Support/Annotations.h"
 #include "llvm/Testing/Support/Error.h"
 #include "llvm/Testing/Support/SupportHelpers.h"
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 
 using namespace clang;
 
 using llvm::Failed;
 using llvm::Succeeded;
 using llvm::ValueIs;
 using tooling::getAssociatedRange;
+using tooling::getExtendedRange;
 using tooling::getExtendedText;
 using tooling::getRangeForEdit;
 using tooling::getText;
+using tooling::maybeExtendRange;
 using tooling::validateEditRange;
 
 namespace {
 
 struct IntLitVisitor : TestVisitor<IntLitVisitor> {
   bool VisitIntegerLiteral(IntegerLiteral *Expr) {
     OnIntLit(Expr, Context);
     return true;

 };
 
 // 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();
 }
 
-MATCHER_P2(EqualsAnnotatedRange, SM, R, "") {
+MATCHER_P2(EqualsAnnotatedRange, Context, R, "") {
   if (arg.getBegin().isMacroID()) {
     *result_listener << "which starts in a macro";
     return false;
   }
   if (arg.getEnd().isMacroID()) {
     *result_listener << "which ends in a macro";
     return false;
   }
 
-  unsigned Begin = SM->getFileOffset(arg.getBegin());
-  unsigned End = SM->getFileOffset(arg.getEnd());
+  CharSourceRange Range = Lexer::getAsCharRange(
+      arg, Context->getSourceManager(), Context->getLangOpts());
+  unsigned Begin = Context->getSourceManager().getFileOffset(Range.getBegin());
+  unsigned End = Context->getSourceManager().getFileOffset(Range.getEnd());
 
-  *result_listener << "which is [" << Begin << ",";
-  if (arg.isTokenRange()) {
-    *result_listener << End << "]";
-    return Begin == R.Begin && End + 1 == R.End;
-  }
-  *result_listener << End << ")";
+  *result_listener << "which is a " << (arg.isTokenRange() ? "Token" : "Char")
+                   << " range [" << Begin << "," << End << ")";
   return Begin == R.Begin && End == R.End;
 }
 
 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)));
 }
 
 // Base class for visitors that expect a single match corresponding to a
 // specific annotated range.
 template <typename T> class AnnotatedCodeVisitor : public TestVisitor<T> {
-  llvm::Annotations Code;
+protected:
   int MatchCount = 0;
+  llvm::Annotations Code;
 
 public:
   AnnotatedCodeVisitor() : Code("$r[[]]") {}
+  // Helper for tests of `getAssociatedRange`.
   bool VisitDeclHelper(Decl *Decl) {
     // Only consider explicit declarations.
     if (Decl->isImplicit())
       return true;
 
     ++MatchCount;
     EXPECT_THAT(getAssociatedRange(*Decl, *this->Context),
-                EqualsAnnotatedRange(&this->Context->getSourceManager(),
-                                     Code.range("r")))
+                EqualsAnnotatedRange(this->Context, Code.range("r")))
         << Code.code();
     return true;
   }
 
   bool runOverAnnotated(llvm::StringRef AnnotatedCode,
                         std::vector<std::string> Args = {}) {
     Code = llvm::Annotations(AnnotatedCode);
     MatchCount = 0;
TEST(SourceCodeTest, getExtendedText) {
   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(SourceCodeTest, maybeExtendRange_TokenRange) {
+  struct ExtendTokenRangeVisitor
+      : AnnotatedCodeVisitor<ExtendTokenRangeVisitor> {
+    bool VisitCallExpr(CallExpr *CE) {
+      ++MatchCount;
+      EXPECT_THAT(getExtendedRange(*CE, tok::TokenKind::semi, *Context),
+                  EqualsAnnotatedRange(Context, Code.range("r")));
+      return true;
+    }
+  };
+
+  ExtendTokenRangeVisitor Visitor;
+  // Extends to include semicolon.
+  Visitor.runOverAnnotated("void f(int x, int y) { $r[[f(x, y);]] }");
+  // Does not extend to include semicolon.
+  Visitor.runOverAnnotated(
+      "int f(int x, int y) { if (0) return $r[[f(x, y)]] + 3; }");
+}
+
+TEST(SourceCodeTest, maybeExtendRange_CharRange) {
+  struct ExtendCharRangeVisitor : AnnotatedCodeVisitor<ExtendCharRangeVisitor> {
+    bool VisitCallExpr(CallExpr *CE) {
+      ++MatchCount;
+      CharSourceRange Call = Lexer::getAsCharRange(CE->getSourceRange(),
+                                                   Context->getSourceManager(),
+                                                   Context->getLangOpts());
+      EXPECT_THAT(maybeExtendRange(Call, tok::TokenKind::semi, *Context),
+                  EqualsAnnotatedRange(Context, Code.range("r")));
+      return true;
+    }
+  };
+  ExtendCharRangeVisitor Visitor;
+  // Extends to include semicolon.
+  Visitor.runOverAnnotated("void f(int x, int y) { $r[[f(x, y);]] }");
+  // Does not extend to include semicolon.
+  Visitor.runOverAnnotated(
+      "int f(int x, int y) { if (0) return $r[[f(x, y)]] + 3; }");
+}
+
 TEST(SourceCodeTest, getAssociatedRange) {
   struct VarDeclsVisitor : AnnotatedCodeVisitor<VarDeclsVisitor> {
     bool VisitVarDecl(VarDecl *Decl) { return VisitDeclHelper(Decl); }
   };
   VarDeclsVisitor Visitor;
 
   // Includes semicolon.
   Visitor.runOverAnnotated("$r[[int x = 4;]]");