[clangd] Copy existing includes in ReplayPreamble

[kadircet]

Summary

ReplayPreamble was just grabbing the reference of IncludeStructure
passed to it and then replayed any includes seen so while exiting
built-in file.

This implies any include seen in built-in files being replayed as part
of preamble, even though they are not. This wasn't an issue until we've
started patching preambles, as includes from built-in files were not
mapped back to main-file.

This patch copies over existing includes at the time of
ReplayPreamble::attach and only replies those to prevent any includes
from the preamble patch getting mixed-in.

[sammccall]

clang-tools-extra/clangd/ParsedAST.cpp

386

it seems like it would be clearer to pass Preamble->preambleIncludes() here, at this point?
(And move within the constructor, so still just one copy)

[kadircet]

• Address comment

clang-tools-extra/clangd/ParsedAST.cpp

static void attach(const IncludeStructure &Includes, CompilerInstance &Clang,
     assert(PP.getPPCallbacks() != ExistingCallbacks &&
            "Expected chaining implementation");
   }
 
 private:
   ReplayPreamble(const IncludeStructure &Includes, PPCallbacks *Delegate,
                  const SourceManager &SM, Preprocessor &PP,
                  const LangOptions &LangOpts, const PreambleBounds &PB)
-      : Includes(Includes), Delegate(Delegate), SM(SM), PP(PP) {
+      : IncludesToReplay(Includes.MainFileIncludes), Delegate(Delegate), SM(SM),
+        PP(PP) {
     // Only tokenize the preamble section of the main file, as we are not
     // interested in the rest of the tokens.
     MainFileTokens = syntax::tokenize(
         syntax::FileRange(SM.getMainFileID(), 0, PB.Size), SM, LangOpts);
   }
 
   // In a normal compile, the preamble traverses the following structure:
   //
void FileChanged(SourceLocation Loc, FileChangeReason Reason,
                    SrcMgr::CharacteristicKind Kind, FileID PrevFID) override {
     // It'd be nice if there was a better way to identify built-in headers...
     if (Reason == FileChangeReason::ExitFile &&
         SM.getBuffer(PrevFID)->getBufferIdentifier() == "<built-in>")
       replay();
   }
 
   void replay() {
-    for (const auto &Inc : Includes.MainFileIncludes) {
+    for (const auto &Inc : IncludesToReplay) {
       const FileEntry *File = nullptr;
       if (Inc.Resolved != "")
         if (auto FE = SM.getFileManager().getFile(Inc.Resolved))
           File = *FE;
 
       // Re-lex the #include directive to find its interesting parts.
       auto HashLoc = SM.getComposedLoc(SM.getMainFileID(), Inc.HashOffset);
       auto HashTok = llvm::partition_point(MainFileTokens,
for (const auto &Inc : IncludesToReplay) {
                               SynthesizedFilenameTok, Inc.FileKind);
       else {
         llvm::SmallString<1> UnusedRecovery;
         Delegate->FileNotFound(WrittenFilename, UnusedRecovery);
       }
     }
   }
 
-  const IncludeStructure &Includes;
+  std::vector<Inclusion> IncludesToReplay;
   PPCallbacks *Delegate;
   const SourceManager &SM;
   Preprocessor &PP;
   std::vector<syntax::Token> MainFileTokens;
 };
 
 } // namespace
 
ParsedAST::build(llvm::StringRef Filename, const ParseInputs &Inputs,
   }
 
   IncludeStructure Includes;
   // If we are using a preamble, copy existing includes.
   if (Preamble) {
     Includes = Preamble->Includes;
     Includes.MainFileIncludes = Patch->preambleIncludes();
     // Replay the preamble includes so that clang-tidy checks can see them.
     ReplayPreamble::attach(Includes, *Clang, Preamble->Preamble.getBounds());
   }
   // Important: collectIncludeStructure is registered *after* ReplayPreamble!
   // Otherwise we would collect the replayed includes again...
   // (We can't *just* use the replayed includes, they don't have Resolved path).
   Clang->getPreprocessor().addPPCallbacks(
       collectIncludeStructureCallback(Clang->getSourceManager(), &Includes));
   // Copy over the macros in the preamble region of the main file, and combine
   // with non-preamble macros below.

clang-tools-extra/clangd/unittests/ParsedASTTests.cpp

for (const auto &R : SIDToRefs.second)
       MacroExpansionPositions.push_back(R.start);
   }
   for (const auto &R : AST.getMacros().UnknownMacros)
     MacroExpansionPositions.push_back(R.start);
   EXPECT_THAT(MacroExpansionPositions,
               testing::UnorderedElementsAreArray(TestCase.points()));
 }
 
+MATCHER_P(WithFileName, Inc, "") { return arg.FileName == Inc; }
+
 TEST(ParsedASTTest, ReplayPreambleForTidyCheckers) {
   struct Inclusion {
     Inclusion(const SourceManager &SM, SourceLocation HashLoc,
               const Token &IncludeTok, llvm::StringRef FileName, bool IsAngled,
               CharSourceRange FilenameRange)
         : HashOffset(SM.getDecomposedLoc(HashLoc).second), IncTok(IncludeTok),
           IncDirective(IncludeTok.getIdentifierInfo()->getName()),
           FileNameOffset(SM.getDecomposedLoc(FilenameRange.getBegin()).second),
EXPECT_EQ(SM.getDecomposedLoc(SkippedFiles[I].location()).second,
               Inc.FileNameOffset);
     // This also contains quotes/angles so increment the range by one from both
     // sides.
     EXPECT_EQ(
         SkippedFiles[I].text(SM),
         Code.substr(FileRange.Begin - 1, FileRange.End - FileRange.Begin + 2));
     EXPECT_EQ(SkippedFiles[I].kind(), tok::header_name);
   }
+
+  // Make sure replay logic works with patched preambles.
+  TU.Code = "";
+  StoreDiags Diags;
+  auto Inputs = TU.inputs();
+  auto CI = buildCompilerInvocation(Inputs, Diags);
+  auto EmptyPreamble =
+      buildPreamble(testPath(TU.Filename), *CI, Inputs, true, nullptr);
+  ASSERT_TRUE(EmptyPreamble);
+  TU.Code = "#include <a.h>";
+  Includes.clear();
+  auto PatchedAST = ParsedAST::build(testPath(TU.Filename), TU.inputs(),
+                                     std::move(CI), {}, EmptyPreamble);
+  ASSERT_TRUE(PatchedAST);
+  // Make sure includes were seen only once.
+  EXPECT_THAT(Includes,
+              ElementsAre(WithFileName(testPath("__preamble_patch__.h")),
+                          WithFileName("a.h")));
 }
 
 TEST(ParsedASTTest, PatchesAdditionalIncludes) {
   llvm::StringLiteral ModifiedContents = R"cpp(
     #include "baz.h"
     #include "foo.h"
     #include "sub/aux.h"
     void bar() {