[clang][Tooling] Add a way to tokenize a FileRange

[sammccall]

clang/lib/Tooling/Syntax/Tokens.cpp

335

Discussed offline, this loop includes an extra token if the truncation is at whitespace between tokens. (Please test this case)

Also the eof comment is confusing.

I think the loop should be rewritten.

346

nit: FileRange(FID, 0, SM.getFileIDSize(FID)) is a lot more direct :-)

[kadircet]

• Address comments

[kadircet]

clang/lib/Tooling/Syntax/Tokens.cpp

335

instead of re-writing the loop I just changed the after-the-loop check to verify latest lexed token to be in range.

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

 /// setting the appropriate token kind (instead of the raw_identifier reported
 /// by lexer in raw mode). This is a very low-level function, most users should
 /// prefer to use TokenCollector. Lexing in raw mode produces wildly different
 /// results from what one might expect when running a C++ frontend, e.g.
 /// preprocessor does not run at all.
 /// The result will *not* have a 'eof' token at the end.
 std::vector<syntax::Token> tokenize(FileID FID, const SourceManager &SM,
                                     const LangOptions &LO);
+/// Similar to one above, instead of whole file tokenizes a part of it. Note
+/// that, the first token might be incomplete if FR.startOffset is not at the
+/// beginning of a token, and the last token returned will start before the
+/// FR.endOffset but might end after it.
+std::vector<syntax::Token>
+tokenize(const FileRange &FR, const SourceManager &SM, const LangOptions &LO);
 
 /// Collects tokens for the main file while running the frontend action. An
 /// instance of this object should be created on
 /// FrontendAction::BeginSourceFile() and the results should be consumed after
 /// FrontendAction::Execute() finishes.
 class TokenCollector {
 public:
   /// Adds the hooks to collect the tokens. Should be called before the

clang/lib/Tooling/Syntax/Tokens.cpp

 }
 
 FileRange syntax::Token::range(const SourceManager &SM,
                                const syntax::Token &First,
                                const syntax::Token &Last) {
   auto F = First.range(SM);
   auto L = Last.range(SM);
   assert(F.file() == L.file() && "tokens from different files");
-  assert((F == L || F.endOffset() <= L.beginOffset()) && "wrong order of tokens");
+  assert((F == L || F.endOffset() <= L.beginOffset()) &&
+         "wrong order of tokens");
   return FileRange(F.file(), F.beginOffset(), L.endOffset());
 }
 
 llvm::raw_ostream &syntax::operator<<(llvm::raw_ostream &OS, const Token &T) {
   return OS << T.str();
 }
 
 FileRange::FileRange(FileID File, unsigned BeginOffset, unsigned EndOffset)
TokenBuffer::macroExpansions(FileID FID) const {
   for (auto Mapping : File.Mappings) {
     const syntax::Token *Token = &Spelled[Mapping.BeginSpelled];
     if (Token->kind() == tok::TokenKind::identifier)
       Expansions.push_back(Token);
   }
   return Expansions;
 }
 
-std::vector<syntax::Token> syntax::tokenize(FileID FID, const SourceManager &SM,
+std::vector<syntax::Token> syntax::tokenize(const FileRange &FR,
+                                            const SourceManager &SM,
                                             const LangOptions &LO) {
   std::vector<syntax::Token> Tokens;
   IdentifierTable Identifiers(LO);
   auto AddToken = [&](clang::Token T) {
     // Fill the proper token kind for keywords, etc.
     if (T.getKind() == tok::raw_identifier && !T.needsCleaning() &&
         !T.hasUCN()) { // FIXME: support needsCleaning and hasUCN cases.
       clang::IdentifierInfo &II = Identifiers.get(T.getRawIdentifier());
       T.setIdentifierInfo(&II);
       T.setKind(II.getTokenID());
     }
     Tokens.push_back(syntax::Token(T));
   };
 
-  Lexer L(FID, SM.getBuffer(FID), SM, LO);
+  auto SrcBuffer = SM.getBufferData(FR.file());
+  Lexer L(SM.getLocForStartOfFile(FR.file()), LO, SrcBuffer.data(),
+          SrcBuffer.data() + FR.beginOffset(),
+          // We can't make BufEnd point to FR.endOffset, as Lexer requires a
+          // null terminated buffer.
+          SrcBuffer.data() + SrcBuffer.size());
 
   clang::Token T;
-  while (!L.LexFromRawLexer(T))
+  while (!L.LexFromRawLexer(T) && L.getCurrentBufferOffset() < FR.endOffset())
     AddToken(T);
-  // 'eof' is only the last token if the input is null-terminated. Never store
-  // it, for consistency.
-  if (T.getKind() != tok::eof)
+  // LexFromRawLexer returns true when it parses the last token of the file, add
+  // it iff it starts within the range we are interested in.
+  if (SM.getFileOffset(T.getLocation()) < FR.endOffset())
     AddToken(T);
   return Tokens;
 }
 
+std::vector<syntax::Token> syntax::tokenize(FileID FID, const SourceManager &SM,
+                                            const LangOptions &LO) {
+  return tokenize(syntax::FileRange(FID, 0, SM.getFileIDSize(FID)), SM, LO);
+}
+
 /// Records information reqired to construct mappings for the token buffer that
 /// we are collecting.
 class TokenCollector::CollectPPExpansions : public PPCallbacks {
 public:
   CollectPPExpansions(TokenCollector &C) : Collector(&C) {}
 
   /// Disabled instance will stop reporting anything to TokenCollector.
   /// This ensures that uses of the preprocessor after TokenCollector::consume()

clang/unittests/Tooling/Syntax/TokensTest.cpp

public:
   // Adds a file to the test VFS.
   void addFile(llvm::StringRef Path, llvm::StringRef Contents) {
     if (!FS->addFile(Path, time_t(),
                      llvm::MemoryBuffer::getMemBufferCopy(Contents))) {
       ADD_FAILURE() << "could not add a file to VFS: " << Path;
     }
   }
 
-  /// Add a new file, run syntax::tokenize() on it and return the results.
+  /// Add a new file, run syntax::tokenize() on the range if any, run it on the
+  /// whole file otherwise and return the results.
   std::vector<syntax::Token> tokenize(llvm::StringRef Text) {
+    llvm::Annotations Annot(Text);
+    auto FID = SourceMgr->createFileID(
+        llvm::MemoryBuffer::getMemBufferCopy(Annot.code()));
     // FIXME: pass proper LangOptions.
+    if (Annot.ranges().empty())
+      return syntax::tokenize(FID, *SourceMgr, LangOptions());
     return syntax::tokenize(
-        SourceMgr->createFileID(llvm::MemoryBuffer::getMemBufferCopy(Text)),
+        syntax::FileRange(FID, Annot.range().Begin, Annot.range().End),
         *SourceMgr, LangOptions());
   }
 
   // Specialized versions of matchers that hide the SourceManager from clients.
   Matcher<syntax::Token> HasText(std::string Text) const {
     return ::HasText(Text, SourceMgr.get());
   }
   Matcher<syntax::Token> RangeIs(llvm::Annotations::Range R) const {
EXPECT_THAT(tokenize("int main() {}"),
                           AllOf(HasText("main"), Kind(tok::identifier)),
                           Kind(tok::l_paren), Kind(tok::r_paren),
                           Kind(tok::l_brace), Kind(tok::r_brace)));
   // Comments are ignored for now.
   EXPECT_THAT(tokenize("/* foo */int a; // more comments"),
               ElementsAre(Kind(tok::kw_int),
                           AllOf(HasText("a"), Kind(tok::identifier)),
                           Kind(tok::semi)));
+  EXPECT_THAT(tokenize("int [[main() {]]}"),
+              ElementsAre(AllOf(HasText("main"), Kind(tok::identifier)),
+                          Kind(tok::l_paren), Kind(tok::r_paren),
+                          Kind(tok::l_brace)));
+  EXPECT_THAT(tokenize("int [[main() {   ]]}"),
+              ElementsAre(AllOf(HasText("main"), Kind(tok::identifier)),
+                          Kind(tok::l_paren), Kind(tok::r_paren),
+                          Kind(tok::l_brace)));
+  // First token is partially parsed, last token is fully included even though
+  // only a part of it is contained in the range.
+  EXPECT_THAT(tokenize("int m[[ain() {ret]]urn 0;}"),
+              ElementsAre(AllOf(HasText("ain"), Kind(tok::identifier)),
+                          Kind(tok::l_paren), Kind(tok::r_paren),
+                          Kind(tok::l_brace), Kind(tok::kw_return)));
 }
 
 TEST_F(TokenCollectorTest, Basic) {
   std::pair</*Input*/ std::string, /*Expected*/ std::string> TestCases[] = {
       {"int main() {}",
        R"(expanded tokens:
   int main ( ) { }
 file './input.cpp'