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

[ioeric]

• Merge with updated D41281
• Fix broken merge

[ioeric]

• Merge with D41289.
• Merge with origin/master

[sammccall]

OK, this is pretty clean now! :-)

clangd/ClangdServer.cpp

139

A comment explaining when this runs and what it does might be nice.

Also this doesn't seem like an ideal long-term solution: rebuilding an index can be slow (less the symbol extraction, and more the rebuilding of index data structures) and we may be able to index on less critical paths.
Probably also worth a comment.

clangd/ClangdUnit.cpp

617

CppFile doesn't need to pass the path, do you want [FileName, ASTCallback](const Context &C, ParsedAST *AST) { ASTCallback(C, FileName, AST); }

clangd/ClangdUnitStore.h

28

synchronously... maybe mention this will block diagnostics and doing expensive things would be bad

clangd/tool/ClangdMain.cpp

98

llvm:🆑:Hidden, if it's experimental?

unittests/clangd/CodeCompleteTests.cpp

521

for this test, I don't see a clear sign that the results come from the index.
Is there a way we know this that you can point out?

For the second test we can tell because there's no #include, but it could use a comment

544

This repeated setup is a bit ugly but I'm planning to pull out a helper for this anyway.

[ioeric]

• Address review comments.

[ioeric]

Thanks for the quick review!

clangd/ClangdUnit.cpp

617

The downside is 1) we need to worry about the life time of FileName 2) we need another type for the new callback, so I am inclined to simply forward the callback.

clangd/ClangdLSPServer.h

 public:
   /// If \p CompileCommandsDir has a value, compile_commands.json will be
   /// loaded only from \p CompileCommandsDir. Otherwise, clangd will look
   /// for compile_commands.json in all parent directories of each file.
   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);
 
   /// 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
   /// class constructor. This method must not be executed more than once for
   /// each instance of ClangdLSPServer.
   ///
   /// \return Wether we received a 'shutdown' request before an 'exit' request
   bool run(std::istream &In);

clangd/ClangdLSPServer.cpp

void ClangdLSPServer::onDocumentHighlight(Ctx C,
 
   reply(C, json::ary(Highlights->Value));
 }
 
 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");
 
   // Set up JSONRPCDispatcher.
   JSONRPCDispatcher Dispatcher([](Context Ctx, const json::Expr &Params) {
     replyError(Ctx, ErrorCode::MethodNotFound, "method not found");
   });

clangd/ClangdServer.h

 
 #include "ClangdUnit.h"
 #include "ClangdUnitStore.h"
 #include "CodeComplete.h"
 #include "DraftStore.h"
 #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"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/StringRef.h"
 #include <condition_variable>
 #include <functional>
 #include <mutex>
public:
   ///
   /// After each parsing request finishes, ClangdServer reports diagnostics to
   /// \p DiagConsumer. Note that a callback to \p DiagConsumer happens on a
   /// worker thread. Therefore, instances of \p DiagConsumer must properly
   /// synchronize access to shared state.
   ///
   /// \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.
   void setRootPath(PathRef RootPath);
 
   /// Add a \p File to the list of tracked C++ files or update the contents if
   /// \p File is already tracked. Also schedules parsing of the AST for it on a
   /// separate thread. When the parsing is complete, DiagConsumer passed in
private:
 
   std::future<Context>
   scheduleCancelRebuild(Context Ctx, std::shared_ptr<CppFile> Resources);
 
   GlobalCompilationDatabase &CDB;
   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.
   llvm::Optional<std::string> RootPath;
   std::shared_ptr<PCHContainerOperations> PCHs;
   bool StorePreamblesInMemory;
   /// Used to serialize diagnostic callbacks.
   /// FIXME(ibiryukov): get rid of an extra map and put all version counters

clangd/ClangdServer.cpp

for (auto &Worker : Workers)
     Worker.join();
 }
 
 ClangdServer::ClangdServer(GlobalCompilationDatabase &CDB,
                            DiagnosticsConsumer &DiagConsumer,
                            FileSystemProvider &FSProvider,
                            unsigned AsyncThreadsCount,
                            bool StorePreamblesInMemory,
+                           bool BuildDynamicSymbolIndex,
                            llvm::Optional<StringRef> ResourceDir)
     : CDB(CDB), DiagConsumer(DiagConsumer), FSProvider(FSProvider),
+      FileIdx(BuildDynamicSymbolIndex ? new FileIndex() : nullptr),
+      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),
       WorkScheduler(AsyncThreadsCount) {}
 
 void ClangdServer::setRootPath(PathRef RootPath) {
   std::string NewRootPath = llvm::sys::path::convert_to_slash(
       RootPath, llvm::sys::path::Style::posix);
void ClangdServer::codeComplete(
   // Remember the current Preamble and use it when async task starts executing.
   // At the point when async task starts executing, we may have a different
   // Preamble in Resources. However, we assume the Preamble that we obtain here
   // is reusable in completion more often.
   std::shared_ptr<const PreambleData> Preamble =
       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.
       [=](Context Ctx, CallbackType Callback) mutable {
         if (!Preamble) {
           // Maybe we built some preamble before processing this request.
           Preamble = Resources->getPossiblyStalePreamble();
         }

clangd/ClangdUnit.h

 private:
   // This wrapper is used as an argument to std::shared_future (and it returns a
   // const ref in get()), but we need to have non-const ref in order to
   // implement some features.
   mutable std::mutex Mutex;
   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> {
 public:
   // We only allow to create CppFile as shared_ptr, because a future returned by
   // deferRebuild will hold references to it.
   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;
   CppFile(CppFile &&) = delete;
 
   /// Cancels a scheduled rebuild, if any, and sets AST and Preamble to nulls.
   /// If a rebuild is in progress, will wait for it to finish.
   void cancelRebuild();
private:
   /// Promise and future for the latests Preamble. Fulfilled during rebuild.
   std::promise<std::shared_ptr<const PreambleData>> PreamblePromise;
   std::shared_future<std::shared_ptr<const PreambleData>> PreambleFuture;
   /// Latest preamble that was built. May be stale, but always available without
   /// waiting for rebuild to finish.
   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.
 SourceLocation getBeginningOfIdentifier(ParsedAST &Unit, const Position &Pos,
                                         const FileEntry *FE);
 
 /// Get definition of symbol at a specified \p Pos.
 std::vector<Location> findDefinitions(const Context &Ctx, ParsedAST &AST,

clangd/ClangdUnit.cpp

ParsedASTWrapper::ParsedASTWrapper(ParsedASTWrapper &&Wrapper)
     : AST(std::move(Wrapper.AST)) {}
 
 ParsedASTWrapper::ParsedASTWrapper(llvm::Optional<ParsedAST> AST)
     : AST(std::move(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 + "] " +
                             llvm::join(this->Command.CommandLine, " "));
 
   std::lock_guard<std::mutex> Lock(Mutex);
   LatestAvailablePreamble = nullptr;
   PreamblePromise.set_value(nullptr);
llvm::Optional<ParsedAST> NewAST;
       SPAN_ATTACH(Tracer, "File", That->FileName);
       NewAST = ParsedAST::Build(Ctx, std::move(CI), std::move(NewPreamble),
                                 std::move(ContentsBuffer), PCHs, VFS);
     }
 
     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();
     }
 
     // Publish the new AST.
     {
       std::lock_guard<std::mutex> Lock(That->Mutex);

clangd/ClangdUnitStore.h

 namespace clang {
 namespace clangd {
 
 class Logger;
 
 /// Thread-safe mapping from FileNames to CppFile.
 class CppFileCollection {
 public:
+  /// \p ASTCallback is called when a file is parsed.
+  explicit CppFileCollection(ASTParsedCallback ASTCallback)
+      : ASTCallback(std::move(ASTCallback)) {}
+
   std::shared_ptr<CppFile>
   getOrCreateFile(PathRef File, PathRef ResourceDir,
                   GlobalCompilationDatabase &CDB, bool StorePreamblesInMemory,
                   std::shared_ptr<PCHContainerOperations> PCHs) {
     std::lock_guard<std::mutex> Lock(Mutex);
 
     auto It = OpenedFiles.find(File);
     if (It == OpenedFiles.end()) {
       auto Command = getCompileCommand(CDB, File, ResourceDir);
 
       It = OpenedFiles
                .try_emplace(File, CppFile::Create(File, std::move(Command),
                                                   StorePreamblesInMemory,
-                                                  std::move(PCHs)))
+                                                  std::move(PCHs), ASTCallback))
                .first;
     }
     return It->second;
   }
 
   struct RecreateResult {
     /// A CppFile, stored in this CppFileCollection for the corresponding
     /// filepath after calling recreateFileIfCompileCommandChanged.
private:
   tooling::CompileCommand getCompileCommand(GlobalCompilationDatabase &CDB,
                                             PathRef File, PathRef ResourceDir);
 
   bool compileCommandsAreEqual(tooling::CompileCommand const &LHS,
                                tooling::CompileCommand const &RHS);
 
   std::mutex Mutex;
   llvm::StringMap<std::shared_ptr<CppFile>> OpenedFiles;
+  ASTParsedCallback ASTCallback;
 };
 } // namespace clangd
 } // namespace clang
 
 #endif

clangd/ClangdUnitStore.cpp

CppFileCollection::recreateFileIfCompileCommandChanged(
 
   RecreateResult Result;
 
   auto It = OpenedFiles.find(File);
   if (It == OpenedFiles.end()) {
     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;
 }
 
 tooling::CompileCommand
 CppFileCollection::getCompileCommand(GlobalCompilationDatabase &CDB,
                                      PathRef File, PathRef ResourceDir) {

clangd/tool/ClangdMain.cpp

static llvm::cl::opt<Path> InputMirrorFile(
     llvm::cl::init(""), llvm::cl::Hidden);
 
 static llvm::cl::opt<Path> TraceFile(
     "trace",
     llvm::cl::desc(
         "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));
+
 int main(int argc, char *argv[]) {
   llvm::cl::ParseCommandLineOptions(argc, argv, "clangd");
 
   if (!RunSynchronously && WorkerThreadsCount == 0) {
     llvm::errs() << "A number of worker threads cannot be 0. Did you mean to "
                     "specify -run-synchronously?";
     return 1;
   }
int main(int argc, char *argv[]) {
   // Change stdin to binary to not lose \r\n on windows.
   llvm::sys::ChangeStdinToBinary();
 
   clangd::CodeCompleteOptions CCOpts;
   CCOpts.EnableSnippets = EnableSnippets;
   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;
 }

unittests/clangd/CodeCompleteTests.cpp

auto Results = completions(R"cpp(
       namespace ns {}
       void f() { ns::X^ }
   )cpp",
                              Opts);
   // Filter is not in the buffer.
   EXPECT_THAT(Results.items, HasEditWithRange("ns::XYZ", {{2, 17}, {2, 19}}));
 }
 
+TEST(CompletionTest, ASTIndexOneFile) {
+  MockFSProvider FS;
+  MockCompilationDatabase CDB;
+  IgnoreDiagnostics DiagConsumer;
+  ClangdServer Server(CDB, DiagConsumer, FS, getDefaultAsyncThreadsCount(),
+                      /*StorePreamblesInMemory=*/true,
+                      /*BuildDynamicSymbolIndex=*/true);
+
+  auto File = getVirtualTestFilePath("foo.cpp");
+  auto Test = parseTextMarker(R"cpp(
+      namespace ns { class XYZ {}; void foo() {} }
+      void f() { ns::^ }
+  )cpp");
+  // Wait to make sure AST is parsed.
+  Server.addDocument(Context::empty(), File, Test.Text).wait();
+  auto Results = Server.codeComplete(Context::empty(), File, Test.MarkerPos, {})
+                     .get()
+                     .second.Value;
+  EXPECT_THAT(Results.items, HasEdit("ns::XYZ", CompletionItemKind::Class));
+  EXPECT_THAT(Results.items, HasEdit("ns::foo", CompletionItemKind::Function));
+}
+
+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;
+  EXPECT_THAT(Results.items, HasEdit("ns::XYZ", CompletionItemKind::Class));
+  EXPECT_THAT(Results.items, HasEdit("ns::foo", CompletionItemKind::Function));
+  EXPECT_THAT(Results.items, HasEdit("ns::XXX", CompletionItemKind::Class));
+  EXPECT_THAT(Results.items,
+              HasEdit("ns::fooooo", CompletionItemKind::Function));
+}
+
 } // namespace
 } // namespace clangd
 } // namespace clang