[clangd] Make in-memory CDB always available as an overlay, refactor.

Summary:
The new implementation is a GlobalCompilationDatabase that overlays a base.
Normally this is the directory-based CDB.
To preserve the behavior of compile_args_from=LSP, the base may be null.

The OverlayCDB is always present, and so the extensions to populate it
are always supported.

It also allows overriding the flags of the fallback command. This is
just unit-tested for now, but the plan is to expose this as an extension
on the initialize message. This addresses use cases like
https://github.com/thomasjo/atom-ide-cpp/issues/16

Reviewers: ilya-biryukov

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

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

llvm-svn: 345970

6 files changed, 118 insertions, 96 deletions

diff --git a/clang-tools-extra/clangd/ClangdLSPServer.cpp b/clang-tools-extra/clangd/ClangdLSPServer.cpp
index 752de60b7ee..2d4a9cf5076 100644
--- a/clang-tools-extra/clangd/ClangdLSPServer.cpp
+++ b/clang-tools-extra/clangd/ClangdLSPServer.cpp
@@ -305,11 +305,12 @@ void ClangdLSPServer::onInitialize(const InitializeParams &Params,
                                       ErrorCode::InvalidRequest));
   if (const auto &Dir = Params.initializationOptions.compilationDatabasePath)
     CompileCommandsDir = Dir;
-  CDB.emplace(UseInMemoryCDB
-                  ? CompilationDB::makeInMemory()
-                  : CompilationDB::makeDirectoryBased(CompileCommandsDir));
-  Server.emplace(CDB->getCDB(), FSProvider,
-                 static_cast<DiagnosticsConsumer &>(*this), ClangdServerOpts);
+  if (UseDirBasedCDB)
+    BaseCDB = llvm::make_unique<DirectoryBasedGlobalCompilationDatabase>(
+        CompileCommandsDir);
+  CDB.emplace(BaseCDB.get());
+  Server.emplace(*CDB, FSProvider, static_cast<DiagnosticsConsumer &>(*this),
+                 ClangdServerOpts);
   applyConfiguration(Params.initializationOptions.ConfigSettings);
 
   CCOpts.EnableSnippets = Params.capabilities.CompletionSnippets;
@@ -658,12 +659,14 @@ void ClangdLSPServer::applyConfiguration(
     /// The opened files need to be reparsed only when some existing
     /// entries are changed.
     PathRef File = Entry.first;
-    if (!CDB->setCompilationCommandForFile(
-            File, tooling::CompileCommand(
-                      std::move(Entry.second.workingDirectory), File,
-                      std::move(Entry.second.compilationCommand),
-                      /*Output=*/"")))
-      ShouldReparseOpenFiles = true;
+    auto Old = CDB->getCompileCommand(File);
+    auto New =
+        tooling::CompileCommand(std::move(Entry.second.workingDirectory), File,
+                                std::move(Entry.second.compilationCommand),
+                                /*Output=*/"");
+    if (Old != New)
+      CDB->setCompileCommand(File, std::move(New));
+    ShouldReparseOpenFiles = true;
   }
   if (ShouldReparseOpenFiles)
     reparseOpenedFiles();
@@ -684,12 +687,12 @@ void ClangdLSPServer::onReference(const ReferenceParams &Params,
 ClangdLSPServer::ClangdLSPServer(class Transport &Transp,
                                  const clangd::CodeCompleteOptions &CCOpts,
                                  Optional<Path> CompileCommandsDir,
-                                 bool ShouldUseInMemoryCDB,
+                                 bool UseDirBasedCDB,
                                  const ClangdServer::Options &Opts)
     : Transp(Transp), MsgHandler(new MessageHandler(*this)), CCOpts(CCOpts),
       SupportedSymbolKinds(defaultSymbolKinds()),
       SupportedCompletionItemKinds(defaultCompletionItemKinds()),
-      UseInMemoryCDB(ShouldUseInMemoryCDB),
+      UseDirBasedCDB(UseDirBasedCDB),
       CompileCommandsDir(std::move(CompileCommandsDir)),
       ClangdServerOpts(Opts) {
   // clang-format off
@@ -783,31 +786,5 @@ void ClangdLSPServer::reparseOpenedFiles() {
                         WantDiagnostics::Auto);
 }
 
-ClangdLSPServer::CompilationDB ClangdLSPServer::CompilationDB::makeInMemory() {
-  return CompilationDB(llvm::make_unique<InMemoryCompilationDb>(),
-                       /*IsDirectoryBased=*/false);
-}
-
-ClangdLSPServer::CompilationDB
-ClangdLSPServer::CompilationDB::makeDirectoryBased(
-    Optional<Path> CompileCommandsDir) {
-  auto CDB = llvm::make_unique<DirectoryBasedGlobalCompilationDatabase>(
-      std::move(CompileCommandsDir));
-  return CompilationDB(std::move(CDB),
-                       /*IsDirectoryBased=*/true);
-}
-
-bool ClangdLSPServer::CompilationDB::setCompilationCommandForFile(
-    PathRef File, tooling::CompileCommand CompilationCommand) {
-  if (IsDirectoryBased) {
-    elog("Trying to set compile command for {0} while using directory-based "
-         "compilation database",
-         File);
-    return false;
-  }
-  return static_cast<InMemoryCompilationDb *>(CDB.get())
-      ->setCompilationCommandForFile(File, std::move(CompilationCommand));
-}
-
 } // namespace clangd
 } // namespace clang
diff --git a/clang-tools-extra/clangd/ClangdLSPServer.h b/clang-tools-extra/clangd/ClangdLSPServer.h
index 12de1150f33..5cfab902494 100644
--- a/clang-tools-extra/clangd/ClangdLSPServer.h
+++ b/clang-tools-extra/clangd/ClangdLSPServer.h
@@ -36,9 +36,11 @@ public:
   /// If \p CompileCommandsDir has a value, compile_commands.json will be
   /// loaded only from \p CompileCommandsDir. Otherwise, clangd will look
   /// for compile_commands.json in all parent directories of each file.
+  /// If UseDirBasedCDB is false, compile commands are not read from disk.
+  // FIXME: Clean up signature around CDBs.
   ClangdLSPServer(Transport &Transp, const clangd::CodeCompleteOptions &CCOpts,
-                  llvm::Optional<Path> CompileCommandsDir,
-                  bool ShouldUseInMemoryCDB, const ClangdServer::Options &Opts);
+                  llvm::Optional<Path> CompileCommandsDir, bool UseDirBasedCDB,
+                  const ClangdServer::Options &Opts);
   ~ClangdLSPServer();
 
   /// Run LSP server loop, communicating with the Transport provided in the
@@ -105,37 +107,6 @@ private:
   /// Caches FixIts per file and diagnostics
   llvm::StringMap<DiagnosticToReplacementMap> FixItsMap;
 
-  /// Encapsulates the directory-based or the in-memory compilation database
-  /// that's used by the LSP server.
-  class CompilationDB {
-  public:
-    static CompilationDB makeInMemory();
-    static CompilationDB
-    makeDirectoryBased(llvm::Optional<Path> CompileCommandsDir);
-
-    /// Sets the compilation command for a particular file.
-    /// Only valid for in-memory CDB, no-op and error log on DirectoryBasedCDB.
-    ///
-    /// \returns True if the File had no compilation command before.
-    bool
-    setCompilationCommandForFile(PathRef File,
-                                 tooling::CompileCommand CompilationCommand);
-
-    /// Returns a CDB that should be used to get compile commands for the
-    /// current instance of ClangdLSPServer.
-    GlobalCompilationDatabase &getCDB() { return *CDB; }
-
-  private:
-    CompilationDB(std::unique_ptr<GlobalCompilationDatabase> CDB,
-                  bool IsDirectoryBased)
-        : CDB(std::move(CDB)), IsDirectoryBased(IsDirectoryBased) {}
-
-    // if IsDirectoryBased is true, an instance of InMemoryCDB.
-    // If IsDirectoryBased is false, an instance of DirectoryBasedCDB.
-    std::unique_ptr<GlobalCompilationDatabase> CDB;
-    bool IsDirectoryBased;
-  };
-
   // Most code should not deal with Transport directly.
   // MessageHandler deals with incoming messages, use call() etc for outgoing.
   clangd::Transport &Transp;
@@ -162,9 +133,11 @@ private:
   DraftStore DraftMgr;
 
   // The CDB is created by the "initialize" LSP method.
-  bool UseInMemoryCDB; // FIXME: make this a capability.
+  bool UseDirBasedCDB;                     // FIXME: make this a capability.
   llvm::Optional<Path> CompileCommandsDir; // FIXME: merge with capability?
-  llvm::Optional<CompilationDB> CDB;
+  std::unique_ptr<GlobalCompilationDatabase> BaseCDB;
+  // CDB is BaseCDB plus any comands overridden via LSP extensions.
+  llvm::Optional<OverlayCDB> CDB;
   // The ClangdServer is created by the "initialize" LSP method.
   // It is destroyed before run() returns, to ensure worker threads exit.
   ClangdServer::Options ClangdServerOpts;
diff --git a/clang-tools-extra/clangd/GlobalCompilationDatabase.cpp b/clang-tools-extra/clangd/GlobalCompilationDatabase.cpp
index acfd5953a16..d02e5e9ff61 100644
--- a/clang-tools-extra/clangd/GlobalCompilationDatabase.cpp
+++ b/clang-tools-extra/clangd/GlobalCompilationDatabase.cpp
@@ -80,22 +80,32 @@ DirectoryBasedGlobalCompilationDatabase::getCDBForFile(PathRef File) const {
 }
 
 Optional<tooling::CompileCommand>
-InMemoryCompilationDb::getCompileCommand(PathRef File) const {
+OverlayCDB::getCompileCommand(PathRef File) const {
+  {
+    std::lock_guard<std::mutex> Lock(Mutex);
+    auto It = Commands.find(File);
+    if (It != Commands.end())
+      return It->second;
+  }
+  return Base ? Base->getCompileCommand(File) : None;
+}
+
+tooling::CompileCommand OverlayCDB::getFallbackCommand(PathRef File) const {
+  auto Cmd = Base ? Base->getFallbackCommand(File)
+                  : GlobalCompilationDatabase::getFallbackCommand(File);
   std::lock_guard<std::mutex> Lock(Mutex);
-  auto It = Commands.find(File);
-  if (It == Commands.end())
-    return None;
-  return It->second;
+  Cmd.CommandLine.insert(Cmd.CommandLine.end(), FallbackFlags.begin(),
+                         FallbackFlags.end());
+  return Cmd;
 }
 
-bool InMemoryCompilationDb::setCompilationCommandForFile(
-    PathRef File, tooling::CompileCommand CompilationCommand) {
+void OverlayCDB::setCompileCommand(
+    PathRef File, llvm::Optional<tooling::CompileCommand> Cmd) {
   std::unique_lock<std::mutex> Lock(Mutex);
-  auto ItInserted = Commands.insert(std::make_pair(File, CompilationCommand));
-  if (ItInserted.second)
-    return true;
-  ItInserted.first->setValue(std::move(CompilationCommand));
-  return false;
+  if (Cmd)
+    Commands[File] = std::move(*Cmd);
+  else
+    Commands.erase(File);
 }
 
 } // namespace clangd
diff --git a/clang-tools-extra/clangd/GlobalCompilationDatabase.h b/clang-tools-extra/clangd/GlobalCompilationDatabase.h
index 692aa956525..8297db1eaf4 100644
--- a/clang-tools-extra/clangd/GlobalCompilationDatabase.h
+++ b/clang-tools-extra/clangd/GlobalCompilationDatabase.h
@@ -74,23 +74,30 @@ private:
   llvm::Optional<Path> CompileCommandsDir;
 };
 
-/// Gets compile args from an in-memory mapping based on a filepath. Typically
-/// used by clients who provide the compile commands themselves.
-class InMemoryCompilationDb : public GlobalCompilationDatabase {
+/// Wraps another compilation database, and supports overriding the commands
+/// using an in-memory mapping.
+class OverlayCDB : public GlobalCompilationDatabase {
 public:
-  /// Gets compile command for \p File from the stored mapping.
+  // Base may be null, in which case no entries are inherited.
+  // FallbackFlags are added to the fallback compile command.
+  OverlayCDB(const GlobalCompilationDatabase *Base,
+             std::vector<std::string> FallbackFlags = {})
+      : Base(Base), FallbackFlags(std::move(FallbackFlags)) {}
+
   llvm::Optional<tooling::CompileCommand>
   getCompileCommand(PathRef File) const override;
+  tooling::CompileCommand getFallbackCommand(PathRef File) const override;
 
-  /// Sets the compilation command for a particular file.
-  ///
-  /// \returns True if the File had no compilation command before.
-  bool setCompilationCommandForFile(PathRef File,
-                                    tooling::CompileCommand CompilationCommand);
+  /// Sets or clears the compilation command for a particular file.
+  void
+  setCompileCommand(PathRef File,
+                    llvm::Optional<tooling::CompileCommand> CompilationCommand);
 
 private:
   mutable std::mutex Mutex;
   llvm::StringMap<tooling::CompileCommand> Commands; /* GUARDED_BY(Mut) */
+  const GlobalCompilationDatabase *Base;
+  std::vector<std::string> FallbackFlags;
 };
 
 } // namespace clangd
diff --git a/clang-tools-extra/clangd/tool/ClangdMain.cpp b/clang-tools-extra/clangd/tool/ClangdMain.cpp
index edf030bf196..bfaec75b035 100644
--- a/clang-tools-extra/clangd/tool/ClangdMain.cpp
+++ b/clang-tools-extra/clangd/tool/ClangdMain.cpp
@@ -322,7 +322,7 @@ int main(int argc, char *argv[]) {
       InputStyle);
   ClangdLSPServer LSPServer(
       *Transport, CCOpts, CompileCommandsDirPath,
-      /*ShouldUseInMemoryCDB=*/CompileArgsFrom == LSPCompileArgs, Opts);
+      /*UseDirBasedCDB=*/CompileArgsFrom == FilesystemCompileArgs, Opts);
   constexpr int NoShutdownRequestErrorCode = 1;
   set_thread_name("clangd.main");
   return LSPServer.run() ? 0 : NoShutdownRequestErrorCode;
diff --git a/clang-tools-extra/unittests/clangd/GlobalCompilationDatabaseTests.cpp b/clang-tools-extra/unittests/clangd/GlobalCompilationDatabaseTests.cpp
index 5e2dfab596c..f938b302b66 100644
--- a/clang-tools-extra/unittests/clangd/GlobalCompilationDatabaseTests.cpp
+++ b/clang-tools-extra/unittests/clangd/GlobalCompilationDatabaseTests.cpp
@@ -33,6 +33,61 @@ TEST(GlobalCompilationDatabaseTest, FallbackCommand) {
                                            testPath("foo/bar.h")));
 }
 
+static tooling::CompileCommand cmd(StringRef File, StringRef Arg) {
+  return tooling::CompileCommand(testRoot(), File, {"clang", Arg, File}, "");
+}
+
+class OverlayCDBTest : public ::testing::Test {
+  class BaseCDB : public GlobalCompilationDatabase {
+  public:
+    Optional<tooling::CompileCommand>
+    getCompileCommand(StringRef File) const override {
+      if (File == testPath("foo.cc"))
+        return cmd(File, "-DA=1");
+      return None;
+    }
+
+    tooling::CompileCommand getFallbackCommand(StringRef File) const override {
+      return cmd(File, "-DA=2");
+    }
+  };
+
+protected:
+  OverlayCDBTest() : Base(llvm::make_unique<BaseCDB>()) {}
+  std::unique_ptr<GlobalCompilationDatabase> Base;
+};
+
+TEST_F(OverlayCDBTest, GetCompileCommand) {
+  OverlayCDB CDB(Base.get());
+  EXPECT_EQ(CDB.getCompileCommand(testPath("foo.cc")),
+            Base->getCompileCommand(testPath("foo.cc")));
+  EXPECT_EQ(CDB.getCompileCommand(testPath("missing.cc")), llvm::None);
+
+  auto Override = cmd(testPath("foo.cc"), "-DA=3");
+  CDB.setCompileCommand(testPath("foo.cc"), Override);
+  EXPECT_EQ(CDB.getCompileCommand(testPath("foo.cc")), Override);
+  EXPECT_EQ(CDB.getCompileCommand(testPath("missing.cc")), llvm::None);
+  CDB.setCompileCommand(testPath("missing.cc"), Override);
+  EXPECT_EQ(CDB.getCompileCommand(testPath("missing.cc")), Override);
+}
+
+TEST_F(OverlayCDBTest, GetFallbackCommand) {
+  OverlayCDB CDB(Base.get(), {"-DA=4"});
+  EXPECT_THAT(CDB.getFallbackCommand(testPath("bar.cc")).CommandLine,
+              ElementsAre("clang", "-DA=2", testPath("bar.cc"), "-DA=4"));
+}
+
+TEST_F(OverlayCDBTest, NoBase) {
+  OverlayCDB CDB(nullptr, {"-DA=6"});
+  EXPECT_EQ(CDB.getCompileCommand(testPath("bar.cc")), None);
+  auto Override = cmd(testPath("bar.cc"), "-DA=5");
+  CDB.setCompileCommand(testPath("bar.cc"), Override);
+  EXPECT_EQ(CDB.getCompileCommand(testPath("bar.cc")), Override);
+
+  EXPECT_THAT(CDB.getFallbackCommand(testPath("foo.cc")).CommandLine,
+              ElementsAre("clang", testPath("foo.cc"), "-DA=6"));
+}
+
 } // namespace
 } // namespace clangd
 } // namespace clang