[clangd] Enable include insertion for static index

Summary:
This enables include insertion by adding canonical includes into
preambledata.

Reviewers: ioeric, ilya-biryukov

Subscribers: javed.absar, MaskRay, jkorous, arphaman, cfe-commits

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

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

10 files changed, 125 insertions, 42 deletions

diff --git a/clangd/ClangdServer.cpp b/clangd/ClangdServer.cpp
index 595e0758..2dcf0fb3 100644
--- a/clangd/ClangdServer.cpp
+++ b/clangd/ClangdServer.cpp
@@ -13,6 +13,7 @@
 #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"
@@ -69,9 +70,10 @@ struct UpdateIndexCallbacks : public ParsingCallbacks {
       : 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 {
diff --git a/clangd/ClangdUnit.cpp b/clangd/ClangdUnit.cpp
index 6b5e1764..ffb96fbf 100644
--- a/clangd/ClangdUnit.cpp
+++ b/clangd/ClangdUnit.cpp
@@ -16,6 +16,7 @@
 #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"
@@ -99,11 +100,16 @@ public:
 
   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 {
@@ -115,10 +121,17 @@ public:
     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;
 };
 
@@ -324,6 +337,17 @@ ParsedAST::build(std::unique_ptr<CompilerInvocation> CI,
   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());
 
@@ -353,7 +377,7 @@ ParsedAST::build(std::unique_ptr<CompilerInvocation> CI,
     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;
@@ -429,21 +453,28 @@ 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);
 }
@@ -510,7 +541,8 @@ buildPreamble(PathRef FileName, CompilerInvocation &CI,
          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;
diff --git a/clangd/ClangdUnit.h b/clangd/ClangdUnit.h
index f4883fd5..d5bdd1ab 100644
--- a/clangd/ClangdUnit.h
+++ b/clangd/ClangdUnit.h
@@ -16,6 +16,7 @@
 #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"
@@ -48,7 +49,8 @@ namespace clangd {
 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;
@@ -59,6 +61,7 @@ struct PreambleData {
   // 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.
@@ -100,13 +103,14 @@ public:
   /// 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.
@@ -125,10 +129,12 @@ private:
   // 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.
diff --git a/clangd/TUScheduler.cpp b/clangd/TUScheduler.cpp
index 14319ab9..c7730827 100644
--- a/clangd/TUScheduler.cpp
+++ b/clangd/TUScheduler.cpp
@@ -45,6 +45,7 @@
 #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"
@@ -385,8 +386,9 @@ void ASTWorker::update(ParseInputs Inputs, WantDiagnostics WantDiags) {
     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);
diff --git a/clangd/TUScheduler.h b/clangd/TUScheduler.h
index cd0c4c78..abcf4aa3 100644
--- a/clangd/TUScheduler.h
+++ b/clangd/TUScheduler.h
@@ -12,6 +12,7 @@
 #include "ClangdUnit.h"
 #include "Function.h"
 #include "Threading.h"
+#include "index/CanonicalIncludes.h"
 #include "llvm/ADT/StringMap.h"
 #include <future>
 
@@ -91,7 +92,8 @@ public:
   /// 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.
diff --git a/clangd/index/FileIndex.cpp b/clangd/index/FileIndex.cpp
index 8ce42b51..bc102d29 100644
--- a/clangd/index/FileIndex.cpp
+++ b/clangd/index/FileIndex.cpp
@@ -10,6 +10,7 @@
 #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"
@@ -28,14 +29,11 @@ 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;
 
@@ -47,7 +45,7 @@ indexSymbols(ASTContext &AST, std::shared_ptr<Preprocessor> PP,
   if (IsIndexMainAST) {
     // We only collect refs when indexing main AST.
     CollectorOpts.RefFilter = RefKind::All;
-  }else {
+  } else {
     IndexOpts.IndexMacrosInPreprocessor = true;
     CollectorOpts.CollectMacro = true;
   }
@@ -72,16 +70,16 @@ indexSymbols(ASTContext &AST, std::shared_ptr<Preprocessor> PP,
 
 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;
 }
@@ -195,8 +193,9 @@ FileIndex::FileIndex(bool UseDex)
       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>());
diff --git a/clangd/index/FileIndex.h b/clangd/index/FileIndex.h
index 74563276..f1e37c33 100644
--- a/clangd/index/FileIndex.h
+++ b/clangd/index/FileIndex.h
@@ -19,6 +19,7 @@
 #include "Index.h"
 #include "MemIndex.h"
 #include "Merge.h"
+#include "index/CanonicalIncludes.h"
 #include "clang/Lex/Preprocessor.h"
 #include <memory>
 
@@ -84,7 +85,8 @@ public:
   /// 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`.
@@ -124,8 +126,8 @@ 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
diff --git a/unittests/clangd/CodeCompleteTests.cpp b/unittests/clangd/CodeCompleteTests.cpp
index ba539fc2..3a0082a2 100644
--- a/unittests/clangd/CodeCompleteTests.cpp
+++ b/unittests/clangd/CodeCompleteTests.cpp
@@ -645,6 +645,36 @@ TEST(CompletionTest, CompletionInPreamble) {
   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;
diff --git a/unittests/clangd/FileIndexTests.cpp b/unittests/clangd/FileIndexTests.cpp
index 898dd64d..0b037a50 100644
--- a/unittests/clangd/FileIndexTests.cpp
+++ b/unittests/clangd/FileIndexTests.cpp
@@ -12,6 +12,7 @@
 #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"
@@ -147,8 +148,8 @@ void update(FileIndex &M, llvm::StringRef Basename, llvm::StringRef Code) {
   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) {
@@ -199,13 +200,16 @@ TEST(FileIndexTest, ClassMembers) {
                                    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) {
@@ -270,10 +274,11 @@ TEST(FileIndexTest, RebuildWithPreamble) {
   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);
 
@@ -358,7 +363,8 @@ TEST(FileIndexTest, ReferencesInMainFileWithPreamble) {
       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,
diff --git a/unittests/clangd/TestTU.cpp b/unittests/clangd/TestTU.cpp
index 8c61bb80..c2f0d897 100644
--- a/unittests/clangd/TestTU.cpp
+++ b/unittests/clangd/TestTU.cpp
@@ -59,13 +59,15 @@ ParsedAST TestTU::build() const {
 
 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);
 }