Searching for tokens including comments

[barcisz]

Summary

This is a small diff that adds an optional argument to the Lexer::findNextToken that allows for searching for the token while including comments (instead of having to
instantiate the lexer separately). This is a prerequisite to D133942, used to insert a statement after another statement while preserving the comments that are on the same line, which is itself a prerequisite to D133956, a new CUDA clang-tidy check

[aaron.ballman]

Can you help me understand the expected use for these change?

clang/lib/Lex/Lexer.cpp

1278

Drive-by change? (We don't typically use top-level const except for occasional structure members.)

[barcisz]

In D133725#3792787, @aaron.ballman wrote:

Can you help me understand the expected use for these change?

Changed the description to answer your question - it's a prerequisite to D133942 and D133956

[aaron.ballman]

Ah, thank you for the extra description, that helped a lot. The changes LGTM, though you should hold off on landing them until we're sure the clang-tidy changes are going in.

clang/include/clang/Lex/Lexer.h

static PreambleBounds ComputePreamble(StringRef Buffer,
                                         const LangOptions &LangOpts,
                                         unsigned MaxLines = 0);
 
   /// Finds the token that comes right after the given location.
   ///
   /// Returns the next token, or none if the location is inside a macro.
   static Optional<Token> findNextToken(SourceLocation Loc,
                                        const SourceManager &SM,
-                                       const LangOptions &LangOpts);
+                                       const LangOptions &LangOpts,
+                                       bool IncludeComments = false);
 
   /// Checks that the given token is the first token that occurs after
   /// the given location (this excludes comments and whitespace). Returns the
   /// location immediately after the specified token. If the token is not found
   /// or the location is inside a macro, the returned source location will be
   /// invalid.
   static SourceLocation findLocationAfterToken(SourceLocation loc,
                                          tok::TokenKind TKind,

clang/lib/Lex/Lexer.cpp

while (true) {
     unsigned NewLineSize = Lexer::getEscapedNewLineSize(AfterEscape);
     if (NewLineSize == 0) return P;
     P = AfterEscape+NewLineSize;
   }
 }
 
 Optional<Token> Lexer::findNextToken(SourceLocation Loc,
                                      const SourceManager &SM,
-                                     const LangOptions &LangOpts) {
+                                     const LangOptions &LangOpts,
+                                     bool IncludeComments) {
   if (Loc.isMacroID()) {
     if (!Lexer::isAtEndOfMacroExpansion(Loc, SM, LangOpts, &Loc))
       return None;
   }
   Loc = Lexer::getLocForEndOfToken(Loc, 0, SM, LangOpts);
 
   // Break down the source location.
   std::pair<FileID, unsigned> LocInfo = SM.getDecomposedLoc(Loc);
 
   // Try to load the file buffer.
   bool InvalidTemp = false;
   StringRef File = SM.getBufferData(LocInfo.first, &InvalidTemp);
   if (InvalidTemp)
     return None;
 
-  const char *TokenBegin = File.data() + LocInfo.second;
+  const char *const TokenBegin = File.data() + LocInfo.second;
 
   // Lex from the start of the given location.
   Lexer lexer(SM.getLocForStartOfFile(LocInfo.first), LangOpts, File.begin(),
                                       TokenBegin, File.end());
+  lexer.SetCommentRetentionState(IncludeComments);
+
   // Find the token.
   Token Tok;
   lexer.LexFromRawLexer(Tok);
   return Tok;
 }
 
 /// Checks that the given token is the first token that occurs after the
 /// given location (this excludes comments and whitespace). Returns the location

clang/unittests/Lex/LexerTest.cpp

if (T->is(tok::eof))
       break;
     GeneratedByNextToken.push_back(getSourceText(*T, *T));
     Loc = T->getLocation();
   }
   EXPECT_THAT(GeneratedByNextToken, ElementsAre("abcd", "=", "0", ";", "int",
                                                 "xyz", "=", "abcd", ";"));
 }
 
+TEST_F(LexerTest, FindNextTokenWithComments) {
+  Lex("int /* type */ abcd = 0; // vardecl\n"
+      "int /*other type*/ xyz = abcd; //other vardecl \n");
+  std::vector<std::string> GeneratedByNextToken;
+  SourceLocation Loc =
+      SourceMgr.getLocForStartOfFile(SourceMgr.getMainFileID());
+  while (true) {
+    auto T = Lexer::findNextToken(Loc, SourceMgr, LangOpts, true);
+    ASSERT_TRUE(T);
+    if (T->is(tok::eof))
+      break;
+    GeneratedByNextToken.push_back(getSourceText(*T, *T));
+    Loc = T->getLocation();
+  }
+  EXPECT_THAT(GeneratedByNextToken, ElementsAre("/* type */", "abcd", "=", "0", ";", "// vardecl", "int", "/*other type*/",
+                                                "xyz", "=", "abcd", ";", "//other vardecl "));
+}
+
 TEST_F(LexerTest, CreatedFIDCountForPredefinedBuffer) {
   TrivialModuleLoader ModLoader;
   auto PP = CreatePP("", ModLoader);
   while (1) {
     Token tok;
     PP->Lex(tok);
     if (tok.is(tok::eof))
       break;