[clangd] Make system header mappings available for PreambleParsedCallback

Summary:
SystemHeaderMappings were added only after takeIncludes call, which
resulted in getting mapping on main file ast updates but not on preamble ast
updates.
Fixes https://github.com/clangd/clangd/issues/8

Reviewers: hokein

Subscribers: ilya-biryukov, ioeric, MaskRay, jkorous, arphaman, cfe-commits

Tags: #clang

Differential Revision: https://reviews.llvm.org/D58029

git-svn-id: https://llvm.org/svn/llvm-project/clang-tools-extra/trunk@353687 91177308-0d34-0410-b5e6-96231b3b80d8

2 files changed, 37 insertions, 5 deletions

diff --git a/clangd/ClangdUnit.cpp b/clangd/ClangdUnit.cpp
index 09ab772d..1689a644 100644
--- a/clangd/ClangdUnit.cpp
+++ b/clangd/ClangdUnit.cpp
@@ -96,14 +96,13 @@ private:
 class CppFilePreambleCallbacks : public PreambleCallbacks {
 public:
   CppFilePreambleCallbacks(PathRef File, PreambleParsedCallback ParsedCallback)
-      : File(File), ParsedCallback(ParsedCallback) {}
+      : File(File), ParsedCallback(ParsedCallback) {
+    addSystemHeadersMapping(&CanonIncludes);
+  }
 
   IncludeStructure takeIncludes() { return std::move(Includes); }
 
-  CanonicalIncludes takeCanonicalIncludes() {
-    addSystemHeadersMapping(&CanonIncludes);
-    return std::move(CanonIncludes);
-  }
+  CanonicalIncludes takeCanonicalIncludes() { return std::move(CanonIncludes); }
 
   void AfterExecute(CompilerInstance &CI) override {
     if (!ParsedCallback)
diff --git a/unittests/clangd/FileIndexTests.cpp b/unittests/clangd/FileIndexTests.cpp
index 0b037a50..fc374bba 100644
--- a/unittests/clangd/FileIndexTests.cpp
+++ b/unittests/clangd/FileIndexTests.cpp
@@ -212,6 +212,39 @@ TEST(FileIndexTest, IncludeCollected) {
               "<the/good/header.h>");
 }
 
+TEST(FileIndexTest, HasSystemHeaderMappingsInPreamble) {
+  FileIndex Index;
+  const std::string Header = R"cpp(
+    class Foo {};
+  )cpp";
+  auto MainFile = testPath("foo.cpp");
+  auto HeaderFile = testPath("bits/alloc_traits.h");
+  std::vector<const char *> Cmd = {"clang", "-xc++", MainFile.c_str(),
+                                   "-include", HeaderFile.c_str()};
+  // Preparse ParseInputs.
+  ParseInputs PI;
+  PI.CompileCommand.Directory = testRoot();
+  PI.CompileCommand.Filename = MainFile;
+  PI.CompileCommand.CommandLine = {Cmd.begin(), Cmd.end()};
+  PI.Contents = "";
+  PI.FS = buildTestFS({{MainFile, ""}, {HeaderFile, Header}});
+
+  // Prepare preamble.
+  auto CI = buildCompilerInvocation(PI);
+  auto PreambleData = buildPreamble(
+      MainFile, *buildCompilerInvocation(PI), /*OldPreamble=*/nullptr,
+      tooling::CompileCommand(), PI, std::make_shared<PCHContainerOperations>(),
+      /*StoreInMemory=*/true,
+      [&](ASTContext &Ctx, std::shared_ptr<Preprocessor> PP,
+          const CanonicalIncludes &Includes) {
+        Index.updatePreamble(MainFile, Ctx, PP, Includes);
+      });
+  auto Symbols = runFuzzyFind(Index, "");
+  EXPECT_THAT(Symbols, ElementsAre(_));
+  EXPECT_THAT(Symbols.begin()->IncludeHeaders.front().IncludeHeader,
+              "<memory>");
+}
+
 TEST(FileIndexTest, TemplateParamsInLabel) {
   auto Source = R"cpp(
 template <class Ty>