[clangd] Implement findOccurrences interface in dynamic index.

[hokein]

Summary

Implement the interface in

• FileIndex
• MemIndex
• MergeIndex

Depends on https://reviews.llvm.org/D50385.

[hokein]

Some numbers of memory usage from running this on my machine:

File	Preamble AST	Main AST
SemaDecl.cpp	14.5MB	108.1KB (symbols) + 16.5KB (occurrences)
CodeComplete.cpp	15.4MB	53.9KB (symbols) + 7.3KB (occurrences)

The memory usage of occurrences is relatively small, I think it is fine to enable it by default.

[ilya-biryukov]

The memory usage looks good. Some NITs and a major consideration wrt to the implementation of merging occurences from dynamic and static index.

clangd/index/FileIndex.cpp

48

There is an implicit assumption (TopLevelDecls != null iff indexing the main file AST in dynamic index) that does not seem quite obvious.
Maybe add two overloads (for indexing the main AST and for indexing the preamble AST) or add an extra parameter to be explicit more about those choices?

80

NIT: maybe remove braces around single-statement branches?

145–149

Why not std::tie(*Slab, *OccurenceSlab) = indexAST(...)?
I assume this should work and give the optimal performance

clangd/index/FileIndex.h

43

Maybe avoid default arguments? Having clients pass nullptr explicitly seems like the right thing to do

105

Maybe use a struct to allow named fields? Would mean a tiny amount of extra code in indexAST, but should make the client code slightly nicer.

clangd/index/MemIndex.h

24

Maybe remove the default parameter? Making the callers specify the occurrences explicitly could make their code more straight-forward to follow.

clangd/index/Merge.cpp

88

Do we want to fix it before checking this in? As discussed offline, this seems to be a limitation that could affect the design (FileOccurrences vs SymbolOccurrence in the symbol interface). In case we want to avoid refactoring the interfaces later.

95

llvm::is_contained seems to be linear time, maybe replace with hash-table lookup.
Or am I missing something?

[ilya-biryukov]

Just noticed I'm not on the reviewers list, sorry for chiming in without invitation. Was really excited about the change :-)

[hokein]

In D51279#1214268, @ilya-biryukov wrote:

Just noticed I'm not on the reviewers list, sorry for chiming in without invitation. Was really excited about the change :-)

Comments are always welcome :)

clangd/index/Merge.cpp

88

Yes, I re-thought it afterwards. We could solve this issue for findOccurrences by redefining the interface.

Unfortunately, we have the same issue for fuzzyFind, lookup too -- we still return the stale symbols from static index if the symbol was deleted in dirty files. I think we should address this issue thoroughly and separately -- we probably want to add a method in dynamic index for query whether a file is in.

[sammccall]

In D51279#1214268, @ilya-biryukov wrote:

Just noticed I'm not on the reviewers list, sorry for chiming in without invitation. Was really excited about the change :-)

fixed :-)

clangd/index/FileIndex.cpp

59–70

combine into one log statement

clangd/index/FileIndex.h

51

This return type seems inconsistent with the design of this class.
The purpose of the shared_ptr<vector<symbol_slab>> in allSymbols is to avoid exposing the concrete storage.
The analogue would I guess be shared_ptr<DenseMap<SymbolID, vector<SymbolOccurrence*>>>.

The cost of creating that does seem a little gross though.

51

allOccurrences
(name should reflect the semantics, not the type)

59

nit: no need for brief unless the "first sentence" heuristic fails. (also, it's misspelled!)

100

update doc

clangd/index/MemIndex.cpp

58–59

we should find a way to avoid having functions implicitly zero out parts of the index like this!

109

This assumes there is no duplication across occurrences, but gives no mechanism through which deduplication could occur (because we're coupled to the underlying storage).

If FileIndex provides the DenseMap<SymbolID, vector<SymbolOccurrence*>> then it could just refrain from providing duplicate pointers.

clangd/index/MemIndex.h

24

(type and name: same comment as on FileIndex)

clangd/index/Merge.cpp

84–103

I'm tempted to say remove this heuristic for now (as you say, we should solve this thoroughly), but the artifacts will be severe, so it's probably correct to include it.

But this comment doesn't explain the heuristic, it just echoes the code.
I'd suggest something like:

// We don't want duplicate occurrences from the static/dynamic indexes,
// and we can't reliably deduplicate because offsets may differ slightly.
// We consider the dynamic index authoritative and report all its occurrences,
// and only report static index occurrences from other files.
// FIXME: this heuristic fails if the dynamic index contains a file, but all
// occurrences were removed (we will report stale ones from static).
// Ultimately we should explicitly check which index has the file instead.

85

SeenFiles -> DynamicIndexFiles?

88

Yes. Please document this limitation on the MergedIndex class too since it applies to everything. It should still be documented here because it's a limitation of the particular heuristic we're using here.

[hokein]

Address review comments.

[hokein]

Minor cleanup.

[hokein]

Thanks for the comments.

clangd/index/FileIndex.cpp

48

This is a hacky way to detect whether we are indexing main AST or preamble AST -- because the current FileIndex::update(PathRef Path, ...) interface doesn't contain this information, and we use it both for PreambleAST and MainAST indexing in the client side (ParsingCallback).

I think we might want two dedicated methods (updateMainAST, and updatePreambleAST) in FileSymbol, and I think we should address it in a follow-up patch. WDYT?

clangd/index/FileIndex.h

51

The cost of creating that does seem a little gross though.

Previously I tried to avoid the cost of iterating/copying all occurrences. After some investigations, it seems that the number of symbol occurrences is small (~500 for SemaDecl.cpp, < 100 for CodeComplete.cpp). The cost should not be too expensive.

105

I think the type of pair explicitly express its meaning well enough? And this function is exposed for unittests, I'm happy to do the change if you insist.

clangd/index/MemIndex.cpp

58–59

Made it explicit.

clangd/index/Merge.cpp

95

Good catch, I thought it had a trait for Map, Set, but it turns out that is_contained is only a std::find wrapper :(

[ilya-biryukov]

clangd/index/FileIndex.cpp

48

Either a dedicated methods or a separate parameter LG.
And it's probably ok to do it with a follow-up change if you want to keep the scope of this patch smaller. That should be a pretty small change overall, so it should be fine either way.

clangd/index/FileIndex.h

105

Not really, that was just a suggestion. Feel free to ignore if the current version looks better to you.

[sammccall]

This basically looks good to go (some fixes needed but they're pretty clear I think let me know if not!)

clangd/index/FileIndex.cpp

61

for this log to be useful, you might want to include the filename.
You can get it from the sourcemanager on the ASTContext.

clangd/index/MemIndex.cpp

55–56

This is still implicitly creating an index with no occurrences. Did you mean to accept a SymbolOccurrencesSlab here?

clangd/index/MemIndex.h

22–23

can you add an explanation to this?

clangd/index/Merge.cpp

94

Unfortunately this is not safe, the backing strings may not live long enough.
This should be a StringSet (by value) instead.

100

or just DynamicIndexFileURIs.count(O.Location.FileURI) and rely on the implicit conversion to bool

clangd/index/SymbolCollector.cpp

227 ↗

(On Diff #163066)

toOccurrenceKind

isn't this just

return Roles & AllOccurrenceKinds with some casts?
(I'd put AllOccurrenceKinds into the header, but it could also be a constant here)

clangd/index/SymbolCollector.h

120 ↗

(On Diff #163066)

move these up next to ReferencedDecls and ReferencedMacros so the comment applies to them too?

Similarly, move SymbolOccurrenceSlab next to the SymbolSlab::Builder? These have strong parallels.

122 ↗

(On Diff #163066)

I'm not sure this works - IIRC can use a SymbolCollector for multiple TUs, with finish() called at the end of each one.
I think you need to (incrementally) build in finish(), and freeze in takeOccurrences().

unittests/clangd/TestTU.h

41 ↗

(On Diff #163066)

We've avoided adding this in the past because it's less readable. Please assign the fields separately instead.

[hokein]

Sorry! Just realised I messed up this patch with D50385 (mostly SymbolCollector changes), all the comments about SymbolCollector are fixed in D50385.

[hokein]

• rebase
• address review comments

[hokein]

Minor cleanup.

clangd/index/MemIndex.cpp

55–56

This is intended, this function only cares about symbol, no occurrences.

unittests/clangd/TestTU.h

41 ↗

(On Diff #163066)

I'm tempted to keep it here, while it's less readable, it does more things, and can make client code more readable (see newly-added tests), otherwise we have to repeat these statements in a few places.

[sammccall]

unittests/clangd/TestTU.h

41 ↗

(On Diff #163066)

That's what I mean, I find the newly-added tests hard to read, because "allcode" isn't clear (particularly when one of the params is a filename) and the order of parameters is not obvious.
I think they would be much clearer with the fields assigned individually as the existing tests do.

[hokein]

address review comments:

• build memindex with symbol slab and occurrence slab
• remove withAllCode in TestTU

clangd/index/FileIndex.h

 /// 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 refereces to snapshots.
 class FileSymbols {
 public:
-  /// \brief Updates all symbols in a file. If \p Slab is nullptr, symbols for
-  /// \p Path will be removed.
-  void update(PathRef Path, std::unique_ptr<SymbolSlab> Slab);
+  /// \brief Updates all symbols and occurrences in a file.
+  /// If \p Slab (Occurrence) is nullptr, symbols (occurrences) for \p Path
+  /// will be removed.
+  void update(PathRef Path, std::unique_ptr<SymbolSlab> Slab,
+              std::unique_ptr<SymbolOccurrenceSlab> Occurrences);
 
   // The shared_ptr keeps the symbols alive
   std::shared_ptr<std::vector<const Symbol *>> allSymbols();
 
+  /// Returns all symbol occurrences for all active files.
+  std::shared_ptr<MemIndex::OccurrenceMap> allOccurrences() const;
+
 private:
   mutable std::mutex Mutex;
 
   /// \brief Stores the latest snapshots for all active files.
   llvm::StringMap<std::shared_ptr<SymbolSlab>> FileToSlabs;
+  /// Stores the latest occurrence slabs for all active files.
+  llvm::StringMap<std::shared_ptr<SymbolOccurrenceSlab>> FileToOccurrenceSlabs;
 };
 
 /// \brief This manages symbls from files and an in-memory index on all symbols.
 class FileIndex : public SymbolIndex {
 public:
   /// If URISchemes is empty, the default schemes in SymbolCollector will be
   /// used.
   FileIndex(std::vector<std::string> URISchemes = {});
public:
   size_t estimateMemoryUsage() const override;
 
 private:
   FileSymbols FSymbols;
   MemIndex Index;
   std::vector<std::string> URISchemes;
 };
 
-/// Retrieves namespace and class level symbols in \p AST.
+/// Retrieves symbols and symbol occurrences in \p AST.
 /// Exposed to assist in unit tests.
 /// If URISchemes is empty, the default schemes in SymbolCollector will be used.
 /// If \p TopLevelDecls is set, only these decls are indexed. Otherwise, all top
 /// level decls obtained from \p AST are indexed.
-SymbolSlab
+std::pair<SymbolSlab, SymbolOccurrenceSlab>
 indexAST(ASTContext &AST, std::shared_ptr<Preprocessor> PP,
          llvm::Optional<llvm::ArrayRef<Decl *>> TopLevelDecls = llvm::None,
          llvm::ArrayRef<std::string> URISchemes = {});
 
 } // namespace clangd
 } // namespace clang
 
 #endif // LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_FILEINDEX_H

clangd/index/FileIndex.cpp

 #include "../Logger.h"
 #include "SymbolCollector.h"
 #include "clang/Index/IndexingAction.h"
 #include "clang/Lex/Preprocessor.h"
 
 namespace clang {
 namespace clangd {
 
-SymbolSlab indexAST(ASTContext &AST, std::shared_ptr<Preprocessor> PP,
-                    llvm::Optional<llvm::ArrayRef<Decl *>> TopLevelDecls,
-                    llvm::ArrayRef<std::string> URISchemes) {
+std::pair<SymbolSlab, SymbolOccurrenceSlab>
+indexAST(ASTContext &AST, std::shared_ptr<Preprocessor> PP,
   SymbolCollector::Options CollectorOpts;
   // FIXME(ioeric): we might also want to collect include headers. We would need
   // to make sure all includes are canonicalized (with CanonicalIncludes), which
   // is not trivial given the current way of collecting symbols: we only have
   // AST at this point, but we also need preprocessor callbacks (e.g.
   // CommentHandler for IWYU pragma) to canonicalize includes.
   CollectorOpts.CollectIncludePath = false;
   CollectorOpts.CountReferences = false;
   if (!URISchemes.empty())
     CollectorOpts.URISchemes = URISchemes;
   CollectorOpts.Origin = SymbolOrigin::Dynamic;
 
-  SymbolCollector Collector(std::move(CollectorOpts));
-  Collector.setPreprocessor(PP);
   index::IndexingOptions IndexOpts;
   // We only need declarations, because we don't count references.
   IndexOpts.SystemSymbolFilter =
       index::IndexingOptions::SystemSymbolFilterKind::DeclarationsOnly;
   IndexOpts.IndexFunctionLocals = false;
 
   std::vector<Decl *> DeclsToIndex;
   if (TopLevelDecls)
     DeclsToIndex.assign(TopLevelDecls->begin(), TopLevelDecls->end());
   else
     DeclsToIndex.assign(AST.getTranslationUnitDecl()->decls().begin(),
                         AST.getTranslationUnitDecl()->decls().end());
 
+  // We only collect occurrences when indexing main AST.
+  // FIXME: this is a hacky way to detect whether we are indexing preamble AST
+  // or main AST, we should make it explicitly.
+  bool IsIndexMainAST = TopLevelDecls.hasValue();
+  if (IsIndexMainAST)
+    CollectorOpts.OccurrenceFilter = AllOccurrenceKinds;
+
+  SymbolCollector Collector(std::move(CollectorOpts));
+  Collector.setPreprocessor(PP);
   index::indexTopLevelDecls(AST, DeclsToIndex, Collector, IndexOpts);
 
-  return Collector.takeSymbols();
+  const auto &SM = AST.getSourceManager();
+  const auto *MainFileEntry = SM.getFileEntryForID(SM.getMainFileID());
+  std::string FileName = MainFileEntry ? MainFileEntry->getName() : "";
+
+  auto Syms = Collector.takeSymbols();
+  auto Occurrences = Collector.takeOccurrences();
+  vlog("index {0}AST for {1}: \n"
+       "  symbol slab: {2} symbols, {3} bytes\n"
+       "  occurrence slab: {4} symbols, {5} bytes",
+       IsIndexMainAST ? "Main" : "Preamble", FileName, Syms.size(),
+       Syms.bytes(), Occurrences.size(), Occurrences.bytes());
+  return {std::move(Syms), std::move(Occurrences)};
 }
 
 FileIndex::FileIndex(std::vector<std::string> URISchemes)
     : URISchemes(std::move(URISchemes)) {}
 
-void FileSymbols::update(PathRef Path, std::unique_ptr<SymbolSlab> Slab) {
+void FileSymbols::update(PathRef Path, std::unique_ptr<SymbolSlab> Slab,
+                         std::unique_ptr<SymbolOccurrenceSlab> Occurrences) {
   std::lock_guard<std::mutex> Lock(Mutex);
   if (!Slab)
     FileToSlabs.erase(Path);
   else
     FileToSlabs[Path] = std::move(Slab);
+  if (!Occurrences)
+    FileToOccurrenceSlabs.erase(Path);
+  else
+    FileToOccurrenceSlabs[Path] = std::move(Occurrences);
 }
 
 std::shared_ptr<std::vector<const Symbol *>> FileSymbols::allSymbols() {
   // The snapshot manages life time of symbol slabs and provides pointers of all
   // symbols in all slabs.
   struct Snapshot {
     std::vector<const Symbol *> Pointers;
     std::vector<std::shared_ptr<SymbolSlab>> KeepAlive;
auto Snap = std::make_shared<Snapshot>();
     }
   }
   auto *Pointers = &Snap->Pointers;
   // Use aliasing constructor to keep the snapshot alive along with the
   // pointers.
   return {std::move(Snap), Pointers};
 }
 
+std::shared_ptr<MemIndex::OccurrenceMap> FileSymbols::allOccurrences() const {
+  // The snapshot manages life time of symbol occurrence slabs and provides
+  // pointers to all occurrences in all occurrence slabs.
+  struct Snapshot {
+    MemIndex::OccurrenceMap Occurrences; // ID => {Occurrence}
+    std::vector<std::shared_ptr<SymbolOccurrenceSlab>> KeepAlive;
+  };
+
+  auto Snap = std::make_shared<Snapshot>();
+  {
+    std::lock_guard<std::mutex> Lock(Mutex);
+
+    for (const auto &FileAndSlab : FileToOccurrenceSlabs) {
+      Snap->KeepAlive.push_back(FileAndSlab.second);
+      for (const auto &IDAndOccurrences : *FileAndSlab.second) {
+        auto &Occurrences = Snap->Occurrences[IDAndOccurrences.first];
+        for (const auto &Occurrence : IDAndOccurrences.second)
+          Occurrences.push_back(&Occurrence);
+      }
+    }
+  }
+
+  return {std::move(Snap), &Snap->Occurrences};
+}
+
 void FileIndex::update(PathRef Path, ASTContext *AST,
                        std::shared_ptr<Preprocessor> PP,
                        llvm::Optional<llvm::ArrayRef<Decl *>> TopLevelDecls) {
   if (!AST) {
-    FSymbols.update(Path, nullptr);
+    FSymbols.update(Path, nullptr, nullptr);
   } else {
     assert(PP);
     auto Slab = llvm::make_unique<SymbolSlab>();
-    *Slab = indexAST(*AST, PP, TopLevelDecls, URISchemes);
-    FSymbols.update(Path, std::move(Slab));
+    auto OccurrenceSlab = llvm::make_unique<SymbolOccurrenceSlab>();
+    auto IndexResults = indexAST(*AST, PP, TopLevelDecls, URISchemes);
+    std::tie(*Slab, *OccurrenceSlab) =
+        indexAST(*AST, PP, TopLevelDecls, URISchemes);
+    FSymbols.update(Path, std::move(Slab), std::move(OccurrenceSlab));
   }
   auto Symbols = FSymbols.allSymbols();
-  Index.build(std::move(Symbols));
+  Index.build(std::move(Symbols), FSymbols.allOccurrences());
 }
 
 bool FileIndex::fuzzyFind(
     const FuzzyFindRequest &Req,
     llvm::function_ref<void(const Symbol &)> Callback) const {
   return Index.fuzzyFind(Req, Callback);
 }
 
 void FileIndex::lookup(
     const LookupRequest &Req,
     llvm::function_ref<void(const Symbol &)> Callback) const {
   Index.lookup(Req, Callback);
 }
 
 void FileIndex::findOccurrences(
     const OccurrencesRequest &Req,
     llvm::function_ref<void(const SymbolOccurrence &)> Callback) const {
-  log("findOccurrences is not implemented.");
+  Index.findOccurrences(Req, Callback);
 }
 
 size_t FileIndex::estimateMemoryUsage() const {
   return Index.estimateMemoryUsage();
 }
 
 } // namespace clangd
 } // namespace clang

clangd/index/Index.h

 class SymbolOccurrenceSlab {
 public:
   using const_iterator =
       llvm::DenseMap<SymbolID, std::vector<SymbolOccurrence>>::const_iterator;
   using iterator = const_iterator;
 
   SymbolOccurrenceSlab() : UniqueStrings(Arena) {}
 
+  static SymbolOccurrenceSlab createEmpty() {
+    SymbolOccurrenceSlab Empty;
+    Empty.freeze();
+    return Empty;
+  }
+
   // Define move semantics for the slab, allowing assignment from an rvalue.
   // Implicit move assignment is deleted by the compiler because
   // StringSaver has a reference type member.
   SymbolOccurrenceSlab(SymbolOccurrenceSlab &&Slab) = default;
   SymbolOccurrenceSlab &operator=(SymbolOccurrenceSlab &&RHS) {
     assert(RHS.Frozen &&
            "SymbolOcucrrenceSlab must be frozen when move assigned!");
     Arena = std::move(RHS.Arena);
     Frozen = true;
     Occurrences = std::move(RHS.Occurrences);
     return *this;
   }
 
   const_iterator begin() const { return Occurrences.begin(); }
   const_iterator end() const { return Occurrences.end(); }
 
+  size_t bytes() const {
+    return sizeof(*this) + Arena.getTotalMemory() + Occurrences.getMemorySize();
+  }
+
+  size_t size() const { return Occurrences.size(); }
+
   // Adds a symbol occurrence.
   // This is a deep copy: underlying FileURI will be owned by the slab.
   void insert(const SymbolID &SymID, const SymbolOccurrence &Occurrence);
 
   llvm::ArrayRef<SymbolOccurrence> find(const SymbolID &ID) const {
     assert(Frozen && "SymbolOccurrenceSlab must be frozen before looking up!");
     auto It = Occurrences.find(ID);
     if (It == Occurrences.end())

clangd/index/MemIndex.h

 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_MEMINDEX_H
 
 #include "Index.h"
 #include <mutex>
 
 namespace clang {
 namespace clangd {
 
-/// \brief This implements an index for a (relatively small) set of symbols that
-/// can be easily managed in memory.
+/// \brief This implements an index for a (relatively small) set of symbols (or
+/// symbol occurrences) that can be easily managed in memory.
 class MemIndex : public SymbolIndex {
 public:
-  /// \brief (Re-)Build index for `Symbols`. All symbol pointers must remain
-  /// accessible as long as `Symbols` is kept alive.
-  void build(std::shared_ptr<std::vector<const Symbol *>> Symbols);
-
-  /// \brief Build index from a symbol slab.
-  static std::unique_ptr<SymbolIndex> build(SymbolSlab Slab);
+  /// Maps from a symbol ID to all corresponding symbol occurrences.
+  /// The map doesn't own occurrence objects.
+  using OccurrenceMap =
+      llvm::DenseMap<SymbolID, std::vector<const SymbolOccurrence *>>;
+
+  /// \brief (Re-)Build index for `Symbols` and update `Occurrences`.
+  /// All symbol pointers and symbol occurrence pointers must remain accessible
+  /// as long as `Symbols` and `Occurrences` are kept alive.
+  void build(std::shared_ptr<std::vector<const Symbol *>> Symbols,
+             std::shared_ptr<OccurrenceMap> Occurrences);
+
+  /// \brief Build index from a symbol slab and a symbol occurrence slab.
+  static std::unique_ptr<SymbolIndex> build(SymbolSlab Symbols,
+                                            SymbolOccurrenceSlab Occurrences);
 
   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;
 
   void findOccurrences(const OccurrencesRequest &Req,
                        llvm::function_ref<void(const SymbolOccurrence &)>
                            Callback) const override;
 
   size_t estimateMemoryUsage() const override;
 
 private:
 
   std::shared_ptr<std::vector<const Symbol *>> Symbols;
   // Index is a set of symbols that are deduplicated by symbol IDs.
   // FIXME: build smarter index structure.
   llvm::DenseMap<SymbolID, const Symbol *> Index;
+  // A map from symbol ID to symbol occurrences, support query by IDs.
+  std::shared_ptr<OccurrenceMap> Occurrences;
   mutable std::mutex Mutex;
 };
 
 // Returns pointers to the symbols in given slab and bundles slab lifetime with
 // returned symbol pointers so that the pointers are never invalid.
 std::shared_ptr<std::vector<const Symbol *>>
 getSymbolsFromSlab(SymbolSlab Slab);
 
 } // namespace clangd
 } // namespace clang
 
 #endif // LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_MEMINDEX_H

clangd/index/MemIndex.cpp

 #include "MemIndex.h"
 #include "../FuzzyMatch.h"
 #include "../Logger.h"
 #include <queue>
 
 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;
 
   // Swap out the old symbols and index.
   {
     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);
 }
 
 bool MemIndex::fuzzyFind(
     const FuzzyFindRequest &Req,
     llvm::function_ref<void(const Symbol &)> Callback) const {
   assert(!StringRef(Req.Query).contains("::") &&
          "There must be no :: in query.");
for (const auto &ID : Req.IDs) {
     auto I = Index.find(ID);
     if (I != Index.end())
       Callback(*I->second);
   }
 }
 
 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 *>>
 getSymbolsFromSlab(SymbolSlab Slab) {
   struct Snapshot {
     SymbolSlab Slab;
     std::vector<const Symbol *> Pointers;
   };

clangd/index/Merge.h

 // The two symbols must have the same ID.
 // Returned symbol may contain data owned by either source.
 Symbol mergeSymbol(const Symbol &L, const Symbol &R);
 
 // mergedIndex returns a composite index based on two provided Indexes:
 //  - the Dynamic index covers few files, but is relatively up-to-date.
 //  - the Static index covers a bigger set of files, but is relatively stale.
 // The returned index attempts to combine results, and avoid duplicates.
+//
+// FIXME: We don't have a mechanism in Index to track deleted symbols and
+// occurrences in dirty files, so the merged index may return stale symbols
+// and occurrences from Static index.
 std::unique_ptr<SymbolIndex> mergeIndex(const SymbolIndex *Dynamic,
                                         const SymbolIndex *Static);
 
 } // namespace clangd
 } // namespace clang
 
 #endif // LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_MERGE_H

clangd/index/Merge.cpp

 //===--- Merge.cpp -----------------------------------------------*- C++-*-===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 
+#include <set>
+
 #include "Merge.h"
 #include "../Logger.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/raw_ostream.h"
 
 namespace clang {
 namespace clangd {
 namespace {
lookup(const LookupRequest &Req,
     for (const auto &ID : RemainingIDs)
       if (const Symbol *Sym = B.find(ID))
         Callback(*Sym);
   }
 
   void findOccurrences(const OccurrencesRequest &Req,
                        llvm::function_ref<void(const SymbolOccurrence &)>
                            Callback) const override {
-    log("findOccurrences is not implemented.");
+    // We don't want duplicated occurrences from the static/dynamic indexes,
+    // and we can't reliably duplicate them because occurrence offsets may
+    // differ slightly.
+    // We consider the dynamic index authoritative and report all its
+    // occurrences, and only report static index occurrences from other files.
+    //
+    // FIXME: The heuristic fails if the dynamic index containts a file, but all
+    // occurrences were removed (we will report stale ones from the static
+    // index). Ultimately we should explicit check which index has the file
+    // instead.
+    std::set<std::string> DynamicIndexFileURIs;
+    Dynamic->findOccurrences(Req, [&](const SymbolOccurrence &O) {
+      DynamicIndexFileURIs.insert(O.Location.FileURI);
+      Callback(O);
+    });
+    Static->findOccurrences(Req, [&](const SymbolOccurrence &O) {
+      if (DynamicIndexFileURIs.count(O.Location.FileURI))
+        return;
+      Callback(O);
+    });
   }
 
   size_t estimateMemoryUsage() const override {
     return Dynamic->estimateMemoryUsage() + Static->estimateMemoryUsage();
   }
 
 private:
   const SymbolIndex *Dynamic, *Static;

clangd/tool/ClangdMain.cpp

if (!Buffer) {
     return nullptr;
   }
   auto Slab = symbolsFromYAML(Buffer.get()->getBuffer());
   SymbolSlab::Builder SymsBuilder;
   for (auto Sym : Slab)
     SymsBuilder.insert(Sym);
 
   return UseDex ? dex::DexIndex::build(std::move(SymsBuilder).build())
-                : MemIndex::build(std::move(SymsBuilder).build());
+                : MemIndex::build(std::move(SymsBuilder).build(),
+                                  SymbolOccurrenceSlab::createEmpty());
 }
 
 } // namespace
 
 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 "

unittests/clangd/CodeCompleteTests.cpp

Matcher<const std::vector<CodeCompletion> &> Has(std::string Name,
   return Contains(AllOf(Named(std::move(Name)), Kind(K)));
 }
 MATCHER(IsDocumented, "") { return !arg.Documentation.empty(); }
 
 std::unique_ptr<SymbolIndex> memIndex(std::vector<Symbol> Symbols) {
   SymbolSlab::Builder Slab;
   for (const auto &Sym : Symbols)
     Slab.insert(Sym);
-  return MemIndex::build(std::move(Slab).build());
+  return MemIndex::build(std::move(Slab).build(),
+                         SymbolOccurrenceSlab::createEmpty());
 }
 
 CodeCompleteResult completions(ClangdServer &Server, StringRef TestCode,
                                Position point,
                                std::vector<Symbol> IndexSymbols = {},
                                clangd::CodeCompleteOptions Opts = {}) {
   std::unique_ptr<SymbolIndex> OverrideIndex;
   if (!IndexSymbols.empty()) {

unittests/clangd/FileIndexTests.cpp

 //===-- FileIndexTests.cpp  ---------------------------*- C++ -*-----------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 
+#include "Annotations.h"
 #include "ClangdUnit.h"
 #include "TestFS.h"
 #include "TestTU.h"
+#include "gmock/gmock.h"
 #include "index/FileIndex.h"
 #include "clang/Frontend/CompilerInvocation.h"
 #include "clang/Frontend/PCHContainerOperations.h"
 #include "clang/Lex/Preprocessor.h"
 #include "clang/Tooling/CompilationDatabase.h"
-#include "gmock/gmock.h"
 #include "gtest/gtest.h"
 
 using testing::UnorderedElementsAre;
+using testing::AllOf;
+
+MATCHER_P(OccurrenceRange, Range, "") {
+  return std::tie(arg.Location.Start.Line, arg.Location.Start.Column,
+                  arg.Location.End.Line, arg.Location.End.Column) ==
+         std::tie(Range.start.line, Range.start.character, Range.end.line,
+                  Range.end.character);
+}
+MATCHER_P(FileURI, F, "") { return arg.Location.FileURI == F; }
 
 namespace clang {
 namespace clangd {
 
 namespace {
 
 Symbol symbol(llvm::StringRef ID) {
   Symbol Sym;
   Sym.ID = SymbolID(ID);
   Sym.Name = ID;
   return Sym;
 }
 
 std::unique_ptr<SymbolSlab> numSlab(int Begin, int End) {
   SymbolSlab::Builder Slab;
   for (int i = Begin; i <= End; i++)
     Slab.insert(symbol(std::to_string(i)));
   return llvm::make_unique<SymbolSlab>(std::move(Slab).build());
 }
 
+std::unique_ptr<SymbolOccurrenceSlab> occurrenceSlab(const SymbolID &ID,
+                                                     llvm::StringRef Path) {
+  auto Slab = llvm::make_unique<SymbolOccurrenceSlab>();
+  SymbolOccurrence Occurrence;
+  Occurrence.Location.FileURI = Path;
+  Slab->insert(ID, Occurrence);
+  return Slab;
+}
+
 std::vector<std::string>
 getSymbolNames(const std::vector<const Symbol *> &Symbols) {
   std::vector<std::string> Names;
   for (const Symbol *Sym : Symbols)
     Names.push_back(Sym->Name);
   return Names;
 }
 
+std::vector<std::string>
+getOccurrencePath(const std::vector<const SymbolOccurrence *> &Occurrences) {
+  std::vector<std::string> Paths;
+  for (const auto *O : Occurrences)
+    Paths.push_back(O->Location.FileURI);
+  return Paths;
+}
+
+std::unique_ptr<SymbolOccurrenceSlab> emptyOccurrence() {
+  auto EmptySlab = llvm::make_unique<SymbolOccurrenceSlab>();
+  EmptySlab->freeze();
+  return EmptySlab;
+}
+
 TEST(FileSymbolsTest, UpdateAndGet) {
   FileSymbols FS;
   EXPECT_THAT(getSymbolNames(*FS.allSymbols()), UnorderedElementsAre());
+  EXPECT_TRUE(FS.allOccurrences()->empty());
 
-  FS.update("f1", numSlab(1, 3));
+  SymbolID ID("1");
+  FS.update("f1", numSlab(1, 3), occurrenceSlab(ID, "f1.cc"));
   EXPECT_THAT(getSymbolNames(*FS.allSymbols()),
               UnorderedElementsAre("1", "2", "3"));
+  auto Occurrences = FS.allOccurrences();
+  EXPECT_THAT(getOccurrencePath((*Occurrences)[ID]),
+              UnorderedElementsAre("f1.cc"));
 }
 
 TEST(FileSymbolsTest, Overlap) {
   FileSymbols FS;
-  FS.update("f1", numSlab(1, 3));
-  FS.update("f2", numSlab(3, 5));
+  FS.update("f1", numSlab(1, 3), emptyOccurrence());
+  FS.update("f2", numSlab(3, 5), emptyOccurrence());
   EXPECT_THAT(getSymbolNames(*FS.allSymbols()),
               UnorderedElementsAre("1", "2", "3", "3", "4", "5"));
 }
 
 TEST(FileSymbolsTest, SnapshotAliveAfterRemove) {
   FileSymbols FS;
 
-  FS.update("f1", numSlab(1, 3));
+  SymbolID ID("1");
+  FS.update("f1", numSlab(1, 3), occurrenceSlab(ID, "f1.cc"));
 
   auto Symbols = FS.allSymbols();
   EXPECT_THAT(getSymbolNames(*Symbols), UnorderedElementsAre("1", "2", "3"));
+  auto Occurrences = FS.allOccurrences();
+  EXPECT_THAT(getOccurrencePath((*Occurrences)[ID]),
+              UnorderedElementsAre("f1.cc"));
 
-  FS.update("f1", nullptr);
+  FS.update("f1", nullptr, nullptr);
   EXPECT_THAT(getSymbolNames(*FS.allSymbols()), UnorderedElementsAre());
   EXPECT_THAT(getSymbolNames(*Symbols), UnorderedElementsAre("1", "2", "3"));
+
+  EXPECT_TRUE(FS.allOccurrences()->empty());
+  EXPECT_THAT(getOccurrencePath((*Occurrences)[ID]),
+              UnorderedElementsAre("f1.cc"));
 }
 
 std::vector<std::string> match(const SymbolIndex &I,
                                const FuzzyFindRequest &Req) {
   std::vector<std::string> Matches;
   I.fuzzyFind(Req, [&](const Symbol &Sym) {
     Matches.push_back((Sym.Scope + Sym.Name).str());
   });
TEST(FileIndexTest, RebuildWithPreamble) {
   Req.Query = "";
   Req.Scopes = {"", "ns_in_header::"};
 
   EXPECT_THAT(
       match(Index, Req),
       UnorderedElementsAre("ns_in_header", "ns_in_header::func_in_header"));
 }
 
+TEST(FileIndexTest, Occurrences) {
+  const char *HeaderCode = "class Foo {};";
+  Annotations MainCode(R"cpp(
+  void f() {
+    $foo[[Foo]] foo;
+  }
+  )cpp");
+
+  auto Foo =
+      findSymbol(TestTU::withHeaderCode(HeaderCode).headerSymbols(), "Foo");
+
+  OccurrencesRequest Request;
+  Request.IDs = {Foo.ID};
+  Request.Filter = SymbolOccurrenceKind::Declaration |
+                   SymbolOccurrenceKind::Definition |
+                   SymbolOccurrenceKind::Reference;
+
+  FileIndex Index(/*URISchemes*/ {"unittest"});
+  // Add test.cc
+  TestTU Test;
+  Test.HeaderCode = HeaderCode;
+  Test.Code = MainCode.code();
+  Test.Filename = "test.cc";
+  auto AST = Test.build();
+  Index.update(Test.Filename, &AST.getASTContext(), AST.getPreprocessorPtr(),
+               AST.getLocalTopLevelDecls());
+  // Add test2.cc
+  TestTU Test2;
+  Test2.HeaderCode = HeaderCode;
+  Test2.Code = MainCode.code();
+  Test2.Filename = "test2.cc";
+  AST = Test2.build();
+  Index.update(Test2.Filename, &AST.getASTContext(), AST.getPreprocessorPtr(),
+               AST.getLocalTopLevelDecls());
+
+  std::vector<SymbolOccurrence> Results;
+  Index.findOccurrences(
+      Request, [&Results](const SymbolOccurrence &O) { Results.push_back(O); });
+
+  EXPECT_THAT(Results,
+              UnorderedElementsAre(AllOf(OccurrenceRange(MainCode.range("foo")),
+                                         FileURI("unittest:///test.cc")),
+                                   AllOf(OccurrenceRange(MainCode.range("foo")),
+                                         FileURI("unittest:///test2.cc"))));
+}
+
 } // namespace
 } // namespace clangd
 } // namespace clang

unittests/clangd/IndexTests.cpp

 //===-- IndexTests.cpp  -------------------------------*- C++ -*-----------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 
+#include "Annotations.h"
 #include "TestIndex.h"
+#include "TestTU.h"
+#include "gmock/gmock.h"
+#include "index/FileIndex.h"
 #include "index/Index.h"
 #include "index/MemIndex.h"
 #include "index/Merge.h"
-#include "gmock/gmock.h"
 #include "gtest/gtest.h"
 
 using testing::Pointee;
 using testing::UnorderedElementsAre;
+using testing::AllOf;
 
 namespace clang {
 namespace clangd {
 namespace {
 
+std::shared_ptr<MemIndex::OccurrenceMap> emptyOccurrences() {
+  return llvm::make_unique<MemIndex::OccurrenceMap>();
+}
+
 MATCHER_P(Named, N, "") { return arg.Name == N; }
+MATCHER_P(OccurrenceRange, Range, "") {
+  return std::tie(arg.Location.Start.Line, arg.Location.Start.Column,
+                  arg.Location.End.Line, arg.Location.End.Column) ==
+         std::tie(Range.start.line, Range.start.character, Range.end.line,
+                  Range.end.character);
+}
+MATCHER_P(FileURI, F, "") { return arg.Location.FileURI == F; }
 
 TEST(SymbolSlab, FindAndIterate) {
   SymbolSlab::Builder B;
   B.insert(symbol("Z"));
   B.insert(symbol("Y"));
   B.insert(symbol("X"));
   EXPECT_EQ(nullptr, B.find(SymbolID("W")));
   for (const char *Sym : {"X", "Y", "Z"})
     EXPECT_THAT(B.find(SymbolID(Sym)), Pointee(Named(Sym)));
 
   SymbolSlab S = std::move(B).build();
   EXPECT_THAT(S, UnorderedElementsAre(Named("X"), Named("Y"), Named("Z")));
   EXPECT_EQ(S.end(), S.find(SymbolID("W")));
   for (const char *Sym : {"X", "Y", "Z"})
     EXPECT_THAT(*S.find(SymbolID(Sym)), Named(Sym));
 }
 
 TEST(MemIndexTest, MemIndexSymbolsRecycled) {
   MemIndex I;
   std::weak_ptr<SlabAndPointers> Symbols;
-  I.build(generateNumSymbols(0, 10, &Symbols));
+  I.build(generateNumSymbols(0, 10, &Symbols), emptyOccurrences());
   FuzzyFindRequest Req;
   Req.Query = "7";
   EXPECT_THAT(match(I, Req), UnorderedElementsAre("7"));
 
   EXPECT_FALSE(Symbols.expired());
   // Release old symbols.
-  I.build(generateNumSymbols(0, 0));
+  I.build(generateNumSymbols(0, 0), emptyOccurrences());
   EXPECT_TRUE(Symbols.expired());
 }
 
 TEST(MemIndexTest, MemIndexDeduplicate) {
   auto Symbols = generateNumSymbols(0, 10);
 
   // Inject some duplicates and make sure we only match the same symbol once.
   auto Sym = symbol("7");
   Symbols->push_back(&Sym);
   Symbols->push_back(&Sym);
   Symbols->push_back(&Sym);
 
   FuzzyFindRequest Req;
   Req.Query = "7";
   MemIndex I;
-  I.build(std::move(Symbols));
+  I.build(std::move(Symbols), emptyOccurrences());
   auto Matches = match(I, Req);
   EXPECT_EQ(Matches.size(), 1u);
 }
 
 TEST(MemIndexTest, MemIndexLimitedNumMatches) {
   MemIndex I;
-  I.build(generateNumSymbols(0, 100));
+  I.build(generateNumSymbols(0, 100), emptyOccurrences());
   FuzzyFindRequest Req;
   Req.Query = "5";
   Req.MaxCandidateCount = 3;
   bool Incomplete;
   auto Matches = match(I, Req, &Incomplete);
   EXPECT_EQ(Matches.size(), Req.MaxCandidateCount);
   EXPECT_TRUE(Incomplete);
 }
 
 TEST(MemIndexTest, FuzzyMatch) {
   MemIndex I;
   I.build(
-      generateSymbols({"LaughingOutLoud", "LionPopulation", "LittleOldLady"}));
+      generateSymbols({"LaughingOutLoud", "LionPopulation", "LittleOldLady"}),
+      emptyOccurrences());
   FuzzyFindRequest Req;
   Req.Query = "lol";
   Req.MaxCandidateCount = 2;
   EXPECT_THAT(match(I, Req),
               UnorderedElementsAre("LaughingOutLoud", "LittleOldLady"));
 }
 
 TEST(MemIndexTest, MatchQualifiedNamesWithoutSpecificScope) {
   MemIndex I;
-  I.build(generateSymbols({"a::y1", "b::y2", "y3"}));
+  I.build(generateSymbols({"a::y1", "b::y2", "y3"}), emptyOccurrences());
   FuzzyFindRequest Req;
   Req.Query = "y";
   EXPECT_THAT(match(I, Req), UnorderedElementsAre("a::y1", "b::y2", "y3"));
 }
 
 TEST(MemIndexTest, MatchQualifiedNamesWithGlobalScope) {
   MemIndex I;
-  I.build(generateSymbols({"a::y1", "b::y2", "y3"}));
+  I.build(generateSymbols({"a::y1", "b::y2", "y3"}), emptyOccurrences());
   FuzzyFindRequest Req;
   Req.Query = "y";
   Req.Scopes = {""};
   EXPECT_THAT(match(I, Req), UnorderedElementsAre("y3"));
 }
 
 TEST(MemIndexTest, MatchQualifiedNamesWithOneScope) {
   MemIndex I;
-  I.build(generateSymbols({"a::y1", "a::y2", "a::x", "b::y2", "y3"}));
+  I.build(generateSymbols({"a::y1", "a::y2", "a::x", "b::y2", "y3"}),
+          emptyOccurrences());
   FuzzyFindRequest Req;
   Req.Query = "y";
   Req.Scopes = {"a::"};
   EXPECT_THAT(match(I, Req), UnorderedElementsAre("a::y1", "a::y2"));
 }
 
 TEST(MemIndexTest, MatchQualifiedNamesWithMultipleScopes) {
   MemIndex I;
-  I.build(generateSymbols({"a::y1", "a::y2", "a::x", "b::y3", "y3"}));
+  I.build(generateSymbols({"a::y1", "a::y2", "a::x", "b::y3", "y3"}),
+          emptyOccurrences());
   FuzzyFindRequest Req;
   Req.Query = "y";
   Req.Scopes = {"a::", "b::"};
   EXPECT_THAT(match(I, Req), UnorderedElementsAre("a::y1", "a::y2", "b::y3"));
 }
 
 TEST(MemIndexTest, NoMatchNestedScopes) {
   MemIndex I;
-  I.build(generateSymbols({"a::y1", "a::b::y2"}));
+  I.build(generateSymbols({"a::y1", "a::b::y2"}), emptyOccurrences());
   FuzzyFindRequest Req;
   Req.Query = "y";
   Req.Scopes = {"a::"};
   EXPECT_THAT(match(I, Req), UnorderedElementsAre("a::y1"));
 }
 
 TEST(MemIndexTest, IgnoreCases) {
   MemIndex I;
-  I.build(generateSymbols({"ns::ABC", "ns::abc"}));
+  I.build(generateSymbols({"ns::ABC", "ns::abc"}), emptyOccurrences());
   FuzzyFindRequest Req;
   Req.Query = "AB";
   Req.Scopes = {"ns::"};
   EXPECT_THAT(match(I, Req), UnorderedElementsAre("ns::ABC", "ns::abc"));
 }
 
 TEST(MemIndexTest, Lookup) {
   MemIndex I;
-  I.build(generateSymbols({"ns::abc", "ns::xyz"}));
+  I.build(generateSymbols({"ns::abc", "ns::xyz"}), emptyOccurrences());
   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(MergeIndexTest, Lookup) {
   MemIndex I, J;
-  I.build(generateSymbols({"ns::A", "ns::B"}));
-  J.build(generateSymbols({"ns::B", "ns::C"}));
+  I.build(generateSymbols({"ns::A", "ns::B"}), emptyOccurrences());
+  J.build(generateSymbols({"ns::B", "ns::C"}), emptyOccurrences());
   EXPECT_THAT(lookup(*mergeIndex(&I, &J), SymbolID("ns::A")),
               UnorderedElementsAre("ns::A"));
   EXPECT_THAT(lookup(*mergeIndex(&I, &J), SymbolID("ns::B")),
               UnorderedElementsAre("ns::B"));
   EXPECT_THAT(lookup(*mergeIndex(&I, &J), SymbolID("ns::C")),
               UnorderedElementsAre("ns::C"));
   EXPECT_THAT(
       lookup(*mergeIndex(&I, &J), {SymbolID("ns::A"), SymbolID("ns::B")}),
       UnorderedElementsAre("ns::A", "ns::B"));
   EXPECT_THAT(
       lookup(*mergeIndex(&I, &J), {SymbolID("ns::A"), SymbolID("ns::C")}),
       UnorderedElementsAre("ns::A", "ns::C"));
   EXPECT_THAT(lookup(*mergeIndex(&I, &J), SymbolID("ns::D")),
               UnorderedElementsAre());
   EXPECT_THAT(lookup(*mergeIndex(&I, &J), {}), UnorderedElementsAre());
 }
 
 TEST(MergeIndexTest, FuzzyFind) {
   MemIndex I, J;
-  I.build(generateSymbols({"ns::A", "ns::B"}));
-  J.build(generateSymbols({"ns::B", "ns::C"}));
+  I.build(generateSymbols({"ns::A", "ns::B"}), emptyOccurrences());
+  J.build(generateSymbols({"ns::B", "ns::C"}), emptyOccurrences());
   FuzzyFindRequest Req;
   Req.Scopes = {"ns::"};
   EXPECT_THAT(match(*mergeIndex(&I, &J), Req),
               UnorderedElementsAre("ns::A", "ns::B", "ns::C"));
 }
 
 TEST(MergeTest, Merge) {
   Symbol L, R;
TEST(MergeTest, PreferSymbolWithDefn) {
 
   R.Definition.FileURI = "file:/right.cpp"; // Now right will be favored.
   M = mergeSymbol(L, R);
   EXPECT_EQ(M.CanonicalDeclaration.FileURI, "file:/right.h");
   EXPECT_EQ(M.Definition.FileURI, "file:/right.cpp");
   EXPECT_EQ(M.Name, "right");
 }
 
+TEST(MergeIndexTest, FindOccurrences) {
+  FileIndex Dyn({"unittest"});
+  FileIndex StaticIndex({"unittest"});
+  auto MergedIndex = mergeIndex(&Dyn, &StaticIndex);
+
+  const char *HeaderCode = "class Foo;";
+  auto HeaderSymbols = TestTU::withHeaderCode("class Foo;").headerSymbols();
+  auto Foo = findSymbol(HeaderSymbols, "Foo");
+
+  // Build dynamic index for test.cc.
+  Annotations Test1Code(R"(class $Foo[[Foo]];)");
+  TestTU Test;
+  Test.HeaderCode = HeaderCode;
+  Test.Code = Test1Code.code();
+  Test.Filename = "test.cc";
+  auto AST = Test.build();
+  Dyn.update(Test.Filename, &AST.getASTContext(), AST.getPreprocessorPtr(),
+             AST.getLocalTopLevelDecls());
+
+  // Build static index for test.cc.
+  Test.HeaderCode = HeaderCode;
+  Test.Code = "// static\nclass Foo {};";
+  Test.Filename = "test.cc";
+  auto StaticAST = Test.build();
+  // Add stale occurrences for test.cc.
+  StaticIndex.update(Test.Filename, &StaticAST.getASTContext(),
+                     StaticAST.getPreprocessorPtr(),
+                     StaticAST.getLocalTopLevelDecls());
+
+  // Add occcurrences for test2.cc
+  Annotations Test2Code(R"(class $Foo[[Foo]] {};)");
+  TestTU Test2;
+  Test2.HeaderCode = HeaderCode;
+  Test2.Code = Test2Code.code();
+  Test2.Filename = "test2.cc";
+  StaticAST = Test2.build();
+  StaticIndex.update(Test2.Filename, &StaticAST.getASTContext(),
+                     StaticAST.getPreprocessorPtr(),
+                     StaticAST.getLocalTopLevelDecls());
+
+  OccurrencesRequest Request;
+  Request.IDs = {Foo.ID};
+  Request.Filter = AllOccurrenceKinds;
+  std::vector<SymbolOccurrence> Results;
+  MergedIndex->findOccurrences(
+      Request, [&](const SymbolOccurrence &O) { Results.push_back(O); });
+
+  EXPECT_THAT(Results, UnorderedElementsAre(
+                           AllOf(OccurrenceRange(Test1Code.range("Foo")),
+                                 FileURI("unittest:///test.cc")),
+                           AllOf(OccurrenceRange(Test2Code.range("Foo")),
+                                 FileURI("unittest:///test2.cc"))));
+}
+
 } // namespace
 } // namespace clangd
 } // namespace clang

unittests/clangd/TestTU.cpp

if (!AST.hasValue()) {
     ADD_FAILURE() << "Failed to build code:\n" << Code;
     llvm_unreachable("Failed to build TestTU!");
   }
   return std::move(*AST);
 }
 
 SymbolSlab TestTU::headerSymbols() const {
   auto AST = build();
-  return indexAST(AST.getASTContext(), AST.getPreprocessorPtr());
+  return indexAST(AST.getASTContext(), AST.getPreprocessorPtr()).first;
 }
 
 std::unique_ptr<SymbolIndex> TestTU::index() const {
-  return MemIndex::build(headerSymbols());
+  // FIXME: we should generate proper occurrences for TestTU.
+  return MemIndex::build(headerSymbols(), SymbolOccurrenceSlab::createEmpty());
 }
 
 const Symbol &findSymbol(const SymbolSlab &Slab, llvm::StringRef QName) {
   const Symbol *Result = nullptr;
   for (const Symbol &S : Slab) {
     if (QName != (S.Scope + S.Name).str())
       continue;
     if (Result) {