3 files changed, 32 insertions, 18 deletions
diff --git a/clang-tools-extra/clangd/index/Serialization.cpp b/clang-tools-extra/clangd/index/Serialization.cpp
index 5e16137f8b2..2ebd2041ff5 100644
--- a/clang-tools-extra/clangd/index/Serialization.cpp
+++ b/clang-tools-extra/clangd/index/Serialization.cpp
@@ -435,8 +435,9 @@ std::unique_ptr<SymbolIndex> loadIndex(llvm::StringRef SymbolFilename,
   }
 
   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)
+             : MemIndex::build(std::move(Symbols), std::move(Refs));
   vlog("Loaded {0} from {1} with estimated memory usage {2}",
        UseDex ? "Dex" : "MemIndex", SymbolFilename,
        Index->estimateMemoryUsage());
diff --git a/clang-tools-extra/clangd/index/dex/Dex.cpp b/clang-tools-extra/clangd/index/dex/Dex.cpp
index 3b236f8918a..08f168dddef 100644
--- a/clang-tools-extra/clangd/index/dex/Dex.cpp
+++ b/clang-tools-extra/clangd/index/dex/Dex.cpp
@@ -24,6 +24,15 @@ 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.
@@ -254,7 +263,10 @@ void Dex::lookup(const LookupRequest &Req,
 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 {
@@ -264,6 +276,7 @@ size_t Dex::estimateMemoryUsage() const {
   Bytes += InvertedIndex.getMemorySize();
   for (const auto &TokenToPostingList : InvertedIndex)
     Bytes += TokenToPostingList.second.bytes();
+  Bytes += Refs.getMemorySize();
   return Bytes + BackingDataSize;
 }
 
diff --git a/clang-tools-extra/clangd/index/dex/Dex.h b/clang-tools-extra/clangd/index/dex/Dex.h
index baa215310fe..d89e6e15bce 100644
--- a/clang-tools-extra/clangd/index/dex/Dex.h
+++ b/clang-tools-extra/clangd/index/dex/Dex.h
@@ -41,9 +41,10 @@ namespace dex {
 // 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.
@@ -53,26 +54,24 @@ public:
     }
     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,
@@ -102,6 +101,7 @@ private:
   /// 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;