[clang] cd9b2e1 - [clang][Syntax] Add spelledTokenAt helper to TokenBuffer

[Kadir Cetinkaya via cfe-commits]

Author: Kadir Cetinkaya
Date: 2020-03-03T14:30:41+01:00
New Revision: cd9b2e18bd69503e8d624d427caa3a0157b34e52

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

LOG: [clang][Syntax] Add spelledTokenAt helper to TokenBuffer

Reviewers: sammccall

Subscribers: cfe-commits

Tags: #clang

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

Added: 
    

Modified: 
    clang/include/clang/Tooling/Syntax/Tokens.h
    clang/lib/Tooling/Syntax/Tokens.cpp
    clang/unittests/Tooling/Syntax/TokensTest.cpp

Removed: 
    


################################################################################
diff  --git a/clang/include/clang/Tooling/Syntax/Tokens.h b/clang/include/clang/Tooling/Syntax/Tokens.h
index 19d120ebbc9f..2ee840074810 100644
--- a/clang/include/clang/Tooling/Syntax/Tokens.h
+++ b/clang/include/clang/Tooling/Syntax/Tokens.h
@@ -245,6 +245,10 @@ class TokenBuffer {
   ///     "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

diff  --git a/clang/lib/Tooling/Syntax/Tokens.cpp b/clang/lib/Tooling/Syntax/Tokens.cpp
index ae5bc687553b..9e12d8b603bf 100644
--- a/clang/lib/Tooling/Syntax/Tokens.cpp
+++ b/clang/lib/Tooling/Syntax/Tokens.cpp
@@ -183,6 +183,16 @@ llvm::ArrayRef<syntax::Token> TokenBuffer::spelledTokens(FileID FID) const {
   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})",

diff  --git a/clang/unittests/Tooling/Syntax/TokensTest.cpp b/clang/unittests/Tooling/Syntax/TokensTest.cpp
index ad0293bc3e07..d4b015393286 100644
--- a/clang/unittests/Tooling/Syntax/TokensTest.cpp
+++ b/clang/unittests/Tooling/Syntax/TokensTest.cpp
@@ -59,6 +59,7 @@ using ::testing::ElementsAre;
 using ::testing::Field;
 using ::testing::Matcher;
 using ::testing::Not;
+using ::testing::Pointee;
 using ::testing::StartsWith;
 
 namespace {
@@ -363,6 +364,12 @@ TEST_F(TokenCollectorTest, Locations) {
                   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) {