[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

115 ↗

(On Diff #260228)

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

265

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

280 ↗

(On Diff #260228)

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

332 ↗

(On Diff #260228)

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.

340 ↗

(On Diff #260228)

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

112 ↗

(On Diff #260228)

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)

115 ↗

(On Diff #260228)

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

280 ↗

(On Diff #260228)

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.

332 ↗

(On Diff #260228)

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.

340 ↗

(On Diff #260228)

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 ↗

(On Diff #267167)

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/Compiler.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_COMPILER_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_COMPILER_H
 
 #include "../clang-tidy/ClangTidyOptions.h"
 #include "GlobalCompilationDatabase.h"
+#include "Headers.h"
 #include "index/Index.h"
 #include "clang/Frontend/CompilerInstance.h"
 #include "clang/Frontend/PrecompiledPreamble.h"
 #include "clang/Tooling/CompilationDatabase.h"
+#include <vector>
 
 namespace clang {
 namespace clangd {
 
 class IgnoreDiagnostics : public DiagnosticConsumer {
 public:
   static void log(DiagnosticsEngine::Level DiagLevel,
                   const clang::Diagnostic &Info);
 
   void HandleDiagnostic(DiagnosticsEngine::Level DiagLevel,
                         const clang::Diagnostic &Info) override;
 };
 
 // Options to run clang e.g. when parsing AST.
 struct ParseOptions {
   tidy::ClangTidyOptions ClangTidyOpts;
   bool SuggestMissingIncludes = false;
   bool BuildRecoveryAST = false;
+  std::vector<Inclusion> AdditionalIncludes;
 };
 
 /// Information required to run clang, e.g. to parse AST or do code completion.
 struct ParseInputs {
   tooling::CompileCommand CompileCommand;
   IntrusiveRefCntPtr<llvm::vfs::FileSystem> FS;
   std::string Contents;
   // Version identifier for Contents, provided by the client and opaque to us.

clang-tools-extra/clangd/ParsedAST.cpp

ParsedAST::build(llvm::StringRef Version,
   if (CI->getLangOpts()->CPlusPlus)
     CI->getLangOpts()->RecoveryAST = Opts.BuildRecoveryAST;
 
   StoreDiags ASTDiags;
   std::string Content = std::string(Buffer->getBuffer());
   std::string Filename =
       std::string(Buffer->getBufferIdentifier()); // Absolute.
 
+  // Process additional includes as implicit includes to ensure AST is built
+  // with necessary symbol information.
+  for (const auto &Inc : Opts.AdditionalIncludes) {
+    // Inc.Written contains quotes or angles, preprocessor requires them to be
+    // stripped.
+    CI->getPreprocessorOpts().Includes.emplace_back(
+        llvm::StringRef(Inc.Written).drop_back().drop_front());
+  }
+
   auto Clang = prepareCompilerInstance(std::move(CI), PreamblePCH,
                                        std::move(Buffer), VFS, ASTDiags);
   if (!Clang)
     return None;
 
   auto Action = std::make_unique<ClangdFrontendAction>();
   const FrontendInputFile &MainInput = Clang->getFrontendOpts().Inputs[0];
   if (!Action->BeginSourceFile(*Clang, MainInput)) {
ParsedAST::build(llvm::StringRef Version,
   // Add diagnostics from the preamble, if any.
   if (Preamble)
     Diags.insert(Diags.end(), Preamble->Diags.begin(), Preamble->Diags.end());
   // Finally, add diagnostics coming from the AST.
   {
     std::vector<Diag> D = ASTDiags.take(CTContext.getPointer());
     Diags.insert(Diags.end(), D.begin(), D.end());
   }
+
+  // Add location information for additional includes.
+  auto &SM = Clang->getSourceManager();
+  auto FID = SM.getMainFileID();
+  const auto *Includer = SM.getFileEntryForID(FID);
+  auto FLoc = SM.getLocForStartOfFile(FID);
+  for (auto &Inc : Opts.AdditionalIncludes) {
+    Includes.MainFileIncludes.push_back(Inc);
+
+    // Try to resolve it using header search.
+    const DirectoryLookup *CurDir = nullptr;
+    auto IncLoc = FLoc.getLocWithOffset(Inc.HashOffset);
+    auto FER = Clang->getPreprocessor().getHeaderSearchInfo().LookupFile(
+        llvm::StringRef(Inc.Written).drop_front().drop_back(), IncLoc,
+        Inc.Written.front() == '<', nullptr, CurDir,
+        {{Includer, Includer->getDir()}}, nullptr, nullptr, nullptr, nullptr,
+        nullptr, nullptr);
+    if (!FER)
+      continue;
+    auto &FE = FER->getFileEntry();
+    Includes.MainFileIncludes.back().Resolved = FE.tryGetRealPathName().str();
+    Includes.recordInclude(Includer->getName(), FE.getName(),
+                           FE.tryGetRealPathName());
+  }
+
   return ParsedAST(Version, std::move(Preamble), std::move(Clang),
                    std::move(Action), std::move(Tokens), std::move(Macros),
                    std::move(ParsedDecls), std::move(Diags),
                    std::move(Includes), std::move(CanonIncludes));
 }
 
 ParsedAST::ParsedAST(ParsedAST &&Other) = default;
 

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.HashOffset, Actual.R, Actual.Written) ==
+         std::tie(Expected.HashOffset, Expected.R, 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, AdditionalIncludes) {
+  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 = R"cpp(
+    #include "baz.h"
+    #include "foo.h"
+    #include "sub/aux.h"
+    void bar() {
+      foo();
+      baz();
+      aux();
+    })cpp";
+  auto ExpectedAST = TU.build();
+
+  // Build preamble with no includes.
+  TU.Code = R"cpp(
+    void bar() {
+      foo();
+      baz();
+      aux();
+    })cpp";
+  StoreDiags Diags;
+  auto Inputs = TU.inputs();
+  auto CI = buildCompilerInvocation(Inputs, Diags);
+  auto Preamble = buildPreamble("foo.cpp", *CI, Inputs, true, nullptr);
+  ASSERT_TRUE(Preamble);
+  EXPECT_THAT(Preamble->Includes.MainFileIncludes, testing::IsEmpty());
+
+  // Now build an AST using additional includes and check that locations are
+  // correctly parsed.
+  TU.Code = R"cpp(
+    #include "baz.h"
+    #include "foo.h"
+    #include "sub/aux.h"
+    void bar() {
+      foo();
+      baz();
+      aux();
+    })cpp";
+  Inputs = TU.inputs();
+  Inputs.Opts.AdditionalIncludes =
+      getPreambleIncludes(TU.Code, ExpectedAST.getLangOpts());
+  auto PatchedAST = buildAST("foo.cpp", std::move(CI), {}, Inputs, Preamble);
+  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