[clang] Mark re-injected tokens appropriately during pragma handling

[kadircet]

Summary

This hides such tokens from TokenWatcher, preventing crashes in clients
trying to match spelled and expanded tokens.

Fixes https://github.com/clangd/clangd/issues/712

[kadircet]

• Add test

[sammccall]

clang/lib/Parse/ParsePragma.cpp

2616

It was surprisingly hard for me to see that the changes here were just marking everything in TokenArray as reinjected.

Could you do that directly with a loop at the end instead? (It's hard to imagine a significant performance difference)

(And the same in all the other modified sites)

2623

as discussed offline, it wasn't clear to me that this was "in time" as we're not trying to affect the behavior of the subsequent PP.Lex call, not the previous one.

Maybe we could extract a function to mark an ArrayRef as reinject with a name that makes this clear.
markAsReinjectedForRelexing() or something...

[kadircet]

• Introduce a helper to mark tokens as reinjected

clang/lib/Parse/ParsePragma.cpp

 //===--- ParsePragma.cpp - Language specific pragma parsing ---------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements the language specific #pragma handlers.
 //
 //===----------------------------------------------------------------------===//
 
 #include "clang/AST/ASTContext.h"
 #include "clang/Basic/PragmaKinds.h"
 #include "clang/Basic/TargetInfo.h"
 #include "clang/Lex/Preprocessor.h"
+#include "clang/Lex/Token.h"
 #include "clang/Parse/LoopHint.h"
 #include "clang/Parse/ParseDiagnostic.h"
 #include "clang/Parse/Parser.h"
 #include "clang/Parse/RAIIObjectsForParser.h"
 #include "clang/Sema/Scope.h"
+#include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/StringSwitch.h"
 using namespace clang;
 
 namespace {
 
 struct PragmaAlignHandler : public PragmaHandler {
   explicit PragmaAlignHandler() : PragmaHandler("align") {}
   void HandlePragma(Preprocessor &PP, PragmaIntroducer Introducer,

 };
 
 struct PragmaMaxTokensTotalHandler : public PragmaHandler {
   PragmaMaxTokensTotalHandler() : PragmaHandler("max_tokens_total") {}
   void HandlePragma(Preprocessor &PP, PragmaIntroducer Introducer,
                     Token &FirstToken) override;
 };
 
+void markAsReinjectedForRelexing(llvm::MutableArrayRef<clang::Token> Toks) {
+  for (auto &T : Toks)
+    T.setFlag(clang::Token::IsReinjected);
+}
 }  // end namespace
 
 void Parser::initializePragmaHandlers() {
   AlignHandler = std::make_unique<PragmaAlignHandler>();
   PP.AddPragmaHandler(AlignHandler.get());
 
   GCCVisibilityHandler = std::make_unique<PragmaGCCVisibilityHandler>();
   PP.AddPragmaHandler("GCC", GCCVisibilityHandler.get());

 
 /// Handle all MS pragmas.  Simply forwards the tokens after inserting
 /// an annotation token.
 void PragmaMSPragma::HandlePragma(Preprocessor &PP,
                                   PragmaIntroducer Introducer, Token &Tok) {
   Token EoF, AnnotTok;
   EoF.startToken();
   EoF.setKind(tok::eof);
   AnnotTok.startToken();
   AnnotTok.setKind(tok::annot_pragma_ms_pragma);
   AnnotTok.setLocation(Tok.getLocation());
   AnnotTok.setAnnotationEndLoc(Tok.getLocation());
   SmallVector<Token, 8> TokenVector;
   // Suck up all of the tokens before the eod.
   for (; Tok.isNot(tok::eod); PP.Lex(Tok)) {
     TokenVector.push_back(Tok);
     AnnotTok.setAnnotationEndLoc(Tok.getLocation());
   }
   // Add a sentinel EoF token to the end of the list.
   TokenVector.push_back(EoF);
   // We must allocate this array with new because EnterTokenStream is going to
   // delete it later.
+  markAsReinjectedForRelexing(TokenVector);
   auto TokenArray = std::make_unique<Token[]>(TokenVector.size());
   std::copy(TokenVector.begin(), TokenVector.end(), TokenArray.get());
   auto Value = new (PP.getPreprocessorAllocator())
       std::pair<std::unique_ptr<Token[]>, size_t>(std::move(TokenArray),
                                                   TokenVector.size());
   AnnotTok.setAnnotationValue(Value);
   PP.EnterToken(AnnotTok, /*IsReinject*/ false);
 }
static bool ParseLoopHintValue(Preprocessor &PP, Token &Tok, Token PragmaName,
   }
 
   Token EOFTok;
   EOFTok.startToken();
   EOFTok.setKind(tok::eof);
   EOFTok.setLocation(Tok.getLocation());
   ValueList.push_back(EOFTok); // Terminates expression for parsing.
 
+  markAsReinjectedForRelexing(ValueList);
   Info.Toks = llvm::makeArrayRef(ValueList).copy(PP.getPreprocessorAllocator());
 
   Info.PragmaName = PragmaName;
   Info.Option = Option;
   return false;
 }
 
 /// Handle the \#pragma clang loop directive.
if ((Info->Action == PragmaAttributeInfo::Push && Tok.isNot(tok::eod)) ||
 
     // Terminate the attribute for parsing.
     Token EOFTok;
     EOFTok.startToken();
     EOFTok.setKind(tok::eof);
     EOFTok.setLocation(EndLoc);
     AttributeTokens.push_back(EOFTok);
 
+    markAsReinjectedForRelexing(AttributeTokens);
     Info->Tokens =
         llvm::makeArrayRef(AttributeTokens).copy(PP.getPreprocessorAllocator());
   }
 
   if (Tok.isNot(tok::eod))
     PP.Diag(Tok.getLocation(), diag::warn_pragma_extra_tokens_at_eol)
         << "clang attribute";
 

clang/unittests/Tooling/Syntax/TokensTest.cpp

TEST_F(TokenBufferTest, ExpandedBySpelled) {
   recordTokens(R"cpp(
     #define EMPTY
     EMPTY int a = 100;
   )cpp");
   EXPECT_THAT(Buffer.expandedForSpelled(findSpelled("EMPTY int").drop_back()),
               IsEmpty());
 }
 
+TEST_F(TokenCollectorTest, Pragmas) {
+  // Tokens coming from concatenations.
+  recordTokens(R"cpp(
+    void foo() {
+      #pragma unroll 4
+      for(int i=0;i<4;++i);
+    }
+  )cpp");
+}
 } // namespace