[clangd] Optionally use dex for the preamble parts of the dynamic index.

[sammccall]

Summary

Reuse the old -use-dex-index experiment flag for this.

To avoid breaking the tests, make Dex deduplicate symbols, addressing an old FIXME.

[hokein]

LGTM.

clangd/ClangdServer.h

79

"too slow" seems confusing, dex is faster, I think here it means too slow to build?

clangd/ClangdServer.h

struct Options {
     ASTRetentionPolicy RetentionPolicy;
 
     /// Cached preambles are potentially large. If false, store them on disk.
     bool StorePreamblesInMemory = true;
 
     /// If true, ClangdServer builds a dynamic in-memory index for symbols in
     /// opened files and uses the index to augment code completion results.
     bool BuildDynamicSymbolIndex = false;
+    /// Use a heavier and faster in-memory index implementation.
+    /// FIXME: we should make this true if it isn't too slow!.
+    bool HeavyweightDynamicSymbolIndex = false;
 
     /// URI schemes to use when building the dynamic index.
     /// If empty, the default schemes in SymbolCollector will be used.
     std::vector<std::string> URISchemes;
 
     /// If set, use this index to augment code completion results.
     SymbolIndex *StaticIndex = nullptr;
 

clangd/ClangdServer.cpp

 
 ClangdServer::ClangdServer(const GlobalCompilationDatabase &CDB,
                            const FileSystemProvider &FSProvider,
                            DiagnosticsConsumer &DiagConsumer,
                            const Options &Opts)
     : CDB(CDB), DiagConsumer(DiagConsumer), FSProvider(FSProvider),
       ResourceDir(Opts.ResourceDir ? Opts.ResourceDir->str()
                                    : getStandardResourceDir()),
-      DynamicIdx(Opts.BuildDynamicSymbolIndex ? new FileIndex(Opts.URISchemes)
-                                              : nullptr),
+      DynamicIdx(Opts.BuildDynamicSymbolIndex
+                     ? new FileIndex(Opts.URISchemes,
+                                     Opts.HeavyweightDynamicSymbolIndex)
       PCHs(std::make_shared<PCHContainerOperations>()),
       // Pass a callback into `WorkScheduler` 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.
       WorkScheduler(Opts.AsyncThreadsCount, Opts.StorePreamblesInMemory,
                     DynamicIdx ? makeUpdateCallbacks(DynamicIdx.get())

clangd/index/Background.cpp

llvm::Error BackgroundIndex::index(tooling::CompileCommand Cmd) {
   IndexedSymbols.update(AbsolutePath,
                         llvm::make_unique<SymbolSlab>(std::move(Symbols)),
                         llvm::make_unique<RefSlab>(std::move(Refs)));
   FileHash[AbsolutePath] = Hash;
 
   // FIXME: this should rebuild once-in-a-while, not after every file.
   //       At that point we should use Dex, too.
   vlog("Rebuilding automatic index");
-  reset(IndexedSymbols.buildMemIndex());
+  reset(IndexedSymbols.buildIndex(IndexType::Light));
   return Error::success();
 }
 
 } // namespace clangd
 } // namespace clang

clangd/index/FileIndex.h

 #include "MemIndex.h"
 #include "Merge.h"
 #include "clang/Lex/Preprocessor.h"
 #include <memory>
 
 namespace clang {
 namespace clangd {
 
+/// Select between in-memory index implementations, which have tradeoffs.
+enum class IndexType {
+  // MemIndex is trivially cheap to build, but has poor query performance.
+  Light,
+  // Dex is relatively expensive to build and uses more memory, but is fast.
+  Heavy,
+};
+
 /// A container of Symbols from several source files. It can be updated
 /// at source-file granularity, replacing all symbols from one file with a new
 /// set.
 ///
 /// This implements a snapshot semantics for symbols in a file. Each update to a
 /// file will create a new snapshot for all symbols in the file. Snapshots are
 /// managed with shared pointers that are shared between this class and the
 /// users. For each file, this class only stores a pointer pointing to the
 /// newest snapshot, and an outdated snapshot is deleted by the last owner of
 /// the snapshot, either this class or the symbol index.
 ///
 /// The snapshot semantics keeps critical sections minimal since we only need
 /// locking when we swap or obtain references to snapshots.
 class FileSymbols {
 public:
   /// Updates all symbols and refs in a file.
   /// If either is nullptr, corresponding data for \p Path will be removed.
   void update(PathRef Path, std::unique_ptr<SymbolSlab> Slab,
               std::unique_ptr<RefSlab> Refs);
 
   // The index keeps the symbols alive.
-  std::unique_ptr<SymbolIndex> buildMemIndex();
+  std::unique_ptr<SymbolIndex>
+  buildIndex(IndexType, ArrayRef<std::string> URISchemes = {});
 
 private:
   mutable std::mutex Mutex;
 
   /// Stores the latest symbol snapshots for all active files.
   llvm::StringMap<std::shared_ptr<SymbolSlab>> FileToSymbols;
   /// Stores the latest ref snapshots for all active files.
   llvm::StringMap<std::shared_ptr<RefSlab>> FileToRefs;
 };
 
 /// This manages symbols from files and an in-memory index on all symbols.
 /// FIXME: Expose an interface to remove files that are closed.
 class FileIndex : public MergedIndex {
 public:
   /// If URISchemes is empty, the default schemes in SymbolCollector will be
   /// used.
-  FileIndex(std::vector<std::string> URISchemes = {});
+  FileIndex(std::vector<std::string> URISchemes = {}, 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);
 
   /// 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.
   std::vector<std::string> URISchemes;
 
   // Contains information from each file's preamble only.
   // These are large, but update fairly infrequently (preambles are stable).
   // Missing information:
   //  - symbol refs (these are always "from the main file")
   //  - definition locations in the main file
   //

clangd/index/FileIndex.cpp

 //===--- FileIndex.cpp - Indexes for files. ------------------------ C++-*-===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 
 #include "FileIndex.h"
 #include "ClangdUnit.h"
 #include "Logger.h"
 #include "SymbolCollector.h"
 #include "index/Index.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 <memory>
 
 namespace clang {
 namespace clangd {
 
void FileSymbols::update(PathRef Path, std::unique_ptr<SymbolSlab> Symbols,
   else
     FileToSymbols[Path] = std::move(Symbols);
   if (!Refs)
     FileToRefs.erase(Path);
   else
     FileToRefs[Path] = std::move(Refs);
 }
 
-std::unique_ptr<SymbolIndex> FileSymbols::buildMemIndex() {
+std::unique_ptr<SymbolIndex>
+FileSymbols::buildIndex(IndexType Type, ArrayRef<std::string> URISchemes) {
   std::vector<std::shared_ptr<SymbolSlab>> SymbolSlabs;
   std::vector<std::shared_ptr<RefSlab>> RefSlabs;
   {
     std::lock_guard<std::mutex> Lock(Mutex);
     for (const auto &FileAndSymbols : FileToSymbols)
       SymbolSlabs.push_back(FileAndSymbols.second);
     for (const auto &FileAndRefs : FileToRefs)
       RefSlabs.push_back(FileAndRefs.second);
FileSymbols::buildIndex(IndexType Type, ArrayRef<std::string> URISchemes) {
 
   size_t StorageSize = RefsStorage.size() * sizeof(Ref);
   for (const auto &Slab : SymbolSlabs)
     StorageSize += Slab->bytes();
   for (const auto &RefSlab : RefSlabs)
     StorageSize += RefSlab->bytes();
 
   // Index must keep the slabs and contiguous ranges alive.
+  switch (Type) {
+  case IndexType::Light:
     return llvm::make_unique<MemIndex>(
         llvm::make_pointee_range(AllSymbols), std::move(AllRefs),
         std::make_tuple(std::move(SymbolSlabs), std::move(RefSlabs),
                         std::move(RefsStorage)),
         StorageSize);
+  case IndexType::Heavy:
+    return llvm::make_unique<dex::Dex>(
+        llvm::make_pointee_range(AllSymbols), std::move(AllRefs),
+        std::make_tuple(std::move(SymbolSlabs), std::move(RefSlabs),
+                        std::move(RefsStorage)),
+        StorageSize, std::move(URISchemes));
+  }
 }
 
-FileIndex::FileIndex(std::vector<std::string> URISchemes)
-    : MergedIndex(&MainFileIndex, &PreambleIndex),
+FileIndex::FileIndex(std::vector<std::string> URISchemes, bool UseDex)
+    : MergedIndex(&MainFileIndex, &PreambleIndex), UseDex(UseDex),
       URISchemes(std::move(URISchemes)),
-      PreambleIndex(PreambleSymbols.buildMemIndex()),
-      MainFileIndex(MainFileSymbols.buildMemIndex()) {}
+      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), URISchemes);
   PreambleSymbols.update(Path,
                          llvm::make_unique<SymbolSlab>(std::move(Symbols)),
                          llvm::make_unique<RefSlab>());
-  PreambleIndex.reset(PreambleSymbols.buildMemIndex());
+  PreambleIndex.reset(PreambleSymbols.buildIndex(
+      UseDex ? IndexType::Heavy : IndexType::Light, URISchemes));
 }
 
 void FileIndex::updateMain(PathRef Path, ParsedAST &AST) {
   auto Contents = indexMainDecls(AST, URISchemes);
   MainFileSymbols.update(
       Path, llvm::make_unique<SymbolSlab>(std::move(Contents.first)),
       llvm::make_unique<RefSlab>(std::move(Contents.second)));
-  MainFileIndex.reset(MainFileSymbols.buildMemIndex());
+  MainFileIndex.reset(MainFileSymbols.buildIndex(IndexType::Light, URISchemes));
 }
 
 } // namespace clangd
 } // namespace clang

clangd/index/dex/Dex.h

Dex(SymbolRange &&Symbols, RefsRange &&Refs,
       llvm::ArrayRef<std::string> Schemes)
       : Corpus(0), URISchemes(Schemes) {
     // If Schemes don't contain any items, fall back to SymbolCollector's
     // default URI schemes.
     if (URISchemes.empty()) {
       SymbolCollector::Options Opts;
       URISchemes = Opts.URISchemes;
     }
+    llvm::DenseSet<SymbolID> SeenIDs;
     for (auto &&Sym : Symbols)
+      if (SeenIDs.insert(Sym.ID).second)
         this->Symbols.push_back(&Sym);
     for (auto &&Ref : Refs)
       this->Refs.try_emplace(Ref.first, Ref.second);
     buildIndex();
   }
   // Symbols and Refs are owned by BackingData, Index takes ownership.
   template <typename SymbolRange, typename RefsRange, typename Payload>
   Dex(SymbolRange &&Symbols, RefsRange &&Refs, Payload &&BackingData,
       size_t BackingDataSize, llvm::ArrayRef<std::string> URISchemes)

clangd/tool/ClangdMain.cpp

 #include <iostream>
 #include <memory>
 #include <string>
 #include <thread>
 
 using namespace clang;
 using namespace clang::clangd;
 
-// FIXME: remove this option when Dex is stable enough.
+// FIXME: remove this option when Dex is cheap enough.
 static llvm::cl::opt<bool>
     UseDex("use-dex-index",
-           llvm::cl::desc("Use experimental Dex static index."),
-           llvm::cl::init(true), llvm::cl::Hidden);
+           llvm::cl::desc("Use experimental Dex dynamic index."),
+           llvm::cl::init(false), llvm::cl::Hidden);
 
 static llvm::cl::opt<Path> CompileCommandsDir(
     "compile-commands-dir",
     llvm::cl::desc("Specify a path to look for compile_commands.json. If path "
                    "is invalid, clangd will look in the current directory and "
                    "parent paths of each source file."));
 
 static llvm::cl::opt<unsigned>
case PCHStorageFlag::Memory:
     break;
   case PCHStorageFlag::Disk:
     Opts.StorePreamblesInMemory = false;
     break;
   }
   if (!ResourceDir.empty())
     Opts.ResourceDir = ResourceDir;
   Opts.BuildDynamicSymbolIndex = EnableIndex;
+  Opts.HeavyweightDynamicSymbolIndex = UseDex;
   std::unique_ptr<SymbolIndex> StaticIdx;
   std::future<void> AsyncIndexLoad; // Block exit while loading the index.
   if (EnableIndex && !IndexFile.empty()) {
     // Load the index asynchronously. Meanwhile SwapIndex returns no results.
     SwapIndex *Placeholder;
     StaticIdx.reset(Placeholder = new SwapIndex(llvm::make_unique<MemIndex>()));
     AsyncIndexLoad = runAsync<void>([Placeholder, &Opts] {
-      if (auto Idx = loadIndex(IndexFile, Opts.URISchemes, UseDex))
+      if (auto Idx = loadIndex(IndexFile, Opts.URISchemes, /*UseDex=*/true))
         Placeholder->reset(std::move(Idx));
     });
     if (RunSynchronously)
       AsyncIndexLoad.wait();
   }
   Opts.StaticIndex = StaticIdx.get();
   Opts.AsyncThreadsCount = WorkerThreadsCount;
 

unittests/clangd/DexTests.cpp

TEST(Dex, FuzzyFind) {
   Req.Query = "";
   Req.Scopes = {};
   EXPECT_THAT(match(*Index, Req),
               UnorderedElementsAre("ns::ABC", "ns::BCD", "::ABC",
                                    "ns::nested::ABC", "other::ABC",
                                    "other::A"));
 }
 
-// FIXME(kbobyrev): This test is different for Dex and MemIndex: while
-// MemIndex manages response deduplication, Dex simply returns all matched
-// symbols which means there might be equivalent symbols in the response.
-// Before drop-in replacement of MemIndex with Dex happens, FileIndex
-// should handle deduplication instead.
 TEST(DexTest, DexDeduplicate) {
   std::vector<Symbol> Symbols = {symbol("1"), symbol("2"), symbol("3"),
                                  symbol("2") /* duplicate */};
   FuzzyFindRequest Req;
   Req.Query = "2";
   Dex I(Symbols, RefSlab(), URISchemes);
-  EXPECT_FALSE(Req.Limit);
-  EXPECT_THAT(match(I, Req), ElementsAre("2", "2"));
+  EXPECT_THAT(match(I, Req), ElementsAre("2"));
 }
 
 TEST(DexTest, DexLimitedNumMatches) {
   auto I = Dex::build(generateNumSymbols(0, 100), RefSlab(), URISchemes);
   FuzzyFindRequest Req;
   Req.Query = "5";
   Req.Limit = 3;
   bool Incomplete;

unittests/clangd/FileIndexTests.cpp

RefSlab getRefs(const SymbolIndex &I, SymbolID ID) {
   Req.IDs = {ID};
   RefSlab::Builder Slab;
   I.refs(Req, [&](const Ref &S) { Slab.insert(ID, S); });
   return std::move(Slab).build();
 }
 
 TEST(FileSymbolsTest, UpdateAndGet) {
   FileSymbols FS;
-  EXPECT_THAT(getSymbolNames(*FS.buildMemIndex()), UnorderedElementsAre());
+  EXPECT_THAT(getSymbolNames(*FS.buildIndex(IndexType::Light)),
+              UnorderedElementsAre());
 
   FS.update("f1", numSlab(1, 3), refSlab(SymbolID("1"), "f1.cc"));
-  EXPECT_THAT(getSymbolNames(*FS.buildMemIndex()),
+  EXPECT_THAT(getSymbolNames(*FS.buildIndex(IndexType::Light)),
               UnorderedElementsAre("1", "2", "3"));
-  EXPECT_THAT(getRefs(*FS.buildMemIndex(), SymbolID("1")),
+  EXPECT_THAT(getRefs(*FS.buildIndex(IndexType::Light), SymbolID("1")),
               RefsAre({FileURI("f1.cc")}));
 }
 
 TEST(FileSymbolsTest, Overlap) {
   FileSymbols FS;
   FS.update("f1", numSlab(1, 3), nullptr);
   FS.update("f2", numSlab(3, 5), nullptr);
-  EXPECT_THAT(getSymbolNames(*FS.buildMemIndex()),
+  for (auto Type : {IndexType::Light, IndexType::Heavy})
+    EXPECT_THAT(getSymbolNames(*FS.buildIndex(Type)),
                 UnorderedElementsAre("1", "2", "3", "4", "5"));
 }
 
 TEST(FileSymbolsTest, SnapshotAliveAfterRemove) {
   FileSymbols FS;
 
   SymbolID ID("1");
   FS.update("f1", numSlab(1, 3), refSlab(ID, "f1.cc"));
 
-  auto Symbols = FS.buildMemIndex();
+  auto Symbols = FS.buildIndex(IndexType::Light);
   EXPECT_THAT(getSymbolNames(*Symbols), UnorderedElementsAre("1", "2", "3"));
   EXPECT_THAT(getRefs(*Symbols, ID), RefsAre({FileURI("f1.cc")}));
 
   FS.update("f1", nullptr, nullptr);
-  auto Empty = FS.buildMemIndex();
+  auto Empty = FS.buildIndex(IndexType::Light);
   EXPECT_THAT(getSymbolNames(*Empty), UnorderedElementsAre());
   EXPECT_THAT(getRefs(*Empty, ID), ElementsAre());
 
   EXPECT_THAT(getSymbolNames(*Symbols), UnorderedElementsAre("1", "2", "3"));
   EXPECT_THAT(getRefs(*Symbols, ID), RefsAre({FileURI("f1.cc")}));
 }
 
 std::vector<std::string> match(const SymbolIndex &I,

unittests/clangd/TestTU.cpp

 
 SymbolSlab TestTU::headerSymbols() const {
   auto AST = build();
   return indexHeaderSymbols(AST.getASTContext(), AST.getPreprocessorPtr());
 }
 
 std::unique_ptr<SymbolIndex> TestTU::index() const {
   auto AST = build();
-  auto Idx = llvm::make_unique<FileIndex>();
+  auto Idx = llvm::make_unique<FileIndex>(
+      /*URISchemes=*/std::vector<std::string>{}, /*UseDex=*/true);
   Idx->updatePreamble(Filename, AST.getASTContext(), AST.getPreprocessorPtr());
   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) {