[clangd] Introduce ProjectAwareIndex

[kadircet]

Summary

An index implementation that can dispatch to a variety of indexes
depending on the file path. Enables clangd to work with multiple indexes in the
same instance, configured via config files.

Depends on D90749, D90746

[kadircet]

This is a WIP patch. Provided as a way to demonstrate how ExternalBlock will be used in an index implementation.

[kadircet]

The index has the following query semantics:

• Lookup: it only queries the associated index. This is to not regress

latency for operations like Hover and Go-To.

• FuzzyFind only queries the associated index when

RestrictForCodeCompletion is set, it queries all otherwise. Again this
is to prevent latency regression for code completion, but make sure we
provide complete results for search operations like WorkspaceSymbol.

• Refs, queries all. As incomplete results for such operations might

result in bad UX for opeartions like rename. Also having an incomplete
set of references might result in bad impressions.

• Relations, queries all. Again this is used by operations like type and

call hierarchy and having incomplete results are useless for such
operations.

[kadircet]

• Rebase

[sammccall]

In D90750#2378987, @kadircet wrote:

The index has the following query semantics:

Rehashing offline conversation:

• These heuristics mostly amount to guessing which high-level operation is being performed. It's a non-obvious place to put this logic, and it's hard to fix without fallout when the guess is wrong.
• These guesses will get stale. e.g. this queries all indexes for relations, but one index for lookup. We'd like to use both for Hover soon.

Our conclusion was we could query one index by default, and allow callers to override this using a context variable.
The arguments for a context variable vs attributes on request:

• it naturally applies to batches of requests
• it's a less invasive change
• it "hides" the option from other indexes

I'm not that sure about this implementation choice, but it also probably doesn't matter much.

[sammccall]

This looks really sensible to me! Nice you managed to reuse Memoize, there's very little ugly caching/locking left.

clang-tools-extra/clangd/index/ProjectAware.cpp

82

I think the cache key needs a bit more consideration.

It's unlikely that a filename will collide with a host/port, but they are currently sitting in the same namespace.

More worryingly, the mountpoint is used in the lambda body but not in the cache key - are we sure this is safe?

The clearest way to enforce this is to make index creation use the cache key only, and not other data from the config. Then you have to include what you're going to use.

82

as discussed offline, allowing the config -> cache key and cache key -> index operations to be injectable is the key to testing I think

101

from memoize:

Concurrent calls for the same key may run the
computation multiple times, but each call will return the same result.

so this side-effect may run multiple times too, resulting in multiple index copies getting leaked. This seems like it may be bad in the case of file indexes.

Why can't the memoize map own the indexes?
This would still result in loading the index multiple times and throwing things away, hmm. Not sure how best to resolve this.

120

consider moving this to unique_ptr<MergedIndex> MergedIndex::create(ArrayRef<unique_ptr<SymbolIndex>>) or so?

Not sure if this is something we'd want to use from ClangdMain...

129

with the helper described above, the Primary ordering could be maybe better expressed as std::stable_partition

clang-tools-extra/clangd/index/ProjectAware.h

31

Is there a reason this class needs to be public?
I think we can expose factories returning SymbolIndex only

[kadircet]

• Address comments
• Only query the associated index
• Use ExternalIndexSpec as the cache key

clang-tools-extra/clangd/index/ProjectAware.cpp

101

so this side-effect may run multiple times too, resulting in multiple index copies getting leaked. This seems like it may be bad in the case of file indexes.

The only way i can see this happening is by synchronizing calls to getIndex.

Why can't the memoize map own the indexes?

Main problem is, it returns a copy to the stored elements, rather than a reference to them. Which might be necessary in general to guarantee thread-safety. Since we can't copy unique_ptrs, I had to store them elsewhere and only keep pointers to it within the memoize map.

So both things considered, maybe we should just give up on memoize and have our own thread-safe cache here?

[kadircet]

ping :)

[sammccall]

clang-tools-extra/clangd/Config.h

88 ↗

(On Diff #305695)

Maybe we could use DenseMap instead? Defining an order on specs seems semantically iffy.

(Maybe one day we'll get to use a better hashtable and DenseMapInfo won't be so gross)

clang-tools-extra/clangd/index/ProjectAware.cpp

101

Yeah given that we know the compute function is going to be fast, we can just run it under the lock. So let's just use a big fat mutex.

138

nit: I think this would be clearer as a named function instead of a lambda

clang-tools-extra/clangd/index/ProjectAware.h

22

nit: "factory" is more common than "generator" for this purpose, probably

26

Mention the default?

The default implementation loads the index asynchronously on the AsyncTaskRunner.
The index will appear empty until loaded.

[kadircet]

• Internalize synchronization of indexstorage rather than using Memoize

[sammccall]

clang-tools-extra/clangd/Config.h

95 ↗

(On Diff #306615)

(nit: static isn't saving anything here I think: cost of constructing == cost of copying)

[kadircet]

• Address comments

clang-tools-extra/clangd/CMakeLists.txt

add_clang_library(clangDaemon
   index/BackgroundQueue.cpp
   index/BackgroundRebuild.cpp
   index/CanonicalIncludes.cpp
   index/FileIndex.cpp
   index/Index.cpp
   index/IndexAction.cpp
   index/MemIndex.cpp
   index/Merge.cpp
+  index/ProjectAware.cpp
   index/Ref.cpp
   index/Relation.cpp
   index/Serialization.cpp
   index/Symbol.cpp
   index/SymbolCollector.cpp
   index/SymbolID.cpp
   index/SymbolLocation.cpp
   index/SymbolOrigin.cpp

clang-tools-extra/clangd/index/ProjectAware.h

+//===--- ProjectAware.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_INDEX_PROJECT_AWARE_H
+#define LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_PROJECT_AWARE_H
+
+#include "Config.h"
+#include "index/Index.h"
+#include "index/MemIndex.h"
+#include "index/Merge.h"
+#include "index/Serialization.h"
+#include "index/remote/Client.h"
+#include "support/Threading.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/StringMap.h"
+#include "llvm/ADT/StringRef.h"
+#include <cstddef>
+#include <memory>
+#include <mutex>
+#include <string>
+#include <vector>
+namespace clang {
+namespace clangd {
+
+class ProjectAwareIndex : public SymbolIndex {
+public:
+  size_t estimateMemoryUsage() const override;
+
+  /// Only queries the associated index with the current context.
+  void lookup(const LookupRequest &Req,
+              llvm::function_ref<void(const Symbol &)> Callback) const override;
+
+  /// Query all indexes while prioritizing the associated one (if any).
+  bool refs(const RefsRequest &Req,
+            llvm::function_ref<void(const Ref &)> Callback) const override;
+
+  /// Queries only the associates index when Req.RestrictForCodeCompletion is
+  /// set, otherwise queries all.
+  bool
+  fuzzyFind(const FuzzyFindRequest &Req,
+            llvm::function_ref<void(const Symbol &)> Callback) const override;
+
+  /// Query all indexes while prioritizing the associated one (if any).
+  void relations(const RelationsRequest &Req,
+                 llvm::function_ref<void(const SymbolID &, const Symbol &)>
+                     Callback) const override;
+
+private:
+  // Returns the index associated with current context, if any.
+  SymbolIndex *getIndex() const;
+  // Adds the Idx into IndexStorage and adjusts the MergedIdx view. Returns a
+  // pointer to the Idx.
+  SymbolIndex *addIndex(std::unique_ptr<SymbolIndex> Idx) const;
+
+  // Storage for all the external indexes.
+  Memoize<llvm::StringMap<SymbolIndex *>> IndexForProject;
+  mutable std::mutex Mu;
+  mutable std::vector<std::unique_ptr<SymbolIndex>> IndexStorage;
+  AsyncTaskRunner Tasks;
+
+  // Used for creating a merged view over all the indexes seen so far to ease
+  // querying them all at once.
+  mutable std::mutex MergedIdxMu;
+  mutable std::vector<std::unique_ptr<SymbolIndex>> MergedIdxStorage;
+  // Points at root of merged view of all smybols.
+  mutable SymbolIndex *MergedIdx = nullptr;
+};
+} // namespace clangd
+} // namespace clang
+
+#endif

clang-tools-extra/clangd/index/ProjectAware.cpp

+//===--- ProjectAware.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
+//
+//===----------------------------------------------------------------------===//
+
+#include "ProjectAware.h"
+#include "index/Index.h"
+#include "index/Merge.h"
+#include "index/Ref.h"
+#include "index/Symbol.h"
+#include "index/SymbolID.h"
+#include "support/Logger.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/StringRef.h"
+#include <memory>
+#include <mutex>
+
+namespace clang {
+namespace clangd {
+size_t ProjectAwareIndex::estimateMemoryUsage() const {
+  size_t Total = 0;
+  std::lock_guard<std::mutex> Lock(Mu);
+  for (auto &Idx : IndexStorage)
+    Total += Idx->estimateMemoryUsage();
+  return Total;
+}
+
+void ProjectAwareIndex::lookup(
+    const LookupRequest &Req,
+    llvm::function_ref<void(const Symbol &)> Callback) const {
+  if (auto *Idx = getIndex())
+    Idx->lookup(Req, Callback);
+}
+
+bool ProjectAwareIndex::refs(
+    const RefsRequest &Req,
+    llvm::function_ref<void(const Ref &)> Callback) const {
+  auto *AssocIdx = getIndex();
+  std::lock_guard<std::mutex> Lock(MergedIdxMu);
+  if (AssocIdx)
+    return MergedIndex(AssocIdx, MergedIdx).refs(Req, Callback);
+  return MergedIdx->refs(Req, Callback);
+}
+
+bool ProjectAwareIndex::fuzzyFind(
+    const FuzzyFindRequest &Req,
+    llvm::function_ref<void(const Symbol &)> Callback) const {
+  auto *AssocIdx = getIndex();
+  if (Req.RestrictForCodeCompletion) {
+    if (AssocIdx)
+      return AssocIdx->fuzzyFind(Req, Callback);
+    return false;
+  }
+  std::lock_guard<std::mutex> Lock(MergedIdxMu);
+  if (AssocIdx)
+    return MergedIndex(AssocIdx, MergedIdx).fuzzyFind(Req, Callback);
+  return MergedIdx->fuzzyFind(Req, Callback);
+}
+
+void ProjectAwareIndex::relations(
+    const RelationsRequest &Req,
+    llvm::function_ref<void(const SymbolID &, const Symbol &)> Callback) const {
+  auto *AssocIdx = getIndex();
+  std::lock_guard<std::mutex> Lock(MergedIdxMu);
+  if (AssocIdx)
+    return MergedIndex(AssocIdx, MergedIdx).relations(Req, Callback);
+  return MergedIdx->relations(Req, Callback);
+}
+
+SymbolIndex *ProjectAwareIndex::getIndex() const {
+  const auto &C = Config::current();
+  if (C.Index.External.MountPoint.empty())
+    return nullptr;
+  if (auto File = C.Index.External.File) {
+    return IndexForProject.get(
+        *File, [File = std::move(*File),
+                MP = llvm::StringRef(C.Index.External.MountPoint), this] {
+          log("Associating {0} with monolithic index at {1}.", MP, File);
+          auto NewIndex =
+              std::make_unique<SwapIndex>(std::make_unique<MemIndex>());
+          SwapIndex *PlaceHolder = NewIndex.get();
+          Tasks.runAsync("Load-index:" + File, [File, PlaceHolder] {
+            if (auto Idx = loadIndex(File, /*UseDex=*/true))
+              PlaceHolder->reset(std::move(Idx));
+          });
+          addIndex(std::move(NewIndex));
+          return PlaceHolder;
+        });
+  }
+  assert(C.Index.External.Server);
+  auto Server = std::move(*C.Index.External.Server);
+  return IndexForProject.get(
+      Server, [Server = std::move(Server),
+               MP = llvm::StringRef(C.Index.External.MountPoint), this] {
+        log("Associating {0} with remote index at {1}.", MP, Server);
+        auto NewIndex = remote::getClient(Server, MP);
+        return addIndex(std::move(NewIndex));
+      });
+}
+
+SymbolIndex *
+ProjectAwareIndex::addIndex(std::unique_ptr<SymbolIndex> Idx) const {
+  SymbolIndex *NewIdx = Idx.get();
+  {
+    std::lock_guard<std::mutex> Lock(Mu);
+    IndexStorage.emplace_back(std::move(Idx));
+  }
+  {
+    std::lock_guard<std::mutex> Lock(MergedIdxMu);
+    if (!MergedIdx)
+      MergedIdx = NewIdx;
+    else {
+      MergedIdxStorage.emplace_back(
+          std::make_unique<MergedIndex>(NewIdx, MergedIdx));
+      MergedIdx = MergedIdxStorage.back().get();
+    }
+  }
+  return NewIdx;
+}
+} // namespace clangd
+} // namespace clang

clang-tools-extra/clangd/unittests/CMakeLists.txt

add_unittest(ClangdUnitTests ClangdTests
   JSONTransportTests.cpp
   LoggerTests.cpp
   LSPClient.cpp
   ModulesTests.cpp
   ParsedASTTests.cpp
   PathMappingTests.cpp
   PreambleTests.cpp
   PrintASTTests.cpp
+  ProjectAwareIndexTests.cpp
   QualityTests.cpp
   RenameTests.cpp
   RIFFTests.cpp
   SelectionTests.cpp
   SemanticHighlightingTests.cpp
   SemanticSelectionTests.cpp
   SerializationTests.cpp
   SourceCodeTests.cpp

clang-tools-extra/clangd/unittests/ProjectAwareIndexTests.cpp

+//===-- ProjectAwareIndexTests.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
+//
+//===----------------------------------------------------------------------===//