[clangd] Make system header mappings available for PreambleParsedCallback

[kadircet]

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

clangd/ClangdUnit.cpp

 
 private:
   std::vector<Decl *> TopLevelDecls;
 };
 
 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)
       return;
     trace::Span Tracer("Running PreambleCallback");
     ParsedCallback(CI.getASTContext(), CI.getPreprocessorPtr(), CanonIncludes);
   }
 

unittests/clangd/FileIndexTests.cpp

update(
       "// IWYU pragma: private, include <the/good/header.h>\nclass string {};");
 
   auto Symbols = runFuzzyFind(M, "");
   EXPECT_THAT(Symbols, ElementsAre(_));
   EXPECT_THAT(Symbols.begin()->IncludeHeaders.front().IncludeHeader,
               "<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>
 class vector {
 };
 
 template <class Ty, class Arg>
 vector<Ty> make_vector(Arg A) {}