1 files changed, 64 insertions, 12 deletions

diff --git a/clang-tools-extra/clangd/index/FileIndex.cpp b/clang-tools-extra/clangd/index/FileIndex.cpp
index 48d05959dae..883faf9f398 100644
--- a/clang-tools-extra/clangd/index/FileIndex.cpp
+++ b/clang-tools-extra/clangd/index/FileIndex.cpp
@@ -16,9 +16,10 @@
 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,
+         llvm::Optional<llvm::ArrayRef<Decl *>> TopLevelDecls,
+         llvm::ArrayRef<std::string> URISchemes) {
   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
@@ -31,8 +32,6 @@ SymbolSlab indexAST(ASTContext &AST, std::shared_ptr<Preprocessor> PP,
     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 =
@@ -46,20 +45,45 @@ SymbolSlab indexAST(ASTContext &AST, std::shared_ptr<Preprocessor> PP,
     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() {
@@ -85,19 +109,47 @@ std::shared_ptr<std::vector<const Symbol *>> FileSymbols::allSymbols() {
   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(
@@ -115,7 +167,7 @@ void FileIndex::lookup(
 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 {