[clangd] Enable include insertion for dynamic index

[kadircet]

Summary

This enables include insertion by adding canonical includes into
preambledata.

[ioeric]

clangd/ClangdUnit.cpp

100 ↗

(On Diff #184485)

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

125 ↗

(On Diff #184485)

This looks like a memory leak? release()?

338 ↗

(On Diff #184485)

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

unittests/clangd/CodeCompleteTests.cpp

2287 ↗

(On Diff #184485)

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

2301 ↗

(On Diff #184485)

nit: just use std::string for readability.

unittests/clangd/FileIndexTests.cpp

207 ↗

(On Diff #184485)

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 ↗

(On Diff #184485)

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

125 ↗

(On Diff #184485)

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

338 ↗

(On Diff #184485)

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 ↗

(On Diff #185030)

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

clang-tools-extra/trunk/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 "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 {

clang-tools-extra/trunk/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>

clang-tools-extra/trunk/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"

 
 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,

clang-tools-extra/trunk/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

clang-tools-extra/trunk/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

clang-tools-extra/trunk/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

clang-tools-extra/trunk/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));
 }
 

clang-tools-extra/trunk/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);
 

clang-tools-extra/trunk/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;

clang-tools-extra/trunk/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())