[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

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)

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),
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 (NewFile) {
+    if (auto *InternalCDB = CDB.lookupCDB(File)) {
+      for (const auto &CDBFile : InternalCDB->getAllFiles()) {
+        if (CDBFile != File && !WorkScheduler.hasFile(CDBFile)) {
+          WithContext WithContext(createProcessingContext(CDBFile));
+          if (Config::current().AST.Build == Config::ASTPolicy::PreBuild) {
+            auto Buffer = llvm::MemoryBuffer::getFile(CDBFile);
+            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/Config.h

struct Config {
 
   /// Style of the codebase.
   struct {
     // Namespaces that should always be fully qualified, meaning no "using"
     // declarations, always spell out the whole name (with or without leading
     // ::). All nested namespaces are affected as well.
     std::vector<std::string> FullyQualifiedNamespaces;
   } Style;
+
+  enum class ASTPolicy { PreBuild, OnDemand };
+  /// Controls AST prebuild behavior.
+  struct {
+    /// Whether this AST should be prebuild.
+    ASTPolicy Build = ASTPolicy::OnDemand;
+  } AST;
 };
 
 } // namespace clangd
 } // namespace clang
 
 #endif

clang-tools-extra/clangd/ConfigCompile.cpp

void compile(Fragment &&F) {
     if (!F.Source.Directory.empty()) {
       FragmentDirectory = llvm::sys::path::convert_to_slash(F.Source.Directory);
       if (FragmentDirectory.back() != '/')
         FragmentDirectory += '/';
     }
     compile(std::move(F.If));
     compile(std::move(F.CompileFlags));
     compile(std::move(F.Index));
+    compile(std::move(F.AST));
   }
 
   void compile(Fragment::IfBlock &&F) {
     if (F.HasUnrecognizedCondition)
       Out.Conditions.push_back([&](const Params &) { return false; });
 
     auto PathMatch = std::make_unique<std::vector<llvm::Regex>>();
     for (auto &Entry : F.PathMatch) {
if (!F.FullyQualifiedNamespaces.empty()) {
                               const Params &, Config &C) {
         C.Style.FullyQualifiedNamespaces.insert(
             C.Style.FullyQualifiedNamespaces.begin(),
             FullyQualifiedNamespaces.begin(), FullyQualifiedNamespaces.end());
       });
     }
   }
 
+  void compile(Fragment::ASTBlock &&F) {
+    if (F.Build) {
+      if (auto Val = compileEnum<Config::ASTPolicy>("Build", **F.Build)
+                         .map("OnDemand", Config::ASTPolicy::OnDemand)
+                         .map("PreBuild", Config::ASTPolicy::PreBuild)
+                         .value()) {
+        Out.Apply.push_back(
+            [Val](const Params &, Config &C) { C.AST.Build = *Val; });
+      }
+    }
+  }
+
   constexpr static llvm::SourceMgr::DiagKind Error = llvm::SourceMgr::DK_Error;
   constexpr static llvm::SourceMgr::DiagKind Warning =
       llvm::SourceMgr::DK_Warning;
   void diag(llvm::SourceMgr::DiagKind Kind, llvm::StringRef Message,
             llvm::SMRange Range) {
     if (Range.isValid() && SourceMgr != nullptr)
       Diagnostic(SourceMgr->GetMessage(Range.Start, Kind, Message, Range));
     else

clang-tools-extra/clangd/ConfigFragment.h

struct Fragment {
   struct StyleBlock {
     // Namespaces that should always be fully qualified, meaning no "using"
     // declarations, always spell out the whole name (with or without leading
     // ::). All nested namespaces are affected as well.
     // Affects availability of the AddUsing tweak.
     std::vector<Located<std::string>> FullyQualifiedNamespaces;
   };
   StyleBlock Style;
+
+  struct ASTBlock {
+    /// Controls whether clangd prebuilds the AST for the current file. When
+    /// opening a file for the first time. clangd iterates over all the files in
+    /// its compilation database and prebuilds the AST for those that have this
+    /// option set to "PreBuild". Legal values are "PreBuild" or "OnDemand".
+    /// Default is "OnDemand".
+    llvm::Optional<Located<std::string>> Build;
+  };
+  ASTBlock AST;
 };
 
 } // namespace config
 } // namespace clangd
 } // namespace clang
 
 #endif

clang-tools-extra/clangd/ConfigYAML.cpp

public:
   // Tries to parse N into F, returning false if it failed and we couldn't
   // meaningfully recover (YAML syntax error, or hard semantic error).
   bool parse(Fragment &F, Node &N) {
     DictParser Dict("Config", this);
     Dict.handle("If", [&](Node &N) { parse(F.If, N); });
     Dict.handle("CompileFlags", [&](Node &N) { parse(F.CompileFlags, N); });
     Dict.handle("Index", [&](Node &N) { parse(F.Index, N); });
     Dict.handle("Style", [&](Node &N) { parse(F.Style, N); });
+    Dict.handle("AST", [&](Node &N) { parse(F.AST, N); });
     Dict.parse(N);
     return !(N.failed() || HadError);
   }
 
 private:
   void parse(Fragment::IfBlock &F, Node &N) {
     DictParser Dict("If", this);
     Dict.unrecognized(
private:
 
   void parse(Fragment::IndexBlock &F, Node &N) {
     DictParser Dict("Index", this);
     Dict.handle("Background",
                 [&](Node &N) { F.Background = scalarValue(N, "Background"); });
     Dict.parse(N);
   }
 
+  void parse(Fragment::ASTBlock &F, Node &N) {
+    DictParser Dict("AST", this);
+    Dict.handle("Build", [&](Node &N) { F.Build = scalarValue(N, "Build"); });
+    Dict.parse(N);
+  }
+
   // Helper for parsing mapping nodes (dictionaries).
   // We don't use YamlIO as we want to control over unknown keys.
   class DictParser {
     llvm::StringRef Description;
     std::vector<std::pair<llvm::StringRef, std::function<void(Node &)>>> Keys;
     std::function<void(llvm::StringRef)> Unknown;
     Parser *Outer;
 

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

auto Task = [=]() mutable {
     }
 
     // Update current inputs so that subsequent reads can see them.
     {
       std::lock_guard<std::mutex> Lock(Mutex);
       FileInputs = Inputs;
     }
 
-    log("ASTWorker building file {0} version {1} with command {2}\n[{3}]\n{4}",
+    log("ASTWorker building file {0} version {1}", FileName, Inputs.Version);
+    vlog("ASTWorker building file {0} version {1} with command {2}\n[{3}]\n{4}",
         FileName, Inputs.Version, Inputs.CompileCommand.Heuristic,
         Inputs.CompileCommand.Directory,
         llvm::join(Inputs.CompileCommand.CommandLine, " "));
 
+    if (!InputsAreTheSame)
+      log("ASTWorker inputs are not the same, cached AST is invalidated");
+
     StoreDiags CompilerInvocationDiagConsumer;
     std::vector<std::string> CC1Args;
     std::unique_ptr<CompilerInvocation> Invocation = buildCompilerInvocation(
         Inputs, CompilerInvocationDiagConsumer, &CC1Args);
     // Log cc1 args even (especially!) if creating invocation failed.
     if (!CC1Args.empty())
       vlog("Driver produced command: cc1 {0}", llvm::join(CC1Args, " "));
     std::vector<Diag> CompilerInvocationDiags =
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/unittests/ClangdTests.cpp

getCompileCommand(PathRef File) const override {
       CanReturnCommand.wait();
       auto FileName = llvm::sys::path::filename(File);
       std::vector<std::string> CommandLine = {"clangd", "-ffreestanding",
                                               std::string(File)};
       return {tooling::CompileCommand(llvm::sys::path::parent_path(File),
                                       FileName, std::move(CommandLine), "")};
     }
 
+    tooling::CompilationDatabase *lookupCDB(PathRef File) const override {
+      return nullptr;
+    }
+
     std::vector<std::string> ExtraClangFlags;
 
   private:
     Notification &CanReturnCommand;
   };
 
   Notification CanReturnCommand;
   DelayedCompilationDatabase CDB(CanReturnCommand);

clang-tools-extra/clangd/unittests/GlobalCompilationDatabaseTests.cpp

class OverlayCDBTest : public ::testing::Test {
   public:
     llvm::Optional<tooling::CompileCommand>
     getCompileCommand(llvm::StringRef File) const override {
       if (File == testPath("foo.cc"))
         return cmd(File, "-DA=1");
       return None;
     }
 
+    tooling::CompilationDatabase *lookupCDB(PathRef File) const override {
+      return nullptr;
+    }
+
     tooling::CompileCommand
     getFallbackCommand(llvm::StringRef File) const override {
       return cmd(File, "-DA=2");
     }
 
     llvm::Optional<ProjectInfo> getProjectInfo(PathRef File) const override {
       return ProjectInfo{testRoot()};
     }

clang-tools-extra/clangd/unittests/TestFS.h

public:
   /// If \p RelPathPrefix is not empty, use that as a prefix in front of the
   /// source file name, instead of using an absolute path.
   MockCompilationDatabase(StringRef Directory = StringRef(),
                           StringRef RelPathPrefix = StringRef());
 
   llvm::Optional<tooling::CompileCommand>
   getCompileCommand(PathRef File) const override;
 
+  tooling::CompilationDatabase *lookupCDB(PathRef File) const override;
+
   llvm::Optional<ProjectInfo> getProjectInfo(PathRef File) const override;
 
   std::vector<std::string> ExtraClangFlags;
 
 private:
   StringRef Directory;
   StringRef RelPathPrefix;
 };

clang-tools-extra/clangd/unittests/TestFS.cpp

MockCompilationDatabase::getCompileCommand(PathRef File) const {
   }
 
   return {tooling::CompileCommand(Directory != llvm::StringRef()
                                       ? Directory
                                       : llvm::sys::path::parent_path(File),
                                   FileName, std::move(CommandLine), "")};
 }
 
+tooling::CompilationDatabase *
+MockCompilationDatabase::lookupCDB(PathRef File) const {
+  return nullptr;
+}
+
 const char *testRoot() {
 #ifdef _WIN32
   return "C:\\clangd-test";
 #else
   return "/clangd-test";
 #endif
 }