[clang][deps] Teach dep directive scanner about _Pragma

[benlangmuir]

Summary

While we cannot handle _Pragma used inside macros, we an handle this at the top level, and it some projects use the _Pragma("once") spelling like that, which was causing spurious failures in the scanner.

Limitations

• Cannot handle #define ONCE _Pragma("once"), same issue as using @import in a macro -- ideally we should diagnose this in obvious cases
• Our LangOpts are currently fixed, so we are not handling u"" strings or R"()" strings that require C11/C++11.

rdar://108629982

clang/include/clang/Lex/Pragma.h

public:
   bool IsEmpty() const { return Handlers.empty(); }
 
   void HandlePragma(Preprocessor &PP, PragmaIntroducer Introducer,
                     Token &Tok) override;
 
   PragmaNamespace *getIfNamespace() override { return this; }
 };
 
+/// Destringize a \c _Pragma("") string according to C11 6.10.9.1:
+/// "The string literal is destringized by deleting any encoding prefix,
+/// deleting the leading and trailing double-quotes, replacing each escape
+/// sequence \" by a double-quote, and replacing each escape sequence \\ by a
+/// single backslash."
+void prepare_PragmaString(SmallVectorImpl<char> &StrVal);
+
 } // namespace clang
 
 #endif // LLVM_CLANG_LEX_PRAGMA_H

clang/lib/Lex/DependencyDirectivesScanner.cpp

 ///
 //===----------------------------------------------------------------------===//
 
 #include "clang/Lex/DependencyDirectivesScanner.h"
 #include "clang/Basic/CharInfo.h"
 #include "clang/Basic/Diagnostic.h"
 #include "clang/Lex/LexDiagnostic.h"
 #include "clang/Lex/Lexer.h"
+#include "clang/Lex/Pragma.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringSwitch.h"
 #include <optional>
 
 using namespace clang;
 using namespace clang::dependency_directives_scan;
Scanner(StringRef Input,
         TheLexer(InputSourceLoc, LangOpts, Input.begin(), Input.begin(),
                  Input.end()) {}
 
   static LangOptions getLangOptsForDepScanning() {
     LangOptions LangOpts;
     // Set the lexer to use 'tok::at' for '@', instead of 'tok::unknown'.
     LangOpts.ObjC = true;
     LangOpts.LineComment = true;
+    // FIXME: we do not enable C11 or C++11, so we are missing u/u8/U"" and
+    // R"()" literals.
     return LangOpts;
   }
 
   /// Lex the provided source and emit the directive tokens.
   ///
   /// \returns True on error.
   bool scan(SmallVectorImpl<Directive> &Directives);
 
 private:
   /// Lexes next token and advances \p First and the \p Lexer.
   [[nodiscard]] dependency_directives_scan::Token &
   lexToken(const char *&First, const char *const End);
 
   dependency_directives_scan::Token &lexIncludeFilename(const char *&First,
                                                         const char *const End);
 
   void skipLine(const char *&First, const char *const End);
   void skipDirective(StringRef Name, const char *&First, const char *const End);
 
+  /// Returns the spelling of a string literal or identifier after performing
+  /// any processing needed to handle \c clang::Token::NeedsCleaning.
+  StringRef cleanStringIfNeeded(const dependency_directives_scan::Token &Tok);
+
   /// Lexes next token and if it is identifier returns its string, otherwise
   /// it skips the current line and returns \p std::nullopt.
   ///
   /// In any case (whatever the token kind) \p First and the \p Lexer will
   /// advance beyond the token.
   [[nodiscard]] std::optional<StringRef>
   tryLexIdentifierOrSkipLine(const char *&First, const char *const End);
 
   /// Used when it is certain that next token is an identifier.
   [[nodiscard]] StringRef lexIdentifier(const char *&First,
                                         const char *const End);
 
   /// Lexes next token and returns true iff it is an identifier that matches \p
   /// Id, otherwise it skips the current line and returns false.
   ///
   /// In any case (whatever the token kind) \p First and the \p Lexer will
   /// advance beyond the token.
   [[nodiscard]] bool isNextIdentifierOrSkipLine(StringRef Id,
                                                 const char *&First,
                                                 const char *const End);
 
+  /// Lexes next token and returns true iff it matches the kind \p K.
+  /// Otherwise it skips the current line and returns false.
+  ///
+  /// In any case (whatever the token kind) \p First and the \p Lexer will
+  /// advance beyond the token.
+  [[nodiscard]] bool isNextTokenOrSkipLine(tok::TokenKind K, const char *&First,
+                                           const char *const End);
+
+  /// Lexes next token and if it is string literal, returns its string.
+  /// Otherwise, it skips the current line and returns \p std::nullopt.
+  ///
+  /// In any case (whatever the token kind) \p First and the \p Lexer will
+  /// advance beyond the token.
+  [[nodiscard]] std::optional<StringRef>
+  tryLexStringLiteralOrSkipLine(const char *&First, const char *const End);
+
   [[nodiscard]] bool scanImpl(const char *First, const char *const End);
   [[nodiscard]] bool lexPPLine(const char *&First, const char *const End);
   [[nodiscard]] bool lexAt(const char *&First, const char *const End);
   [[nodiscard]] bool lexModule(const char *&First, const char *const End);
   [[nodiscard]] bool lexDefine(const char *HashLoc, const char *&First,
                                const char *const End);
   [[nodiscard]] bool lexPragma(const char *&First, const char *const End);
+  [[nodiscard]] bool lex_Pragma(const char *&First, const char *const End);
   [[nodiscard]] bool lexEndif(const char *&First, const char *const End);
   [[nodiscard]] bool lexDefault(DirectiveKind Kind, const char *&First,
                                 const char *const End);
   [[nodiscard]] bool lexModuleDirectiveBody(DirectiveKind Kind,
                                             const char *&First,
                                             const char *const End);
   void lexPPDirectiveBody(const char *&First, const char *const End);
 

 void Scanner::lexPPDirectiveBody(const char *&First, const char *const End) {
   while (true) {
     const dependency_directives_scan::Token &Tok = lexToken(First, End);
     if (Tok.is(tok::eod))
       break;
   }
 }
 
-[[nodiscard]] std::optional<StringRef>
-Scanner::tryLexIdentifierOrSkipLine(const char *&First, const char *const End) {
-  const dependency_directives_scan::Token &Tok = lexToken(First, End);
-  if (Tok.isNot(tok::raw_identifier)) {
-    if (!Tok.is(tok::eod))
-      skipLine(First, End);
-    return std::nullopt;
-  }
-
+StringRef
+Scanner::cleanStringIfNeeded(const dependency_directives_scan::Token &Tok) {
   bool NeedsCleaning = Tok.Flags & clang::Token::NeedsCleaning;
   if (LLVM_LIKELY(!NeedsCleaning))
     return Input.slice(Tok.Offset, Tok.getEnd());
 
   SmallString<64> Spelling;
   Spelling.resize(Tok.Length);
 
+  // FIXME: C++11 raw string literals need special handling (see getSpellingSlow
+  // in the Lexer). Currently we cannot see them due to our LangOpts.
+
   unsigned SpellingLength = 0;
   const char *BufPtr = Input.begin() + Tok.Offset;
   const char *AfterIdent = Input.begin() + Tok.getEnd();
   while (BufPtr < AfterIdent) {
     unsigned Size;
     Spelling[SpellingLength++] =
         Lexer::getCharAndSizeNoWarn(BufPtr, Size, LangOpts);
     BufPtr += Size;
   }
 
   return SplitIds.try_emplace(StringRef(Spelling.begin(), SpellingLength), 0)
       .first->first();
 }
 
+std::optional<StringRef>
+Scanner::tryLexIdentifierOrSkipLine(const char *&First, const char *const End) {
+  const dependency_directives_scan::Token &Tok = lexToken(First, End);
+  if (Tok.isNot(tok::raw_identifier)) {
+    if (!Tok.is(tok::eod))
+      skipLine(First, End);
+    return std::nullopt;
+  }
+
+  return cleanStringIfNeeded(Tok);
+}
+
 StringRef Scanner::lexIdentifier(const char *&First, const char *const End) {
   std::optional<StringRef> Id = tryLexIdentifierOrSkipLine(First, End);
   assert(Id && "expected identifier token");
   return *Id;
 }
 
 bool Scanner::isNextIdentifierOrSkipLine(StringRef Id, const char *&First,
                                          const char *const End) {
   if (std::optional<StringRef> FoundId =
           tryLexIdentifierOrSkipLine(First, End)) {
     if (*FoundId == Id)
       return true;
     skipLine(First, End);
   }
   return false;
 }
 
+bool Scanner::isNextTokenOrSkipLine(tok::TokenKind K, const char *&First,
+                                    const char *const End) {
+  const dependency_directives_scan::Token &Tok = lexToken(First, End);
+  if (Tok.is(K))
+    return true;
+  skipLine(First, End);
+  return false;
+}
+
+std::optional<StringRef>
+Scanner::tryLexStringLiteralOrSkipLine(const char *&First,
+                                       const char *const End) {
+  const dependency_directives_scan::Token &Tok = lexToken(First, End);
+  if (!tok::isStringLiteral(Tok.Kind)) {
+    if (!Tok.is(tok::eod))
+      skipLine(First, End);
+    return std::nullopt;
+  }
+
+  return cleanStringIfNeeded(Tok);
+}
+
 bool Scanner::lexAt(const char *&First, const char *const End) {
   // Handle "@import".
 
   // Lex '@'.
   const dependency_directives_scan::Token &AtTok = lexToken(First, End);
   assert(AtTok.is(tok::at));
   (void)AtTok;
 
bool Scanner::lexModule(const char *&First, const char *const End) {
   if (Id == "module")
     Kind = Export ? cxx_export_module_decl : cxx_module_decl;
   else
     Kind = Export ? cxx_export_import_decl : cxx_import_decl;
 
   return lexModuleDirectiveBody(Kind, First, End);
 }
 
+bool Scanner::lex_Pragma(const char *&First, const char *const End) {
+  if (!isNextTokenOrSkipLine(tok::l_paren, First, End))
+    return false;
+
+  std::optional<StringRef> Str = tryLexStringLiteralOrSkipLine(First, End);
+
+  if (!Str || !isNextTokenOrSkipLine(tok::r_paren, First, End))
+    return false;
+
+  SmallString<64> Buffer(*Str);
+  prepare_PragmaString(Buffer);
+
+  // Use a new scanner instance since the tokens will be inside the allocated
+  // string. We should already have captured all the relevant tokens in the
+  // current scanner.
+  SmallVector<dependency_directives_scan::Token> DiscardTokens;
+  const char *Begin = Buffer.c_str();
+  Scanner PragmaScanner{StringRef(Begin, Buffer.size()), DiscardTokens, Diags,
+                        InputSourceLoc};
+
+  PragmaScanner.TheLexer.setParsingPreprocessorDirective(true);
+  if (PragmaScanner.lexPragma(Begin, Buffer.end()))
+    return true;
+
+  DirectiveKind K = PragmaScanner.topDirective();
+  if (K == pp_none) {
+    skipLine(First, End);
+    return false;
+  }
+
+  assert(Begin == Buffer.end());
+  pushDirective(K);
+  return false;
+}
+
 bool Scanner::lexPragma(const char *&First, const char *const End) {
   std::optional<StringRef> FoundId = tryLexIdentifierOrSkipLine(First, End);
   if (!FoundId)
     return false;
 
   StringRef Id = *FoundId;
   auto Kind = llvm::StringSwitch<DirectiveKind>(Id)
                   .Case("once", pp_pragma_once)

 
 static bool isStartOfRelevantLine(char First) {
   switch (First) {
   case '#':
   case '@':
   case 'i':
   case 'e':
   case 'm':
+  case '_':
     return true;
   }
   return false;
 }
 
 bool Scanner::lexPPLine(const char *&First, const char *const End) {
   assert(First != End);
 
bool Scanner::lexPPLine(const char *&First, const char *const End) {
 
   // Handle "@import".
   if (*First == '@')
     return lexAt(First, End);
 
   if (*First == 'i' || *First == 'e' || *First == 'm')
     return lexModule(First, End);
 
+  if (*First == '_') {
+    if (isNextIdentifierOrSkipLine("_Pragma", First, End))
+      return lex_Pragma(First, End);
+    return false;
+  }
+
   // Handle preprocessing directives.
 
   TheLexer.setParsingPreprocessorDirective(true);
   auto ScEx2 = make_scope_exit(
       [&]() { TheLexer.setParsingPreprocessorDirective(false); });
 
   // Lex '#'.
   const dependency_directives_scan::Token &HashTok = lexToken(First, End);

clang/lib/Lex/Pragma.cpp

void Preprocessor::Handle_Pragma(Token &Tok) {
   // If we're expanding a macro argument, put the tokens back.
   if (InMacroArgPreExpansion) {
     Toks.revert();
     return;
   }
 
   SourceLocation RParenLoc = Tok.getLocation();
   bool Invalid = false;
-  std::string StrVal = getSpelling(StrTok, &Invalid);
+  SmallString<64> StrVal;
+  StrVal.resize(StrTok.getLength());
+  StringRef StrValRef = getSpelling(StrTok, StrVal, &Invalid);
   if (Invalid) {
     Diag(PragmaLoc, diag::err__Pragma_malformed);
     return;
   }
 
-  // The _Pragma is lexically sound.  Destringize according to C11 6.10.9.1:
-  // "The string literal is destringized by deleting any encoding prefix,
-  // deleting the leading and trailing double-quotes, replacing each escape
-  // sequence \" by a double-quote, and replacing each escape sequence \\ by a
-  // single backslash."
+  assert(StrValRef.size() <= StrVal.size());
+
+  // If the token was spelled somewhere else, copy it.
+  if (StrValRef.begin() != StrVal.begin())
+    StrVal.assign(StrValRef);
+  // Truncate if necessary.
+  else if (StrValRef.size() != StrVal.size())
+    StrVal.resize(StrValRef.size());
+
+  // The _Pragma is lexically sound.  Destringize according to C11 6.10.9.1.
+  prepare_PragmaString(StrVal);
+
+  // Plop the string (including the newline and trailing null) into a buffer
+  // where we can lex it.
+  Token TmpTok;
+  TmpTok.startToken();
+  CreateString(StrVal, TmpTok);
+  SourceLocation TokLoc = TmpTok.getLocation();
+
+  // Make and enter a lexer object so that we lex and expand the tokens just
+  // like any others.
+  Lexer *TL = Lexer::Create_PragmaLexer(TokLoc, PragmaLoc, RParenLoc,
+                                        StrVal.size(), *this);
+
+  EnterSourceFileWithLexer(TL, nullptr);
+
+  // With everything set up, lex this as a #pragma directive.
+  HandlePragmaDirective({PIK__Pragma, PragmaLoc});
+
+  // Finally, return whatever came after the pragma directive.
+  return Lex(Tok);
+}
+
+void clang::prepare_PragmaString(SmallVectorImpl<char> &StrVal) {
   if (StrVal[0] == 'L' || StrVal[0] == 'U' ||
       (StrVal[0] == 'u' && StrVal[1] != '8'))
     StrVal.erase(StrVal.begin());
   else if (StrVal[0] == 'u')
     StrVal.erase(StrVal.begin(), StrVal.begin() + 2);
 
   if (StrVal[0] == 'R') {
     // FIXME: C++11 does not specify how to handle raw-string-literals here.
     // We strip off the 'R', the quotes, the d-char-sequences, and the parens.
     assert(StrVal[1] == '"' && StrVal[StrVal.size() - 1] == '"' &&
            "Invalid raw string token!");
 
     // Measure the length of the d-char-sequence.
     unsigned NumDChars = 0;
     while (StrVal[2 + NumDChars] != '(') {
       assert(NumDChars < (StrVal.size() - 5) / 2 &&
              "Invalid raw string token!");
       ++NumDChars;
     }
     assert(StrVal[StrVal.size() - 2 - NumDChars] == ')');
 
     // Remove 'R " d-char-sequence' and 'd-char-sequence "'. We'll replace the
     // parens below.
-    StrVal.erase(0, 2 + NumDChars);
-    StrVal.erase(StrVal.size() - 1 - NumDChars);
+    StrVal.erase(StrVal.begin(), StrVal.begin() + 2 + NumDChars);
+    StrVal.erase(StrVal.end() - 1 - NumDChars, StrVal.end());
   } else {
     assert(StrVal[0] == '"' && StrVal[StrVal.size()-1] == '"' &&
            "Invalid string token!");
 
     // Remove escaped quotes and escapes.
     unsigned ResultPos = 1;
     for (size_t i = 1, e = StrVal.size() - 1; i != e; ++i) {
       // Skip escapes.  \\ -> '\' and \" -> '"'.
       if (StrVal[i] == '\\' && i + 1 < e &&
           (StrVal[i + 1] == '\\' || StrVal[i + 1] == '"'))
         ++i;
       StrVal[ResultPos++] = StrVal[i];
     }
     StrVal.erase(StrVal.begin() + ResultPos, StrVal.end() - 1);
   }
 
   // Remove the front quote, replacing it with a space, so that the pragma
   // contents appear to have a space before them.
   StrVal[0] = ' ';
 
   // Replace the terminating quote with a \n.
-  StrVal[StrVal.size()-1] = '\n';
-
-  // Plop the string (including the newline and trailing null) into a buffer
-  // where we can lex it.
-  Token TmpTok;
-  TmpTok.startToken();
-  CreateString(StrVal, TmpTok);
-  SourceLocation TokLoc = TmpTok.getLocation();
-
-  // Make and enter a lexer object so that we lex and expand the tokens just
-  // like any others.
-  Lexer *TL = Lexer::Create_PragmaLexer(TokLoc, PragmaLoc, RParenLoc,
-                                        StrVal.size(), *this);
-
-  EnterSourceFileWithLexer(TL, nullptr);
-
-  // With everything set up, lex this as a #pragma directive.
-  HandlePragmaDirective({PIK__Pragma, PragmaLoc});
-
-  // Finally, return whatever came after the pragma directive.
-  return Lex(Tok);
+  StrVal[StrVal.size() - 1] = '\n';
 }
 
 /// HandleMicrosoft__pragma - Like Handle_Pragma except the pragma text
 /// is not enclosed within a string literal.
 void Preprocessor::HandleMicrosoft__pragma(Token &Tok) {
   // During macro pre-expansion, check the syntax now but put the tokens back
   // into the token stream for later consumption. Same as Handle_Pragma.
   TokenCollector Toks = {*this, InMacroArgPreExpansion, {}, Tok};

clang/test/ClangScanDeps/_Pragma-once.c

+// Test scanning deps works with _Pragma syntax when not inside a macro.
+
+// RUN: rm -rf %t
+// RUN: split-file %s %t
+// RUN: sed "s|DIR|%/t|g" %t/cdb.json.template > %t/cdb.json
+
+// RUN: clang-scan-deps -compilation-database %t/cdb.json -j 1
+
+//--- cdb.json.template
+[{
+  "directory": "DIR",
+  "command": "clang -fsyntax-only DIR/tu.c",
+  "file": "DIR/tu.c"
+}]
+
+//--- a.h
+_Pragma("once")
+#include "b.h"
+
+//--- b.h
+#include "a.h"
+
+//--- tu.c
+#include "a.h"

clang/unittests/Lex/DependencyDirectivesScannerTest.cpp

ASSERT_FALSE(minimizeSourceToDependencyDirectives(
       "#pragma clang module impor\n", Out));
   EXPECT_STREQ("<TokBeforeEOF>\n", Out.data());
 
   ASSERT_FALSE(minimizeSourceToDependencyDirectives(
       "#pragma clang module import\n", Out));
   EXPECT_STREQ("#pragma clang module import\n", Out.data());
 }
 
+TEST(MinimizeSourceToDependencyDirectivesTest, UnderscorePragma) {
+  SmallVector<char, 128> Out;
+
+  ASSERT_FALSE(minimizeSourceToDependencyDirectives(R"(_)", Out));
+  EXPECT_STREQ("<TokBeforeEOF>\n", Out.data());
+  ASSERT_FALSE(minimizeSourceToDependencyDirectives(R"(_Pragma)", Out));
+  EXPECT_STREQ("<TokBeforeEOF>\n", Out.data());
+  ASSERT_FALSE(minimizeSourceToDependencyDirectives(R"(_Pragma()", Out));
+  EXPECT_STREQ("<TokBeforeEOF>\n", Out.data());
+  ASSERT_FALSE(minimizeSourceToDependencyDirectives(R"(_Pragma())", Out));
+  EXPECT_STREQ("<TokBeforeEOF>\n", Out.data());
+  ASSERT_FALSE(minimizeSourceToDependencyDirectives(R"(_Pragma(")", Out));
+  EXPECT_STREQ("<TokBeforeEOF>\n", Out.data());
+  ASSERT_FALSE(minimizeSourceToDependencyDirectives(R"(_Pragma("A"))", Out));
+  EXPECT_STREQ("<TokBeforeEOF>\n", Out.data());
+
+  ASSERT_FALSE(minimizeSourceToDependencyDirectives(
+      R"x(_Pragma("push_macro(\"MACRO\")"))x", Out));
+  EXPECT_STREQ(R"x(_Pragma("push_macro(\"MACRO\")"))x"
+               "\n",
+               Out.data());
+
+  ASSERT_FALSE(minimizeSourceToDependencyDirectives(
+      R"x(_Pragma("pop_macro(\"MACRO\")"))x", Out));
+  EXPECT_STREQ(R"x(_Pragma("pop_macro(\"MACRO\")"))x"
+               "\n",
+               Out.data());
+
+  ASSERT_FALSE(minimizeSourceToDependencyDirectives(
+      R"x(_Pragma("include_alias(\"A\", \"B\")"))x", Out));
+  EXPECT_STREQ(R"x(_Pragma("include_alias(\"A\", \"B\")"))x"
+               "\n",
+               Out.data());
+
+  ASSERT_FALSE(minimizeSourceToDependencyDirectives(
+      R"x(_Pragma("include_alias(<A>, <B>)"))x", Out));
+  EXPECT_STREQ(R"x(_Pragma("include_alias(<A>, <B>)"))x"
+               "\n",
+               Out.data());
+
+  ASSERT_FALSE(
+      minimizeSourceToDependencyDirectives(R"(_Pragma("clang"))", Out));
+  EXPECT_STREQ("<TokBeforeEOF>\n", Out.data());
+
+  ASSERT_FALSE(
+      minimizeSourceToDependencyDirectives(R"(_Pragma("clang module"))", Out));
+  EXPECT_STREQ("<TokBeforeEOF>\n", Out.data());
+
+  ASSERT_FALSE(minimizeSourceToDependencyDirectives(
+      R"(_Pragma("clang module impor"))", Out));
+  EXPECT_STREQ("<TokBeforeEOF>\n", Out.data());
+
+  ASSERT_FALSE(minimizeSourceToDependencyDirectives(
+      R"(_Pragma("clang module import"))", Out));
+  EXPECT_STREQ(R"(_Pragma("clang module import"))"
+               "\n",
+               Out.data());
+
+  ASSERT_FALSE(minimizeSourceToDependencyDirectives(
+      R"(_Pragma("clang \
+  module \
+  import"))",
+      Out));
+  EXPECT_STREQ(R"(_Pragma("clang \
+  module \
+  import"))"
+               "\n",
+               Out.data());
+
+  ASSERT_FALSE(minimizeSourceToDependencyDirectives(
+      R"(_Pragma(L"clang module import"))", Out));
+  EXPECT_STREQ(R"(_Pragma(L"clang module import"))"
+               "\n",
+               Out.data());
+
+  // FIXME: u"" strings depend on using C11 language mode
+  ASSERT_FALSE(minimizeSourceToDependencyDirectives(
+      R"(_Pragma(u"clang module import"))", Out));
+  EXPECT_STREQ("<TokBeforeEOF>\n", Out.data());
+
+  // FIXME: R"()" strings depend on using C++ 11 language mode
+  ASSERT_FALSE(minimizeSourceToDependencyDirectives(
+      R"(_Pragma(R"abc(clang module import)abc"))", Out));
+  EXPECT_STREQ("<TokBeforeEOF>\n", Out.data());
+}
+
 TEST(MinimizeSourceToDependencyDirectivesTest, Include) {
   SmallVector<char, 128> Out;
 
   ASSERT_FALSE(minimizeSourceToDependencyDirectives("#include \"A\"\n", Out));
   EXPECT_STREQ("#include \"A\"\n", Out.data());
 
   ASSERT_FALSE(minimizeSourceToDependencyDirectives("#include <A>\n", Out));
   EXPECT_STREQ("#include <A>\n", Out.data());
TEST(MinimizeSourceToDependencyDirectivesTest, PragmaOnce) {
   SmallVector<char, 128> Out;
   SmallVector<dependency_directives_scan::Token, 4> Tokens;
   SmallVector<Directive, 4> Directives;
 
   StringRef Source = R"(// comment
 #pragma once
 // another comment
 #include <test.h>
+_Pragma("once")
 )";
   ASSERT_FALSE(
       minimizeSourceToDependencyDirectives(Source, Out, Tokens, Directives));
-  EXPECT_STREQ("#pragma once\n#include <test.h>\n", Out.data());
-  ASSERT_EQ(Directives.size(), 3u);
+  EXPECT_STREQ("#pragma once\n#include <test.h>\n_Pragma(\"once\")\n",
+               Out.data());
+  ASSERT_EQ(Directives.size(), 4u);
   EXPECT_EQ(Directives[0].Kind, dependency_directives_scan::pp_pragma_once);
+  EXPECT_EQ(Directives[2].Kind, dependency_directives_scan::pp_pragma_once);
 
   Source = R"(// comment
     #pragma once extra tokens
     // another comment
     #include <test.h>
+    _Pragma("once") extra tokens
     )";
   ASSERT_FALSE(minimizeSourceToDependencyDirectives(Source, Out));
-  EXPECT_STREQ("#pragma once extra tokens\n#include <test.h>\n", Out.data());
+  EXPECT_STREQ("#pragma once extra tokens\n#include "
+               "<test.h>\n_Pragma(\"once\")<TokBeforeEOF>\n",
+               Out.data());
 }
 
 TEST(MinimizeSourceToDependencyDirectivesTest,
      SkipLineStringCharLiteralsUntilNewline) {
   SmallVector<char, 128> Out;
 
   StringRef Source = R"(#if NEVER_ENABLED
     #define why(fmt, ...) #error don't try me