[clangd] Make FileIndex aware of the main file

[ilya-biryukov]

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.

[ilya-biryukov]

This is still WIP, but wanted to share anyway the progress anyway.
Need to fix the FIXMEs currently added to the code and add tests.

clangd/index/FileIndex.cpp

60

We need to be more careful here now that we have to slabs in the list. Should only erase when both are removed.

66

This should definitely be done outside the lock, only the actual swap of the symbol tables should be there.

[ilya-biryukov]

As discussed offline, we could instead have a two-layer scheme for dynamic index. A top layer will contain main file symbols, second layer will contain preamble symbols.

[ilya-biryukov]

• Reverted changes to FileIndex, use merged index instead

[ilya-biryukov]

Still missing tests, but otherwise should be good to review.

[hokein]

clangd/index/FileIndex.h

70

Any strong reason to unify the interface (and indexAST)?

It seems to me that we do different things on PreambleAST and MainAST:

• on PreambleAST, we only index Symbols for code completion.
• on MainAST, we index Symbols (to collect other information missing from PreambleAST, e.g. definition location), and collect symbol references.

Therefore, we'd need to set different options for SymbolCollector accordingly. Having a single interface would make the implementation tricky (checking TopLevelDecls to see whether this is a PreambleASt).

[ioeric]

clangd/ClangdServer.cpp

71

Maybe provide noop implementations for ParsingCallbacks::onPreambleAST() and ParsingCallbacks::onMainAST by default?

clangd/ClangdServer.h

215

Can we move this into ClangdServer.cpp as an implementation helper?

226

nit: s/FileIdx/DymIdx? (also need to update the comment above)

clangd/index/FileIndex.cpp

46

nit: I'd try avoiding modify the input argument if possible. Maybe set Storage properly according to TopLevelDecls?

[ilya-biryukov]

• Move DynamicIndex into ClangdServer.cpp
• Rename FileIdx to DynamicIdx
• Avoid modifying input args

[ilya-biryukov]

clangd/ClangdServer.cpp

71

I'll address this in a parent revision.

clangd/ClangdServer.h

226

Went with DynamicIdx.

[ilya-biryukov]

• Use noop callbacks from the parent revision

[ilya-biryukov]

All the issues should be addressed at this point

clangd/ClangdServer.cpp

71

Done

[ilya-biryukov]

clangd/index/FileIndex.h

70

The only reason is to keep the changes to FileIndex minimal in this change. We'll have to change more code to cover those differences, but I'm not sure I have enough context on what's going to be different for this change, a change that will plug xrefs into clangd is probably a better candidate for those.

[ioeric]

lgtm

[ilya-biryukov]

• Rebase onto head

clangd/ClangdServer.h

public:
   /// if compilation arguments changed on calls to forceReparse().
   ///
   /// After each parsing request finishes, ClangdServer reports diagnostics to
   /// \p DiagConsumer. Note that a callback to \p DiagConsumer happens on a
   /// worker thread. Therefore, instances of \p DiagConsumer must properly
   /// 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);
 
   /// Add a \p File to the list of tracked C++ files or update the contents if
   /// \p File is already tracked. Also schedules parsing of the AST for it on a
   /// separate thread. When the parsing is complete, DiagConsumer passed in
   /// constructor will receive onDiagnosticsReady callback.

 
 private:
   /// FIXME: This stats several files to find a .clang-format file. I/O can be
   /// slow. Think of a way to cache this.
   llvm::Expected<tooling::Replacements>
   formatCode(llvm::StringRef Code, PathRef File,
              ArrayRef<tooling::Range> Ranges);
 
+  class DynamicIndex;
   typedef uint64_t DocVersion;
 
   void consumeDiagnostics(PathRef File, DocVersion Version,
                           std::vector<Diag> Diags);
 
   tooling::CompileCommand getCompileCommand(PathRef File);
 
   GlobalCompilationDatabase &CDB;
   DiagnosticsConsumer &DiagConsumer;
   FileSystemProvider &FSProvider;
 
   /// Used to synchronize diagnostic responses for added and removed files.
   llvm::StringMap<DocVersion> InternalVersion;
 
   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;
   std::shared_ptr<PCHContainerOperations> PCHs;
   /// Used to serialize diagnostic callbacks.
   /// FIXME(ibiryukov): get rid of an extra map and put all version counters
   /// into CppFile.
   std::mutex DiagnosticsMutex;

clangd/ClangdServer.cpp

public:
 
   void handle(tooling::AtomicChanges SourceReplacements) override {
     assert(!Result.hasValue());
     Result = std::move(SourceReplacements);
   }
 
   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!
   Opts.StorePreamblesInMemory = true;
   Opts.AsyncThreadsCount = 4; // Consistent!
   return Opts;
 }
 
 ClangdServer::ClangdServer(GlobalCompilationDatabase &CDB,
                            FileSystemProvider &FSProvider,
                            DiagnosticsConsumer &DiagConsumer,
                            const Options &Opts)
     : 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)
       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
       // is parsed.
       // 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);
   if (!Status)
     elog("Failed to get status for RootPath {0}: {1}", RootPath,
          Status.getError().message());
   else if (Status->isDirectory())
     this->RootPath = RootPath;

clangd/index/FileIndex.h

public:
   /// If URISchemes is empty, the default schemes in SymbolCollector will be
   /// used.
   FileIndex(std::vector<std::string> URISchemes = {});
 
   /// \brief Update symbols in \p Path with symbols in \p AST. If \p AST is
   /// 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,
             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:
   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.
-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,
 
 } // namespace clangd
 } // namespace clang
 
 #endif // LLVM_CLANG_TOOLS_EXTRA_CLANGD_INDEX_FILEINDEX_H

clangd/index/FileIndex.cpp

 //===--- FileIndex.cpp - Indexes for files. ------------------------ C++-*-===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 
 #include "FileIndex.h"
-#include "SymbolCollector.h"
 #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) {
   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<const Decl *> TopLevelDecls(
-      AST.getTranslationUnitDecl()->decls().begin(),
+  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, TopLevelDecls, Collector, IndexOpts);
+
+  index::indexTopLevelDecls(AST, DeclsToIndex, Collector, IndexOpts);
 
   return Collector.takeSymbols();
 }
 
 FileIndex::FileIndex(std::vector<std::string> URISchemes)
     : URISchemes(std::move(URISchemes)) {}
 
 void FileSymbols::update(PathRef Path, std::unique_ptr<SymbolSlab> Slab) {
   std::lock_guard<std::mutex> Lock(Mutex);
   if (!Slab)
     FileToSlabs.erase(Path);
   else
     FileToSlabs[Path] = std::move(Slab);
 }
 
 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};
 }
 
 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();
   Index.build(std::move(Symbols));
 }
 
 bool FileIndex::fuzzyFind(
     const FuzzyFindRequest &Req,