[clangd] Add AST config to prebuild ASTs

[DaanDeMeyer]

Summary

The new AST config allows prebuilding the ASTs for files in a
compilation database when a single file in that compilation database
is opened by the client. This reduces latency when opening subsequent
files in the same compilation database. This matters in large
codebases where building the AST can take multiple seconds.

This is still missing tests and more documentation but I wanted to get some feedback first before continuing.

[DaanDeMeyer]

Switched from all or nothing command line option to per file config.

[DaanDeMeyer]

Fixed formatting

[DaanDeMeyer]

Fix tests

[DaanDeMeyer]

Fix all the bugs after doing some proper testing.

One thing we don't handle yet is header files since those don't appear in the compilation database. Maybe we need to add prebuilding for those as well?

[DaanDeMeyer]

Abandoning as the end result is too slow even when including a limited number of files.

clang-tools-extra/clangd/ClangdServer.h

struct Options {
     std::vector<std::string> QueryDriverGlobs;
 
     /// Enable notification-based semantic highlighting.
     bool TheiaSemanticHighlighting = false;
 
     /// Enable preview of FoldingRanges feature.
     bool FoldingRanges = false;
 
+    /// Prebuild ASTs for files in the compilation database.
+    bool PrebuildASTs = false;
+
     explicit operator TUScheduler::Options() const;
   };
   // Sensible default options for use in tests.
   // Features like indexing must be enabled if desired.
   static Options optsForTest();
 
   /// Creates a new ClangdServer instance.
   ///
private:
   /// This includes the current configuration (see Options::ConfigProvider).
   /// The empty string means no particular file is the target.
   /// Rather than called by each feature, this is exposed to the components
   /// that control worker threads, like TUScheduler and BackgroundIndex.
   /// This means it's OK to do some IO here, and it cuts across all features.
   Context createProcessingContext(PathRef) const;
   config::Provider *ConfigProvider = nullptr;
 
+  const GlobalCompilationDatabase &CDB;
   const ThreadsafeFS &TFS;
 
   Path ResourceDir;
   // The index used to look up symbols. This could be:
   //   - null (all index functionality is optional)
   //   - the dynamic index owned by ClangdServer (DynamicIdx)
   //   - the static index passed to the constructor
   //   - a merged view of a static and dynamic index (MergedIndex)
private:
   // If this is true, suggest include insertion fixes for diagnostic errors that
   // can be caused by missing includes (e.g. member access in incomplete type).
   bool SuggestMissingIncludes = false;
 
   // If true, preserve expressions in AST for broken code.
   bool BuildRecoveryAST = true;
   // If true, preserve the type for recovery AST.
   bool PreserveRecoveryASTType = false;
+  // If true, prebuild ASTs for all files in a compilation database.
+  bool PrebuildASTs = false;
 
   // GUARDED_BY(CachedCompletionFuzzyFindRequestMutex)
   llvm::StringMap<llvm::Optional<FuzzyFindRequest>>
       CachedCompletionFuzzyFindRequestByFile;
   mutable std::mutex CachedCompletionFuzzyFindRequestMutex;
 
   llvm::Optional<std::string> WorkspaceRoot;
   // WorkScheduler has to be the last member, because its destructor has to be

clang-tools-extra/clangd/ClangdServer.cpp

ClangdServer::Options::operator TUScheduler::Options() const {
   Opts.UpdateDebounce = UpdateDebounce;
   Opts.AsyncPreambleBuilds = AsyncPreambleBuilds;
   return Opts;
 }
 
 ClangdServer::ClangdServer(const GlobalCompilationDatabase &CDB,
                            const ThreadsafeFS &TFS, const Options &Opts,
                            Callbacks *Callbacks)
-    : ConfigProvider(Opts.ConfigProvider), TFS(TFS),
+    : ConfigProvider(Opts.ConfigProvider), CDB(CDB), TFS(TFS),
       DynamicIdx(Opts.BuildDynamicSymbolIndex
                      ? new FileIndex(Opts.HeavyweightDynamicSymbolIndex,
                                      Opts.CollectMainFileRefs)
                      : nullptr),
       GetClangTidyOptions(Opts.GetClangTidyOptions),
       SuggestMissingIncludes(Opts.SuggestMissingIncludes),
       BuildRecoveryAST(Opts.BuildRecoveryAST),
       PreserveRecoveryASTType(Opts.PreserveRecoveryASTType),
-      WorkspaceRoot(Opts.WorkspaceRoot),
+      PrebuildASTs(Opts.PrebuildASTs), WorkspaceRoot(Opts.WorkspaceRoot),
       // 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(
           CDB,
           [&, this] {
void ClangdServer::addDocument(PathRef File, llvm::StringRef Contents,
   Opts.SuggestMissingIncludes = SuggestMissingIncludes;
 
   // Compile command is set asynchronously during update, as it can be slow.
   ParseInputs Inputs;
   Inputs.TFS = &TFS;
   Inputs.Contents = std::string(Contents);
   Inputs.Version = Version.str();
   Inputs.ForceRebuild = ForceRebuild;
-  Inputs.Opts = std::move(Opts);
+  Inputs.Opts = Opts;
   Inputs.Index = Index;
   Inputs.Opts.BuildRecoveryAST = BuildRecoveryAST;
   Inputs.Opts.PreserveRecoveryASTType = PreserveRecoveryASTType;
   bool NewFile = WorkScheduler.update(File, Inputs, WantDiags);
   // If we loaded Foo.h, we want to make sure Foo.cpp is indexed.
   if (NewFile && BackgroundIdx)
     BackgroundIdx->boostRelated(File);
+
+  // If --prebuild-asts is enabled, loop over all the files in the compilation
+  // database and schedule an AST update for them if they haven't yet been added
+  // to the WorkScheduler before.
+  if (NewFile && PrebuildASTs) {
+    if (auto *InternalCDB = CDB.lookupCDB(File)) {
+      for (const auto &CDBFile : InternalCDB->getAllFiles()) {
+        if (CDBFile != File && !WorkScheduler.hasFile(File)) {
+          auto Buffer = llvm::MemoryBuffer::getFile(File);
+          Inputs.Contents = std::string(Buffer->get()->getBuffer());
+          WorkScheduler.update(CDBFile, Inputs, WantDiagnostics::No);
+          if (BackgroundIdx)
+            BackgroundIdx->boostRelated(CDBFile);
+        }
+      }
+    }
+  }
 }
 
 void ClangdServer::removeDocument(PathRef File) { WorkScheduler.remove(File); }
 
 void ClangdServer::codeComplete(PathRef File, Position Pos,
                                 const clangd::CodeCompleteOptions &Opts,
                                 Callback<CodeCompleteResult> CB) {
   // Copy completion options for passing them to async task handler.

clang-tools-extra/clangd/GlobalCompilationDatabase.h

 class GlobalCompilationDatabase {
 public:
   virtual ~GlobalCompilationDatabase() = default;
 
   /// If there are any known-good commands for building this file, returns one.
   virtual llvm::Optional<tooling::CompileCommand>
   getCompileCommand(PathRef File) const = 0;
 
+  virtual tooling::CompilationDatabase *lookupCDB(PathRef File) const = 0;
+
   /// Finds the closest project to \p File.
   virtual llvm::Optional<ProjectInfo> getProjectInfo(PathRef File) const {
     return llvm::None;
   }
 
   /// Makes a guess at how to build a file.
   /// The default implementation just runs clang on the file.
   /// Clangd should treat the results as unreliable.
public:
   ~DirectoryBasedGlobalCompilationDatabase() override;
 
   /// Scans File's parents looking for compilation databases.
   /// Any extra flags will be added.
   /// Might trigger OnCommandChanged, if CDB wasn't broadcasted yet.
   llvm::Optional<tooling::CompileCommand>
   getCompileCommand(PathRef File) const override;
 
+  virtual tooling::CompilationDatabase *lookupCDB(PathRef File) const override;
+
   /// Returns the path to first directory containing a compilation database in
   /// \p File's parents.
   llvm::Optional<ProjectInfo> getProjectInfo(PathRef File) const override;
 
 private:
   /// Caches compilation databases loaded from directories.
   struct CachedCDB {
     std::string Path; // Not case-folded.
public:
   // FallbackFlags are added to the fallback compile command.
   // Adjuster is applied to all commands, fallback or not.
   OverlayCDB(const GlobalCompilationDatabase *Base,
              std::vector<std::string> FallbackFlags = {},
              tooling::ArgumentsAdjuster Adjuster = nullptr);
 
   llvm::Optional<tooling::CompileCommand>
   getCompileCommand(PathRef File) const override;
+
+  tooling::CompilationDatabase *lookupCDB(PathRef File) const override;
+
   tooling::CompileCommand getFallbackCommand(PathRef File) const override;
   /// Project info is gathered purely from the inner compilation database to
   /// ensure consistency.
   llvm::Optional<ProjectInfo> getProjectInfo(PathRef File) const override;
 
   /// Sets or clears the compilation command for a particular file.
   void
   setCompileCommand(PathRef File,

clang-tools-extra/clangd/GlobalCompilationDatabase.cpp

DirectoryBasedGlobalCompilationDatabase(
         llvm::Optional<Path> CompileCommandsDir)
     : CompileCommandsDir(std::move(CompileCommandsDir)) {}
 
 DirectoryBasedGlobalCompilationDatabase::
     ~DirectoryBasedGlobalCompilationDatabase() = default;
 
 llvm::Optional<tooling::CompileCommand>
 DirectoryBasedGlobalCompilationDatabase::getCompileCommand(PathRef File) const {
+  auto *CDB = lookupCDB(File);
+  if (!CDB) {
+    return llvm::None;
+  }
+
+  auto Candidates = CDB->getCompileCommands(File);
+  if (!Candidates.empty())
+    return std::move(Candidates.front());
+
+  return None;
+}
+
+tooling::CompilationDatabase *
+DirectoryBasedGlobalCompilationDatabase::lookupCDB(PathRef File) const {
   CDBLookupRequest Req;
   Req.FileName = File;
   Req.ShouldBroadcast = true;
 
   auto Res = lookupCDB(Req);
   if (!Res) {
     log("Failed to find compilation database for {0}", File);
-    return llvm::None;
+    return nullptr;
   }
 
-  auto Candidates = Res->CDB->getCompileCommands(File);
-  if (!Candidates.empty())
-    return std::move(Candidates.front());
-
-  return None;
+  return Res->CDB;
 }
 
 // For platforms where paths are case-insensitive (but case-preserving),
 // we need to do case-insensitive comparisons and use lowercase keys.
 // FIXME: Make Path a real class with desired semantics instead.
 //        This class is not the only place this problem exists.
 // FIXME: Mac filesystems default to case-insensitive, but may be sensitive.
 
if (!Cmd && Base)
     Cmd = Base->getCompileCommand(File);
   if (!Cmd)
     return llvm::None;
   if (ArgsAdjuster)
     Cmd->CommandLine = ArgsAdjuster(Cmd->CommandLine, Cmd->Filename);
   return Cmd;
 }
 
+tooling::CompilationDatabase *OverlayCDB::lookupCDB(PathRef File) const {
+  return Base ? Base->lookupCDB(File) : nullptr;
+}
+
 tooling::CompileCommand OverlayCDB::getFallbackCommand(PathRef File) const {
   auto Cmd = Base ? Base->getFallbackCommand(File)
                   : GlobalCompilationDatabase::getFallbackCommand(File);
   std::lock_guard<std::mutex> Lock(Mutex);
   Cmd.CommandLine.insert(Cmd.CommandLine.end(), FallbackFlags.begin(),
                          FallbackFlags.end());
   if (ArgsAdjuster)
     Cmd.CommandLine = ArgsAdjuster(Cmd.CommandLine, Cmd.Filename);

clang-tools-extra/clangd/QueryDriverDatabase.cpp

std::vector<std::string> SystemIncludes =
         QueriedDrivers.get(/*Key=*/(Driver + ":" + Lang).str(), [&] {
           return extractSystemIncludes(Driver, Lang, Cmd->CommandLine,
                                        QueryDriverRegex);
         });
 
     return addSystemIncludes(*Cmd, SystemIncludes);
   }
 
+  tooling::CompilationDatabase *lookupCDB(PathRef File) const override {
+    return Base->lookupCDB(File);
+  }
+
   llvm::Optional<ProjectInfo> getProjectInfo(PathRef File) const override {
     return Base->getProjectInfo(File);
   }
 
 private:
   // Caches includes extracted from a driver. Key is driver:lang.
   Memoize<llvm::StringMap<std::vector<std::string>>> QueriedDrivers;
   llvm::Regex QueryDriverRegex;

clang-tools-extra/clangd/TUScheduler.h

public:
   /// The compile command in \p Inputs is ignored; worker queries CDB to get
   /// the actual compile command.
   /// If diagnostics are requested (Yes), and the context is cancelled
   /// before they are prepared, they may be skipped if eventual-consistency
   /// permits it (i.e. WantDiagnostics is downgraded to Auto).
   /// Returns true if the file was not previously tracked.
   bool update(PathRef File, ParseInputs Inputs, WantDiagnostics WD);
 
+  bool hasFile(PathRef File);
+
   /// Remove \p File from the list of tracked files and schedule removal of its
   /// resources. Pending diagnostics for closed files may not be delivered, even
   /// if requested with WantDiags::Auto or WantDiags::Yes.
   void remove(PathRef File);
 
   /// Returns a snapshot of all file buffer contents, per last update().
   llvm::StringMap<std::string> getAllFileContents() const;
 

clang-tools-extra/clangd/TUScheduler.cpp

FD = std::unique_ptr<FileData>(
         new FileData{Inputs.Contents, std::move(Worker)});
   } else {
     FD->Contents = Inputs.Contents;
   }
   FD->Worker->update(std::move(Inputs), WantDiags);
   return NewFile;
 }
 
+bool TUScheduler::hasFile(PathRef File) { return Files[File] == nullptr; }
+
 void TUScheduler::remove(PathRef File) {
   bool Removed = Files.erase(File);
   if (!Removed)
     elog("Trying to remove file from TUScheduler that is not tracked: {0}",
          File);
 }
 
 llvm::StringMap<std::string> TUScheduler::getAllFileContents() const {

clang-tools-extra/clangd/tool/ClangdMain.cpp

 
 opt<bool> CollectMainFileRefs{
     "collect-main-file-refs",
     cat(Misc),
     desc("Store references to main-file-only symbols in the index"),
     init(false),
 };
 
+opt<bool> PrebuildASTs{
+    "prebuild-asts",
+    cat(Misc),
+    desc("Prebuild ASTs for all files in the compilation database"),
+    init(false),
+};
+
 #if CLANGD_ENABLE_REMOTE
 opt<std::string> RemoteIndexAddress{
     "remote-index-address",
     cat(Features),
     desc("Address of the remote index server"),
 };
 
 // FIXME(kirillbobyrev): Should this be the location of compile_commands.json?
if (TweakList.getNumOccurrences())
       return llvm::is_contained(TweakList, T.id());
     return true;
   };
   if (ForceOffsetEncoding != OffsetEncoding::UnsupportedEncoding)
     Opts.Encoding = ForceOffsetEncoding;
 
   // Shall we allow to customize the file limit?
   Opts.Rename.AllowCrossFile = CrossFileRename;
+  Opts.PrebuildASTs = PrebuildASTs;
 
   if (CheckFile.getNumOccurrences()) {
     llvm::SmallString<256> Path;
     llvm::sys::fs::real_path(CheckFile, Path, /*expand_tilde=*/true);
     log("Entering check mode (no LSP server)");
     return check(Path, TFS, Opts)
                ? 0
                : static_cast<int>(ErrorResultCode::CheckFailed);