[clangd] Make FileIndex aware of the main file

Summary:
It was previously only indexing the preamble decls. The new
implementation will index both the preamble and the main AST and
report both sets of symbols, preferring the ones from the main AST
whenever the symbol is present in both.
The symbols in the main AST slab always store all information
available in the preamble symbols, possibly adding more,
e.g. definition locations.

Reviewers: hokein, ioeric

Reviewed By: ioeric

Subscribers: kadircet, MaskRay, jkorous, arphaman, cfe-commits

Differential Revision: https://reviews.llvm.org/D50889

llvm-svn: 340404

4 files changed, 64 insertions, 33 deletions

diff --git a/clang-tools-extra/clangd/ClangdServer.cpp b/clang-tools-extra/clangd/ClangdServer.cpp
index efaa687af37..0a3a26c82d6 100644
--- a/clang-tools-extra/clangd/ClangdServer.cpp
+++ b/clang-tools-extra/clangd/ClangdServer.cpp
@@ -65,27 +65,39 @@ public:
 
   Optional<Expected<tooling::AtomicChanges>> Result;
 };
+} // namespace
 
-class UpdateFileIndex : public ParsingCallbacks {
+/// Manages dynamic index for open files. Each file might contribute two sets
+/// of symbols to the dynamic index: symbols from the preamble and symbols
+/// from the file itself. Those have different lifetimes and we merge results
+/// from both
+class ClangdServer::DynamicIndex : public ParsingCallbacks {
 public:
-  UpdateFileIndex(FileIndex *FileIdx) : FileIdx(FileIdx) {}
+  DynamicIndex(std::vector<std::string> URISchemes)
+      : PreambleIdx(URISchemes), MainFileIdx(URISchemes),
+        MergedIndex(mergeIndex(&MainFileIdx, &PreambleIdx)) {}
+
+  SymbolIndex &index() const { return *MergedIndex; }
 
   void onPreambleAST(PathRef Path, ASTContext &Ctx,
                      std::shared_ptr<clang::Preprocessor> PP) override {
-    if (FileIdx)
-      FileIdx->update(Path, &Ctx, std::move(PP));
+    PreambleIdx.update(Path, &Ctx, PP, /*TopLevelDecls=*/llvm::None);
   }
 
   void onMainAST(PathRef Path, ParsedAST &AST) override {
-    // FIXME: merge results from the main file into the index too.
+
+    MainFileIdx.update(Path, &AST.getASTContext(), AST.getPreprocessorPtr(),
+                       AST.getLocalTopLevelDecls());
   }
 
 private:
-  FileIndex *FileIdx;
+  FileIndex PreambleIdx;
+  FileIndex MainFileIdx;
+  /// Merged view into both indexes. Merges are performed in a similar manner
+  /// to the merges of dynamic and static index.
+  std::unique_ptr<SymbolIndex> MergedIndex;
 };
 
-} // namespace
-
 ClangdServer::Options ClangdServer::optsForTest() {
   ClangdServer::Options Opts;
   Opts.UpdateDebounce = std::chrono::steady_clock::duration::zero(); // Faster!
@@ -101,9 +113,9 @@ ClangdServer::ClangdServer(GlobalCompilationDatabase &CDB,
     : CDB(CDB), DiagConsumer(DiagConsumer), FSProvider(FSProvider),
       ResourceDir(Opts.ResourceDir ? Opts.ResourceDir->str()
                                    : getStandardResourceDir()),
-      FileIdx(Opts.BuildDynamicSymbolIndex ? new FileIndex(Opts.URISchemes)
-                                           : nullptr),
-      FileIdxUpdater(llvm::make_unique<UpdateFileIndex>(FileIdx.get())),
+      DynamicIdx(Opts.BuildDynamicSymbolIndex
+                     ? new DynamicIndex(Opts.URISchemes)
+                     : nullptr),
       PCHs(std::make_shared<PCHContainerOperations>()),
       // Pass a callback into `WorkScheduler` to extract symbols from a newly
       // parsed file and rebuild the file index synchronously each time an AST
@@ -111,19 +123,23 @@ ClangdServer::ClangdServer(GlobalCompilationDatabase &CDB,
       // FIXME(ioeric): this can be slow and we may be able to index on less
       // critical paths.
       WorkScheduler(Opts.AsyncThreadsCount, Opts.StorePreamblesInMemory,
-                    *FileIdxUpdater, Opts.UpdateDebounce,
-                    Opts.RetentionPolicy) {
-  if (FileIdx && Opts.StaticIndex) {
-    MergedIndex = mergeIndex(FileIdx.get(), Opts.StaticIndex);
+                    DynamicIdx ? *DynamicIdx : noopParsingCallbacks(),
+                    Opts.UpdateDebounce, Opts.RetentionPolicy) {
+  if (DynamicIdx && Opts.StaticIndex) {
+    MergedIndex = mergeIndex(&DynamicIdx->index(), Opts.StaticIndex);
     Index = MergedIndex.get();
-  } else if (FileIdx)
-    Index = FileIdx.get();
+  } else if (DynamicIdx)
+    Index = &DynamicIdx->index();
   else if (Opts.StaticIndex)
     Index = Opts.StaticIndex;
   else
     Index = nullptr;
 }
 
+// Destructor has to be in .cpp file to see the definition of
+// ClangdServer::DynamicIndex.
+ClangdServer::~ClangdServer() = default;
+
 void ClangdServer::setRootPath(PathRef RootPath) {
   auto FS = FSProvider.getFileSystem();
   auto Status = FS->status(RootPath);
diff --git a/clang-tools-extra/clangd/ClangdServer.h b/clang-tools-extra/clangd/ClangdServer.h
index ebe1390ff38..5613553810a 100644
--- a/clang-tools-extra/clangd/ClangdServer.h
+++ b/clang-tools-extra/clangd/ClangdServer.h
@@ -99,6 +99,7 @@ public:
   /// synchronize access to shared state.
   ClangdServer(GlobalCompilationDatabase &CDB, FileSystemProvider &FSProvider,
                DiagnosticsConsumer &DiagConsumer, const Options &Opts);
+  ~ClangdServer();
 
   /// Set the root path of the workspace.
   void setRootPath(PathRef RootPath);
@@ -200,6 +201,7 @@ private:
   formatCode(llvm::StringRef Code, PathRef File,
              ArrayRef<tooling::Range> Ranges);
 
+  class DynamicIndex;
   typedef uint64_t DocVersion;
 
   void consumeDiagnostics(PathRef File, DocVersion Version,
@@ -217,15 +219,14 @@ private:
   Path ResourceDir;
   // The index used to look up symbols. This could be:
   //   - null (all index functionality is optional)
-  //   - the dynamic index owned by ClangdServer (FileIdx)
+  //   - the dynamic index owned by ClangdServer (DynamicIdx)
   //   - the static index passed to the constructor
   //   - a merged view of a static and dynamic index (MergedIndex)
   SymbolIndex *Index;
-  // If present, an up-to-date of symbols in open files. Read via Index.
-  std::unique_ptr<FileIndex> FileIdx;
-  /// Callbacks responsible for updating FileIdx.
-  std::unique_ptr<ParsingCallbacks> FileIdxUpdater;
-  // If present, a merged view of FileIdx and an external index. Read via Index.
+  /// If present, an up-to-date of symbols in open files. Read via Index.
+  std::unique_ptr<DynamicIndex> DynamicIdx;
+  // If present, a merged view of DynamicIdx and an external index. Read via
+  // Index.
   std::unique_ptr<SymbolIndex> MergedIndex;
   // If set, this represents the workspace path.
   llvm::Optional<std::string> RootPath;
diff --git a/clang-tools-extra/clangd/index/FileIndex.cpp b/clang-tools-extra/clangd/index/FileIndex.cpp
index 8ee5018eda0..3d8e17bfb2f 100644
--- a/clang-tools-extra/clangd/index/FileIndex.cpp
+++ b/clang-tools-extra/clangd/index/FileIndex.cpp
@@ -8,8 +8,8 @@
 //===----------------------------------------------------------------------===//
 
 #include "FileIndex.h"
-#include "SymbolCollector.h"
 #include "../Logger.h"
+#include "SymbolCollector.h"
 #include "clang/Index/IndexingAction.h"
 #include "clang/Lex/Preprocessor.h"
 
@@ -17,6 +17,7 @@ namespace clang {
 namespace clangd {
 
 SymbolSlab 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
@@ -38,10 +39,14 @@ SymbolSlab indexAST(ASTContext &AST, std::shared_ptr<Preprocessor> PP,
       index::IndexingOptions::SystemSymbolFilterKind::DeclarationsOnly;
   IndexOpts.IndexFunctionLocals = false;
 
-  std::vector<const Decl *> TopLevelDecls(
-      AST.getTranslationUnitDecl()->decls().begin(),
-      AST.getTranslationUnitDecl()->decls().end());
-  index::indexTopLevelDecls(AST, TopLevelDecls, Collector, IndexOpts);
+  std::vector<Decl *> DeclsToIndex;
+  if (TopLevelDecls)
+    DeclsToIndex.assign(TopLevelDecls->begin(), TopLevelDecls->end());
+  else
+    DeclsToIndex.assign(AST.getTranslationUnitDecl()->decls().begin(),
+                        AST.getTranslationUnitDecl()->decls().end());
+
+  index::indexTopLevelDecls(AST, DeclsToIndex, Collector, IndexOpts);
 
   return Collector.takeSymbols();
 }
@@ -81,13 +86,14 @@ std::shared_ptr<std::vector<const Symbol *>> FileSymbols::allSymbols() {
 }
 
 void FileIndex::update(PathRef Path, ASTContext *AST,
-                       std::shared_ptr<Preprocessor> PP) {
+                       std::shared_ptr<Preprocessor> PP,
+                       llvm::Optional<llvm::ArrayRef<Decl *>> TopLevelDecls) {
   if (!AST) {
     FSymbols.update(Path, nullptr);
   } else {
     assert(PP);
     auto Slab = llvm::make_unique<SymbolSlab>();
-    *Slab = indexAST(*AST, PP, URISchemes);
+    *Slab = indexAST(*AST, PP, TopLevelDecls, URISchemes);
     FSymbols.update(Path, std::move(Slab));
   }
   auto Symbols = FSymbols.allSymbols();
diff --git a/clang-tools-extra/clangd/index/FileIndex.h b/clang-tools-extra/clangd/index/FileIndex.h
index 59bf58c6fef..6f24f9d2c23 100644
--- a/clang-tools-extra/clangd/index/FileIndex.h
+++ b/clang-tools-extra/clangd/index/FileIndex.h
@@ -64,7 +64,11 @@ public:
   /// nullptr, this removes all symbols in the file.
   /// If \p AST is not null, \p PP cannot be null and it should be the
   /// preprocessor that was used to build \p AST.
-  void update(PathRef Path, ASTContext *AST, std::shared_ptr<Preprocessor> PP);
+  /// If \p TopLevelDecls is set, only these decls are indexed. Otherwise, all
+  /// top level decls obtained from \p AST are indexed.
+  void
+  update(PathRef Path, ASTContext *AST, std::shared_ptr<Preprocessor> PP,
+         llvm::Optional<llvm::ArrayRef<Decl *>> TopLevelDecls = llvm::None);
 
   bool
   fuzzyFind(const FuzzyFindRequest &Req,
@@ -86,8 +90,12 @@ private:
 /// 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.
-SymbolSlab indexAST(ASTContext &AST, std::shared_ptr<Preprocessor> PP,
-                    llvm::ArrayRef<std::string> URISchemes = {});
+/// If \p TopLevelDecls is set, only these decls are indexed. Otherwise, all top
+/// level decls obtained from \p AST are indexed.
+SymbolSlab
+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