[clangd] Handle additional includes while parsing ASTs

[kadircet]

Summary

Depends on D77392.

Enables building ASTs with stale preambles by handling additional preamble
includes. Sets the correct location information for those imaginary includes so
that features like gotodef/documentlink keeps functioning propoerly.

[kadircet]

this doesn't handle source locations correctly in presence of deleted headers.

[kadircet]

• Use PreamblePatch
• Add more tests
• Handle deleted includes

[kadircet]

this is ready for review now, to explain the flow:

While creating a patch we also record remaining includes from Baseline left in Modified, as the new ones will be put into the patch already.
That way ParsedAST can keep track of latest state of the preamble + main file includes, with the patched includes ending up in the latter.

It doesn't modify include depths, as it is only used by code completion. If need be we can prune the tree to get rid of deleted includes.

[kadircet]

clang-tools-extra/clangd/Preamble.h

117

i am not happy about this interface. I was thinking about a std::vector<Inclusion> remainingBaselineIncludes() const but didn't go for it to keep the concept of patching.

I suppose the former is more intuitive though, WDYT?

[sammccall]

This is excellent, a pretty nice model in the end!
Good job puzzling it out, when we discussed the idea in the past I imagined this part would be a pile of hacks.

Thanks for the offline help understanding this, too...

clang-tools-extra/clangd/ParsedAST.cpp

273

As discussed offline, might be clearer not to have the patch exist unless the preamble does.

clang-tools-extra/clangd/ParsedAST.h

51 ↗

(On Diff #260228)

ah, this is going to conflict with https://github.com/llvm/llvm-project/commit/6880c4dfa3981ec26d44b5f4844b641342a4e3e8
Sorry about that, but that change also ParseInputs available here so I think you can just drop your changes.

clang-tools-extra/clangd/Preamble.cpp

279

BTW do we want to check isPreambleCompatible and bail out early if it is, or is that the caller's job?

337

This explicitly looks at the existing includes based on spelling, and not whether the include was resolved during scanning or not.

One implication is that if you insert an #include that was already transitively included (which IIUC will hit the preamble VFS cache and thus be resolved) then we're not going to take advantage of this to record its resolution now.

Instead we're going to parse it along with the mainfile, and... well, I *hope* everything works out:

• do we pay for the stat, or is the preamble's stat cache still in effect?
• if the file is include-guarded, we're not going to read or parse it I guess?
• if it's not include-guarded but #imported, then again we won't re-read or re-parse it because we use the same directive?
• if the file isn't include-guarded, I guess we must read it (because preamble doesn't contain the actual buffer) and then parse it, but such is life.

344–345

nice to have a comment similar to the one on line 333 for this case

// Include is new in the modified preamble.
// Inject it into the patch and use #line to set the presumed location to where it is spelled

clang-tools-extra/clangd/Preamble.h

114

This comment explains what but not why, and the why is pretty non-obvious here.
Assuming we switch to returning a vector, I'd call it preambleIncludes() and say something like:

Returns #include directives from the \c Modified preamble that were resolved using the \c Baseline preamble.
This covers the new locations of inclusions that were moved around, but not inclusions of new files.
Those will be recorded when parsing the main file: the includes in the injected section will be resolved back to their spelled positions in the main file using the presumed-location mechanism.

("injected section" isn't a term we've defined anywhere, it'd be good to give this idea some name in the class comment as it's the main mechanism behind preamble patching)

117

Looking at the callsite, I think I like returning a vector better. buildAST is copyng a data structure specifically to be patched, and the patch is just overwriting it. The interface looks good but doesn't really reflect what's going on.

To make the vector thing work cleanly, I think we should have it always overwrite. This implies it needs to have the preamble's includes, so the default constructor needs to go away. I'd suggest adding a static factory PreamblePatch::unmodified(const PreambleData&) and hiding all constructors.
This means we can't have a PreamblePatch instance with no preamble, which means dealing with Optional in places but it also seems clearer.

[kadircet]

clang-tools-extra/clangd/Preamble.cpp

279

I am planning to leave it to the caller. Eventually this PreamblePatch should be generated by TUScheduler, in which we can check for staleness and issue either an unmodified patch or create a patch and use it until new preamble is ready or patch is invalidated.

337

This explicitly looks at the existing includes based on spelling, and not whether the include was resolved during scanning or not.

Right. Scanning won't resolve anything as we pass an empty FS to preprocessor.

Instead we're going to parse it along with the mainfile, and... well, I *hope* everything works out:
do we pay for the stat, or is the preamble's stat cache still in effect?

It depends on the user of PreamblePatch, while building an AST we build preprocessor with cached fs coming from preamble.

if the file is include-guarded, we're not going to read or parse it I guess?

Yes that's the case. We'll only see the include directive.

if it's not include-guarded but #imported, then again we won't re-read or re-parse it because we use the same directive?

Yes again.

if the file isn't include-guarded, I guess we must read it (because preamble doesn't contain the actual buffer) and then parse it, but such is life.

Unfortunately yes again, but as you said i am not sure if that's a case we should try and recover from as i can only see two use cases:

• You include the header in a different preprocessor state and intentionally want it to parse again, hopefully we'll handle this after macro patching.
• You deleted the include from a transitively included file and inserted it back into this file again. Well, now the hell breaks loose until we build the preamble again. In the worst case AST will turn into a garbage and we won't be able to serve anything but code completions, which is the state of the world without preamble patching, we are just consuming more cpu now.

344–345

handled in D78743

[kadircet]

• Address comments

[sammccall]

Sorry for the delay. I think this is good though I have a nagging feeling I don't understand all the implications yet :-)

[kadircet]

• Rebase

[kadircet]

• Fix windows build bots in the face of #import directives

[jkorous]

Hi @kadircet!
I am investigating a failure of PatchesAdditionalIncludes test that we got internally. It's a failed assert in ReplayPreamble::replay.
Our clangd source code is for all practical purposes identical to upstream version but not so with clang source. Specifically what we do is that in CompilerInstance::createPreprocessor we always add one particular callback.
This means that when in the test we are calling buildPreamble for TestTU TU with empty buffer we never hit the early return in ReplayPreamble::attach() (ParsedAST.cpp:124) like upstream version does and proceed to create a new ReplayPreamble object with PreambleBounds of size() == 0 which leads to ReplayPreamble::MainFileTokens to be empty and later we hit failed assert in ReplayPreamble::replay about assert(HashTok != MainFileTokens.end() && ...).

Now, while I can just tweak either ReplayPreamble::attach() or remove the PP callback in the test internally I am wondering whether you consider empty preamble & PP with callbacks valid use-case of ReplayPreamble and worth a fix.

This is where we are creating the empty MainFileTokens:

* frame #0: 0x0000000103337649 ClangdTests` clang::clangd::(anonymous namespace)::ReplayPreamble::ReplayPreamble(this=0x0000000114f45da0, Includes=0x00007ffeefbfc678, Delegate=0x0000000114f3bd80, SM=0x0000000114f3dd80, PP=0x0000000115850218, LangOpts=0x0000000114f38eb0, PB=0x00007ffeefbfc658)  + 169 at ParsedAST.cpp:142
  frame #1: 0x000000010333756d ClangdTests` clang::clangd::(anonymous namespace)::ReplayPreamble::ReplayPreamble(this=0x0000000114f45da0, Includes=0x00007ffeefbfc678, Delegate=0x0000000114f3bd80, SM=0x0000000114f3dd80, PP=0x0000000115850218, LangOpts=0x0000000114f38eb0, PB=0x00007ffeefbfc658)  + 77 at ParsedAST.cpp:139
  frame #2: 0x0000000103334fa7 ClangdTests` clang::clangd::(anonymous namespace)::ReplayPreamble::attach(Includes=0x00007ffeefbfc678, Clang=0x0000000114f33ea0, PB=0x00007ffeefbfc658)  + 183 at ParsedAST.cpp:126
  frame #3: 0x0000000103334189 ClangdTests` clang::clangd::ParsedAST::build(Filename=(Data = "/clangd-test/foo.cpp", Length = 20), Inputs=0x00007ffeefbfdf98, CI=nullptr, CompilerInvocationDiags=ArrayRef<clang::clangd::Diag> @ 0x00007ffeefbfd830, Preamble=std::__1::shared_ptr<const clang::clangd::PreambleData>::element_type @ 0x0000000114f3e728 strong=2 weak=1)  + 3897 at ParsedAST.cpp:385
  frame #4: 0x00000001006f1032 ClangdTests` clang::clangd::(anonymous namespace)::ParsedASTTest_PatchesAdditionalIncludes_Test::TestBody(this=0x0000000114f04090)  + 1778 at ParsedASTTests.cpp:477

This is the failed assert:

* frame #4: 0x0000000103337a0c ClangdTests` clang::clangd::(anonymous namespace)::ReplayPreamble::replay(this=0x0000000114f45da0)  + 556 at ParsedAST.cpp:182
  frame #5: 0x000000010333777c ClangdTests` clang::clangd::(anonymous namespace)::ReplayPreamble::FileChanged(this=0x0000000114f45da0, Loc=(ID = 74), Reason=ExitFile, Kind=C_User, PrevFID=(ID = 2))  + 156 at ParsedAST.cpp:166
  frame #6: 0x0000000101b7900a ClangdTests` clang::PPChainedCallbacks::FileChanged(this=0x0000000116204080, Loc=(ID = 74), Reason=ExitFile, FileType=C_User, PrevFID=(ID = 2))  + 90 at PPCallbacks.h:390
  frame #7: 0x0000000101b79046 ClangdTests` clang::PPChainedCallbacks::FileChanged(this=0x00000001162040c0, Loc=(ID = 74), Reason=ExitFile, FileType=C_User, PrevFID=(ID = 2))  + 150 at PPCallbacks.h:391
  frame #8: 0x0000000101b79046 ClangdTests` clang::PPChainedCallbacks::FileChanged(this=0x0000000116204100, Loc=(ID = 74), Reason=ExitFile, FileType=C_User, PrevFID=(ID = 2))  + 150 at PPCallbacks.h:391
  frame #9: 0x0000000101b79046 ClangdTests` clang::PPChainedCallbacks::FileChanged(this=0x0000000116204160, Loc=(ID = 74), Reason=ExitFile, FileType=C_User, PrevFID=(ID = 2))  + 150 at PPCallbacks.h:391
  frame #10: 0x0000000101b79046 ClangdTests` clang::PPChainedCallbacks::FileChanged(this=0x0000000116205480, Loc=(ID = 74), Reason=ExitFile, FileType=C_User, PrevFID=(ID = 2))  + 150 at PPCallbacks.h:391
  frame #11: 0x0000000101b9a590 ClangdTests` clang::Preprocessor::HandleEndOfFile(this=0x0000000115850218, Result=0x0000000116808210, isEndOfMacro=false)  + 3904 at PPLexerChange.cpp:469
  frame #12: 0x0000000101b38713 ClangdTests` clang::Lexer::LexEndOfFile(this=0x0000000116206330, Result=0x0000000116808210, CurPtr="")  + 931 at Lexer.cpp:2833
  frame #13: 0x0000000101b39c44 ClangdTests` clang::Lexer::LexTokenInternal(this=0x0000000116206330, Result=0x0000000116808210, TokAtPhysicalStartOfLine=true)  + 420 at Lexer.cpp:3265
  frame #14: 0x0000000101b382f8 ClangdTests` clang::Lexer::Lex(this=0x0000000116206330, Result=0x0000000116808210)  + 216 at Lexer.cpp:3216
  frame #15: 0x0000000101bd7396 ClangdTests` clang::Preprocessor::Lex(this=0x0000000115850218, Result=0x0000000116808210)  + 118 at Preprocessor.cpp:900
  frame #16: 0x0000000101b75ef1 ClangdTests` clang::Preprocessor::CachingLex(this=0x0000000115850218, Result=0x0000000116808210)  + 289 at PPCaching.cpp:63
  frame #17: 0x0000000101bd73d5 ClangdTests` clang::Preprocessor::Lex(this=0x0000000115850218, Result=0x0000000116808210)  + 181 at Preprocessor.cpp:906
  frame #18: 0x0000000104e3c21c ClangdTests` clang::Parser::TryConsumeToken(this=0x0000000116808200, Expected=semi)  + 204 at Parser.h:489
  frame #19: 0x0000000104e3bf1a ClangdTests` clang::Parser::ExpectAndConsumeSemi(this=0x0000000116808200, DiagID=1526)  + 42 at Parser.cpp:162
  frame #20: 0x0000000104d5f0e0 ClangdTests` clang::Parser::ParseDeclGroup(this=0x0000000116808200, DS=0x00007ffeefbfb6e8, Context=FileContext, DeclEnd=0x0000000000000000, FRI=0x0000000000000000)  + 3168 at ParseDecl.cpp:2243
  frame #21: 0x0000000104e42c48 ClangdTests` clang::Parser::ParseDeclOrFunctionDefInternal(this=0x0000000116808200, attrs=0x00007ffeefbfbc80, DS=0x00007ffeefbfb6e8, AS=AS_none)  + 1704 at Parser.cpp:1109
  frame #22: 0x0000000104e42170 ClangdTests` clang::Parser::ParseDeclarationOrFunctionDefinition(this=0x0000000116808200, attrs=0x00007ffeefbfbc80, DS=0x0000000000000000, AS=AS_none)  + 144 at Parser.cpp:1125
  frame #23: 0x0000000104e411fa ClangdTests` clang::Parser::ParseExternalDeclaration(this=0x0000000116808200, attrs=0x00007ffeefbfbc80, DS=0x0000000000000000)  + 3642 at Parser.cpp:945
  frame #24: 0x0000000104e3f254 ClangdTests` clang::Parser::ParseTopLevelDecl(this=0x0000000116808200, Result=0x00007ffeefbfbde8, IsFirstDecl=false)  + 1828 at Parser.cpp:693
  frame #25: 0x0000000104d465e3 ClangdTests` clang::ParseAST(S=0x000000011680c200, PrintStats=false, SkipFunctionBodies=false)  + 595 at ParseAST.cpp:158
  frame #26: 0x0000000101992662 ClangdTests` clang::ASTFrontendAction::ExecuteAction(this=0x0000000114f349a0)  + 322 at FrontendAction.cpp:1066
  frame #27: 0x0000000101991b78 ClangdTests` clang::FrontendAction::Execute(this=0x0000000114f349a0)  + 136 at FrontendAction.cpp:959
  frame #28: 0x00000001033343ab ClangdTests` clang::clangd::ParsedAST::build(Filename=(Data = "/clangd-test/foo.cpp", Length = 20), Inputs=0x00007ffeefbfdf98, CI=nullptr, CompilerInvocationDiags=ArrayRef<clang::clangd::Diag> @ 0x00007ffeefbfd830, Preamble=std::__1::shared_ptr<const clang::clangd::PreambleData>::element_type @ 0x0000000114f3e728 strong=2 weak=1)  + 4443 at ParsedAST.cpp:415
  frame #29: 0x00000001006f1032 ClangdTests` clang::clangd::(anonymous namespace)::ParsedASTTest_PatchesAdditionalIncludes_Test::TestBody(this=0x0000000114f04090)  + 1778 at ParsedASTTests.cpp:477

[uabelho]

clang-tools-extra/clangd/CodeComplete.cpp

1034

Hi!

When compiling with gcc 7.4 I see a warning that I think originates from this line:

[4032/4668] Building CXX object tools/clang/tools/extra/clangd/CMakeFiles/obj.clangDaemon.dir/CodeComplete.cpp.o
In file included from /data/repo/master/llvm/include/llvm/ADT/STLExtras.h:19:0,
                 from /data/repo/master/llvm/include/llvm/ADT/StringRef.h:12,
                 from /data/repo/master/clang-tools-extra/clangd/URI.h:12,
                 from /data/repo/master/clang-tools-extra/clangd/Protocol.h:26,
                 from /data/repo/master/clang-tools-extra/clangd/Headers.h:12,
                 from /data/repo/master/clang-tools-extra/clangd/CodeComplete.h:18,
                 from /data/repo/master/clang-tools-extra/clangd/CodeComplete.cpp:20:
/data/repo/master/llvm/include/llvm/ADT/Optional.h: In instantiation of 'void llvm::optional_detail::OptionalStorage<T, <anonymous> >::emplace(Args&& ...) [with Args = {const clang::clangd::PreamblePatch}; T = const clang::clangd::PreamblePatch; bool <anonymous> = false]':
/data/repo/master/llvm/include/llvm/ADT/Optional.h:55:14:   required from 'llvm::optional_detail::OptionalStorage<T, <anonymous> >::OptionalStorage(llvm::optional_detail::OptionalStorage<T, <anonymous> >&&) [with T = const clang::clangd::PreamblePatch; bool <anonymous> = false]'
/data/repo/master/llvm/include/llvm/ADT/Optional.h:228:3:   required from here
/data/repo/master/llvm/include/llvm/ADT/Optional.h:89:12: warning: cast from type 'const clang::clangd::PreamblePatch*' to type 'void*' casts away qualifiers [-Wcast-qual]
     ::new ((void *)std::addressof(value)) T(std::forward<Args>(args)...);
            ^~~~~~~~~~~~~~~~~~~~~~~~~~~~~

clang-tools-extra/clangd/CodeComplete.cpp

private:
   CodeCompletionTUInfo CCTUInfo;
   const SymbolIndex *Index;
 }; // SignatureHelpCollector
 
 struct SemaCompleteInput {
   PathRef FileName;
   const tooling::CompileCommand &Command;
   const PreambleData &Preamble;
-  const PreamblePatch &Patch;
+  llvm::Optional<const PreamblePatch> Patch;
   llvm::StringRef Contents;
   size_t Offset;
   llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> VFS;
 };
 
 void loadMainFilePreambleMacros(const Preprocessor &PP,
                                 const PreambleData &Preamble) {
   // The ExternalPreprocessorSource has our macros, if we know where to look.
bool semaCodeComplete(std::unique_ptr<CodeCompleteConsumer> Consumer,
   // correctness/performance tradeoff: building without a preamble is slow, and
   // completion is latency-sensitive.
   // However, if we're completing *inside* the preamble section of the draft,
   // overriding the preamble will break sema completion. Fortunately we can just
   // skip all includes in this case; these completions are really simple.
   PreambleBounds PreambleRegion =
       ComputePreambleBounds(*CI->getLangOpts(), ContentsBuffer.get(), 0);
   bool CompletingInPreamble = PreambleRegion.Size > Input.Offset;
-  Input.Patch.apply(*CI);
+  if (Input.Patch)
+    Input.Patch->apply(*CI);
   // NOTE: we must call BeginSourceFile after prepareCompilerInstance. Otherwise
   // the remapped buffers do not get freed.
   auto Clang = prepareCompilerInstance(
       std::move(CI), !CompletingInPreamble ? &Input.Preamble.Preamble : nullptr,
       std::move(ContentsBuffer), std::move(VFS), IgnoreDiags);
   Clang->getPreprocessorOpts().SingleFileParseMode = CompletingInPreamble;
   Clang->setCodeCompletionConsumer(Consumer.release());
 
codeComplete(PathRef FileName, const tooling::CompileCommand &Command,
   auto Flow = CodeCompleteFlow(
       FileName, Preamble ? Preamble->Includes : IncludeStructure(),
       SpecFuzzyFind, Opts);
   return (!Preamble || Opts.RunParser == CodeCompleteOptions::NeverParse)
              ? std::move(Flow).runWithoutSema(Contents, *Offset, VFS)
              : std::move(Flow).run({FileName, Command, *Preamble,
                                     // We want to serve code completions with
                                     // low latency, so don't bother patching.
-                                    PreamblePatch(), Contents, *Offset, VFS});
+                                    /*PreamblePatch=*/llvm::None, Contents,
+                                    *Offset, VFS});
 }
 
 SignatureHelp signatureHelp(PathRef FileName,
                             const tooling::CompileCommand &Command,
                             const PreambleData &Preamble,
                             llvm::StringRef Contents, Position Pos,
                             llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> VFS,
                             const SymbolIndex *Index) {
   auto Offset = positionToOffset(Contents, Pos);
   if (!Offset) {
     elog("Signature help position was invalid {0}", Offset.takeError());
     return SignatureHelp();
   }
   SignatureHelp Result;
   clang::CodeCompleteOptions Options;
   Options.IncludeGlobals = false;
   Options.IncludeMacros = false;
   Options.IncludeCodePatterns = false;
   Options.IncludeBriefComments = false;
 
   ParseInputs PI;
   PI.CompileCommand = Command;
   PI.Contents = Contents.str();
   PI.FS = std::move(VFS);
-  auto PP = PreamblePatch::create(FileName, PI, Preamble);
   semaCodeComplete(
       std::make_unique<SignatureHelpCollector>(Options, Index, Result), Options,
-      {FileName, Command, Preamble, PP, Contents, *Offset, std::move(PI.FS)});
+      {FileName, Command, Preamble,
+       PreamblePatch::create(FileName, PI, Preamble), Contents, *Offset,
+       std::move(PI.FS)});
   return Result;
 }
 
 bool isIndexedForCodeCompletion(const NamedDecl &ND, ASTContext &ASTCtx) {
   auto InTopLevelScope = [](const NamedDecl &ND) {
     switch (ND.getDeclContext()->getDeclKind()) {
     case Decl::TranslationUnit:
     case Decl::Namespace:

clang-tools-extra/clangd/ParsedAST.cpp

 //===--- ParsedAST.cpp -------------------------------------------*- C++-*-===//
 //
 // 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 "ParsedAST.h"
 #include "../clang-tidy/ClangTidyDiagnosticConsumer.h"
 #include "../clang-tidy/ClangTidyModuleRegistry.h"
 #include "AST.h"
 #include "Compiler.h"
 #include "Diagnostics.h"
 #include "Headers.h"
 #include "IncludeFixer.h"
+#include "Preamble.h"
 #include "SourceCode.h"
 #include "index/CanonicalIncludes.h"
 #include "index/Index.h"
 #include "support/Logger.h"
 #include "support/Trace.h"
 #include "clang/AST/ASTContext.h"
 #include "clang/AST/Decl.h"
 #include "clang/Basic/LangOptions.h"

 #include "clang/Tooling/Syntax/Tokens.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/raw_ostream.h"
 #include <algorithm>
 #include <memory>
+#include <vector>
 
 // Force the linker to link in Clang-tidy modules.
 // clangd doesn't support the static analyzer.
 #define CLANG_TIDY_DISABLE_STATIC_ANALYZER_CHECKS
 #include "../clang-tidy/ClangTidyForceLinker.h"
 
 namespace clang {
 namespace clangd {
const PrecompiledPreamble *PreamblePCH =
       Preamble ? &Preamble->Preamble : nullptr;
 
   // This is on-by-default in windows to allow parsing SDK headers, but it
   // breaks many features. Disable it for the main-file (not preamble).
   CI->getLangOpts()->DelayedTemplateParsing = false;
 
   StoreDiags ASTDiags;
 
+  llvm::Optional<PreamblePatch> Patch;
+  if (Preamble) {
+    Patch = PreamblePatch::create(Filename, Inputs, *Preamble);
+    Patch->apply(*CI);
+  }
   auto Clang = prepareCompilerInstance(
       std::move(CI), PreamblePCH,
       llvm::MemoryBuffer::getMemBufferCopy(Inputs.Contents, Filename), VFS,
       ASTDiags);
   if (!Clang)
     return None;
 
   auto Action = std::make_unique<ClangdFrontendAction>();
FixIncludes.emplace(Filename, Inserter, *Inputs.Index,
                         /*IndexRequestLimit=*/5);
     ASTDiags.contributeFixes([&FixIncludes](DiagnosticsEngine::Level DiagLevl,
                                             const clang::Diagnostic &Info) {
       return FixIncludes->fix(DiagLevl, Info);
     });
     Clang->setExternalSemaSource(FixIncludes->unresolvedNameRecorder());
   }
 
-  // Copy over the includes from the preamble, then combine with the
-  // non-preamble includes below.
-  auto Includes = Preamble ? Preamble->Includes : IncludeStructure{};
+  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.
-  if (Preamble)
     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.
   MainFileMacros Macros;

clang-tools-extra/clangd/Preamble.h

 /// return true when some missing headers are now available.
 /// FIXME: Should return more information about the delta between \p Preamble
 /// and \p Inputs, e.g. new headers.
 bool isPreambleCompatible(const PreambleData &Preamble,
                           const ParseInputs &Inputs, PathRef FileName,
                           const CompilerInvocation &CI);
 
 /// Stores information required to parse a TU using a (possibly stale) Baseline
-/// preamble. Updates compiler invocation to approximately reflect additions to
-/// the preamble section of Modified contents, e.g. new include directives.
+/// preamble. Later on this information can be injected into the main file by
+/// updating compiler invocation with \c apply. This injected section
+/// approximately reflects additions to the preamble in Modified contents, e.g.
+/// new include directives.
 class PreamblePatch {
 public:
-  // With an empty patch, the preamble is used verbatim.
-  PreamblePatch() = default;
+  /// \p Preamble is used verbatim.
+  static PreamblePatch unmodified(const PreambleData &Preamble);
   /// Builds a patch that contains new PP directives introduced to the preamble
   /// section of \p Modified compared to \p Baseline.
   /// FIXME: This only handles include directives, we should at least handle
   /// define/undef.
   static PreamblePatch create(llvm::StringRef FileName,
                               const ParseInputs &Modified,
                               const PreambleData &Baseline);
   /// Adjusts CI (which compiles the modified inputs) to be used with the
   /// baseline preamble. This is done by inserting an artifical include to the
   /// \p CI that contains new directives calculated in create.
   void apply(CompilerInvocation &CI) const;
 
+  /// Returns #include directives from the \c Modified preamble that were
+  /// resolved using the \c Baseline preamble. This covers the new locations of
+  /// inclusions that were moved around, but not inclusions of new files. Those
+  /// will be recorded when parsing the main file: the includes in the injected
+  /// section will be resolved back to their spelled positions in the main file
+  /// using the presumed-location mechanism.
+  std::vector<Inclusion> preambleIncludes() const;
+
 private:
+  PreamblePatch() = default;
   std::string PatchContents;
   std::string PatchFileName;
+  /// Includes that are present in both \p Baseline and \p Modified. Used for
+  /// patching includes of baseline preamble.
+  std::vector<Inclusion> PreambleIncludes;
 };
 
 } // namespace clangd
 } // namespace clang
 
 #endif // LLVM_CLANG_TOOLS_EXTRA_CLANGD_PREAMBLE_H

clang-tools-extra/clangd/Preamble.cpp

 #include "clang/Frontend/CompilerInvocation.h"
 #include "clang/Frontend/FrontendActions.h"
 #include "clang/Lex/Lexer.h"
 #include "clang/Lex/PPCallbacks.h"
 #include "clang/Lex/Preprocessor.h"
 #include "clang/Lex/PreprocessorOptions.h"
 #include "clang/Tooling/CompilationDatabase.h"
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/IntrusiveRefCntPtr.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/StringSet.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/ErrorHandling.h"
for (char C : Text) {
     }
     OS << C;
   }
 }
 
 PreamblePatch PreamblePatch::create(llvm::StringRef FileName,
                                     const ParseInputs &Modified,
                                     const PreambleData &Baseline) {
+  assert(llvm::sys::path::is_absolute(FileName) && "relative FileName!");
   // First scan the include directives in Baseline and Modified. These will be
   // used to figure out newly added directives in Modified. Scanning can fail,
   // the code just bails out and creates an empty patch in such cases, as:
   // - If scanning for Baseline fails, no knowledge of existing includes hence
   //   patch will contain all the includes in Modified. Leading to rebuild of
   //   whole preamble, which is terribly slow.
   // - If scanning for Modified fails, cannot figure out newly added ones so
   //   there's nothing to do but generate an empty patch.
auto ModifiedIncludes = scanPreambleIncludes(
       Modified.CompileCommand);
   if (!ModifiedIncludes) {
     elog("Failed to scan includes for modified contents of {0}: {1}", FileName,
          ModifiedIncludes.takeError());
     return {};
   }
   // No patch needed if includes are equal.
   if (*BaselineIncludes == *ModifiedIncludes)
-    return {};
+    return PreamblePatch::unmodified(Baseline);
 
   PreamblePatch PP;
   // This shouldn't coincide with any real file name.
   llvm::SmallString<128> PatchName;
   llvm::sys::path::append(PatchName, llvm::sys::path::parent_path(FileName),
                           "__preamble_patch__.h");
   PP.PatchFileName = PatchName.str().str();
 
   // We are only interested in newly added includes, record the ones in Baseline
   // for exclusion.
-  llvm::DenseSet<std::pair<tok::PPKeywordKind, llvm::StringRef>>
+  llvm::DenseMap<std::pair<tok::PPKeywordKind, llvm::StringRef>,
+                 /*Resolved=*/llvm::StringRef>
       ExistingIncludes;
+  for (const auto &Inc : Baseline.Includes.MainFileIncludes)
+    ExistingIncludes[{Inc.Directive, Inc.Written}] = Inc.Resolved;
+  // There might be includes coming from disabled regions, record these for
+  // exclusion too. note that we don't have resolved paths for those.
   for (const auto &Inc : *BaselineIncludes)
-    ExistingIncludes.insert({Inc.Directive, Inc.Written});
+    ExistingIncludes.try_emplace({Inc.Directive, Inc.Written});
   // Calculate extra includes that needs to be inserted.
   llvm::raw_string_ostream Patch(PP.PatchContents);
   // Set default filename for subsequent #line directives
   Patch << "#line 0 \"";
   // FileName part of a line directive is subject to backslash escaping, which
   // might lead to problems on windows especially.
   escapeBackslashAndQuotes(FileName, Patch);
   Patch << "\"\n";
-  for (const auto &Inc : *ModifiedIncludes) {
-    if (ExistingIncludes.count({Inc.Directive, Inc.Written}))
+  for (auto &Inc : *ModifiedIncludes) {
+    auto It = ExistingIncludes.find({Inc.Directive, Inc.Written});
+    // Include already present in the baseline preamble. Set resolved path and
+    // put into preamble includes.
+    if (It != ExistingIncludes.end()) {
+      Inc.Resolved = It->second.str();
+      PP.PreambleIncludes.push_back(Inc);
       continue;
+    }
     // Include is new in the modified preamble. Inject it into the patch and use
     // #line to set the presumed location to where it is spelled.
     auto LineCol = offsetToClangLineColumn(Modified.Contents, Inc.HashOffset);
     Patch << llvm::formatv("#line {0}\n", LineCol.first);
     Patch << llvm::formatv("#{0} {1}\n", spellingForIncDirective(Inc.Directive),
                            Inc.Written);
   }
   Patch.flush();
 
   // FIXME: Handle more directives, e.g. define/undef.
auto PatchBuffer =
       llvm::MemoryBuffer::getMemBufferCopy(PatchContents, PatchFileName);
   // CI will take care of the lifetime of the buffer.
   PPOpts.addRemappedFile(PatchFileName, PatchBuffer.release());
   // The patch will be parsed after loading the preamble ast and before parsing
   // the main file.
   PPOpts.Includes.push_back(PatchFileName);
 }
 
+std::vector<Inclusion> PreamblePatch::preambleIncludes() const {
+  return PreambleIncludes;
+}
+
+PreamblePatch PreamblePatch::unmodified(const PreambleData &Preamble) {
+  PreamblePatch PP;
+  PP.PreambleIncludes = Preamble.Includes.MainFileIncludes;
+  return PP;
+}
+
 } // namespace clangd
 } // namespace clang

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

TEST_F(DocumentSymbolsTest, UsingDirectives) {
   addFile(FilePath, Source.code());
   EXPECT_THAT(getSymbols(FilePath),
               ElementsAre(WithName("ns"), WithName("ns_alias"),
                           WithName("using namespace ::ns"),
                           WithName("using namespace ns_alias")));
 }
 
 TEST_F(DocumentSymbolsTest, TempSpecs) {
-  addFile("foo.cpp", R"cpp(
+  std::string FilePath = testPath("foo.cpp");
+  addFile(FilePath, R"cpp(
       template <typename T, typename U, int X = 5> class Foo {};
       template <typename T> class Foo<int, T> {};
       template <> class Foo<bool, int> {};
       template <> class Foo<bool, int, 3> {};
       )cpp");
   // Foo is higher ranked because of exact name match.
   EXPECT_THAT(
-      getSymbols("foo.cpp"),
+      getSymbols(FilePath),
       UnorderedElementsAre(
           AllOf(WithName("Foo"), WithKind(SymbolKind::Class)),
           AllOf(WithName("Foo<int, T>"), WithKind(SymbolKind::Class)),
           AllOf(WithName("Foo<bool, int>"), WithKind(SymbolKind::Class)),
           AllOf(WithName("Foo<bool, int, 3>"), WithKind(SymbolKind::Class))));
 }
 
 TEST_F(DocumentSymbolsTest, Qualifiers) {
-  addFile("foo.cpp", R"cpp(
+  std::string FilePath = testPath("foo.cpp");
+  addFile(FilePath, R"cpp(
     namespace foo { namespace bar {
       struct Cls;
 
       int func1();
       int func2();
       int func3();
       int func4();
     }}
 
     struct foo::bar::Cls { };
 
     int foo::bar::func1() { return 10; }
     int ::foo::bar::func2() { return 20; }
 
     using namespace foo;
     int bar::func3() { return 30; }
 
     namespace alias = foo::bar;
     int ::alias::func4() { return 40; }
   )cpp");
 
   // All the qualifiers should be preserved exactly as written.
-  EXPECT_THAT(getSymbols("foo.cpp"),
+  EXPECT_THAT(getSymbols(FilePath),
               UnorderedElementsAre(
                   WithName("foo"), WithName("foo::bar::Cls"),
                   WithName("foo::bar::func1"), WithName("::foo::bar::func2"),
                   WithName("using namespace foo"), WithName("bar::func3"),
                   WithName("alias"), WithName("::alias::func4")));
 }
 
 TEST_F(DocumentSymbolsTest, QualifiersWithTemplateArgs) {
-  addFile("foo.cpp", R"cpp(
+  std::string FilePath = testPath("foo.cpp");
+  addFile(FilePath, R"cpp(
       template <typename T, typename U = double> class Foo;
 
       template <>
       class Foo<int, double> {
         int method1();
         int method2();
         int method3();
       };
 
       using int_type = int;
 
       // Typedefs should be preserved!
       int Foo<int_type, double>::method1() { return 10; }
 
       // Default arguments should not be shown!
       int Foo<int>::method2() { return 20; }
 
       using Foo_type = Foo<int>;
       // If the whole type is aliased, this should be preserved too!
       int Foo_type::method3() { return 30; }
       )cpp");
   EXPECT_THAT(
-      getSymbols("foo.cpp"),
+      getSymbols(FilePath),
       UnorderedElementsAre(WithName("Foo"), WithName("Foo<int, double>"),
                            WithName("int_type"),
                            WithName("Foo<int_type, double>::method1"),
                            WithName("Foo<int>::method2"), WithName("Foo_type"),
                            WithName("Foo_type::method3")));
 }
 
 } // namespace clangd
 } // namespace clang

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

 //===----------------------------------------------------------------------===//
 
 #include "../../clang-tidy/ClangTidyModule.h"
 #include "../../clang-tidy/ClangTidyModuleRegistry.h"
 #include "AST.h"
 #include "Annotations.h"
 #include "Compiler.h"
 #include "Diagnostics.h"
+#include "Headers.h"
 #include "ParsedAST.h"
+#include "Preamble.h"
 #include "SourceCode.h"
 #include "TestFS.h"
 #include "TestTU.h"
 #include "clang/AST/DeclTemplate.h"
 #include "clang/Basic/SourceLocation.h"
 #include "clang/Basic/SourceManager.h"
 #include "clang/Basic/TokenKinds.h"
 #include "clang/Lex/PPCallbacks.h"
 #include "clang/Lex/Token.h"
 #include "clang/Tooling/Syntax/Tokens.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/ScopedPrinter.h"
 #include "gmock/gmock-matchers.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 
 namespace clang {
 namespace clangd {
if (const NamedDecl *ND = dyn_cast<NamedDecl>(arg))
     return printTemplateSpecializationArgs(*ND) == ArgName;
   return false;
 }
 
 MATCHER_P(RangeIs, R, "") {
   return arg.beginOffset() == R.Begin && arg.endOffset() == R.End;
 }
 
+MATCHER(EqInc, "") {
+  Inclusion Actual = testing::get<0>(arg);
+  Inclusion Expected = testing::get<1>(arg);
+  return std::tie(Actual.HashLine, Actual.Written) ==
+         std::tie(Expected.HashLine, Expected.Written);
+}
+
 TEST(ParsedASTTest, TopLevelDecls) {
   TestTU TU;
   TU.HeaderCode = R"(
     int header1();
     int header2;
   )";
   TU.Code = "int main();";
   auto AST = TU.build();
for (size_t I = 0; I < Includes.size(); ++I) {
     // 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);
   }
 }
 
+TEST(ParsedASTTest, PatchesAdditionalIncludes) {
+  llvm::StringLiteral ModifiedContents = R"cpp(
+    #include "baz.h"
+    #include "foo.h"
+    #include "sub/aux.h"
+    void bar() {
+      foo();
+      baz();
+      aux();
+    })cpp";
+  // Build expected ast with symbols coming from headers.
+  TestTU TU;
+  TU.Filename = "foo.cpp";
+  TU.AdditionalFiles["foo.h"] = "void foo();";
+  TU.AdditionalFiles["sub/baz.h"] = "void baz();";
+  TU.AdditionalFiles["sub/aux.h"] = "void aux();";
+  TU.ExtraArgs = {"-I" + testPath("sub")};
+  TU.Code = ModifiedContents.str();
+  auto ExpectedAST = TU.build();
+
+  // Build preamble with no includes.
+  TU.Code = "";
+  StoreDiags Diags;
+  auto Inputs = TU.inputs();
+  auto CI = buildCompilerInvocation(Inputs, Diags);
+  auto EmptyPreamble =
+      buildPreamble(testPath("foo.cpp"), *CI, Inputs, true, nullptr);
+  ASSERT_TRUE(EmptyPreamble);
+  EXPECT_THAT(EmptyPreamble->Includes.MainFileIncludes, testing::IsEmpty());
+
+  // Now build an AST using empty preamble and ensure patched includes worked.
+  TU.Code = ModifiedContents.str();
+  Inputs = TU.inputs();
+  auto PatchedAST = ParsedAST::build(testPath("foo.cpp"), Inputs, std::move(CI),
+                                     {}, EmptyPreamble);
+  ASSERT_TRUE(PatchedAST);
+  ASSERT_TRUE(PatchedAST->getDiagnostics().empty());
+
+  // Ensure source location information is correct, including resolved paths.
+  EXPECT_THAT(PatchedAST->getIncludeStructure().MainFileIncludes,
+              testing::Pointwise(
+                  EqInc(), ExpectedAST.getIncludeStructure().MainFileIncludes));
+  auto StringMapToVector = [](const llvm::StringMap<unsigned> SM) {
+    std::vector<std::pair<std::string, unsigned>> Res;
+    for (const auto &E : SM)
+      Res.push_back({E.first().str(), E.second});
+    llvm::sort(Res);
+    return Res;
+  };
+  // Ensure file proximity signals are correct.
+  EXPECT_EQ(StringMapToVector(PatchedAST->getIncludeStructure().includeDepth(
+                testPath("foo.cpp"))),
+            StringMapToVector(ExpectedAST.getIncludeStructure().includeDepth(
+                testPath("foo.cpp"))));
+}
+
+TEST(ParsedASTTest, PatchesDeletedIncludes) {
+  TestTU TU;
+  TU.Filename = "foo.cpp";
+  TU.Code = "";
+  auto ExpectedAST = TU.build();
+
+  // Build preamble with no includes.
+  TU.Code = R"cpp(#include <foo.h>)cpp";
+  StoreDiags Diags;
+  auto Inputs = TU.inputs();
+  auto CI = buildCompilerInvocation(Inputs, Diags);
+  auto BaselinePreamble =
+      buildPreamble(testPath("foo.cpp"), *CI, Inputs, true, nullptr);
+  ASSERT_TRUE(BaselinePreamble);
+  EXPECT_THAT(BaselinePreamble->Includes.MainFileIncludes,
+              ElementsAre(testing::Field(&Inclusion::Written, "<foo.h>")));
+
+  // Now build an AST using additional includes and check that locations are
+  // correctly parsed.
+  TU.Code = "";
+  Inputs = TU.inputs();
+  auto PatchedAST = ParsedAST::build(testPath("foo.cpp"), Inputs, std::move(CI),
+                                     {}, BaselinePreamble);
+  ASSERT_TRUE(PatchedAST);
+
+  // Ensure source location information is correct.
+  EXPECT_THAT(PatchedAST->getIncludeStructure().MainFileIncludes,
+              testing::Pointwise(
+                  EqInc(), ExpectedAST.getIncludeStructure().MainFileIncludes));
+  auto StringMapToVector = [](const llvm::StringMap<unsigned> SM) {
+    std::vector<std::pair<std::string, unsigned>> Res;
+    for (const auto &E : SM)
+      Res.push_back({E.first().str(), E.second});
+    llvm::sort(Res);
+    return Res;
+  };
+  // Ensure file proximity signals are correct.
+  EXPECT_EQ(StringMapToVector(PatchedAST->getIncludeStructure().includeDepth(
+                testPath("foo.cpp"))),
+            StringMapToVector(ExpectedAST.getIncludeStructure().includeDepth(
+                testPath("foo.cpp"))));
+}
+
 } // namespace
 } // namespace clangd
 } // namespace clang

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

 //===--- PreambleTests.cpp --------------------------------------*- C++ -*-===//
 //
 // 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 "Annotations.h"
 #include "Compiler.h"
+#include "Headers.h"
 #include "Preamble.h"
 #include "TestFS.h"
 #include "TestTU.h"
 #include "clang/Frontend/PrecompiledPreamble.h"
 #include "clang/Lex/PreprocessorOptions.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/VirtualFileSystem.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 #include <clang/Frontend/FrontendActions.h>
+#include <memory>
 #include <string>
 #include <vector>
 
 using testing::Field;
 
 namespace clang {
 namespace clangd {
 namespace {
 
+MATCHER_P2(Distance, File, D, "") {
+  return arg.first() == File && arg.second == D;
+}
+
+std::shared_ptr<const PreambleData>
+createPreamble(llvm::StringRef Contents = "") {
+  auto TU = TestTU::withCode(Contents);
+  // ms-compatibility changes meaning of #import, make sure it is turned off.
+  TU.ExtraArgs = {"-fno-ms-compatibility"};
+  TU.Filename = "preamble.cpp";
+  auto PI = TU.inputs();
+  IgnoreDiagnostics Diags;
+  auto CI = buildCompilerInvocation(PI, Diags);
+  if (!CI) {
+    ADD_FAILURE() << "failed to build compiler invocation";
+    return nullptr;
+  }
+  if (auto Preamble = buildPreamble(TU.Filename, *CI, PI, true, nullptr))
+    return Preamble;
+  ADD_FAILURE() << "failed to build preamble";
+  return nullptr;
+}
+
 // Builds a preamble for BaselineContents, patches it for ModifiedContents and
 // returns the includes in the patch.
 IncludeStructure
 collectPatchedIncludes(llvm::StringRef ModifiedContents,
                        llvm::StringRef BaselineContents,
                        llvm::StringRef MainFileName = "main.cpp") {
-  std::string MainFile = testPath(MainFileName);
-  ParseInputs PI;
-  PI.FS = new llvm::vfs::InMemoryFileSystem;
-  MockCompilationDatabase CDB;
+  auto BaselinePreamble = createPreamble(BaselineContents);
+  // Create the patch.
+  auto TU = TestTU::withCode(ModifiedContents);
+  TU.Filename = MainFileName.str();
   // ms-compatibility changes meaning of #import, make sure it is turned off.
-  CDB.ExtraClangFlags.push_back("-fno-ms-compatibility");
-  PI.CompileCommand = CDB.getCompileCommand(MainFile).getValue();
-  // Create invocation
+  TU.ExtraArgs = {"-fno-ms-compatibility"};
+  auto PI = TU.inputs();
+  auto PP = PreamblePatch::create(testPath(TU.Filename), PI, *BaselinePreamble);
+  // Collect patch contents.
   IgnoreDiagnostics Diags;
   auto CI = buildCompilerInvocation(PI, Diags);
-  assert(CI && "failed to create compiler invocation");
-  // Build baseline preamble.
-  PI.Contents = BaselineContents.str();
-  PI.Version = "baseline preamble";
-  auto BaselinePreamble = buildPreamble(MainFile, *CI, PI, true, nullptr);
-  assert(BaselinePreamble && "failed to build baseline preamble");
-  // Create the patch.
-  PI.Contents = ModifiedContents.str();
-  PI.Version = "modified contents";
-  auto PP = PreamblePatch::create(MainFile, PI, *BaselinePreamble);
-  // Collect patch contents.
   PP.apply(*CI);
-  llvm::StringRef PatchContents;
-  for (const auto &Rempaped : CI->getPreprocessorOpts().RemappedFileBuffers) {
-    if (Rempaped.first == testPath("__preamble_patch__.h")) {
-      PatchContents = Rempaped.second->getBuffer();
-      break;
-    }
-  }
-  // Run preprocessor over the modified contents with patched Invocation to and
-  // BaselinePreamble to collect includes in the patch. We trim the input to
-  // only preamble section to not collect includes in the mainfile.
+  // Run preprocessor over the modified contents with patched Invocation. We
+  // provide a preamble and trim contents to ensure only the implicit header
+  // introduced by the patch is parsed and nothing else.
+  // We don't run PP directly over the patch cotents to test production
+  // behaviour.
   auto Bounds = Lexer::ComputePreamble(ModifiedContents, *CI->getLangOpts());
   auto Clang =
       prepareCompilerInstance(std::move(CI), &BaselinePreamble->Preamble,
                               llvm::MemoryBuffer::getMemBufferCopy(
                                   ModifiedContents.slice(0, Bounds.Size).str()),
                               PI.FS, Diags);
-  Clang->getPreprocessorOpts().ImplicitPCHInclude.clear();
   PreprocessOnlyAction Action;
   if (!Action.BeginSourceFile(*Clang, Clang->getFrontendOpts().Inputs[0])) {
     ADD_FAILURE() << "failed begin source file";
     return {};
   }
   IncludeStructure Includes;
   Clang->getPreprocessor().addPPCallbacks(
       collectIncludeStructureCallback(Clang->getSourceManager(), &Includes));

 }
 
 TEST(PreamblePatchTest, MainFileIsEscaped) {
   auto Includes = collectPatchedIncludes("#include <a.h>", "", "file\"name.cpp")
                       .MainFileIncludes;
   EXPECT_THAT(Includes, ElementsAre(AllOf(Field(&Inclusion::Written, "<a.h>"),
                                           Field(&Inclusion::HashLine, 0))));
 }
+
+TEST(PreamblePatchTest, PatchesPreambleIncludes) {
+  IgnoreDiagnostics Diags;
+  auto TU = TestTU::withCode(R"cpp(
+    #include "a.h"
+    #include "c.h"
+  )cpp");
+  TU.AdditionalFiles["a.h"] = "#include \"b.h\"";
+  TU.AdditionalFiles["b.h"] = "";
+  TU.AdditionalFiles["c.h"] = "";
+  auto PI = TU.inputs();
+  auto BaselinePreamble = buildPreamble(
+      TU.Filename, *buildCompilerInvocation(PI, Diags), PI, true, nullptr);
+  // We drop c.h from modified and add a new header. Since the latter is patched
+  // we should only get a.h in preamble includes.
+  TU.Code = R"cpp(
+    #include "a.h"
+    #include "b.h"
+  )cpp";
+  auto PP = PreamblePatch::create(testPath(TU.Filename), TU.inputs(),
+                                  *BaselinePreamble);
+  // Only a.h should exists in the preamble, as c.h has been dropped and b.h was
+  // newly introduced.
+  EXPECT_THAT(PP.preambleIncludes(),
+              ElementsAre(AllOf(Field(&Inclusion::Written, "\"a.h\""),
+                                Field(&Inclusion::Resolved, testPath("a.h")))));
+}
 } // namespace
 } // namespace clangd
 } // namespace clang