[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?

74

NIT: maybe remove braces around single-statement branches?

126

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

clangd/index/FileIndex.h

45

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

100

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

27

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

58

combine into one log statement

clangd/index/FileIndex.h

50

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.

50

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

57

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

95

update doc

clangd/index/MemIndex.cpp

36

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

86

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

27

(type and name: same comment as on FileIndex)

clangd/index/Merge.cpp

84

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

50

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.

100

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

36

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

100

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

60

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

34

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

clangd/index/MemIndex.h

25–27

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

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

34

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

unittests/clangd/TestTU.h

41

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

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

 /// 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);
+  void update(PathRef Path, std::unique_ptr<SymbolSlab> Slab,
+              std::unique_ptr<SymbolOccurrenceSlab> Occurrences = nullptr);
 
   // The shared_ptr keeps the symbols alive
   std::shared_ptr<std::vector<const Symbol *>> allSymbols();
 
+  std::vector<std::shared_ptr<SymbolOccurrenceSlab>> allOccurrenceSlabs() const;
+
 private:
   mutable std::mutex Mutex;
 
   /// \brief Stores the latest snapshots for all active files.
   llvm::StringMap<std::shared_ptr<SymbolSlab>> FileToSlabs;
+  /// \breif 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 = {});
void findOccurrences(const OccurrencesRequest &Req,
                        llvm::function_ref<void(const SymbolOccurrence &)>
                            Callback) const override;
 private:
   FileSymbols FSymbols;
   MemIndex Index;
   std::vector<std::string> URISchemes;
 };
 
 /// Retrieves namespace and class level symbols 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());
 
+  if (TopLevelDecls) { // index main AST, set occurrence flag.
+    CollectorOpts.OccurrenceFilter = SymbolOccurrenceKind::Declaration |
+                                     SymbolOccurrenceKind::Definition |
+                                     SymbolOccurrenceKind::Reference;
+  }
+
+  SymbolCollector Collector(std::move(CollectorOpts));
+  Collector.setPreprocessor(PP);
   index::indexTopLevelDecls(AST, DeclsToIndex, Collector, IndexOpts);
 
-  return Collector.takeSymbols();
+  vlog("index for AST: ");
+  auto Syms = Collector.takeSymbols();
+  auto Occurrences = Collector.takeOccurrences();
+  vlog("  symbol slab: {0} symbols, {1} bytes", Syms.size(), Syms.bytes());
+  vlog("  occurrence slab: {0} symbols, {1} 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)
+  if (!Slab) {
     FileToSlabs.erase(Path);
-  else
+  } 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>();
   {
     std::lock_guard<std::mutex> Lock(Mutex);
 
     for (const auto &FileAndSlab : FileToSlabs) {
       Snap->KeepAlive.push_back(FileAndSlab.second);
       for (const auto &Iter : *FileAndSlab.second)
         Snap->Pointers.push_back(&Iter);
     }
   }
   auto *Pointers = &Snap->Pointers;
   // Use aliasing constructor to keep the snapshot alive along with the
   // pointers.
   return {std::move(Snap), Pointers};
 }
 
+std::vector<std::shared_ptr<SymbolOccurrenceSlab>>
+FileSymbols::allOccurrenceSlabs() const{
+  std::vector<std::shared_ptr<SymbolOccurrenceSlab>> Slabs;
+  std::lock_guard<std::mutex> Lock(Mutex);
+
+  for (const auto &FileAndSlab : FileToOccurrenceSlabs)
+    Slabs.push_back(FileAndSlab.second);
+  return Slabs;
+}
+
 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 IndexResults = indexAST(*AST, PP, TopLevelDecls, URISchemes);
+    *Slab = std::move(IndexResults.first);
+    auto OccurrenceSlab = llvm::make_unique<SymbolOccurrenceSlab>();
+    *OccurrenceSlab = std::move(IndexResults.second);
+    FSymbols.update(Path, std::move(Slab), std::move(OccurrenceSlab));
   }
   auto Symbols = FSymbols.allSymbols();
-  Index.build(std::move(Symbols));
+  Index.build(std::move(Symbols), FSymbols.allOccurrenceSlabs());
 }
 
 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);
 }
 
 } // namespace clangd
 } // namespace clang

clangd/index/Index.h

SymbolOccurrenceSlab& operator=(SymbolOccurrenceSlab&& RHS) {
      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 {
      auto It = Occurrences.find(ID);
      if (It == Occurrences.end())
        return {};

clangd/index/MemIndex.h

 namespace clangd {
 
 /// \brief This implements an index for a (relatively small) set of symbols 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);
+  void
+  build(std::shared_ptr<std::vector<const Symbol *>> Symbols,
+        std::vector<std::shared_ptr<SymbolOccurrenceSlab>> OccurrenceSlabs = {});
 
   /// \brief Build index from a symbol slab.
   static std::unique_ptr<SymbolIndex> build(SymbolSlab Slab);
 
   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;
 
 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;
+
+  std::vector<std::shared_ptr<SymbolOccurrenceSlab>> OccurrenceSlabs;
   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) {
+void MemIndex::build(
+    std::shared_ptr<std::vector<const Symbol *>> Syms,
+    std::vector<std::shared_ptr<SymbolOccurrenceSlab>> Occurrences) {
   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.
+    OccurrenceSlabs = std::move(Occurrences);
   }
 }
 
 std::unique_ptr<SymbolIndex> MemIndex::build(SymbolSlab Slab) {
   auto Idx = llvm::make_unique<MemIndex>();
-  Idx->build(getSymbolsFromSlab(std::move(Slab)));
+  Idx->build(getSymbolsFromSlab(std::move(Slab)), {});
   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.");
+  for (const auto &Slab : OccurrenceSlabs) {
+    for (const auto &ID : Req.IDs) {
+      for (const auto &Occurrence : Slab->find(ID)) {
+        if (static_cast<int>(Req.Filter & Occurrence.Kind)) {
+          Callback(Occurrence);
+        }
+      }
+    }
+  }
 }
 
 std::shared_ptr<std::vector<const Symbol *>>
 getSymbolsFromSlab(SymbolSlab Slab) {
   struct Snapshot {
     SymbolSlab Slab;
     std::vector<const Symbol *> Pointers;
   };

clangd/index/Merge.cpp

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.");
+    // Store all seen files by dynamic index.
+    llvm::DenseSet<llvm::StringRef> SeenFiles;
+    // Emit all dynamic occurrences, and static occurrences that are not
+    // in seen files.
+    // FIXME: If a file has been updated, and there are no occurrences indexed
+    // in dynamic index, stale results are still returned (from static index).
+    Dynamic->findOccurrences(Req, [&](const SymbolOccurrence &O) {
+      SeenFiles.insert(O.Location.FileURI);
+      Callback(O);
+    });
+    Static->findOccurrences(Req, [&](const SymbolOccurrence& O) {
+      if (llvm::is_contained(SeenFiles, O.Location.FileURI))
+        return;
+      Callback(O);
+    });
   }
 
 private:
   const SymbolIndex *Dynamic, *Static;
 };
 } // namespace
 
 Symbol
 mergeSymbol(const Symbol &L, const Symbol &R, Symbol::Details *Scratch) {
   assert(L.ID == R.ID);

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 "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;
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
+  auto AST =
+      TestTU::withAllCode(HeaderCode, MainCode.code(), "test.cc").build();
+  Index.update("test.cc", &AST.getASTContext(), AST.getPreprocessorPtr(),
+               AST.getLocalTopLevelDecls());
+  // Add test2.cc
+  AST = TestTU::withAllCode(HeaderCode, MainCode.code(), "test2.cc").build();
+  Index.update("test2.cc", &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 "TestIndex.h"
+#include "TestTU.h"
+#include "Annotations.h"
 #include "index/Index.h"
 #include "index/MemIndex.h"
 #include "index/Merge.h"
+#include "index/FileIndex.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 
 using testing::Pointee;
 using testing::UnorderedElementsAre;
+using testing::AllOf;
 
 namespace clang {
 namespace clangd {
 namespace {
 
 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"})
TEST(MergeTest, PreferSymbolWithDefn) {
 
   R.Definition.FileURI = "file:/right.cpp"; // Now right will be favored.
   M = mergeSymbol(L, R, &Scratch);
   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.
+  Annotations Test1Code(R"(class $Foo[[Foo]];)");
+  auto AST =
+      TestTU::withAllCode(HeaderCode, Test1Code.code(), "test.cc").build();
+  Dyn.update("test.cc", &AST.getASTContext(), AST.getPreprocessorPtr(),
+             AST.getLocalTopLevelDecls());
+
+  // Build static index.
+  auto StaticAST =
+      TestTU::withAllCode(HeaderCode, "// static\nclass Foo {};", "test.cc")
+          .build();
+  // Add stale occurrences for test.cc.
+  StaticIndex.update("test.cc", &StaticAST.getASTContext(),
+                     StaticAST.getPreprocessorPtr(),
+                     StaticAST.getLocalTopLevelDecls());
+
+  // Add occcurrences for test2.cc
+  Annotations Test2Code(R"(class $Foo[[Foo]] {};)");
+  StaticAST =
+      TestTU::withAllCode(HeaderCode, Test2Code.code(), "test2.cc").build();
+  StaticIndex.update("test2.cc", &StaticAST.getASTContext(),
+                     StaticAST.getPreprocessorPtr(),
+                     StaticAST.getLocalTopLevelDecls());
+
+  OccurrencesRequest Request;
+  Request.IDs = {Foo.ID};
+  Request.Filter = SymbolOccurrenceKind::Declaration |
+                   SymbolOccurrenceKind::Definition |
+                   SymbolOccurrenceKind::Reference;
+  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.h

struct TestTU {
   }
 
   static TestTU withHeaderCode(llvm::StringRef HeaderCode) {
     TestTU TU;
     TU.HeaderCode = HeaderCode;
     return TU;
   }
 
+  static TestTU withAllCode(llvm::StringRef HeaderCode, llvm::StringRef Code,
+                            llvm::StringRef Filename = "") {
+    TestTU TU;
+    TU.HeaderCode = HeaderCode;
+    TU.Code = Code;
+    if (!Filename.empty())
+     TU.Filename = Filename;
+    return TU;
+  }
+
   // The code to be compiled.
   std::string Code;
   std::string Filename = "TestTU.cpp";
 
   // Define contents of a header which will be implicitly included by Code.
   std::string HeaderCode;
   std::string HeaderFilename = "TestTU.h";
 

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());
 }
 
 const Symbol &findSymbol(const SymbolSlab &Slab, llvm::StringRef QName) {
   const Symbol *Result = nullptr;