[clangd] Enable include insertion for dynamic index

[kadircet]

Summary

This enables include insertion by adding canonical includes into
preambledata.

[ioeric]

clangd/ClangdUnit.cpp

100

Does this have to own the IWYUHandler? Could we create one when getCommentHandler is called?

124

This looks like a memory leak? release()?

340

I think so? A header file (with IWYU pragma) can also be the main file.

unittests/clangd/CodeCompleteTests.cpp

2317

nit: maybe put this close to test cases that involve preamble. e.g. IndexSuppressesPreambleCompletions

2331

nit: just use std::string for readability.

unittests/clangd/FileIndexTests.cpp

207

Are we sure the comment will be included in the preamble if there is no #include block?

[kadircet]

• Address comments.
• Call addSystemHeaderMappings when we are building ast without preamble.

clangd/ClangdUnit.cpp

100

The class needs to own it, since preprocessor doesn't take the ownership. But we can move the creation to getcommenthandler.

124

Preprocessor doesn't take the ownership, so this seems to be correct ?

340

Also call addSystemHeadersMapping in case there is no preamble.

[ioeric]

lgtm.

Note that this will enable include insertions for everyone (even without index switched on). Not sure if we should provide options.

clangd/ClangdServer.cpp

16

nit: is this used?

[kadircet]

I don't think so, since the only reason this already wasn't enabled was because there was a fixme ? Also we want people to use -background-index by default in future so I think it should be OK

[kadircet]

• Delete unnecessary include

clangd/ClangdServer.cpp

 //===--- ClangdServer.cpp - Main clangd server code --------------*- 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 "ClangdServer.h"
 #include "ClangdUnit.h"
 #include "CodeComplete.h"
 #include "FindSymbols.h"
 #include "Headers.h"
 #include "SourceCode.h"
 #include "Trace.h"
+#include "XRefs.h"
+#include "index/CanonicalIncludes.h"
 #include "index/FileIndex.h"
 #include "index/Merge.h"
 #include "refactor/Tweak.h"
 #include "clang/Format/Format.h"
 #include "clang/Frontend/CompilerInstance.h"
 #include "clang/Frontend/CompilerInvocation.h"
 #include "clang/Lex/Preprocessor.h"
 #include "clang/Tooling/CompilationDatabase.h"

 };
 
 // Update the FileIndex with new ASTs and plumb the diagnostics responses.
 struct UpdateIndexCallbacks : public ParsingCallbacks {
   UpdateIndexCallbacks(FileIndex *FIndex, DiagnosticsConsumer &DiagConsumer)
       : FIndex(FIndex), DiagConsumer(DiagConsumer) {}
 
   void onPreambleAST(PathRef Path, ASTContext &Ctx,
-                     std::shared_ptr<clang::Preprocessor> PP) override {
+                     std::shared_ptr<clang::Preprocessor> PP,
+                     const CanonicalIncludes &CanonIncludes) override {
     if (FIndex)
-      FIndex->updatePreamble(Path, Ctx, std::move(PP));
+      FIndex->updatePreamble(Path, Ctx, std::move(PP), CanonIncludes);
   }
 
   void onMainAST(PathRef Path, ParsedAST &AST) override {
     if (FIndex)
       FIndex->updateMain(Path, AST);
   }
 
   void onDiagnostics(PathRef File, std::vector<Diag> Diags) override {

clangd/ClangdUnit.h

 
 #include "Compiler.h"
 #include "Diagnostics.h"
 #include "FS.h"
 #include "Function.h"
 #include "Headers.h"
 #include "Path.h"
 #include "Protocol.h"
+#include "index/CanonicalIncludes.h"
 #include "index/Index.h"
 #include "clang/Frontend/FrontendAction.h"
 #include "clang/Frontend/PrecompiledPreamble.h"
 #include "clang/Lex/Preprocessor.h"
 #include "clang/Serialization/ASTBitCodes.h"
 #include "clang/Tooling/CompilationDatabase.h"
 #include "clang/Tooling/Core/Replacement.h"
 #include <memory>

 } // namespace tooling
 
 namespace clangd {
 
 // Stores Preamble and associated data.
 struct PreambleData {
   PreambleData(PrecompiledPreamble Preamble, std::vector<Diag> Diags,
                IncludeStructure Includes,
-               std::unique_ptr<PreambleFileStatusCache> StatCache);
+               std::unique_ptr<PreambleFileStatusCache> StatCache,
+               CanonicalIncludes CanonIncludes);
 
   tooling::CompileCommand CompileCommand;
   PrecompiledPreamble Preamble;
   std::vector<Diag> Diags;
   // Processes like code completions and go-to-definitions will need #include
   // information, and their compile action skips preamble range.
   IncludeStructure Includes;
   // Cache of FS operations performed when building the preamble.
   // When reusing a preamble, this cache can be consumed to save IO.
   std::unique_ptr<PreambleFileStatusCache> StatCache;
+  CanonicalIncludes CanonIncludes;
 };
 
 /// Stores and provides access to parsed AST.
 class ParsedAST {
 public:
   /// Attempts to run Clang and store parsed AST. If \p Preamble is non-null
   /// it is reused during parsing.
   static llvm::Optional<ParsedAST>
public:
   ArrayRef<Decl *> getLocalTopLevelDecls();
 
   const std::vector<Diag> &getDiagnostics() const;
 
   /// Returns the esitmated size of the AST and the accessory structures, in
   /// bytes. Does not include the size of the preamble.
   std::size_t getUsedBytes() const;
   const IncludeStructure &getIncludeStructure() const;
+  const CanonicalIncludes &getCanonicalIncludes() const;
 
 private:
   ParsedAST(std::shared_ptr<const PreambleData> Preamble,
             std::unique_ptr<CompilerInstance> Clang,
             std::unique_ptr<FrontendAction> Action,
             std::vector<Decl *> LocalTopLevelDecls, std::vector<Diag> Diags,
-            IncludeStructure Includes);
+            IncludeStructure Includes, CanonicalIncludes CanonIncludes);
 
   // In-memory preambles must outlive the AST, it is important that this member
   // goes before Clang and Action.
   std::shared_ptr<const PreambleData> Preamble;
   // We store an "incomplete" FrontendAction (i.e. no EndSourceFile was called
   // on it) and CompilerInstance used to run it. That way we don't have to do
   // complex memory management of all Clang structures on our own. (They are
   // stored in CompilerInstance and cleaned up by
   // FrontendAction.EndSourceFile).
   std::unique_ptr<CompilerInstance> Clang;
   std::unique_ptr<FrontendAction> Action;
 
   // Data, stored after parsing.
   std::vector<Diag> Diags;
   // Top-level decls inside the current file. Not that this does not include
   // top-level decls from the preamble.
   std::vector<Decl *> LocalTopLevelDecls;
   IncludeStructure Includes;
+  CanonicalIncludes CanonIncludes;
 };
 
 using PreambleParsedCallback =
-    std::function<void(ASTContext &, std::shared_ptr<clang::Preprocessor>)>;
+    std::function<void(ASTContext &, std::shared_ptr<clang::Preprocessor>,
+                       const CanonicalIncludes &)>;
 
 /// Rebuild the preamble for the new inputs unless the old one can be reused.
 /// If \p OldPreamble can be reused, it is returned unchanged.
 /// If \p OldPreamble is null, always builds the preamble.
 /// If \p PreambleCallback is set, it will be run on top of the AST while
 /// building the preamble. Note that if the old preamble was reused, no AST is
 /// built and, therefore, the callback will not be executed.
 std::shared_ptr<const PreambleData>

clangd/ClangdUnit.cpp

 #include "../clang-tidy/ClangTidyModuleRegistry.h"
 #include "Compiler.h"
 #include "Diagnostics.h"
 #include "Headers.h"
 #include "IncludeFixer.h"
 #include "Logger.h"
 #include "SourceCode.h"
 #include "Trace.h"
+#include "index/CanonicalIncludes.h"
 #include "index/Index.h"
 #include "clang/AST/ASTContext.h"
 #include "clang/Basic/LangOptions.h"
 #include "clang/Frontend/CompilerInstance.h"
 #include "clang/Frontend/CompilerInvocation.h"
 #include "clang/Frontend/FrontendActions.h"
 #include "clang/Frontend/Utils.h"
 #include "clang/Index/IndexDataConsumer.h"

 private:
   std::vector<Decl *> TopLevelDecls;
 };
 
 class CppFilePreambleCallbacks : public PreambleCallbacks {
 public:
   CppFilePreambleCallbacks(PathRef File, PreambleParsedCallback ParsedCallback)
       : File(File), ParsedCallback(ParsedCallback) {}
 
   IncludeStructure takeIncludes() { return std::move(Includes); }
 
+  CanonicalIncludes takeCanonicalIncludes() {
+    addSystemHeadersMapping(&CanonIncludes);
+    return std::move(CanonIncludes);
+  }
+
   void AfterExecute(CompilerInstance &CI) override {
     if (!ParsedCallback)
       return;
     trace::Span Tracer("Running PreambleCallback");
-    ParsedCallback(CI.getASTContext(), CI.getPreprocessorPtr());
+    ParsedCallback(CI.getASTContext(), CI.getPreprocessorPtr(), CanonIncludes);
   }
 
   void BeforeExecute(CompilerInstance &CI) override {
     SourceMgr = &CI.getSourceManager();
   }
 
   std::unique_ptr<PPCallbacks> createPPCallbacks() override {
     assert(SourceMgr && "SourceMgr must be set at this point");
     return collectIncludeStructureCallback(*SourceMgr, &Includes);
   }
 
+  CommentHandler *getCommentHandler() override {
+    IWYUHandler = collectIWYUHeaderMaps(&CanonIncludes);
+    return IWYUHandler.get();
+  }
+
 private:
   PathRef File;
   PreambleParsedCallback ParsedCallback;
   IncludeStructure Includes;
+  CanonicalIncludes CanonIncludes;
+  std::unique_ptr<CommentHandler> IWYUHandler = nullptr;
   SourceManager *SourceMgr = nullptr;
 };
 
 // When using a preamble, only preprocessor events outside its bounds are seen.
 // This is almost what we want: replaying transitive preprocessing wastes time.
 // However this confuses clang-tidy checks: they don't see any #includes!
 // So we replay the *non-transitive* #includes that appear in the main-file.
 // It would be nice to replay other events (macro definitions, ifdefs etc) but
ParsedAST::build(std::unique_ptr<CompilerInvocation> CI,
   if (Preamble)
     ReplayPreamble::attach(Includes, *Clang);
   // 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 includes from the preamble, then combine with the
+  // non-preamble includes below.
+  CanonicalIncludes CanonIncludes;
+  if (Preamble)
+    CanonIncludes = Preamble->CanonIncludes;
+  else
+    addSystemHeadersMapping(&CanonIncludes);
+  std::unique_ptr<CommentHandler> IWYUHandler =
+      collectIWYUHeaderMaps(&CanonIncludes);
+  Clang->getPreprocessor().addCommentHandler(IWYUHandler.get());
+
   if (!Action->Execute())
     log("Execute() failed when building AST for {0}", MainInput.getFile());
 
   std::vector<Decl *> ParsedDecls = Action->takeTopLevelDecls();
   // AST traversals should exclude the preamble, to avoid performance cliffs.
   Clang->getASTContext().setTraversalScope(ParsedDecls);
   {
     // Run the AST-dependent part of the clang-tidy checks.
ParsedAST::build(std::unique_ptr<CompilerInvocation> CI,
   Clang->getPreprocessor().EndSourceFile();
 
   std::vector<Diag> Diags = ASTDiags.take();
   // Add diagnostics from the preamble, if any.
   if (Preamble)
     Diags.insert(Diags.begin(), Preamble->Diags.begin(), Preamble->Diags.end());
   return ParsedAST(std::move(Preamble), std::move(Clang), std::move(Action),
                    std::move(ParsedDecls), std::move(Diags),
-                   std::move(Includes));
+                   std::move(Includes), std::move(CanonIncludes));
 }
 
 ParsedAST::ParsedAST(ParsedAST &&Other) = default;
 
 ParsedAST &ParsedAST::operator=(ParsedAST &&Other) = default;
 
 ParsedAST::~ParsedAST() {
   if (Action) {
std::size_t ParsedAST::getUsedBytes() const {
 
   return Total;
 }
 
 const IncludeStructure &ParsedAST::getIncludeStructure() const {
   return Includes;
 }
 
+const CanonicalIncludes &ParsedAST::getCanonicalIncludes() const {
+  return CanonIncludes;
+}
+
 PreambleData::PreambleData(PrecompiledPreamble Preamble,
                            std::vector<Diag> Diags, IncludeStructure Includes,
-                           std::unique_ptr<PreambleFileStatusCache> StatCache)
+                           std::unique_ptr<PreambleFileStatusCache> StatCache,
+                           CanonicalIncludes CanonIncludes)
     : Preamble(std::move(Preamble)), Diags(std::move(Diags)),
-      Includes(std::move(Includes)), StatCache(std::move(StatCache)) {}
+      Includes(std::move(Includes)), StatCache(std::move(StatCache)),
+      CanonIncludes(std::move(CanonIncludes)) {}
 
 ParsedAST::ParsedAST(std::shared_ptr<const PreambleData> Preamble,
                      std::unique_ptr<CompilerInstance> Clang,
                      std::unique_ptr<FrontendAction> Action,
                      std::vector<Decl *> LocalTopLevelDecls,
-                     std::vector<Diag> Diags, IncludeStructure Includes)
+                     std::vector<Diag> Diags, IncludeStructure Includes,
+                     CanonicalIncludes CanonIncludes)
     : Preamble(std::move(Preamble)), Clang(std::move(Clang)),
       Action(std::move(Action)), Diags(std::move(Diags)),
       LocalTopLevelDecls(std::move(LocalTopLevelDecls)),
-      Includes(std::move(Includes)) {
+      Includes(std::move(Includes)), CanonIncludes(std::move(CanonIncludes)) {
   assert(this->Clang);
   assert(this->Action);
 }
 
 std::shared_ptr<const PreambleData>
 buildPreamble(PathRef FileName, CompilerInvocation &CI,
               std::shared_ptr<const PreambleData> OldPreamble,
               const tooling::CompileCommand &OldCompileCommand,
buildPreamble(PathRef FileName, CompilerInvocation &CI,
   // bodies.
   CI.getFrontendOpts().SkipFunctionBodies = false;
 
   if (BuiltPreamble) {
     vlog("Built preamble of size {0} for file {1}", BuiltPreamble->getSize(),
          FileName);
     return std::make_shared<PreambleData>(
         std::move(*BuiltPreamble), PreambleDiagnostics.take(),
-        SerializedDeclsCollector.takeIncludes(), std::move(StatCache));
+        SerializedDeclsCollector.takeIncludes(), std::move(StatCache),
+        SerializedDeclsCollector.takeCanonicalIncludes());
   } else {
     elog("Could not build a preamble for file {0}", FileName);
     return nullptr;
   }
 }
 
 llvm::Optional<ParsedAST>
 buildAST(PathRef FileName, std::unique_ptr<CompilerInvocation> Invocation,

clangd/TUScheduler.h

 //===--- TUScheduler.h -------------------------------------------*-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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_TUSCHEDULER_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_TUSCHEDULER_H
 
 #include "ClangdUnit.h"
 #include "Function.h"
 #include "Threading.h"
+#include "index/CanonicalIncludes.h"
 #include "llvm/ADT/StringMap.h"
 #include <future>
 
 namespace clang {
 namespace clangd {
 
 /// Returns a number of a default async threads to use for TUScheduler.
 /// Returned value is always >= 1 (i.e. will not cause requests to be processed

 class ParsingCallbacks {
 public:
   virtual ~ParsingCallbacks() = default;
 
   /// Called on the AST that was built for emitting the preamble. The built AST
   /// contains only AST nodes from the #include directives at the start of the
   /// file. AST node in the current file should be observed on onMainAST call.
   virtual void onPreambleAST(PathRef Path, ASTContext &Ctx,
-                             std::shared_ptr<clang::Preprocessor> PP) {}
+                             std::shared_ptr<clang::Preprocessor> PP,
+                             const CanonicalIncludes &) {}
   /// Called on the AST built for the file itself. Note that preamble AST nodes
   /// are not deserialized and should be processed in the onPreambleAST call
   /// instead.
   /// The \p AST always contains all AST nodes for the main file itself, and
   /// only a portion of the AST nodes deserialized from the preamble. Note that
   /// some nodes from the preamble may have been deserialized and may also be
   /// accessed from the main file AST, e.g. redecls of functions from preamble,
   /// etc. Clients are expected to process only the AST nodes from the main file

clangd/TUScheduler.cpp

 //   pending inputs and reset the timer.
 // - If any reads of the AST are scheduled, we start building the AST
 //   immediately.
 
 #include "TUScheduler.h"
 #include "Cancellation.h"
 #include "Logger.h"
 #include "Trace.h"
+#include "index/CanonicalIncludes.h"
 #include "clang/Frontend/CompilerInvocation.h"
 #include "clang/Frontend/PCHContainerOperations.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/Support/Errc.h"
 #include "llvm/Support/Path.h"
 #include <algorithm>
 #include <memory>
 #include <queue>
if (!Invocation) {
       return;
     }
 
     std::shared_ptr<const PreambleData> OldPreamble =
         getPossiblyStalePreamble();
     std::shared_ptr<const PreambleData> NewPreamble = buildPreamble(
         FileName, *Invocation, OldPreamble, OldCommand, Inputs, PCHs,
         StorePreambleInMemory,
-        [this](ASTContext &Ctx, std::shared_ptr<clang::Preprocessor> PP) {
-          Callbacks.onPreambleAST(FileName, Ctx, std::move(PP));
+        [this](ASTContext &Ctx, std::shared_ptr<clang::Preprocessor> PP,
+               const CanonicalIncludes &CanonIncludes) {
+          Callbacks.onPreambleAST(FileName, Ctx, std::move(PP), CanonIncludes);
         });
 
     bool CanReuseAST = InputsAreTheSame && (OldPreamble == NewPreamble);
     {
       std::lock_guard<std::mutex> Lock(Mutex);
       LastBuiltPreamble = NewPreamble;
     }
     // Before doing the expensive AST reparse, we want to release our reference

clangd/index/FileIndex.h

 
 #ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_FILEINDEX_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_FILEINDEX_H
 
 #include "ClangdUnit.h"
 #include "Index.h"
 #include "MemIndex.h"
 #include "Merge.h"
+#include "index/CanonicalIncludes.h"
 #include "clang/Lex/Preprocessor.h"
 #include <memory>
 
 namespace clang {
 namespace clangd {
 
 /// Select between in-memory index implementations, which have tradeoffs.
 enum class IndexType {

 /// FIXME: Expose an interface to remove files that are closed.
 class FileIndex : public MergedIndex {
 public:
   FileIndex(bool UseDex = true);
 
   /// Update preamble symbols of file \p Path with all declarations in \p AST
   /// and macros in \p PP.
   void updatePreamble(PathRef Path, ASTContext &AST,
-                      std::shared_ptr<Preprocessor> PP);
+                      std::shared_ptr<Preprocessor> PP,
+                      const CanonicalIncludes &Includes);
 
   /// Update symbols and references from main file \p Path with
   /// `indexMainDecls`.
   void updateMain(PathRef Path, ParsedAST &AST);
 
 private:
   bool UseDex; // FIXME: this should be always on.
 

 
 /// Retrieves symbols and refs of local top level decls in \p AST (i.e.
 /// `AST.getLocalTopLevelDecls()`).
 /// Exposed to assist in unit tests.
 std::pair<SymbolSlab, RefSlab> indexMainDecls(ParsedAST &AST);
 
 /// Idex declarations from \p AST and macros from \p PP that are declared in
 /// included headers.
-SymbolSlab indexHeaderSymbols(ASTContext &AST,
-                              std::shared_ptr<Preprocessor> PP);
+SymbolSlab indexHeaderSymbols(ASTContext &AST, std::shared_ptr<Preprocessor> PP,
+                              const CanonicalIncludes &Includes);
 
 } // namespace clangd
 } // namespace clang
 
 #endif // LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_FILEINDEX_H

clangd/index/FileIndex.cpp

 //===--- FileIndex.cpp - Indexes for files. ------------------------ 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 "FileIndex.h"
 #include "ClangdUnit.h"
 #include "Logger.h"
 #include "SymbolCollector.h"
+#include "index/CanonicalIncludes.h"
 #include "index/Index.h"
 #include "index/MemIndex.h"
 #include "index/Merge.h"
 #include "index/dex/Dex.h"
 #include "clang/Index/IndexingAction.h"
 #include "clang/Lex/MacroInfo.h"
 #include "clang/Lex/Preprocessor.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/StringRef.h"
 #include <memory>
 
 namespace clang {
 namespace clangd {
 
 static std::pair<SymbolSlab, RefSlab>
 indexSymbols(ASTContext &AST, std::shared_ptr<Preprocessor> PP,
-             llvm::ArrayRef<Decl *> DeclsToIndex, bool IsIndexMainAST) {
+             llvm::ArrayRef<Decl *> DeclsToIndex,
+             const CanonicalIncludes &Includes, bool IsIndexMainAST) {
   SymbolCollector::Options CollectorOpts;
-  // FIXME(ioeric): we might also want to collect include headers. We would need
-  // to make sure all includes are canonicalized (with CanonicalIncludes), which
-  // is not trivial given the current way of collecting symbols: we only have
-  // AST at this point, but we also need preprocessor callbacks (e.g.
-  // CommentHandler for IWYU pragma) to canonicalize includes.
-  CollectorOpts.CollectIncludePath = false;
+  CollectorOpts.CollectIncludePath = true;
+  CollectorOpts.Includes = &Includes;
   CollectorOpts.CountReferences = false;
   CollectorOpts.Origin = SymbolOrigin::Dynamic;
 
   index::IndexingOptions IndexOpts;
   // We only need declarations, because we don't count references.
   IndexOpts.SystemSymbolFilter =
       index::IndexingOptions::SystemSymbolFilterKind::DeclarationsOnly;
   IndexOpts.IndexFunctionLocals = false;
   if (IsIndexMainAST) {
     // We only collect refs when indexing main AST.
     CollectorOpts.RefFilter = RefKind::All;
-  }else {
+  } else {
     IndexOpts.IndexMacrosInPreprocessor = true;
     CollectorOpts.CollectMacro = true;
   }
 
   SymbolCollector Collector(std::move(CollectorOpts));
   Collector.setPreprocessor(PP);
   index::indexTopLevelDecls(AST, *PP, DeclsToIndex, Collector, IndexOpts);
 
   const auto &SM = AST.getSourceManager();
   const auto *MainFileEntry = SM.getFileEntryForID(SM.getMainFileID());
   std::string FileName = MainFileEntry ? MainFileEntry->getName() : "";
 
   auto Syms = Collector.takeSymbols();
   auto Refs = Collector.takeRefs();
   vlog("index AST for {0} (main={1}): \n"
        "  symbol slab: {2} symbols, {3} bytes\n"
        "  ref slab: {4} symbols, {5} refs, {6} bytes",
        FileName, IsIndexMainAST, Syms.size(), Syms.bytes(), Refs.size(),
        Refs.numRefs(), Refs.bytes());
   return {std::move(Syms), std::move(Refs)};
 }
 
 std::pair<SymbolSlab, RefSlab> indexMainDecls(ParsedAST &AST) {
   return indexSymbols(AST.getASTContext(), AST.getPreprocessorPtr(),
-                      AST.getLocalTopLevelDecls(),
+                      AST.getLocalTopLevelDecls(), AST.getCanonicalIncludes(),
                       /*IsIndexMainAST=*/true);
 }
 
-SymbolSlab indexHeaderSymbols(ASTContext &AST,
-                              std::shared_ptr<Preprocessor> PP) {
+SymbolSlab indexHeaderSymbols(ASTContext &AST, std::shared_ptr<Preprocessor> PP,
+                              const CanonicalIncludes &Includes) {
   std::vector<Decl *> DeclsToIndex(
       AST.getTranslationUnitDecl()->decls().begin(),
       AST.getTranslationUnitDecl()->decls().end());
-  return indexSymbols(AST, std::move(PP), DeclsToIndex,
+  return indexSymbols(AST, std::move(PP), DeclsToIndex, Includes,
                       /*IsIndexMainAST=*/false)
       .first;
 }
 
 void FileSymbols::update(PathRef Path, std::unique_ptr<SymbolSlab> Symbols,
                          std::unique_ptr<RefSlab> Refs) {
   std::lock_guard<std::mutex> Lock(Mutex);
   if (!Symbols)

 }
 
 FileIndex::FileIndex(bool UseDex)
     : MergedIndex(&MainFileIndex, &PreambleIndex), UseDex(UseDex),
       PreambleIndex(llvm::make_unique<MemIndex>()),
       MainFileIndex(llvm::make_unique<MemIndex>()) {}
 
 void FileIndex::updatePreamble(PathRef Path, ASTContext &AST,
-                               std::shared_ptr<Preprocessor> PP) {
-  auto Symbols = indexHeaderSymbols(AST, std::move(PP));
+                               std::shared_ptr<Preprocessor> PP,
+                               const CanonicalIncludes &Includes) {
+  auto Symbols = indexHeaderSymbols(AST, std::move(PP), Includes);
   PreambleSymbols.update(Path,
                          llvm::make_unique<SymbolSlab>(std::move(Symbols)),
                          llvm::make_unique<RefSlab>());
   PreambleIndex.reset(
       PreambleSymbols.buildIndex(UseDex ? IndexType::Heavy : IndexType::Light,
                                  DuplicateHandling::PickOne));
 }
 

unittests/clangd/CodeCompleteTests.cpp

auto Results = completions(R"cpp(
     #include <bar.h>
     int foo() {}
     #endif
     )cpp")
                      .Completions;
   EXPECT_THAT(Results, ElementsAre(Named("ifndef")));
 }
 
+TEST(CompletionTest, DynamicIndexIncludeInsertion) {
+  MockFSProvider FS;
+  MockCompilationDatabase CDB;
+  IgnoreDiagnostics DiagConsumer;
+  ClangdServer::Options Opts = ClangdServer::optsForTest();
+  Opts.BuildDynamicSymbolIndex = true;
+  ClangdServer Server(CDB, FS, DiagConsumer, Opts);
+
+  FS.Files[testPath("foo_header.h")] = R"cpp(
+    struct Foo {
+       // Member doc
+       int foo();
+    };
+  )cpp";
+  const std::string FileContent(R"cpp(
+    #include "foo_header.h"
+    int Foo::foo() {
+      return 42;
+    }
+  )cpp");
+  Server.addDocument(testPath("foo_impl.cpp"), FileContent);
+  // Wait for the dynamic index being built.
+  ASSERT_TRUE(Server.blockUntilIdleForTest());
+  EXPECT_THAT(
+      completions(Server, "Foo^ foo;").Completions,
+      ElementsAre(AllOf(Named("Foo"),
+                        HasInclude('"' + testPath("foo_header.h") + '"'),
+                        InsertInclude())));
+}
+
 TEST(CompletionTest, DynamicIndexMultiFile) {
   MockFSProvider FS;
   MockCompilationDatabase CDB;
   IgnoreDiagnostics DiagConsumer;
   auto Opts = ClangdServer::optsForTest();
   Opts.BuildDynamicSymbolIndex = true;
   ClangdServer Server(CDB, FS, DiagConsumer, Opts);
 
int main() {
       for (auto [loopVar] : y ) { // y has to be unresolved.
         int z = loopV^;
       }
     }
   )cpp");
   EXPECT_THAT(R.Completions, ElementsAre(Named("loopVar")));
 }
 
 } // namespace
 } // namespace clangd
 } // namespace clang

unittests/clangd/FileIndexTests.cpp

 //===-- FileIndexTests.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 "AST.h"
 #include "Annotations.h"
 #include "ClangdUnit.h"
 #include "SyncAPI.h"
 #include "TestFS.h"
 #include "TestTU.h"
+#include "index/CanonicalIncludes.h"
 #include "index/FileIndex.h"
 #include "index/Index.h"
 #include "clang/Frontend/CompilerInvocation.h"
 #include "clang/Frontend/PCHContainerOperations.h"
 #include "clang/Frontend/Utils.h"
 #include "clang/Index/IndexSymbol.h"
 #include "clang/Lex/Preprocessor.h"
 #include "clang/Tooling/CompilationDatabase.h"

 
 // Adds Basename.cpp, which includes Basename.h, which contains Code.
 void update(FileIndex &M, llvm::StringRef Basename, llvm::StringRef Code) {
   TestTU File;
   File.Filename = (Basename + ".cpp").str();
   File.HeaderFilename = (Basename + ".h").str();
   File.HeaderCode = Code;
   auto AST = File.build();
-  M.updatePreamble(File.Filename, AST.getASTContext(),
-                   AST.getPreprocessorPtr());
+  M.updatePreamble(File.Filename, AST.getASTContext(), AST.getPreprocessorPtr(),
+                   AST.getCanonicalIncludes());
 }
 
 TEST(FileIndexTest, CustomizedURIScheme) {
   FileIndex M;
   update(M, "f", "class string {};");
 
   EXPECT_THAT(runFuzzyFind(M, ""), ElementsAre(DeclURI("unittest:///f.h")));
 }
TEST(FileIndexTest, ClassMembers) {
   FileIndex M;
   update(M, "f1", "class X { static int m1; int m2; static void f(); };");
 
   EXPECT_THAT(runFuzzyFind(M, ""),
               UnorderedElementsAre(QName("X"), QName("X::m1"), QName("X::m2"),
                                    QName("X::f")));
 }
 
-TEST(FileIndexTest, NoIncludeCollected) {
+TEST(FileIndexTest, IncludeCollected) {
   FileIndex M;
-  update(M, "f", "class string {};");
+  update(
+      M, "f",
+      "// IWYU pragma: private, include <the/good/header.h>\nclass string {};");
 
   auto Symbols = runFuzzyFind(M, "");
   EXPECT_THAT(Symbols, ElementsAre(_));
-  EXPECT_THAT(Symbols.begin()->IncludeHeaders, IsEmpty());
+  EXPECT_THAT(Symbols.begin()->IncludeHeaders.front().IncludeHeader,
+              "<the/good/header.h>");
 }
 
 TEST(FileIndexTest, TemplateParamsInLabel) {
   auto Source = R"cpp(
 template <class Ty>
 class vector {
 };
 
TEST(FileIndexTest, RebuildWithPreamble) {
   // Rebuild the file.
   auto CI = buildCompilerInvocation(PI);
 
   FileIndex Index;
   bool IndexUpdated = false;
   buildPreamble(
       FooCpp, *CI, /*OldPreamble=*/nullptr, tooling::CompileCommand(), PI,
       std::make_shared<PCHContainerOperations>(), /*StoreInMemory=*/true,
-      [&](ASTContext &Ctx, std::shared_ptr<Preprocessor> PP) {
+      [&](ASTContext &Ctx, std::shared_ptr<Preprocessor> PP,
+          const CanonicalIncludes &CanonIncludes) {
         EXPECT_FALSE(IndexUpdated) << "Expected only a single index update";
         IndexUpdated = true;
-        Index.updatePreamble(FooCpp, Ctx, std::move(PP));
+        Index.updatePreamble(FooCpp, Ctx, std::move(PP), CanonIncludes);
       });
   ASSERT_TRUE(IndexUpdated);
 
   // Check the index contains symbols from the preamble, but not from the main
   // file.
   FuzzyFindRequest Req;
   Req.Query = "";
   Req.Scopes = {"", "ns_in_header::"};
TEST(FileIndexTest, ReferencesInMainFileWithPreamble) {
   PI.FS = buildTestFS({{MainFile, Main.code()}, {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) {});
+      [&](ASTContext &Ctx, std::shared_ptr<Preprocessor> PP,
+          const CanonicalIncludes &) {});
   // Build AST for main file with preamble.
   auto AST =
       ParsedAST::build(createInvocationFromCommandLine(Cmd), PreambleData,
                        llvm::MemoryBuffer::getMemBufferCopy(Main.code()),
                        std::make_shared<PCHContainerOperations>(), PI.FS,
                        /*Index=*/nullptr, ParseOptions());
   ASSERT_TRUE(AST);
   FileIndex Index;

unittests/clangd/TestTU.cpp

if (!AST.hasValue()) {
     ADD_FAILURE() << "Failed to build code:\n" << Code;
     llvm_unreachable("Failed to build TestTU!");
   }
   return std::move(*AST);
 }
 
 SymbolSlab TestTU::headerSymbols() const {
   auto AST = build();
-  return indexHeaderSymbols(AST.getASTContext(), AST.getPreprocessorPtr());
+  return indexHeaderSymbols(AST.getASTContext(), AST.getPreprocessorPtr(),
+                            AST.getCanonicalIncludes());
 }
 
 std::unique_ptr<SymbolIndex> TestTU::index() const {
   auto AST = build();
   auto Idx = llvm::make_unique<FileIndex>(/*UseDex=*/true);
-  Idx->updatePreamble(Filename, AST.getASTContext(), AST.getPreprocessorPtr());
+  Idx->updatePreamble(Filename, AST.getASTContext(), AST.getPreprocessorPtr(),
+                      AST.getCanonicalIncludes());
   Idx->updateMain(Filename, AST);
   return std::move(Idx);
 }
 
 const Symbol &findSymbol(const SymbolSlab &Slab, llvm::StringRef QName) {
   const Symbol *Result = nullptr;
   for (const Symbol &S : Slab) {
     if (QName != (S.Scope + S.Name).str())