[clangd] Build dynamic index and use it for code completion.

Reviewers: sammccall

Reviewed By: sammccall

Subscribers: klimek, ilya-biryukov, cfe-commits

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

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

10 files changed, 99 insertions, 15 deletions

diff --git a/clangd/ClangdLSPServer.cpp b/clangd/ClangdLSPServer.cpp
index fe330bf5..e680b121 100644
--- a/clangd/ClangdLSPServer.cpp
+++ b/clangd/ClangdLSPServer.cpp
@@ -282,10 +282,11 @@ ClangdLSPServer::ClangdLSPServer(JSONOutput &Out, unsigned AsyncThreadsCount,
                                  bool StorePreamblesInMemory,
                                  const clangd::CodeCompleteOptions &CCOpts,
                                  llvm::Optional<StringRef> ResourceDir,
-                                 llvm::Optional<Path> CompileCommandsDir)
+                                 llvm::Optional<Path> CompileCommandsDir,
+                                 bool BuildDynamicSymbolIndex)
     : Out(Out), CDB(std::move(CompileCommandsDir)), CCOpts(CCOpts),
       Server(CDB, /*DiagConsumer=*/*this, FSProvider, AsyncThreadsCount,
-             StorePreamblesInMemory, ResourceDir) {}
+             StorePreamblesInMemory, BuildDynamicSymbolIndex, ResourceDir) {}
 
 bool ClangdLSPServer::run(std::istream &In) {
   assert(!IsDone && "Run was called before");
diff --git a/clangd/ClangdLSPServer.h b/clangd/ClangdLSPServer.h
index 8441fee6..ed055251 100644
--- a/clangd/ClangdLSPServer.h
+++ b/clangd/ClangdLSPServer.h
@@ -34,7 +34,8 @@ public:
                   bool StorePreamblesInMemory,
                   const clangd::CodeCompleteOptions &CCOpts,
                   llvm::Optional<StringRef> ResourceDir,
-                  llvm::Optional<Path> CompileCommandsDir);
+                  llvm::Optional<Path> CompileCommandsDir,
+                  bool BuildDynamicSymbolIndex);
 
   /// Run LSP server loop, receiving input for it from \p In. \p In must be
   /// opened in binary mode. Output will be written using Out variable passed to
diff --git a/clangd/ClangdServer.cpp b/clangd/ClangdServer.cpp
index 129fc074..e7c9c7a5 100644
--- a/clangd/ClangdServer.cpp
+++ b/clangd/ClangdServer.cpp
@@ -134,8 +134,19 @@ ClangdServer::ClangdServer(GlobalCompilationDatabase &CDB,
                            FileSystemProvider &FSProvider,
                            unsigned AsyncThreadsCount,
                            bool StorePreamblesInMemory,
+                           bool BuildDynamicSymbolIndex,
                            llvm::Optional<StringRef> ResourceDir)
     : CDB(CDB), DiagConsumer(DiagConsumer), FSProvider(FSProvider),
+      FileIdx(BuildDynamicSymbolIndex ? new FileIndex() : nullptr),
+      // Pass a callback into `Units` to extract symbols from a newly parsed
+      // file and rebuild the file index synchronously each time an AST is
+      // parsed.
+      // FIXME(ioeric): this can be slow and we may be able to index on less
+      // critical paths.
+      Units(FileIdx
+                ? [this](const Context &Ctx, PathRef Path,
+                         ParsedAST *AST) { FileIdx->update(Ctx, Path, AST); }
+                : ASTParsedCallback()),
       ResourceDir(ResourceDir ? ResourceDir->str() : getStandardResourceDir()),
       PCHs(std::make_shared<PCHContainerOperations>()),
       StorePreamblesInMemory(StorePreamblesInMemory),
@@ -238,6 +249,8 @@ void ClangdServer::codeComplete(
       Resources->getPossiblyStalePreamble();
   // Copy completion options for passing them to async task handler.
   auto CodeCompleteOpts = Opts;
+  if (FileIdx)
+    CodeCompleteOpts.Index = FileIdx.get();
   // A task that will be run asynchronously.
   auto Task =
       // 'mutable' to reassign Preamble variable.
diff --git a/clangd/ClangdServer.h b/clangd/ClangdServer.h
index f9c8f92c..746e6d03 100644
--- a/clangd/ClangdServer.h
+++ b/clangd/ClangdServer.h
@@ -17,6 +17,7 @@
 #include "Function.h"
 #include "GlobalCompilationDatabase.h"
 #include "Protocol.h"
+#include "index/FileIndex.h"
 #include "clang/Tooling/CompilationDatabase.h"
 #include "clang/Tooling/Core/Replacement.h"
 #include "llvm/ADT/IntrusiveRefCntPtr.h"
@@ -195,10 +196,15 @@ public:
   ///
   /// \p StorePreamblesInMemory defines whether the Preambles generated by
   /// clangd are stored in-memory or on disk.
+  ///
+  /// If \p BuildDynamicSymbolIndex is true, ClangdServer builds a dynamic
+  /// in-memory index for symbols in all opened files and uses the index to
+  /// augment code completion results.
   ClangdServer(GlobalCompilationDatabase &CDB,
                DiagnosticsConsumer &DiagConsumer,
                FileSystemProvider &FSProvider, unsigned AsyncThreadsCount,
                bool StorePreamblesInMemory,
+               bool BuildDynamicSymbolIndex = false,
                llvm::Optional<StringRef> ResourceDir = llvm::None);
 
   /// Set the root path of the workspace.
@@ -330,6 +336,8 @@ private:
   DiagnosticsConsumer &DiagConsumer;
   FileSystemProvider &FSProvider;
   DraftStore DraftMgr;
+  /// If set, this manages index for symbols in opened files.
+  std::unique_ptr<FileIndex> FileIdx;
   CppFileCollection Units;
   std::string ResourceDir;
   // If set, this represents the workspace path.
diff --git a/clangd/ClangdUnit.cpp b/clangd/ClangdUnit.cpp
index f735e09f..b6bf3f75 100644
--- a/clangd/ClangdUnit.cpp
+++ b/clangd/ClangdUnit.cpp
@@ -357,17 +357,21 @@ ParsedASTWrapper::ParsedASTWrapper(llvm::Optional<ParsedAST> AST)
 std::shared_ptr<CppFile>
 CppFile::Create(PathRef FileName, tooling::CompileCommand Command,
                 bool StorePreamblesInMemory,
-                std::shared_ptr<PCHContainerOperations> PCHs) {
-  return std::shared_ptr<CppFile>(new CppFile(
-      FileName, std::move(Command), StorePreamblesInMemory, std::move(PCHs)));
+                std::shared_ptr<PCHContainerOperations> PCHs,
+                ASTParsedCallback ASTCallback) {
+  return std::shared_ptr<CppFile>(
+      new CppFile(FileName, std::move(Command), StorePreamblesInMemory,
+                  std::move(PCHs), std::move(ASTCallback)));
 }
 
 CppFile::CppFile(PathRef FileName, tooling::CompileCommand Command,
                  bool StorePreamblesInMemory,
-                 std::shared_ptr<PCHContainerOperations> PCHs)
+                 std::shared_ptr<PCHContainerOperations> PCHs,
+                 ASTParsedCallback ASTCallback)
     : FileName(FileName), Command(std::move(Command)),
       StorePreamblesInMemory(StorePreamblesInMemory), RebuildCounter(0),
-      RebuildInProgress(false), PCHs(std::move(PCHs)) {
+      RebuildInProgress(false), PCHs(std::move(PCHs)),
+      ASTCallback(std::move(ASTCallback)) {
   // FIXME(ibiryukov): we should pass a proper Context here.
   log(Context::empty(), "Opened file " + FileName + " with command [" +
                             this->Command.Directory + "] " +
@@ -570,6 +574,8 @@ CppFile::deferRebuild(StringRef NewContents,
     if (NewAST) {
       Diagnostics.insert(Diagnostics.end(), NewAST->getDiagnostics().begin(),
                          NewAST->getDiagnostics().end());
+      if (That->ASTCallback)
+        That->ASTCallback(Ctx, That->FileName, NewAST.getPointer());
     } else {
       // Don't report even Preamble diagnostics if we coulnd't build AST.
       Diagnostics.clear();
diff --git a/clangd/ClangdUnit.h b/clangd/ClangdUnit.h
index 47902c54..9ec115d1 100644
--- a/clangd/ClangdUnit.h
+++ b/clangd/ClangdUnit.h
@@ -136,6 +136,9 @@ private:
   mutable llvm::Optional<ParsedAST> AST;
 };
 
+using ASTParsedCallback =
+    std::function<void(const Context &Ctx, PathRef Path, ParsedAST *)>;
+
 /// Manages resources, required by clangd. Allows to rebuild file with new
 /// contents, and provides AST and Preamble for it.
 class CppFile : public std::enable_shared_from_this<CppFile> {
@@ -145,12 +148,14 @@ public:
   static std::shared_ptr<CppFile>
   Create(PathRef FileName, tooling::CompileCommand Command,
          bool StorePreamblesInMemory,
-         std::shared_ptr<PCHContainerOperations> PCHs);
+         std::shared_ptr<PCHContainerOperations> PCHs,
+         ASTParsedCallback ASTCallback);
 
 private:
   CppFile(PathRef FileName, tooling::CompileCommand Command,
           bool StorePreamblesInMemory,
-          std::shared_ptr<PCHContainerOperations> PCHs);
+          std::shared_ptr<PCHContainerOperations> PCHs,
+          ASTParsedCallback ASTCallback);
 
 public:
   CppFile(CppFile const &) = delete;
@@ -252,6 +257,9 @@ private:
   std::shared_ptr<const PreambleData> LatestAvailablePreamble;
   /// Utility class, required by clang.
   std::shared_ptr<PCHContainerOperations> PCHs;
+  /// This is called after the file is parsed. This can be nullptr if there is
+  /// no callback.
+  ASTParsedCallback ASTCallback;
 };
 
 /// Get the beginning SourceLocation at a specified \p Pos.
diff --git a/clangd/ClangdUnitStore.cpp b/clangd/ClangdUnitStore.cpp
index c92ae59a..402fa4ba 100644
--- a/clangd/ClangdUnitStore.cpp
+++ b/clangd/ClangdUnitStore.cpp
@@ -41,13 +41,14 @@ CppFileCollection::recreateFileIfCompileCommandChanged(
     It = OpenedFiles
              .try_emplace(File, CppFile::Create(File, std::move(NewCommand),
                                                 StorePreamblesInMemory,
-                                                std::move(PCHs)))
+                                                std::move(PCHs), ASTCallback))
              .first;
   } else if (!compileCommandsAreEqual(It->second->getCompileCommand(),
                                       NewCommand)) {
     Result.RemovedFile = std::move(It->second);
-    It->second = CppFile::Create(File, std::move(NewCommand),
-                                 StorePreamblesInMemory, std::move(PCHs));
+    It->second =
+        CppFile::Create(File, std::move(NewCommand), StorePreamblesInMemory,
+                        std::move(PCHs), ASTCallback);
   }
   Result.FileInCollection = It->second;
   return Result;
diff --git a/clangd/ClangdUnitStore.h b/clangd/ClangdUnitStore.h
index ab31cb94..4b363847 100644
--- a/clangd/ClangdUnitStore.h
+++ b/clangd/ClangdUnitStore.h
@@ -25,6 +25,11 @@ class Logger;
 /// Thread-safe mapping from FileNames to CppFile.
 class CppFileCollection {
 public:
+  /// \p ASTCallback is called when a file is parsed synchronously. This should
+  /// not be expensive since it blocks diagnostics.
+  explicit CppFileCollection(ASTParsedCallback ASTCallback)
+      : ASTCallback(std::move(ASTCallback)) {}
+
   std::shared_ptr<CppFile>
   getOrCreateFile(PathRef File, PathRef ResourceDir,
                   GlobalCompilationDatabase &CDB, bool StorePreamblesInMemory,
@@ -38,7 +43,7 @@ public:
       It = OpenedFiles
                .try_emplace(File, CppFile::Create(File, std::move(Command),
                                                   StorePreamblesInMemory,
-                                                  std::move(PCHs)))
+                                                  std::move(PCHs), ASTCallback))
                .first;
     }
     return It->second;
@@ -85,6 +90,7 @@ private:
 
   std::mutex Mutex;
   llvm::StringMap<std::shared_ptr<CppFile>> OpenedFiles;
+  ASTParsedCallback ASTCallback;
 };
 } // namespace clangd
 } // namespace clang
diff --git a/clangd/tool/ClangdMain.cpp b/clangd/tool/ClangdMain.cpp
index 4c91fe86..c63a01a3 100644
--- a/clangd/tool/ClangdMain.cpp
+++ b/clangd/tool/ClangdMain.cpp
@@ -90,6 +90,13 @@ static llvm::cl::opt<Path> TraceFile(
         "Trace internal events and timestamps in chrome://tracing JSON format"),
     llvm::cl::init(""), llvm::cl::Hidden);
 
+static llvm::cl::opt<bool> EnableIndexBasedCompletion(
+    "enable-index-based-completion",
+    llvm::cl::desc(
+        "Enable index-based global code completion (experimental). Clangd will "
+        "use index built from symbols in opened files"),
+    llvm::cl::init(false), llvm::cl::Hidden);
+
 int main(int argc, char *argv[]) {
   llvm::cl::ParseCommandLineOptions(argc, argv, "clangd");
 
@@ -180,7 +187,8 @@ int main(int argc, char *argv[]) {
   CCOpts.IncludeIneligibleResults = IncludeIneligibleResults;
   // Initialize and run ClangdLSPServer.
   ClangdLSPServer LSPServer(Out, WorkerThreadsCount, StorePreamblesInMemory,
-                            CCOpts, ResourceDirRef, CompileCommandsDirPath);
+                            CCOpts, ResourceDirRef, CompileCommandsDirPath,
+                            EnableIndexBasedCompletion);
   constexpr int NoShutdownRequestErrorCode = 1;
   llvm::set_thread_name("clangd.main");
   return LSPServer.run(std::cin) ? 0 : NoShutdownRequestErrorCode;
diff --git a/unittests/clangd/CodeCompleteTests.cpp b/unittests/clangd/CodeCompleteTests.cpp
index bbe0c8d1..312725fc 100644
--- a/unittests/clangd/CodeCompleteTests.cpp
+++ b/unittests/clangd/CodeCompleteTests.cpp
@@ -558,6 +558,38 @@ TEST(CompletionTest, FullyQualifiedScope) {
   EXPECT_THAT(Results.items, Has("XYZ", CompletionItemKind::Class));
 }
 
+TEST(CompletionTest, ASTIndexMultiFile) {
+  MockFSProvider FS;
+  MockCompilationDatabase CDB;
+  IgnoreDiagnostics DiagConsumer;
+  ClangdServer Server(CDB, DiagConsumer, FS, getDefaultAsyncThreadsCount(),
+                      /*StorePreamblesInMemory=*/true,
+                      /*BuildDynamicSymbolIndex=*/true);
+
+  Server
+      .addDocument(Context::empty(), getVirtualTestFilePath("foo.cpp"), R"cpp(
+      namespace ns { class XYZ {}; void foo() {} }
+  )cpp")
+      .wait();
+
+  auto File = getVirtualTestFilePath("bar.cpp");
+  auto Test = parseTextMarker(R"cpp(
+      namespace ns { class XXX {}; void fooooo() {} }
+      void f() { ns::^ }
+  )cpp");
+  Server.addDocument(Context::empty(), File, Test.Text).wait();
+
+  auto Results = Server.codeComplete(Context::empty(), File, Test.MarkerPos, {})
+                     .get()
+                     .second.Value;
+  // "XYZ" and "foo" are not included in the file being completed but are still
+  // visible through the index.
+  EXPECT_THAT(Results.items, Has("XYZ", CompletionItemKind::Class));
+  EXPECT_THAT(Results.items, Has("foo", CompletionItemKind::Function));
+  EXPECT_THAT(Results.items, Has("XXX", CompletionItemKind::Class));
+  EXPECT_THAT(Results.items, Has("fooooo", CompletionItemKind::Function));
+}
+
 } // namespace
 } // namespace clangd
 } // namespace clang