[clangd] Collect symbol occurrences in SymbolCollector

[hokein]

Summary

SymbolCollector will be used for two cases:

• collect Symbol type only, used for indexing preamble AST.
• collect Symbol and SymbolOccurrences, used for indexing main AST.

For finding local references from the AST, we will implement it in other ways.

[ioeric]

2 high-level questions:

1. What's the reason for having a separate SymbolOccurrenceSlab? Could store occurrences as extra payload of Symbol?

2. Could we merge SymbolOccurrenceCollector into the existing SymbolCollector? They look a lot alike. Having another index data consumer seems like more overhead on the user side.

[hokein]

In D50385#1191914, @ioeric wrote:

2 high-level questions:

1. What's the reason for having a separate SymbolOccurrenceSlab? Could store occurrences as extra payload of Symbol?

Storing occurrences in Symbol structure is easy to misuse by users IMO -- if we go through this way, we will end up having a getOccurrences-like method in Symbol structure. Once users get the Symbol instance, it is natural for them to call getOccurrences to get all occurrences of the symbol. However this getOccurrences method doesn't do what users expected (just returning an incomplete set of results or empty). To query the symbol occurrences, we should always use index interface.

Therefore, I think we should try to avoid these confusions in the design.

2. Could we merge SymbolOccurrenceCollector into the existing SymbolCollector? They look a lot alike. Having another index data consumer seems like more overhead on the user side.

The SymbolOccurrenceCollector has many responsibilities (collecting declaration, definition, code completion information etc), and the code is growing complex now. Merging the SymbolOccurrenceCollector to it will make it more complicated -- we will introduce more option flags like collect-symbol-only, collect-occurrence-only to configure it for our different use cases (we need to the implementation detail clearly in order to make a correct option for SymbolCollector). And I can foresee these two collectors might be run at different point (runWithPreamble vs runWithAST) in dynamic index.

They might use same facilities, but we could always share them.

[ioeric]

In D50385#1193545, @hokein wrote:

In D50385#1191914, @ioeric wrote:

2 high-level questions:

1. What's the reason for having a separate SymbolOccurrenceSlab? Could store occurrences as extra payload of Symbol?

Storing occurrences in Symbol structure is easy to misuse by users IMO -- if we go through this way, we will end up having a getOccurrences-like method in Symbol structure. Once users get the Symbol instance, it is natural for them to call getOccurrences to get all occurrences of the symbol. However this getOccurrences method doesn't do what users expected (just returning an incomplete set of results or empty). To query the symbol occurrences, we should always use index interface.

Therefore, I think we should try to avoid these confusions in the design.

Hmm, I think this is the same for other symbol payload e.g. definition can be missing for a symbol. And it seems to me that the concern is on the SymbolSlab level: if a slab is for a single TU, users should expect missing information; if a slab is merged from all TUs, then users can expect "complete" information. I think it's reasonable to assume that users of SymbolSlab are aware of this. I think it's probably not worth the overhead of maintaining and using two separate slabs.

2. Could we merge SymbolOccurrenceCollector into the existing SymbolCollector? They look a lot alike. Having another index data consumer seems like more overhead on the user side.

The SymbolOccurrenceCollector has many responsibilities (collecting declaration, definition, code completion information etc), and the code is growing complex now. Merging the SymbolOccurrenceCollector to it will make it more

Although the existing SymbolCollector supports different options, I think it still has a pretty well-defined responsibility: gather information about symbols. IMO, cross-reference is one of the property of symbol, and I don't see strong reasons to keep them separated.

complicated -- we will introduce more option flags like collect-symbol-only, collect-occurrence-only to configure it for our different use cases (we need to the implementation detail clearly in order to make a correct option for SymbolCollector).

I think these options are reasonable if they turn out to be necessary. And making the SymbolCollector more complicated doesn't seem to be a problem if we are indeed doing more complicated work, but I don't think this would turn into a big problem as logic of xrefs seems pretty isolated. Conversely, I think implementing xrefs in a separate class would likely to cause more duplicate and maintenance, e.g. two sets of options, two sets of initializations or life-time tracking of collectors (they look a lot alike), the same boilerplate factory code in tests, passing around two collectors in user code.

And I can foresee these two collectors might be run at different point (runWithPreamble vs runWithAST) in dynamic index.

With some options, this should be a problem I think?

[ilya-biryukov]

In D50385#1193600, @ioeric wrote:

In D50385#1193545, @hokein wrote:

In D50385#1191914, @ioeric wrote:

2 high-level questions:

1. What's the reason for having a separate SymbolOccurrenceSlab? Could store occurrences as extra payload of Symbol?

Storing occurrences in Symbol structure is easy to misuse by users IMO -- if we go through this way, we will end up having a getOccurrences-like method in Symbol structure. Once users get the Symbol instance, it is natural for them to call getOccurrences to get all occurrences of the symbol. However this getOccurrences method doesn't do what users expected (just returning an incomplete set of results or empty). To query the symbol occurrences, we should always use index interface.

Therefore, I think we should try to avoid these confusions in the design.

Hmm, I think this is the same for other symbol payload e.g. definition can be missing for a symbol. And it seems to me that the concern is on the SymbolSlab level: if a slab is for a single TU, users should expect missing information; if a slab is merged from all TUs, then users can expect "complete" information. I think it's reasonable to assume that users of SymbolSlab are aware of this. I think it's probably not worth the overhead of maintaining and using two separate slabs.

I think it's reasonable to keep occurrences away from Symbol's Detail field. Stashing them together is only fine for the collector API, having any way to directly access occurrences through Symbol will be totally confusing for all the other users.
E.g., the Index::lookup() will not provide occurrences in the Symbol instances it returns, and if the accessors for those will be there it will only add confusion. So +1 to keeping them out of the Symbol class.

On the other hand, SymbolSlab feels like a perfectly reasonable place to store the occurrences in addition to the symbols themselves and it feels we should reuse its memory arena for storing any strings we need to allocate, etc.

2. Could we merge SymbolOccurrenceCollector into the existing SymbolCollector? They look a lot alike. Having another index data consumer seems like more overhead on the user side.

The SymbolOccurrenceCollector has many responsibilities (collecting declaration, definition, code completion information etc), and the code is growing complex now. Merging the SymbolOccurrenceCollector to it will make it more

Although the existing SymbolCollector supports different options, I think it still has a pretty well-defined responsibility: gather information about symbols. IMO, cross-reference is one of the property of symbol, and I don't see strong reasons to keep them separated.

complicated -- we will introduce more option flags like collect-symbol-only, collect-occurrence-only to configure it for our different use cases (we need to the implementation detail clearly in order to make a correct option for SymbolCollector).

I think these options are reasonable if they turn out to be necessary. And making the SymbolCollector more complicated doesn't seem to be a problem if we are indeed doing more complicated work, but I don't think this would turn into a big problem as logic of xrefs seems pretty isolated. Conversely, I think implementing xrefs in a separate class would likely to cause more duplicate and maintenance, e.g. two sets of options, two sets of initializations or life-time tracking of collectors (they look a lot alike), the same boilerplate factory code in tests, passing around two collectors in user code.

And I can foresee these two collectors might be run at different point (runWithPreamble vs runWithAST) in dynamic index.

With some options, this should be a problem I think?

+1 to merging into the SymbolCollector. Keeping the responsibilities separate inside a single class should be easy, e.g. something like that should be simple enough:

SymbolCollector::handleDeclOccurence(args) {
  this->processForSymbol(args); // handles keeping the Symbol structure up-to-date, i.e. adds definition locations, etc.
  this->processForOccurrences(args); // appends occurrences to a list of xrefs.
};

The main advantage that we get is less clang-specific boilerplate. The less IndexDataConsumers, FrontendActionFactorys, FrontendActions we create, the more focused and concise our code is.
And in that case, SymbolCollector is already handling those responsibilities for us and reusing looks like a good idea.

[hokein]

Hmm, I think this is the same for other symbol payload e.g. definition can be missing for a symbol. And it seems to me that the concern is on the SymbolSlab level: if a slab is for a single TU, users should expect missing information; if a slab is merged from all TUs, then users can expect "complete" information. I think it's reasonable to assume that users of SymbolSlab are aware of this. I think it's probably not worth the overhead of maintaining and using two separate slabs.

My concerns of storing occurrences as an extra payload of Symbol are:

• SymbolSlab is more like an implementation detail. Users of SymbolIndex are not aware of it, they only get Symbol objects, so it easily confuses users if they see any occurrence-related interface/member in Symbol. And we will write a looong comment explaining its correct behavior. It'd be better if we avoid this confusion in the API level.
• The fields in Symbol structure are symbol properties, and could be stored in memory. However, occurrences are not, we can't guarantee that.
• It seems that we are coupling ID, Symbol, SymbolOccurrence together: in the index implementation, we will go through ID=>Symbol=>Occurrences rather than ID=>Occurrences.

I think these options are reasonable if they turn out to be necessary.

I think they are necessary. For collecting all occurrences for local symbols from the AST, we only need symbol occurrence information, other information (e.g. declaration&definition location, #include) should be discarded; Index for code completion should not collect symbol occurrences.

And making the SymbolCollector more complicated doesn't seem to be a problem if we are indeed doing more complicated work, but I don't think this would turn into a big problem as logic of xrefs seems pretty isolated.

If xrefs is quite isolated, I think it is a good signal to have a dedicated class handling it.

I think implementing xrefs in a separate class would likely to cause more duplicate and maintenance, e.g. two sets of options, two sets of initializations or life-time tracking of collectors (they look a lot alike), the same boilerplate factory code in tests, passing around two collectors in user code.

Merging xrefs to SymbolCollector couldn't avoid these problems, I think it is a matter of where we put these code:

• different initialization of SymbolCollector for different use cases (e.g. setting different flags in SymbolCollectorOptions).
• for dynamic index, index for xrefs and code completion would be triggered at different point: index for xrefs should happen when AST is ready; index for code completion happens when Preamble is ready; we might end up with two slabs instances in the dynamic index (1 symbol slab + 1 occurrence slab vs. 2 symbol slabs).

The duplication is mainly about AST frontend action boilerplate code. To eliminate it, we could do some refactorings:

• get rid of the clang ast action code in SymbolCollector, and SymbolOccurrenceCollector
• introduce an IndexSymbol which is a subclass index::IndexDataConsumer
• the IndexSymbol has two mode (indexing symbol or indexing occurrence), and dispatch ast information to SymbolCollector/SymbolOccurrenceCollector.

[hokein]

Update the patch based on our offline discussion

• only one single clang intefaces implementation, and move finding references to current symbol collector;
• store references in SymbolSlab;

[ilya-biryukov]

clangd/index/Index.cpp

140

NIT: remove the lambda? using < is the default.

146

NIT: remove the lambda? Using == is the default.

159

Is this used for debugging?
In that case maybe consider having a user-readable representation instead of the number?

clangd/index/Index.h

46

NIT: having friend decls inside the classes themselves might prove to be more readable.
Not opposed to the current one too, feel free to ignore.

292

Maybe add a comment or remove the empty line?

293

Any store occurences in a file-centric manner?
E.g.

/// Occurences inside a single file.
class FileOccurences {
  StringRef File;
  vector<pair<Point, OccurenceKind>> Locations;
};

// ....
DenseMap<SymbolID, vector<FileOccurences>>  SymbolOccurences;

As discussed previously, this representation is better suited for both merging and serialization.

clangd/index/SymbolCollector.cpp

272 ↗

(On Diff #161927)

NIT: maybe use early exits and inverted conditions to keep the nesting down?

321 ↗

(On Diff #161927)

If we any Options here, why have an extra CollectorSymbolOptions?

clangd/index/SymbolCollector.h

59 ↗

(On Diff #161927)

Could you elaborate on what this option will be used for? How do we know in advance which symbols we're interested in?

[hokein]

Address review comments.

[hokein]

Add one more comment.

[hokein]

clangd/index/Index.h

46

These operator implementations seem not as much interesting as members in the structure, putting them to the structure probably adds some noise to readers.

293

The file-centric manner doesn't seem to suite our current model: whenever we update the index for the main AST, we just replace the symbol slab with the new one; and for index merging, we only use the index findOccurrences interfaces.

It would save some memory usage of StringRef File, but AFAIK, the memory usage of current model is relatively small (comparing with the SymbolSlab for code completion) since we only store occurrences in main file (~50KB for CodeComplete.cpp).

I'd leave it as it is now, and we could revisit it later.

clangd/index/SymbolCollector.h

59 ↗

(On Diff #161927)

This is used for finding references in the AST as a part of the xref implementation, basically the workflow would be:

1. find SymbolIDs of the symbol under the cursor, using DeclarationAndMacrosFinder
2. run symbol collector to find all occurrences in the main AST with all SymbolIDs in #1
3. query the index, to get more occurrences
4. merge them

[ilya-biryukov]

clangd/index/Index.h

46

Ok, LG outside too

293

Isn't the merging model different for the occurrences?
We would actually have to drop all references from the older index when merging if the new one contains locations in the same file.
If the merge if file-centric, the file-based representation makes more sense in the first place.

Apart from simpler merging the code, the file-based representation also buys us more efficient serialization for the static index, arguably efficient enough to stash all the occurrences even into our YAML index.

Postponing till later is also fine, but I'm not sure it buys us much now. These arguments only apply if we think the file-centric approach is a the right final design, though.

clangd/index/SymbolCollector.h

59 ↗

(On Diff #161927)

Can we instead find all the occurences in DeclarationAndMacrosFinder directly?
Extra run of SymbolCollector means another AST traversal, which is slow by itself, and SymbolCollector s designed for a much more hairy problem, its interface is just not nicely suited for things like only occurrences. The latter seems to be a simpler problem, and we can have a simpler interface to solve it (possibly shared between SymbolCollector and DeclarationAndMacrosFinder). WDYT?

[ioeric]

clangd/index/SymbolCollector.h

74 ↗

(On Diff #161972)

Use llvm::Optional?

[ioeric]

clangd/index/SymbolCollector.cpp

241 ↗

(On Diff #161972)

I don't see a strong reason for the separation of CollectOccurrence and CollectSymbol. There are some pieceis that are only used by one of them, but they seem cheap enough to ignore? Intuitively, it seems to me reference collection could just be a member function of SymbolCollector.

[sammccall]

clangd/index/Index.h

269

As discussed offline: the merge of occurrences into SymbolSlab seems problematic to me.
On the consumer side, we have a separation between Symbol APIs and SymbolOccurrence APIs - they don't really interact.
The Symbol type can often only be used with SymbolSlab, and so including occurrences drags them into the mess for consumers that don't care about them.

For producers (index implementations), they will usually have both and they may want to share arena storage. But this probably doesn't matter much, and if it does we can use another mechanism (like allowing SymbolSlabBuilder and SymbolOccurrenceSlab to share UniqueStringSaver)

clangd/index/SymbolCollector.cpp

439 ↗

(On Diff #161972)

note that here we've done basically all the work needed to record the occurrence.

If you add a DenseMap<Decl*, {SourceLocation, SymbolRole}> then you'll have enough info at the end to fill in the occurrences, like we do with referenceddecls -> references.

clangd/index/SymbolCollector.h

59 ↗

(On Diff #161927)

Yeah, I don't think we need this.
For "find references in the AST" we have an implementation in XRefs for highlights which we don't need to share.

40 ↗

(On Diff #161972)

Not sure this split is justified.
if IDs goes away (see below), all that's left can be represented in a SymbolOccurenceKind filter (which is 0 to collect no occurrences)

76 ↗

(On Diff #161972)

collecting symbols doesn't actually need to be optional I think - it's the core responsibility of this class, and "find occurrences of a decl in an ast" can be implemented more easily in other ways

[hokein]

Update the patch based on our new discussion

• SymbolOccurrenceSlab for storing underlying occurrence data
• reuse SymbolCollector to collect symbol occurrences

[sammccall]

This looks pretty good!

clangd/index/Index.h

375

assert frozen?
looking up in a non-frozen array is probably a mistake.
if we choose to optimize this, it probably won't be possible.

376

return Occurrences.lookup(ID)?

clangd/index/SymbolCollector.cpp

227 ↗

(On Diff #162611)

nit: toOccurrenceKind

229 ↗

(On Diff #162611)

If you want to filter out the unsupported bits, maybe just add an explicit AllOccurrenceKinds constant to the header file, and return AllOccurrenceKinds & Roles here? (plus casts)

328 ↗

(On Diff #162611)

just compute the spelling loc once and reuse?

329 ↗

(On Diff #162611)

you get the spelling loc on the previous line to check for mainfile - so surely we should be using spelling loc here?

455 ↗

(On Diff #162611)

nit: const auto& for clarity since we're not mutating

460 ↗

(On Diff #162611)

so this seems maybe gratuitously inefficient, we're copying the filename then going through the URI conversion dance for each reference - even though the filename is the same for each.

consider splitting out part of getTokenLocation into getTokenRange(SymbolLocation&) and only calling that here.

clangd/index/SymbolCollector.h

54 ↗

(On Diff #162611)

this should be next to OccurrenceFilter, they're very closely related (the name mismatch is a little unfortunate)

121 ↗

(On Diff #162611)

please move next to ReferencedDecls/ReferencedMacros so the comment applies to this too

unittests/clangd/SymbolCollectorTests.cpp

466 ↗

(On Diff #162611)

this is cute - if possible, consider adding a matcher factory function for readability here, so you can write EXPECT_THAT(..., HaveRanges(Main.ranges("foo"))

[hokein]

Address review comments and fix code style.

[hokein]

Minor cleanup.

[hokein]

clangd/index/Index.h

376

The DenseMap::lookup returns a copy of Value (vector) which doesn't suit our use case :( -- we will return an ArrayRef which stores an reference of a local vector object.

unittests/clangd/SymbolCollectorTests.cpp

466 ↗

(On Diff #162611)

Wrapped this into HaveRanges.

[hokein]

Address review comments in D51279.

[hokein]

Committed in rL341208.

clangd/CMakeLists.txt

add_clang_library(clangDaemon
   XRefs.cpp
 
   index/CanonicalIncludes.cpp
   index/FileIndex.cpp
   index/Index.cpp
   index/MemIndex.cpp
   index/Merge.cpp
   index/SymbolCollector.cpp
+  index/SymbolOccurrenceCollector.cpp
   index/SymbolYAML.cpp
 
   index/dex/Iterator.cpp
   index/dex/Trigram.cpp
 
   LINK_LIBS
   clangAST
   clangASTMatchers

clangd/index/Index.h

 
 #include "clang/Index/IndexSymbol.h"
 #include "clang/Lex/Lexer.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/Hashing.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/StringExtras.h"
+#include "llvm/Support/StringSaver.h"
 #include <array>
 #include <string>
 
 namespace clang {
 namespace clangd {
 
 struct SymbolLocation {
   // Specify a position (Line, Column) of symbol. Using Line/Column allows us to
   // build LSP responses without reading the file content.
   struct Position {
     uint32_t Line = 0; // 0-based
     // Using UTF-16 code units.
     uint32_t Column = 0; // 0-based
-    bool operator==(const Position& P) const {
-      return Line == P.Line && Column == P.Column;
-    }
   };
 
   // The URI of the source file where a symbol occurs.
   llvm::StringRef FileURI;
 
   /// The symbol range, using half-open range [Start, End).
   Position Start;
   Position End;
 
   explicit operator bool() const { return !FileURI.empty(); }
-  bool operator==(const SymbolLocation& Loc) const {
-    return std::tie(FileURI, Start, End) ==
-           std::tie(Loc.FileURI, Loc.Start, Loc.End);
-  }
 };
+inline bool operator==(const SymbolLocation::Position &L,
+                       const SymbolLocation::Position &R){
+  return std::tie(L.Line, L.Column) == std::tie(R.Line, R.Column);
+}
+inline bool operator<(const SymbolLocation::Position &L,
+                      const SymbolLocation::Position &R){
+  return std::tie(L.Line, L.Column) < std::tie(R.Line, R.Column);
+}
+inline bool operator==(const SymbolLocation&L,
+                       const SymbolLocation&R){
+  return std::tie(L.FileURI, L.Start, L.End) ==
+         std::tie(R.FileURI, R.Start, R.End);
+}
+inline bool operator<(const SymbolLocation&L,
+                      const SymbolLocation&R){
+  return std::tie(L.FileURI, L.Start, L.End) <
+         std::tie(R.FileURI, R.Start, R.End);
+}
 llvm::raw_ostream &operator<<(llvm::raw_ostream &, const SymbolLocation &);
 
 // The class identifies a particular C++ symbol (class, function, method, etc).
 //
 // As USRs (Unified Symbol Resolution) could be large, especially for functions
 // with long type arguments, SymbolID is using 160-bits SHA1(USR) values to
 // guarantee the uniqueness of symbols while using a relatively small amount of
 // memory (vs storing USRs directly).
public:
   size_t size() const { return Symbols.size(); }
   // Estimates the total memory usage.
   size_t bytes() const {
     return sizeof(*this) + Arena.getTotalMemory() +
            Symbols.capacity() * sizeof(Symbol);
   }
 
   // SymbolSlab::Builder is a mutable container that can 'freeze' to SymbolSlab.
   // The frozen SymbolSlab will use less memory.
   class Builder {
   public:
     // Adds a symbol, overwriting any existing one with the same ID.
     // This is a deep copy: underlying strings will be owned by the slab.
     void insert(const Symbol &S);
 
     // Returns the symbol with an ID, if it exists. Valid until next insert().
     const Symbol *find(const SymbolID &ID) {
       auto I = SymbolIndex.find(ID);
       return I == SymbolIndex.end() ? nullptr : &Symbols[I->second];
     }
 
     // Consumes the builder to finalize the slab.
     SymbolSlab build() &&;
 
   private:
     llvm::BumpPtrAllocator Arena;
     // Intern table for strings. Contents are on the arena.
     llvm::DenseSet<llvm::StringRef> Strings;
     std::vector<Symbol> Symbols;
     // Values are indices into Symbols vector.
     llvm::DenseMap<SymbolID, size_t> SymbolIndex;
   };
 
 private:
   SymbolSlab(llvm::BumpPtrAllocator Arena, std::vector<Symbol> Symbols)
       : Arena(std::move(Arena)), Symbols(std::move(Symbols)) {}
 
   llvm::BumpPtrAllocator Arena; // Owns Symbol data that the Symbols do not.
   std::vector<Symbol> Symbols;  // Sorted by SymbolID to allow lookup.
 };
 
inline SymbolOccurrenceKind &operator|=(SymbolOccurrenceKind &L,
                                         SymbolOccurrenceKind R) {
   return L = L | R;
 }
 inline SymbolOccurrenceKind operator&(SymbolOccurrenceKind A,
                                       SymbolOccurrenceKind B) {
   return static_cast<SymbolOccurrenceKind>(static_cast<uint8_t>(A) &
                                            static_cast<uint8_t>(B));
 }
+raw_ostream &operator<<(raw_ostream &OS, SymbolOccurrenceKind K);
 
 // Represents a symbol occurrence in the source file. It could be a
 // declaration/definition/reference occurrence.
 //
 // WARNING: Location does not own the underlying data - Copies are shallow.
 struct SymbolOccurrence {
   // The location of the occurrence.
   SymbolLocation Location;
   SymbolOccurrenceKind Kind = SymbolOccurrenceKind::Unknown;
 };
+inline bool operator<(const SymbolOccurrence &L, const SymbolOccurrence &R) {
+  return std::tie(L.Location, L.Kind) < std::tie(R.Location, R.Kind);
+}
+inline bool operator==(const SymbolOccurrence &L, const SymbolOccurrence &R) {
+  return std::tie(L.Location, L.Kind) == std::tie(R.Location, R.Kind);
+}
+llvm::raw_ostream &operator<<(llvm::raw_ostream &OS,
+                              const SymbolOccurrence &Occurrence);
+
+// An efficient structure of storing large set of symbol occurrences in memory.
+// Filenames are deduplicated.
+class SymbolOccurrenceSlab {
+ public:
+   using const_iterator =
+       llvm::DenseMap<SymbolID, std::vector<SymbolOccurrence>>::const_iterator;
+   using iterator = const_iterator;
+
+   SymbolOccurrenceSlab() : UniqueStrings(Arena) {}
+
+   const_iterator begin() const { return Occurrences.begin(); }
+   const_iterator end() const { return Occurrences.end(); }
+
+   // 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 {};
+     return It->second;
+  }
+
+  void freeze();
+
+private:
+  bool Frozen = false;
+  llvm::BumpPtrAllocator Arena;
+  llvm::UniqueStringSaver UniqueStrings;
+  llvm::DenseMap<SymbolID, std::vector<SymbolOccurrence>> Occurrences;
+};
 
 struct FuzzyFindRequest {
   /// \brief A query string for the fuzzy find. This is matched against symbols'
   /// un-qualified identifiers and should not contain qualifiers like "::".
   std::string Query;
   /// \brief If this is non-empty, symbols must be in at least one of the scopes
   /// (e.g. namespaces) excluding nested scopes. For example, if a scope "xyz::"
   /// is provided, the matched symbols must be defined in namespace xyz but not

clangd/index/Index.cpp

SymbolSlab SymbolSlab::Builder::build() && {
   // We may have unused strings from overwritten symbols. Build a new arena.
   BumpPtrAllocator NewArena;
   DenseSet<StringRef> Strings;
   for (auto &S : Symbols)
     own(S, Strings, NewArena);
   return SymbolSlab(std::move(NewArena), std::move(Symbols));
 }
 
+raw_ostream &operator<<(raw_ostream &OS, SymbolOccurrenceKind K) {
+  OS << static_cast<unsigned>(K);
+  return OS;
+}
+
+llvm::raw_ostream &operator<<(llvm::raw_ostream &OS,
+                              const SymbolOccurrence &Occurrence) {
+  OS << Occurrence.Location << ":" << Occurrence.Kind;
+  return OS;
+}
+
+void SymbolOccurrenceSlab::insert(const SymbolID &SymID,
+                                  const SymbolOccurrence &Occurrence) {
+  assert(!Frozen &&
+         "Can't insert a symbol occurrence after the slab has been frozen!");
+  auto& SymOccurrences = Occurrences[SymID];
+  SymOccurrences.push_back(Occurrence);
+  SymOccurrences.back().Location.FileURI =
+      UniqueStrings.save(Occurrence.Location.FileURI);
+}
+
+void SymbolOccurrenceSlab::freeze() {
+  // We may have duplicated symbol occurrences. Deduplicate them.
+  for (auto &IDAndOccurrence : Occurrences) {
+    auto &Occurrence = IDAndOccurrence.getSecond();
+    std::sort(Occurrence.begin(), Occurrence.end(),
+              [](const SymbolOccurrence &L, const SymbolOccurrence &R) {
+                return L < R;
+              });
+    Occurrence.erase(
+        std::unique(Occurrence.begin(), Occurrence.end(),
+                    [](const SymbolOccurrence &L, const SymbolOccurrence &R) {
+                      return L == R;
+                    }),
+        Occurrence.end());
+  }
+  Frozen = true;
+}
+
 } // namespace clangd
 } // namespace clang

clangd/index/SymbolOccurrenceCollector.h

+#ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_SYMBOL_OCCURRENCE_COLLECTOR_H
+#define LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_SYMBOL_OCCURRENCE_COLLECTOR_H
+
+#include "Index.h"
+#include "llvm/ADT/DenseSet.h"
+#include "clang/AST/ASTContext.h"
+#include "clang/AST/Decl.h"
+#include "clang/Index/IndexDataConsumer.h"
+#include "clang/Index/IndexSymbol.h"
+
+namespace clang {
+namespace clangd {
+
+// A collector collects all symbol occurrences from the AST.
+class SymbolOccurrenceCollector: public index::IndexDataConsumer {
+public:
+  struct Options {
+    // The symbol occurrence kind that will be collected.
+    SymbolOccurrenceKind Filter;
+    // A whitelist symbols which will be collected.
+    // If none, all symbol occurrences will be collected.
+    llvm::Optional<llvm::DenseSet<SymbolID>> IDs = llvm::None;
+  };
+
+  SymbolOccurrenceCollector(Options Opts)
+      : Opts(Opts), Occurrences(new SymbolOccurrenceSlab()) {}
+
+  void initialize(ASTContext &Ctx) override {
+    ASTCtx = &Ctx;
+  }
+
+  bool
+  handleDeclOccurence(const Decl *D, index::SymbolRoleSet Roles,
+                      ArrayRef<index::SymbolRelation> Relations,
+                      SourceLocation Loc,
+                      index::IndexDataConsumer::ASTNodeInfo ASTNode) override;
+
+  void finish() override {
+    Occurrences->freeze();
+  }
+
+  std::unique_ptr<SymbolOccurrenceSlab> takeOccurrences() {
+    return std::move(Occurrences);
+  }
+
+private:
+  ASTContext *ASTCtx;
+  Options Opts;
+  std::unique_ptr<SymbolOccurrenceSlab> Occurrences;
+};
+
+} // namespace clangd
+} // namespace clang
+
+#endif // LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_SYMBOL_OCCURRENCE_COLLECTOR_H

clangd/index/SymbolOccurrenceCollector.cpp

+//===--- SymbolOccurrenceCollector.cpp ---------------------------*- C++-*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#include "SymbolOccurrenceCollector.h"
+#include "clang/Index/USRGeneration.h"
+#include "../AST.h"
+#include "../SourceCode.h"
+#include "llvm/ADT/STLExtras.h"
+
+namespace clang {
+namespace clangd {
+namespace {
+llvm::Optional<SymbolLocation>
+getTokenLocation(SourceLocation TokLoc, const ASTContext* ASTCtx,
+                 std::string &FileURIStorage) {
+  const auto& SM = ASTCtx->getSourceManager();
+  auto TokenLength = clang::Lexer::MeasureTokenLength(TokLoc, SM,
+                                                      ASTCtx->getLangOpts());
+
+  auto CreatePosition = [&SM](SourceLocation Loc) {
+    auto LSPLoc = sourceLocToPosition(SM, Loc);
+    SymbolLocation::Position Pos;
+    Pos.Line = LSPLoc.line;
+    Pos.Column = LSPLoc.character;
+    return Pos;
+  };
+
+  SymbolLocation Result;
+  Result.Start = CreatePosition(TokLoc);
+  auto EndLoc = TokLoc.getLocWithOffset(TokenLength);
+  Result.End = CreatePosition(EndLoc);
+
+  const auto* F = SM.getFileEntryForID(SM.getFileID(TokLoc));
+  if (!F)
+    return llvm::None;
+  auto FilePath = getAbsoluteFilePath(F, SM);
+  if (!FilePath)
+    return llvm::None;
+  // FIXME: support custom URIs.
+  FileURIStorage = URI::createFile(*FilePath).toString();
+  Result.FileURI = FileURIStorage;
+  return Result;
+}
+
+SymbolOccurrenceKind ToOccurrenceKind(index::SymbolRoleSet Roles) {
+  SymbolOccurrenceKind Kind;
+  for (auto Mask : {SymbolOccurrenceKind::Declaration,
+                    SymbolOccurrenceKind::Definition,
+                    SymbolOccurrenceKind::Reference}) {
+    if (Roles & static_cast<unsigned>(Mask))
+      Kind |= Mask;
+  }
+  return Kind;
+}
+} // namespace
+
+bool SymbolOccurrenceCollector::handleDeclOccurence(
+    const Decl *D, index::SymbolRoleSet Roles,
+    ArrayRef<index::SymbolRelation> Relations, SourceLocation Loc,
+    index::IndexDataConsumer::ASTNodeInfo ASTNode) {
+  if (D->isImplicit())
+    return true;
+
+  std::string FileURI;
+  auto AddOccurrence = [&](SourceLocation L, const SymbolID& ID) {
+    if (auto Location = getTokenLocation(Loc, ASTCtx, FileURI)) {
+      SymbolOccurrence Occurrence;
+      Occurrence.Location = *Location;
+      Occurrence.Kind = ToOccurrenceKind(Roles);
+      Occurrences->insert(ID, Occurrence);
+    }
+  };
+  if (static_cast<unsigned>(Opts.Filter) & Roles) {
+    if (auto ID = getSymbolID(D)) {
+      if (!Opts.IDs || llvm::is_contained(*Opts.IDs, *ID))
+        AddOccurrence(Loc, *ID);
+    }
+  }
+  return true;
+}
+
+} // namespace clangd
+} // namespace clang

unittests/clangd/CMakeLists.txt

add_extra_unittest(ClangdTests
   FindSymbolsTests.cpp
   FuzzyMatchTests.cpp
   GlobalCompilationDatabaseTests.cpp
   HeadersTests.cpp
   IndexTests.cpp
   QualityTests.cpp
   SourceCodeTests.cpp
   SymbolCollectorTests.cpp
+  SymbolOccurrenceCollectorTests.cpp
   SyncAPI.cpp
   TUSchedulerTests.cpp
   TestFS.cpp
   TestTU.cpp
   ThreadingTests.cpp
   TraceTests.cpp
   URITests.cpp
   XRefsTests.cpp

unittests/clangd/SymbolOccurrenceCollectorTests.cpp

+//===-- SymbolOccurrenceCollectorTests.cpp  ---------------------*- C++ -*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#include "index/SymbolOccurrenceCollector.h"
+#include "Annotations.h"
+#include "TestFS.h"
+#include "TestTU.h"
+#include "clang/Index/IndexingAction.h"
+#include "clang/Basic/FileManager.h"
+#include "clang/Basic/FileSystemOptions.h"
+#include "clang/Basic/VirtualFileSystem.h"
+#include "clang/Frontend/CompilerInstance.h"
+#include "clang/Index/IndexingAction.h"
+#include "clang/Tooling/Tooling.h"
+#include "llvm/ADT/IntrusiveRefCntPtr.h"
+#include "llvm/ADT/StringRef.h"
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+
+#include <memory>
+#include <string>
+
+MATCHER(OccurrenceRange, "") {
+  const clang::clangd::SymbolOccurrence& Pos = testing::get<0>(arg);
+  const clang::clangd::Range& Range = testing::get<1>(arg);
+  return std::tie(Pos.Location.Start.Line,
+                  Pos.Location.Start.Column,
+                  Pos.Location.End.Line,
+                  Pos.Location.End.Column) ==
+         std::tie(Range.start.line, Range.start.character, Range.end.line,
+                  Range.end.character);
+}
+
+namespace clang {
+namespace clangd {
+namespace {
+
+class SymbolIndexActionFactory : public tooling::FrontendActionFactory {
+public:
+  SymbolIndexActionFactory() = default;
+
+  clang::FrontendAction *create() override {
+    index::IndexingOptions IndexOpts;
+    IndexOpts.SystemSymbolFilter =
+        index::IndexingOptions::SystemSymbolFilterKind::All;
+    IndexOpts.IndexFunctionLocals = true;
+    SymbolOccurrenceKind Filter = SymbolOccurrenceKind::Declaration |
+                                  SymbolOccurrenceKind::Definition |
+                                  SymbolOccurrenceKind::Reference;
+    SymbolOccurrenceCollector::Options Opts;
+    Opts.Filter = Filter;
+    Opts.IDs = llvm::None;
+    Collector = std::make_shared<SymbolOccurrenceCollector>(Opts);
+    return index::createIndexingAction(Collector, IndexOpts, nullptr).release();
+  }
+
+  std::shared_ptr<SymbolOccurrenceCollector> Collector;
+};
+
+class OccurrenceCollectorTest : public ::testing::Test {
+public:
+  OccurrenceCollectorTest()
+      : InMemoryFileSystem(new vfs::InMemoryFileSystem),
+        TestHeaderName(testPath("symbol.h")),
+        TestFileName(testPath("symbol.cc")) {
+    TestHeaderURI = URI::createFile(TestHeaderName).toString();
+    TestFileURI = URI::createFile(TestFileName).toString();
+  }
+
+  bool collectOccurrences(StringRef HeaderCode, StringRef MainCode,
+                          const std::vector<std::string> &ExtraArgs = {}) {
+    llvm::IntrusiveRefCntPtr<FileManager> Files(
+        new FileManager(FileSystemOptions(), InMemoryFileSystem));
+
+    auto Factory = llvm::make_unique<SymbolIndexActionFactory>();
+
+    std::vector<std::string> Args = {
+        "symbol_occurrence_collector", "-fsyntax-only", "-xc++",
+        "-std=c++11",       "-include",      TestHeaderName};
+    Args.insert(Args.end(), ExtraArgs.begin(), ExtraArgs.end());
+    // This allows to override the "-xc++" with something else, i.e.
+    // -xobjective-c++.
+    Args.push_back(TestFileName);
+
+    tooling::ToolInvocation Invocation(
+        Args,
+        Factory->create(), Files.get(),
+        std::make_shared<PCHContainerOperations>());
+
+    InMemoryFileSystem->addFile(TestHeaderName, 0,
+                                llvm::MemoryBuffer::getMemBuffer(HeaderCode));
+    InMemoryFileSystem->addFile(TestFileName, 0,
+                                llvm::MemoryBuffer::getMemBuffer(MainCode));
+    Invocation.run();
+    Occurrences = Factory->Collector->takeOccurrences();
+    return true;
+  }
+
+protected:
+  llvm::IntrusiveRefCntPtr<vfs::InMemoryFileSystem> InMemoryFileSystem;
+  std::string TestHeaderName;
+  std::string TestHeaderURI;
+  std::string TestFileName;
+  std::string TestFileURI;
+  std::unique_ptr<SymbolOccurrenceSlab> Occurrences;
+};
+
+std::vector<Range> operator+(const std::vector<Range> &L,
+                             const std::vector<Range> &R) {
+  std::vector<Range> Result = L;
+  Result.insert(Result.end(), R.begin(), R.end());
+  return Result;
+}
+
+TEST_F(OccurrenceCollectorTest, Reference) {
+  Annotations Header(R"(
+  class $foo[[Foo]] {
+  public:
+    $foo[[Foo]]() {}
+    $foo[[Foo]](int);
+  };
+  class $bar[[Bar]];
+  void $func[[func]]();
+  )");
+  Annotations Main(R"(
+  class $bar[[Bar]] {};
+
+  void $func[[func]]();
+
+  void fff() {
+    $foo[[Foo]] foo;
+    $bar[[Bar]] bar;
+    $func[[func]]();
+    int abc = 0;
+    $foo[[Foo]] foo2 = abc;
+  }
+  )");
+  collectOccurrences(Header.code(), Main.code());
+  auto H = TestTU::withHeaderCode(Header.code());
+  auto Symbols = H.headerSymbols();
+  auto Foo = findSymbol(Symbols, "Foo");
+  auto Bar = findSymbol(Symbols, "Bar");
+  auto Func = findSymbol(Symbols, "func");
+
+  EXPECT_THAT(
+      Occurrences->find(Foo.ID),
+      testing::UnorderedPointwise(OccurrenceRange(),
+                                  Header.ranges("foo") + Main.ranges("foo")));
+  EXPECT_THAT(
+      Occurrences->find(Bar.ID),
+      testing::UnorderedPointwise(OccurrenceRange(),
+                                  Header.ranges("bar") + Main.ranges("bar")));
+
+  EXPECT_THAT(
+      Occurrences->find(Func.ID),
+      testing::UnorderedPointwise(OccurrenceRange(),
+                                  Header.ranges("func") + Main.ranges("func")));
+}
+
+} // namespace
+} // namespace clangd
+} // namespace clang