[clang][Syntax] Add spelledTokenAt helper to TokenBuffer

[sammccall]

Thanks for doing this!

clang/include/clang/Tooling/Syntax/Tokens.h

public:
   /// Lexed tokens of a file before preprocessing. E.g. for the following input
   ///     #define DECL(name) int name = 10
   ///     DECL(a);
   /// spelledTokens() returns
   ///    {"#", "define", "DECL", "(", "name", ")", "int", "name", "=", "10",
   ///     "DECL", "(", "a", ")", ";"}
   llvm::ArrayRef<syntax::Token> spelledTokens(FileID FID) const;
 
+  /// Returns the spelled Token starting at Loc, if there are no such tokens
+  /// returns nullptr.
+  const syntax::Token *spelledTokenAt(SourceLocation Loc) const;
+
   /// Get all tokens that expand a macro in \p FID. For the following input
   ///     #define FOO B
   ///     #define FOO2(X) int X
   ///     FOO2(XY)
   ///     int B;
   ///     FOO;
   /// macroExpansions() returns {"FOO2", "FOO"} (from line 3 and 5
   /// respecitvely).

clang/lib/Tooling/Syntax/Tokens.cpp

 }
 
 llvm::ArrayRef<syntax::Token> TokenBuffer::spelledTokens(FileID FID) const {
   auto It = Files.find(FID);
   assert(It != Files.end());
   return It->second.SpelledTokens;
 }
 
+const syntax::Token *TokenBuffer::spelledTokenAt(SourceLocation Loc) const {
+  assert(Loc.isFileID());
+  const auto *Tok = llvm::partition_point(
+      spelledTokens(SourceMgr->getFileID(Loc)),
+      [&](const syntax::Token &Tok) { return Tok.location() < Loc; });
+  if (!Tok || Tok->location() != Loc)
+    return nullptr;
+  return Tok;
+}
+
 std::string TokenBuffer::Mapping::str() const {
   return std::string(
       llvm::formatv("spelled tokens: [{0},{1}), expanded tokens: [{2},{3})",
                     BeginSpelled, EndSpelled, BeginExpanded, EndExpanded));
 }
 
 llvm::Optional<llvm::ArrayRef<syntax::Token>>
 TokenBuffer::spelledForExpanded(llvm::ArrayRef<syntax::Token> Expanded) const {

clang/unittests/Tooling/Syntax/TokensTest.cpp

 
 using llvm::ValueIs;
 using ::testing::AllOf;
 using ::testing::Contains;
 using ::testing::ElementsAre;
 using ::testing::Field;
 using ::testing::Matcher;
 using ::testing::Not;
+using ::testing::Pointee;
 using ::testing::StartsWith;
 
 namespace {
 // Checks the passed ArrayRef<T> has the same begin() and end() iterators as the
 // argument.
 MATCHER_P(SameRange, A, "") {
   return A.begin() == arg.begin() && A.end() == arg.end();
 }
TEST_F(TokenCollectorTest, Locations) {
   // Check spelled tokens.
   EXPECT_THAT(
       Buffer.spelledTokens(SourceMgr->getMainFileID()),
       ElementsAre(AllOf(Kind(tok::kw_int), RangeIs(Code.range("r1"))),
                   AllOf(Kind(tok::identifier), RangeIs(Code.range("r2"))),
                   AllOf(Kind(tok::equal), RangeIs(Code.range("r3"))),
                   AllOf(Kind(tok::string_literal), RangeIs(Code.range("r4"))),
                   AllOf(Kind(tok::semi), RangeIs(Code.range("r5")))));
+
+  auto StartLoc = SourceMgr->getLocForStartOfFile(SourceMgr->getMainFileID());
+  for (auto &R : Code.ranges()) {
+    EXPECT_THAT(Buffer.spelledTokenAt(StartLoc.getLocWithOffset(R.Begin)),
+                Pointee(RangeIs(R)));
+  }
 }
 
 TEST_F(TokenCollectorTest, MacroDirectives) {
   // Macro directives are not stored anywhere at the moment.
   std::string Code = R"cpp(
     #define FOO a
     #include "unresolved_file.h"
     #undef FOO