[clangd] Support refs() in dex. Largely cloned from MemIndex.

[hokein]

clangd/index/dex/Dex.h

60–61

nit: this comment is stale too.

unittests/clangd/DexTests.cpp

630

The test seems doesn't verify any thing?

[sammccall]

unittests/clangd/DexTests.cpp

630

Oops, thanks!

[sammccall]

Fix test and comment, and sync

[hokein]

LGTM

clangd/index/Serialization.cpp

if (auto I = readIndexFile(Buffer->get()->getBuffer())) {
         Symbols = std::move(*I->Symbols);
     } else {
       llvm::errs() << "Bad Index: " << llvm::toString(I.takeError()) << "\n";
       return nullptr;
     }
   }
 
   trace::Span Tracer("BuildIndex");
-  auto Index = UseDex ? dex::Dex::build(std::move(Symbols), URISchemes)
-                      : MemIndex::build(std::move(Symbols), std::move(Refs));
+  auto Index =
+      UseDex ? dex::Dex::build(std::move(Symbols), std::move(Refs), URISchemes)
   vlog("Loaded {0} from {1} with estimated memory usage {2}",
        UseDex ? "Dex" : "MemIndex", SymbolFilename,
        Index->estimateMemoryUsage());
   return Index;
 }
 
 } // namespace clangd
 } // namespace clang

clangd/index/dex/Dex.h

 /// In-memory Dex trigram-based index implementation.
 // FIXME(kbobyrev): Introduce serialization and deserialization of the symbol
 // index so that it can be loaded from the disk. Since static index is not
 // changed frequently, it's safe to assume that it has to be built only once
 // (when the clangd process starts). Therefore, it can be easier to store built
 // index on disk and then load it if available.
 class Dex : public SymbolIndex {
 public:
-  // All symbols must outlive this index.
-  template <typename Range>
-  Dex(Range &&Symbols, llvm::ArrayRef<std::string> Schemes)
+  // All data must outlive this index.
+  template <typename SymbolRange, typename RefsRange>
+  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;
     }
     for (auto &&Sym : Symbols)
       this->Symbols.push_back(&Sym);
+    for (auto &&Ref : Refs)
+      this->Refs.try_emplace(Ref.first, Ref.second);
     buildIndex();
   }
-  // Symbols are owned by BackingData, Index takes ownership.
-  template <typename Range, typename Payload>
-  Dex(Range &&Symbols, Payload &&BackingData, size_t BackingDataSize,
-      llvm::ArrayRef<std::string> URISchemes)
-      : Dex(std::forward<Range>(Symbols), URISchemes) {
+  // 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)
+      : Dex(std::forward<SymbolRange>(Symbols), std::forward<RefsRange>(Refs),
+            URISchemes) {
     KeepAlive = std::shared_ptr<void>(
         std::make_shared<Payload>(std::move(BackingData)), nullptr);
     this->BackingDataSize = BackingDataSize;
   }
 
-  /// Builds an index from a slab. The index takes ownership of the slab.
+  /// Builds an index from slabs. The index takes ownership of the slab.
   static std::unique_ptr<SymbolIndex>
-  build(SymbolSlab Slab, llvm::ArrayRef<std::string> URISchemes) {
-    // Store Slab size before it is moved.
-    const auto BackingDataSize = Slab.bytes();
-    return llvm::make_unique<Dex>(Slab, std::move(Slab), BackingDataSize,
-                                  URISchemes);
-  }
+  build(SymbolSlab, RefSlab, llvm::ArrayRef<std::string> URISchemes);
 
   bool
   fuzzyFind(const FuzzyFindRequest &Req,
             llvm::function_ref<void(const Symbol &)> Callback) const override;
 
   void lookup(const LookupRequest &Req,
               llvm::function_ref<void(const Symbol &)> Callback) const override;
 
private:
   /// Inverted index is a mapping from the search token to the posting list,
   /// which contains all items which can be characterized by such search token.
   /// For example, if the search token is scope "std::", the corresponding
   /// posting list would contain all indices of symbols defined in namespace
   /// std. Inverted index is used to retrieve posting lists which are processed
   /// during the fuzzyFind process.
   llvm::DenseMap<Token, PostingList> InvertedIndex;
   dex::Corpus Corpus;
+  llvm::DenseMap<SymbolID, llvm::ArrayRef<Ref>> Refs;
   std::shared_ptr<void> KeepAlive; // poor man's move-only std::any
   // Size of memory retained by KeepAlive.
   size_t BackingDataSize = 0;
 
   std::vector<std::string> URISchemes;
 };
 
 /// Returns Search Token for a number of parent directories of given Path.

clangd/index/dex/Dex.cpp

 #include "llvm/Support/ScopedPrinter.h"
 #include <algorithm>
 #include <queue>
 
 namespace clang {
 namespace clangd {
 namespace dex {
 
+std::unique_ptr<SymbolIndex>
+Dex::build(SymbolSlab Symbols, RefSlab Refs,
+           llvm::ArrayRef<std::string> URISchemes) {
+  auto Size = Symbols.bytes() + Refs.bytes();
+  auto Data = std::make_pair(std::move(Symbols), std::move(Refs));
+  return llvm::make_unique<Dex>(Data.first, Data.second, std::move(Data), Size,
+                                std::move(URISchemes));
+}
+
 namespace {
 
 // Mark symbols which are can be used for code completion.
 const Token RestrictedForCodeCompletion =
     Token(Token::Kind::Sentinel, "Restricted For Code Completion");
 
 // Returns the tokens which are given symbol's characteristics. Currently, the
 // generated tokens only contain fuzzy matching trigrams and symbol's scope,
for (const auto &ID : Req.IDs) {
     if (I != LookupTable.end())
       Callback(*I->second);
   }
 }
 
 void Dex::refs(const RefsRequest &Req,
                llvm::function_ref<void(const Ref &)> Callback) const {
   trace::Span Tracer("Dex refs");
-  log("refs is not implemented.");
+  for (const auto &ID : Req.IDs)
+    for (const auto &Ref : Refs.lookup(ID))
+      if (static_cast<int>(Req.Filter & Ref.Kind))
+        Callback(Ref);
 }
 
 size_t Dex::estimateMemoryUsage() const {
   size_t Bytes = Symbols.size() * sizeof(const Symbol *);
   Bytes += SymbolQuality.size() * sizeof(float);
   Bytes += LookupTable.getMemorySize();
   Bytes += InvertedIndex.getMemorySize();
   for (const auto &TokenToPostingList : InvertedIndex)
     Bytes += TokenToPostingList.second.bytes();
+  Bytes += Refs.getMemorySize();
   return Bytes + BackingDataSize;
 }
 
 std::vector<std::string> generateProximityURIs(llvm::StringRef URIPath) {
   std::vector<std::string> Result;
   auto ParsedURI = URI::parse(URIPath);
   assert(ParsedURI &&
          "Non-empty argument of generateProximityURIs() should be a valid "

clangd/indexer/IndexerMain.cpp

static llvm::cl::opt<std::string> AssumedHeaderDir(
     llvm::cl::init(""));
 
 static llvm::cl::opt<IndexFileFormat>
     Format("format", llvm::cl::desc("Format of the index to be written"),
            llvm::cl::values(clEnumValN(IndexFileFormat::YAML, "yaml",
                                        "human-readable YAML format"),
                             clEnumValN(IndexFileFormat::RIFF, "binary",
                                        "binary RIFF format")),
-           llvm::cl::init(IndexFileFormat::YAML));
+           llvm::cl::init(IndexFileFormat::RIFF));
 
 class IndexActionFactory : public tooling::FrontendActionFactory {
 public:
   IndexActionFactory(IndexFileIn &Result) : Result(Result) {}
 
   clang::FrontendAction *create() override {
     SymbolCollector::Options Opts;
     Opts.FallbackDir = AssumedHeaderDir;

unittests/clangd/DexTests.cpp

EXPECT_THAT(generateProximityURIs("unittest:///a/b/c.h"),
                           "unittest:///a", "unittest:///"));
 }
 
 //===----------------------------------------------------------------------===//
 // Index tests.
 //===----------------------------------------------------------------------===//
 
 TEST(Dex, Lookup) {
-  auto I = Dex::build(generateSymbols({"ns::abc", "ns::xyz"}), URISchemes);
+  auto I = Dex::build(generateSymbols({"ns::abc", "ns::xyz"}), RefSlab(),
+                      URISchemes);
   EXPECT_THAT(lookup(*I, SymbolID("ns::abc")), UnorderedElementsAre("ns::abc"));
   EXPECT_THAT(lookup(*I, {SymbolID("ns::abc"), SymbolID("ns::xyz")}),
               UnorderedElementsAre("ns::abc", "ns::xyz"));
   EXPECT_THAT(lookup(*I, {SymbolID("ns::nonono"), SymbolID("ns::xyz")}),
               UnorderedElementsAre("ns::xyz"));
   EXPECT_THAT(lookup(*I, SymbolID("ns::nonono")), UnorderedElementsAre());
 }
 
 TEST(Dex, FuzzyFind) {
   auto Index =
       Dex::build(generateSymbols({"ns::ABC", "ns::BCD", "::ABC",
                                   "ns::nested::ABC", "other::ABC", "other::A"}),
-                 URISchemes);
+                 RefSlab(), URISchemes);
   FuzzyFindRequest Req;
   Req.Query = "ABC";
   Req.Scopes = {"ns::"};
   EXPECT_THAT(match(*Index, Req), UnorderedElementsAre("ns::ABC"));
   Req.Scopes = {"ns::", "ns::nested::"};
   EXPECT_THAT(match(*Index, Req),
               UnorderedElementsAre("ns::ABC", "ns::nested::ABC"));
   Req.Query = "A";

 // 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, URISchemes);
+  Dex I(Symbols, RefSlab(), URISchemes);
   EXPECT_FALSE(Req.Limit);
   EXPECT_THAT(match(I, Req), ElementsAre("2", "2"));
 }
 
 TEST(DexTest, DexLimitedNumMatches) {
-  auto I = Dex::build(generateNumSymbols(0, 100), URISchemes);
+  auto I = Dex::build(generateNumSymbols(0, 100), RefSlab(), URISchemes);
   FuzzyFindRequest Req;
   Req.Query = "5";
   Req.Limit = 3;
   bool Incomplete;
   auto Matches = match(*I, Req, &Incomplete);
   EXPECT_TRUE(Req.Limit);
   EXPECT_EQ(Matches.size(), *Req.Limit);
   EXPECT_TRUE(Incomplete);
 }
 
 TEST(DexTest, FuzzyMatch) {
   auto I = Dex::build(
       generateSymbols({"LaughingOutLoud", "LionPopulation", "LittleOldLady"}),
-      URISchemes);
+      RefSlab(), URISchemes);
   FuzzyFindRequest Req;
   Req.Query = "lol";
   Req.Limit = 2;
   EXPECT_THAT(match(*I, Req),
               UnorderedElementsAre("LaughingOutLoud", "LittleOldLady"));
 }
 
 TEST(DexTest, MatchQualifiedNamesWithoutSpecificScope) {
-  auto I = Dex::build(generateSymbols({"a::y1", "b::y2", "y3"}), URISchemes);
+  auto I = Dex::build(generateSymbols({"a::y1", "b::y2", "y3"}), RefSlab(),
+                      URISchemes);
   FuzzyFindRequest Req;
   Req.Query = "y";
   EXPECT_THAT(match(*I, Req), UnorderedElementsAre("a::y1", "b::y2", "y3"));
 }
 
 TEST(DexTest, MatchQualifiedNamesWithGlobalScope) {
-  auto I = Dex::build(generateSymbols({"a::y1", "b::y2", "y3"}), URISchemes);
+  auto I = Dex::build(generateSymbols({"a::y1", "b::y2", "y3"}), RefSlab(),
+                      URISchemes);
   FuzzyFindRequest Req;
   Req.Query = "y";
   Req.Scopes = {""};
   EXPECT_THAT(match(*I, Req), UnorderedElementsAre("y3"));
 }
 
 TEST(DexTest, MatchQualifiedNamesWithOneScope) {
-  auto I = Dex::build(
-      generateSymbols({"a::y1", "a::y2", "a::x", "b::y2", "y3"}), URISchemes);
+  auto I =
+      Dex::build(generateSymbols({"a::y1", "a::y2", "a::x", "b::y2", "y3"}),
+                 RefSlab(), URISchemes);
   FuzzyFindRequest Req;
   Req.Query = "y";
   Req.Scopes = {"a::"};
   EXPECT_THAT(match(*I, Req), UnorderedElementsAre("a::y1", "a::y2"));
 }
 
 TEST(DexTest, MatchQualifiedNamesWithMultipleScopes) {
   auto I = Dex::build(
-      generateSymbols({"a::y1", "a::y2", "a::x", "b::y3", "y3"}), URISchemes);
+      generateSymbols({"a::y1", "a::y2", "a::x", "b::y3", "y3"}), RefSlab(), URISchemes);
   FuzzyFindRequest Req;
   Req.Query = "y";
   Req.Scopes = {"a::", "b::"};
   EXPECT_THAT(match(*I, Req), UnorderedElementsAre("a::y1", "a::y2", "b::y3"));
 }
 
 TEST(DexTest, NoMatchNestedScopes) {
-  auto I = Dex::build(generateSymbols({"a::y1", "a::b::y2"}), URISchemes);
+  auto I = Dex::build(generateSymbols({"a::y1", "a::b::y2"}), RefSlab(), URISchemes);
   FuzzyFindRequest Req;
   Req.Query = "y";
   Req.Scopes = {"a::"};
   EXPECT_THAT(match(*I, Req), UnorderedElementsAre("a::y1"));
 }
 
 TEST(DexTest, WildcardScope) {
   auto I =
-      Dex::build(generateSymbols({"a::y1", "a::b::y2", "c::y3"}), URISchemes);
+      Dex::build(generateSymbols({"a::y1", "a::b::y2", "c::y3"}), RefSlab(), URISchemes);
   FuzzyFindRequest Req;
   Req.Query = "y";
   Req.Scopes = {"a::"};
   Req.AnyScope = true;
   EXPECT_THAT(match(*I, Req),
               UnorderedElementsAre("a::y1", "a::b::y2", "c::y3"));
 }
 
 TEST(DexTest, IgnoreCases) {
-  auto I = Dex::build(generateSymbols({"ns::ABC", "ns::abc"}), URISchemes);
+  auto I = Dex::build(generateSymbols({"ns::ABC", "ns::abc"}), RefSlab(), URISchemes);
   FuzzyFindRequest Req;
   Req.Query = "AB";
   Req.Scopes = {"ns::"};
   EXPECT_THAT(match(*I, Req), UnorderedElementsAre("ns::ABC", "ns::abc"));
 }
 
 TEST(DexTest, Lookup) {
-  auto I = Dex::build(generateSymbols({"ns::abc", "ns::xyz"}), URISchemes);
+  auto I = Dex::build(generateSymbols({"ns::abc", "ns::xyz"}), RefSlab(), URISchemes);
   EXPECT_THAT(lookup(*I, SymbolID("ns::abc")), UnorderedElementsAre("ns::abc"));
   EXPECT_THAT(lookup(*I, {SymbolID("ns::abc"), SymbolID("ns::xyz")}),
               UnorderedElementsAre("ns::abc", "ns::xyz"));
   EXPECT_THAT(lookup(*I, {SymbolID("ns::nonono"), SymbolID("ns::xyz")}),
               UnorderedElementsAre("ns::xyz"));
   EXPECT_THAT(lookup(*I, SymbolID("ns::nonono")), UnorderedElementsAre());
 }
 
 TEST(DexTest, SymbolIndexOptionsFilter) {
   auto CodeCompletionSymbol = symbol("Completion");
   auto NonCodeCompletionSymbol = symbol("NoCompletion");
   CodeCompletionSymbol.Flags = Symbol::SymbolFlag::IndexedForCodeCompletion;
   NonCodeCompletionSymbol.Flags = Symbol::SymbolFlag::None;
   std::vector<Symbol> Symbols{CodeCompletionSymbol, NonCodeCompletionSymbol};
-  Dex I(Symbols, URISchemes);
+  Dex I(Symbols, RefSlab(), URISchemes);
   FuzzyFindRequest Req;
   Req.RestrictForCodeCompletion = false;
   EXPECT_THAT(match(I, Req), ElementsAre("Completion", "NoCompletion"));
   Req.RestrictForCodeCompletion = true;
   EXPECT_THAT(match(I, Req), ElementsAre("Completion"));
 }
 
 TEST(DexTest, ProximityPathsBoosting) {
   auto RootSymbol = symbol("root::abc");
   RootSymbol.CanonicalDeclaration.FileURI = "unittest:///file.h";
   auto CloseSymbol = symbol("close::abc");
   CloseSymbol.CanonicalDeclaration.FileURI = "unittest:///a/b/c/d/e/f/file.h";
 
   std::vector<Symbol> Symbols{CloseSymbol, RootSymbol};
-  Dex I(Symbols, URISchemes);
+  Dex I(Symbols, RefSlab(), URISchemes);
 
   FuzzyFindRequest Req;
   Req.Query = "abc";
   // The best candidate can change depending on the proximity paths.
   Req.Limit = 1;
 
   // FuzzyFind request comes from the file which is far from the root: expect
   // CloseSymbol to come out.
   Req.ProximityPaths = {testPath("a/b/c/d/e/f/file.h")};
   EXPECT_THAT(match(I, Req), ElementsAre("close::abc"));
 
   // FuzzyFind request comes from the file which is close to the root: expect
   // RootSymbol to come out.
   Req.ProximityPaths = {testPath("file.h")};
   EXPECT_THAT(match(I, Req), ElementsAre("root::abc"));
 }
 
+TEST(DexTests, Refs) {
+  DenseMap<SymbolID, std::vector<Ref>> Refs;
+  auto AddRef = [&](const Symbol& Sym, StringRef Filename, RefKind Kind) {
+    auto& SymbolRefs = Refs[Sym.ID];
+    SymbolRefs.emplace_back();
+    SymbolRefs.back().Kind = Kind;
+    SymbolRefs.back().Location.FileURI = Filename;
+  };
+  auto Foo = symbol("foo");
+  auto Bar = symbol("bar");
+  AddRef(Foo, "foo.h", RefKind::Declaration);
+  AddRef(Foo, "reffoo.h", RefKind::Reference);
+  AddRef(Bar, "bar.h", RefKind::Declaration);
+
+  std::vector<std::string> Files;
+  RefsRequest Req;
+  Req.IDs.insert(Foo.ID);
+  Req.Filter = RefKind::Declaration | RefKind::Definition;
+  Dex(std::vector<Symbol>{Foo, Bar}, Refs, {}).refs(Req, [&](const Ref &R) {
+    Files.push_back(R.Location.FileURI);
+  });
+
+  EXPECT_THAT(Files, ElementsAre("foo.h"));
+}
+
 } // namespace
 } // namespace dex
 } // namespace clangd
 } // namespace clang