[Lex/DependencyDirectivesScanner] Keep track of the presence of tokens between the last scanned directive and EOF

[akyrtzi]

Summary

Directive dependency_directives_scan::tokens_present_before_eof is introduced to indicate there were tokens present before
the last scanned dependency directive and EOF.
This is useful to ensure we correctly identify the macro guards when lexing using the dependency directives.

[akyrtzi]

Remove a couple of unused local variables.

[benlangmuir]

clang/include/clang/Lex/DependencyDirectivesScanner.h

131

Why would you want to print without this? It seems important for correctness of the output. I would have expected we would always print it.

clang/lib/Lex/DependencyDirectivesScanner.cpp

868

How about "TokBeforeEOF"? I think "BEOF" is too cryptic.

[akyrtzi]

clang/include/clang/Lex/DependencyDirectivesScanner.h

131

The -print-dependency-directives-minimized-source clang option is using this function, and if you print the sources with "<TokBEOF>" at the end then the source text will not be parsable.

But the tests that pass -print-dependency-directives-minimized-source don't try to parse the code, so I think I can switch the default to be true for printing the "tokens-before-eof marker" and if a caller has a need to ignore it it can pass false for the parameter.

clang/lib/Lex/DependencyDirectivesScanner.cpp

868

SGTM 👍

[benlangmuir]

clang/include/clang/Lex/DependencyDirectivesScanner.h

131

If someone uses -print-dependency-directives-minimized-source and creates a minimized file for each header, it will "parse", but it won't behave correctly for multiple includes, right? My preference is we don't allow printing this without the TokBEOF. If we care about making it parse, we should add a real token of some kind -- maybe there is a no-op #pragma or something?

[benlangmuir]

clang/include/clang/Lex/DependencyDirectivesScanner.h

131

Oops, missed half my comment. Meant to also add:

But the tests that pass -print-dependency-directives-minimized-source don't try to parse the code, so I think I can switch the default to be true for printing the "tokens-before-eof marker" and if a caller has a need to ignore it it can pass false for the parameter.

SGTM

[akyrtzi]

clang/include/clang/Lex/DependencyDirectivesScanner.h

131

Ok, you convinced me that we should just always print it. This function is for testing purposes only anyway, if later on we actually have a need to parse back the minimized source we can re-evaluate what to do with that directive marker.

[akyrtzi]

Always print tokens_present_before_eof and print it as "<TokBeforeEOF>" for clarity.

[akyrtzi]

@benlangmuir see latest diff.

[benlangmuir]

You forgot to remove the \param PrintMarkerForTokensBeforeEOF ... from the doc comment. Otherwise LGTM.

[akyrtzi]

Remove leftover doc-comment parameter.

clang/include/clang/Lex/DependencyDirectivesScanner.h

enum DirectiveKind : uint8_t {
   pp_elifndef,
   pp_else,
   pp_endif,
   decl_at_import,
   cxx_module_decl,
   cxx_import_decl,
   cxx_export_module_decl,
   cxx_export_import_decl,
+  /// Indicates that there are tokens present between the last scanned directive
+  /// and eof. The \p Directive::Tokens array will be empty for this kind.
+  tokens_present_before_eof,
   pp_eof,
 };
 
 /// Represents a directive that's lexed as part of the dependency directives
 /// scanning. It's used to track various preprocessor directives that could
 /// potentially have an effect on the dependencies.
 struct Directive {
   ArrayRef<Token> Tokens;

 ///
 /// \param Source The original source text that the dependency directives were
 /// scanned from.
 /// \param Directives The previously scanned dependency
 /// directives.
 /// \param OS the stream to print the dependency directives on.
 ///
 /// This is used primarily for testing purposes, during dependency scanning the
 /// \p Lexer uses the tokens directly, not their printed version.
 void printDependencyDirectivesAsSource(
     StringRef Source,
     ArrayRef<dependency_directives_scan::Directive> Directives,
     llvm::raw_ostream &OS);
 
 } // end namespace clang
 
 #endif // LLVM_CLANG_LEX_DEPENDENCYDIRECTIVESSCANNER_H

clang/lib/Lex/DependencyDirectivesScanner.cpp

 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);
+
   /// Lexes next token and if it is identifier returns its string, otherwise
   /// it skips the current line and returns \p None.
   ///
   /// In any case (whatever the token kind) \p First and the \p Lexer will
   /// advance beyond the token.
   [[nodiscard]] Optional<StringRef>
   tryLexIdentifierOrSkipLine(const char *&First, const char *const End);
 
private:
   bool reportError(const char *CurPtr, unsigned Err);
 
   StringMap<char> SplitIds;
   StringRef Input;
   SmallVectorImpl<dependency_directives_scan::Token> &Tokens;
   DiagnosticsEngine *Diags;
   SourceLocation InputSourceLoc;
 
+  const char *LastTokenPtr = nullptr;
   /// Keeps track of the tokens for the currently lexed directive. Once a
   /// directive is fully lexed and "committed" then the tokens get appended to
   /// \p Tokens and \p CurDirToks is cleared for the next directive.
   SmallVector<dependency_directives_scan::Token, 32> CurDirToks;
   /// The directives that were lexed along with the number of tokens that each
   /// directive contains. The tokens of all the directives are kept in \p Tokens
   /// vector, in the same order as the directives order in \p DirsWithToks.
   SmallVector<DirectiveWithTokens, 64> DirsWithToks;
static bool isQuoteCppDigitSeparator(const char *const Start,
   if (Prev == '8' && (Cur - 1 != Start) && *(Cur - 2) == 'u')
     return false;
   if (!isPreprocessingNumberBody(Prev))
     return false;
   // The next character should be a valid identifier body character.
   return (Cur + 1) < End && isAsciiIdentifierContinue(*(Cur + 1));
 }
 
-static void skipLine(const char *&First, const char *const End) {
+void Scanner::skipLine(const char *&First, const char *const End) {
   for (;;) {
     assert(First <= End);
     if (First == End)
       return;
 
     if (isVerticalWhitespace(*First)) {
       skipNewline(First, End);
       return;
     }
     const char *Start = First;
     while (First != End && !isVerticalWhitespace(*First)) {
       // Iterate over strings correctly to avoid comments and newlines.
       if (*First == '"' ||
           (*First == '\'' && !isQuoteCppDigitSeparator(Start, First, End))) {
+        LastTokenPtr = First;
         if (isRawStringLiteral(Start, First))
           skipRawString(First, End);
         else
           skipString(First, End);
         continue;
       }
 
       // Iterate over comments correctly.
       if (*First != '/' || End - First < 2) {
+        LastTokenPtr = First;
         ++First;
         continue;
       }
 
       if (First[1] == '/') {
         // "//...".
         skipLineComment(First, End);
         continue;
       }
 
       if (First[1] != '*') {
+        LastTokenPtr = First;
         ++First;
         continue;
       }
 
       // "/*...*/".
       skipBlockComment(First, End);
     }
     if (First == End)
       return;
 
     // Skip over the newline.
     unsigned Len = skipNewline(First, End);
     if (!wasLineContinuation(First, Len)) // Continue past line-continuations.
       break;
   }
 }
 
-static void skipDirective(StringRef Name, const char *&First,
+void Scanner::skipDirective(StringRef Name, const char *&First,
                             const char *const End) {
   if (llvm::StringSwitch<bool>(Name)
           .Case("warning", true)
           .Case("error", true)
           .Default(false))
     // Do not process quotes or comments.
     skipToNewlineRaw(First, End);
   else
     skipLine(First, End);
if (First == End)
     return false;
 
   if (!isStartOfRelevantLine(*First)) {
     skipLine(First, End);
     assert(First <= End);
     return false;
   }
 
+  LastTokenPtr = First;
+
   TheLexer.seek(getOffsetAt(First), /*IsAtStartOfLine*/ true);
 
   auto ScEx1 = make_scope_exit([&]() {
     /// Clear Scanner's CurDirToks before returning, in case we didn't push a
     /// new directive.
     CurDirToks.clear();
   });
 
bool Scanner::scanImpl(const char *First, const char *const End) {
   return false;
 }
 
 bool Scanner::scan(SmallVectorImpl<Directive> &Directives) {
   bool Error = scanImpl(Input.begin(), Input.end());
 
   if (!Error) {
     // Add an EOF on success.
+    if (LastTokenPtr &&
+        (Tokens.empty() || LastTokenPtr > Input.begin() + Tokens.back().Offset))
+      pushDirective(tokens_present_before_eof);
     pushDirective(pp_eof);
   }
 
   ArrayRef<dependency_directives_scan::Token> RemainingTokens = Tokens;
   for (const DirectiveWithTokens &DirWithToks : DirsWithToks) {
     assert(RemainingTokens.size() >= DirWithToks.NumTokens);
     Directives.emplace_back(DirWithToks.Kind,
                             RemainingTokens.take_front(DirWithToks.NumTokens));
if (Prev == tok::r_paren &&
       return true;
     if (Prev == tok::comma &&
         Tok.isOneOf(tok::l_paren, tok::string_literal, tok::less))
       return true;
     return false;
   };
 
   for (const dependency_directives_scan::Directive &Directive : Directives) {
+    if (Directive.Kind == tokens_present_before_eof)
+      OS << "<TokBeforeEOF>";
     Optional<tok::TokenKind> PrevTokenKind;
     for (const dependency_directives_scan::Token &Tok : Directive.Tokens) {
       if (PrevTokenKind && needsSpaceSeparator(*PrevTokenKind, Tok))
         OS << ' ';
       PrevTokenKind = Tok.Kind;
       OS << Source.slice(Tok.Offset, Tok.getEnd());
     }
   }
 }

clang/lib/Lex/Lexer.cpp

 bool Lexer::LexDependencyDirectiveToken(Token &Result) {
   assert(isDependencyDirectivesLexer());
 
   using namespace dependency_directives_scan;
 
   while (NextDepDirectiveTokenIndex == DepDirectives.front().Tokens.size()) {
     if (DepDirectives.front().Kind == pp_eof)
       return LexEndOfFile(Result, BufferEnd);
+    if (DepDirectives.front().Kind == tokens_present_before_eof)
+      MIOpt.ReadToken();
     NextDepDirectiveTokenIndex = 0;
     DepDirectives = DepDirectives.drop_front();
   }
 
   const dependency_directives_scan::Token &DDTok =
       DepDirectives.front().Tokens[NextDepDirectiveTokenIndex++];
   if (NextDepDirectiveTokenIndex > 1 || DDTok.Kind != tok::hash) {
     // Read something other than a preprocessor directive hash.
do {
     case pp_pragma_pop_macro:
     case pp_pragma_include_alias:
     case pp_include_next:
     case decl_at_import:
     case cxx_module_decl:
     case cxx_import_decl:
     case cxx_export_module_decl:
     case cxx_export_import_decl:
+    case tokens_present_before_eof:
       break;
     case pp_if:
     case pp_ifdef:
     case pp_ifndef:
       ++NestedIfs;
       break;
     case pp_elif:
     case pp_elifdef:

clang/unittests/Lex/CMakeLists.txt

 set(LLVM_LINK_COMPONENTS
   Support
   )
 
 add_clang_unittest(LexTests
   DependencyDirectivesScannerTest.cpp
   HeaderMapTest.cpp
   HeaderSearchTest.cpp
   LexerTest.cpp
   PPCallbacksTest.cpp
   PPConditionalDirectiveRecordTest.cpp
+  PPDependencyDirectivesTest.cpp
   PPMemoryAllocationsTest.cpp
   )
 
 clang_target_link_libraries(LexTests
   PRIVATE
   clangAST
   clangBasic
   clangLex
   clangParse
   clangSema
+
+  LLVMTestingSupport
   )

clang/unittests/Lex/DependencyDirectivesScannerTest.cpp

ASSERT_FALSE(
       minimizeSourceToDependencyDirectives("", Out, Tokens, Directives));
   EXPECT_TRUE(Out.empty());
   EXPECT_TRUE(Tokens.empty());
   ASSERT_EQ(1u, Directives.size());
   ASSERT_EQ(pp_eof, Directives.back().Kind);
 
   ASSERT_FALSE(minimizeSourceToDependencyDirectives("abc def\nxyz", Out, Tokens,
                                                     Directives));
-  EXPECT_TRUE(Out.empty());
+  EXPECT_STREQ("<TokBeforeEOF>\n", Out.data());
   EXPECT_TRUE(Tokens.empty());
-  ASSERT_EQ(1u, Directives.size());
-  ASSERT_EQ(pp_eof, Directives.back().Kind);
+  ASSERT_EQ(2u, Directives.size());
+  EXPECT_EQ(tokens_present_before_eof, Directives[0].Kind);
+  EXPECT_EQ(pp_eof, Directives[1].Kind);
 }
 
 TEST(MinimizeSourceToDependencyDirectivesTest, AllTokens) {
   SmallVector<char, 128> Out;
   SmallVector<dependency_directives_scan::Token, 4> Tokens;
   SmallVector<Directive, 4> Directives;
 
   ASSERT_FALSE(
EXPECT_STREQ("#ifdef A\n"
                "#endif\n",
                Out.data());
 }
 
 TEST(MinimizeSourceToDependencyDirectivesTest, Pragma) {
   SmallVector<char, 128> Out;
 
   ASSERT_FALSE(minimizeSourceToDependencyDirectives("#pragma A\n", Out));
-  EXPECT_STREQ("", Out.data());
+  EXPECT_STREQ("<TokBeforeEOF>\n", Out.data());
 
   ASSERT_FALSE(minimizeSourceToDependencyDirectives(
       "#pragma push_macro(\"MACRO\")\n", Out));
   EXPECT_STREQ("#pragma push_macro(\"MACRO\")\n", Out.data());
 
   ASSERT_FALSE(minimizeSourceToDependencyDirectives(
       "#pragma pop_macro(\"MACRO\")\n", Out));
   EXPECT_STREQ("#pragma pop_macro(\"MACRO\")\n", Out.data());
 
   ASSERT_FALSE(minimizeSourceToDependencyDirectives(
       "#pragma include_alias(\"A\", \"B\")\n", Out));
   EXPECT_STREQ("#pragma include_alias(\"A\", \"B\")\n", Out.data());
 
   ASSERT_FALSE(minimizeSourceToDependencyDirectives(
       "#pragma include_alias(<A>, <B>)\n", Out));
   EXPECT_STREQ("#pragma include_alias(<A>, <B>)\n", Out.data());
 
   ASSERT_FALSE(minimizeSourceToDependencyDirectives("#pragma clang\n", Out));
-  EXPECT_STREQ("", Out.data());
+  EXPECT_STREQ("<TokBeforeEOF>\n", Out.data());
 
   ASSERT_FALSE(
       minimizeSourceToDependencyDirectives("#pragma clang module\n", Out));
-  EXPECT_STREQ("", Out.data());
+  EXPECT_STREQ("<TokBeforeEOF>\n", Out.data());
 
   ASSERT_FALSE(minimizeSourceToDependencyDirectives(
       "#pragma clang module impor\n", Out));
-  EXPECT_STREQ("", Out.data());
+  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, Include) {
   SmallVector<char, 128> Out;
for (auto Source : {
     EXPECT_STREQ("#include <t.h>\n", Out.data());
   }
 
   for (auto Source : {
            "#warning \\\n#include <t.h>\n",
            "#error \\\n#include <t.h>\n",
        }) {
     ASSERT_FALSE(minimizeSourceToDependencyDirectives(Source, Out));
-    EXPECT_STREQ("", Out.data());
+    EXPECT_STREQ("<TokBeforeEOF>\n", Out.data());
   }
 
   for (auto Source : {
            "#if MACRO\n#warning '\n#endif\n",
            "#if MACRO\n#warning \"\n#endif\n",
            "#if MACRO\n#warning /*\n#endif\n",
            "#if MACRO\n#error '\n#endif\n",
            "#if MACRO\n#error \"\n#endif\n",
TEST(MinimizeSourceToDependencyDirectivesTest,
   StringRef Source = R"(#if NEVER_ENABLED
     #define why(fmt, ...) #error don't try me
     #endif
 
     void foo();
 )";
   ASSERT_FALSE(minimizeSourceToDependencyDirectives(Source, Out));
   EXPECT_STREQ(
-      "#if NEVER_ENABLED\n#define why(fmt,...) #error don't try me\n#endif\n",
+      "#if NEVER_ENABLED\n#define why(fmt,...) #error don't try me\n#endif\n"
+      "<TokBeforeEOF>\n",
       Out.data());
 
   Source = R"(#if NEVER_ENABLED
       #define why(fmt, ...) "quote dropped
       #endif
 
       void foo();
   )";
   ASSERT_FALSE(minimizeSourceToDependencyDirectives(Source, Out));
   EXPECT_STREQ(
-      "#if NEVER_ENABLED\n#define why(fmt,...) \"quote dropped\n#endif\n",
+      "#if NEVER_ENABLED\n#define why(fmt,...) \"quote dropped\n#endif\n"
+      "<TokBeforeEOF>\n",
       Out.data());
 }
 
 TEST(MinimizeSourceToDependencyDirectivesTest,
      SupportWhitespaceBeforeLineContinuation) {
   SmallVector<char, 128> Out;
 
   ASSERT_FALSE(minimizeSourceToDependencyDirectives("#define FOO(BAR) \\\n"
                                                     "  #BAR\\\n"
                                                     "  baz\n",
                                                     Out));
   EXPECT_STREQ("#define FOO(BAR) #BAR baz\n", Out.data());
 }
 
 TEST(MinimizeSourceToDependencyDirectivesTest,
      SupportWhitespaceBeforeLineContinuationInStringSkipping) {
   SmallVector<char, 128> Out;
 
   StringRef Source = "#define X '\\ \t\nx'\nvoid foo() {}";
   ASSERT_FALSE(minimizeSourceToDependencyDirectives(Source, Out));
-  EXPECT_STREQ("#define X '\\ \t\nx'\n", Out.data());
+  EXPECT_STREQ("#define X '\\ \t\nx'\n<TokBeforeEOF>\n", Out.data());
 
   Source = "#define X \"\\ \r\nx\"\nvoid foo() {}";
   ASSERT_FALSE(minimizeSourceToDependencyDirectives(Source, Out));
-  EXPECT_STREQ("#define X \"\\ \r\nx\"\n", Out.data());
+  EXPECT_STREQ("#define X \"\\ \r\nx\"\n<TokBeforeEOF>\n", Out.data());
 
   Source = "#define X \"\\ \r\nx\n#include <x>\n";
   ASSERT_FALSE(minimizeSourceToDependencyDirectives(Source, Out));
   EXPECT_STREQ("#define X\"\\ \r\nx\n#include <x>\n", Out.data());
 }
 
 TEST(MinimizeSourceToDependencyDirectivesTest, CxxModules) {
   SmallVector<char, 128> Out;
ort \
     }
     )";
   ASSERT_FALSE(
       minimizeSourceToDependencyDirectives(Source, Out, Tokens, Directives));
   EXPECT_STREQ("#include \"textual-header.h\"\nexport module m;"
                "exp\\\nort import:l[[rename]];"
                "import<<=3;import a b d e d e f e;"
                "import foo[[no_unique_address]];import foo();"
-               "import f(:sefse);import f(->a=3);\n",
+               "import f(:sefse);import f(->a=3);"
+               "<TokBeforeEOF>\n",
                Out.data());
-  ASSERT_EQ(Directives.size(), 10u);
+  ASSERT_EQ(Directives.size(), 11u);
   EXPECT_EQ(Directives[0].Kind, pp_include);
   EXPECT_EQ(Directives[1].Kind, cxx_export_module_decl);
 }
 
+TEST(MinimizeSourceToDependencyDirectivesTest, TokensBeforeEOF) {
+  SmallString<128> Out;
+
+  StringRef Source = R"(
+    #define A
+    #ifdef B
+    int x;
+    #endif
+    )";
+  ASSERT_FALSE(minimizeSourceToDependencyDirectives(Source, Out));
+  EXPECT_STREQ("#define A\n<TokBeforeEOF>\n", Out.data());
+
+  Source = R"(
+    #ifndef A
+    #define A
+    #endif // some comment
+
+    // other comment
+    )";
+  ASSERT_FALSE(minimizeSourceToDependencyDirectives(Source, Out));
+  EXPECT_STREQ("#ifndef A\n#define A\n#endif\n", Out.data());
+
+  Source = R"(
+    #ifndef A
+    #define A
+    #endif /* some comment
+
+    */
+    )";
+  ASSERT_FALSE(minimizeSourceToDependencyDirectives(Source, Out));
+  EXPECT_STREQ("#ifndef A\n#define A\n#endif\n", Out.data());
+
+  Source = R"(
+    #ifndef A
+    #define A
+    #endif /* some comment
+
+    */
+    int x;
+    )";
+  ASSERT_FALSE(minimizeSourceToDependencyDirectives(Source, Out));
+  EXPECT_STREQ("#ifndef A\n#define A\n#endif\n<TokBeforeEOF>\n", Out.data());
+}
+
 } // end anonymous namespace

clang/unittests/Lex/PPDependencyDirectivesTest.cpp

+//===- unittests/Lex/PPDependencyDirectivesTest.cpp -------------------------=//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "clang/Basic/Diagnostic.h"
+#include "clang/Basic/DiagnosticOptions.h"
+#include "clang/Basic/FileManager.h"
+#include "clang/Basic/LangOptions.h"
+#include "clang/Basic/SourceManager.h"
+#include "clang/Basic/TargetInfo.h"
+#include "clang/Basic/TargetOptions.h"
+#include "clang/Lex/DependencyDirectivesScanner.h"
+#include "clang/Lex/HeaderSearch.h"
+#include "clang/Lex/HeaderSearchOptions.h"
+#include "clang/Lex/ModuleLoader.h"
+#include "clang/Lex/Preprocessor.h"
+#include "clang/Lex/PreprocessorOptions.h"
+#include "llvm/Testing/Support/Error.h"
+#include "gtest/gtest.h"
+
+using namespace clang;
+
+namespace {
+
+// The test fixture.
+class PPDependencyDirectivesTest : public ::testing::Test {
+protected:
+  PPDependencyDirectivesTest()
+      : FileMgr(FileMgrOpts), DiagID(new DiagnosticIDs()),
+        Diags(DiagID, new DiagnosticOptions, new IgnoringDiagConsumer()),
+        SourceMgr(Diags, FileMgr), TargetOpts(new TargetOptions) {
+    TargetOpts->Triple = "x86_64-apple-macos12";
+    Target = TargetInfo::CreateTargetInfo(Diags, TargetOpts);
+  }
+
+  FileSystemOptions FileMgrOpts;
+  FileManager FileMgr;
+  IntrusiveRefCntPtr<DiagnosticIDs> DiagID;
+  DiagnosticsEngine Diags;
+  SourceManager SourceMgr;
+  LangOptions LangOpts;
+  std::shared_ptr<TargetOptions> TargetOpts;
+  IntrusiveRefCntPtr<TargetInfo> Target;
+};
+
+class IncludeCollector : public PPCallbacks {
+public:
+  Preprocessor &PP;
+  SmallVectorImpl<StringRef> &IncludedFiles;
+
+  IncludeCollector(Preprocessor &PP, SmallVectorImpl<StringRef> &IncludedFiles)
+      : PP(PP), IncludedFiles(IncludedFiles) {}
+
+  void LexedFileChanged(FileID FID, LexedFileChangeReason Reason,
+                        SrcMgr::CharacteristicKind FileType, FileID PrevFID,
+                        SourceLocation Loc) override {
+    if (Reason != LexedFileChangeReason::EnterFile)
+      return;
+    if (FID == PP.getPredefinesFileID())
+      return;
+    StringRef Filename =
+        PP.getSourceManager().getSLocEntry(FID).getFile().getName();
+    IncludedFiles.push_back(Filename);
+  }
+};
+
+TEST_F(PPDependencyDirectivesTest, MacroGuard) {
+  // "head1.h" has a macro guard and should only be included once.
+  // "head2.h" and "head3.h" have tokens following the macro check, they should
+  // be included multiple times.
+
+  auto VFS = new llvm::vfs::InMemoryFileSystem();
+  VFS->addFile(
+      "head1.h", 0,
+      llvm::MemoryBuffer::getMemBuffer("#ifndef H1_H\n#define H1_H\n#endif\n"));
+  VFS->addFile(
+      "head2.h", 0,
+      llvm::MemoryBuffer::getMemBuffer("#ifndef H2_H\n#define H2_H\n#endif\n\n"
+                                       "extern int foo;\n"));
+  VFS->addFile("head3.h", 0,
+               llvm::MemoryBuffer::getMemBuffer(
+                   "#ifndef H3_H\n#define H3_H\n#endif\n\n"
+                   "#ifdef SOMEMAC\nextern int foo;\n#endif\n"));
+  VFS->addFile("main.c", 0,
+               llvm::MemoryBuffer::getMemBuffer(
+                   "#include \"head1.h\"\n#include \"head1.h\"\n"
+                   "#include \"head2.h\"\n#include \"head2.h\"\n"
+                   "#include \"head3.h\"\n#include \"head3.h\"\n"));
+  FileMgr.setVirtualFileSystem(VFS);
+
+  Optional<FileEntryRef> FE;
+  ASSERT_THAT_ERROR(FileMgr.getFileRef("main.c").moveInto(FE),
+                    llvm::Succeeded());
+  SourceMgr.setMainFileID(
+      SourceMgr.createFileID(*FE, SourceLocation(), SrcMgr::C_User));
+
+  struct DepDirectives {
+    SmallVector<dependency_directives_scan::Token> Tokens;
+    SmallVector<dependency_directives_scan::Directive> Directives;
+  };
+  SmallVector<std::unique_ptr<DepDirectives>> DepDirectivesObjects;
+
+  auto getDependencyDirectives = [&](FileEntryRef File)
+      -> Optional<ArrayRef<dependency_directives_scan::Directive>> {
+    DepDirectivesObjects.push_back(std::make_unique<DepDirectives>());
+    StringRef Input = (*FileMgr.getBufferForFile(File))->getBuffer();
+    bool Err = scanSourceForDependencyDirectives(
+        Input, DepDirectivesObjects.back()->Tokens,
+        DepDirectivesObjects.back()->Directives);
+    EXPECT_FALSE(Err);
+    return llvm::makeArrayRef(DepDirectivesObjects.back()->Directives);
+  };
+
+  auto PPOpts = std::make_shared<PreprocessorOptions>();
+  PPOpts->DependencyDirectivesForFile = [&](FileEntryRef File)
+      -> Optional<ArrayRef<dependency_directives_scan::Directive>> {
+    return getDependencyDirectives(File);
+  };
+
+  TrivialModuleLoader ModLoader;
+  HeaderSearch HeaderInfo(std::make_shared<HeaderSearchOptions>(), SourceMgr,
+                          Diags, LangOpts, Target.get());
+  Preprocessor PP(PPOpts, Diags, LangOpts, SourceMgr, HeaderInfo, ModLoader,
+                  /*IILookup =*/nullptr,
+                  /*OwnsHeaderSearch =*/false);
+  PP.Initialize(*Target);
+
+  SmallVector<StringRef> IncludedFiles;
+  PP.addPPCallbacks(std::make_unique<IncludeCollector>(PP, IncludedFiles));
+  PP.EnterMainSourceFile();
+  while (true) {
+    Token tok;
+    PP.Lex(tok);
+    if (tok.is(tok::eof))
+      break;
+  }
+
+  SmallVector<StringRef> ExpectedIncludes{
+      "main.c", "./head1.h", "./head2.h", "./head2.h", "./head3.h", "./head3.h",
+  };
+  EXPECT_EQ(IncludedFiles, ExpectedIncludes);
+}
+
+} // anonymous namespace