1 files changed, 36 insertions, 4 deletions

diff --git a/clang-tools-extra/clangd/index/MemIndex.cpp b/clang-tools-extra/clangd/index/MemIndex.cpp
index 19a64ad9525..8d314d1febd 100644
--- a/clang-tools-extra/clangd/index/MemIndex.cpp
+++ b/clang-tools-extra/clangd/index/MemIndex.cpp
@@ -15,7 +15,27 @@
 namespace clang {
 namespace clangd {
 
-void MemIndex::build(std::shared_ptr<std::vector<const Symbol *>> Syms) {
+static std::shared_ptr<MemIndex::OccurrenceMap>
+getOccurrencesFromSlab(SymbolOccurrenceSlab OccurrencesSlab) {
+  struct Snapshot {
+    SymbolOccurrenceSlab Slab;
+    MemIndex::OccurrenceMap Occurrences;
+  };
+
+  auto Snap = std::make_shared<Snapshot>();
+  Snap->Slab = std::move(OccurrencesSlab);
+  for (const auto &IDAndOccurrences : Snap->Slab) {
+    auto &Occurrences = Snap->Occurrences[IDAndOccurrences.first];
+    for (const auto &Occurrence : IDAndOccurrences.second)
+      Occurrences.push_back(&Occurrence);
+  }
+  return {std::move(Snap), &Snap->Occurrences};
+}
+
+void MemIndex::build(std::shared_ptr<std::vector<const Symbol *>> Syms,
+                     std::shared_ptr<OccurrenceMap> AllOccurrences) {
+  assert(Syms && "Syms must be set when build MemIndex");
+  assert(AllOccurrences && "Occurrences must be set when build MemIndex");
   llvm::DenseMap<SymbolID, const Symbol *> TempIndex;
   for (const Symbol *Sym : *Syms)
     TempIndex[Sym->ID] = Sym;
@@ -25,15 +45,18 @@ void MemIndex::build(std::shared_ptr<std::vector<const Symbol *>> Syms) {
     std::lock_guard<std::mutex> Lock(Mutex);
     Index = std::move(TempIndex);
     Symbols = std::move(Syms); // Relase old symbols.
+    Occurrences = std::move(AllOccurrences);
   }
 
   vlog("Built MemIndex with estimated memory usage {0} bytes.",
        estimateMemoryUsage());
 }
 
-std::unique_ptr<SymbolIndex> MemIndex::build(SymbolSlab Slab) {
+std::unique_ptr<SymbolIndex> MemIndex::build(SymbolSlab Symbols,
+                                             SymbolOccurrenceSlab Occurrences) {
   auto Idx = llvm::make_unique<MemIndex>();
-  Idx->build(getSymbolsFromSlab(std::move(Slab)));
+  Idx->build(getSymbolsFromSlab(std::move(Symbols)),
+             getOccurrencesFromSlab(std::move(Occurrences)));
   return std::move(Idx);
 }
 
@@ -84,7 +107,16 @@ void MemIndex::lookup(const LookupRequest &Req,
 void MemIndex::findOccurrences(
     const OccurrencesRequest &Req,
     llvm::function_ref<void(const SymbolOccurrence &)> Callback) const {
-  log("findOccurrences is not implemented.");
+  std::lock_guard<std::mutex> Lock(Mutex);
+  for (const auto &ReqID : Req.IDs) {
+    auto FoundOccurrences = Occurrences->find(ReqID);
+    if (FoundOccurrences == Occurrences->end())
+      continue;
+    for (const auto *O : FoundOccurrences->second) {
+      if (static_cast<int>(Req.Filter & O->Kind))
+        Callback(*O);
+    }
+  }
 }
 
 std::shared_ptr<std::vector<const Symbol *>>