[clangd] Introduce ASTHooks to FeatureModules

These can be invoked at different stages while building an AST to let
FeatureModules implement features on top of it. The patch also
introduces a sawDiagnostic hook, which can mutate the final clangd::Diag
while reading a clang::Diagnostic.

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

Author: kadircet

Diff for: clang-tools-extra/clangd/ClangdServer.cpp

@@ -231,6 +231,7 @@ void ClangdServer::addDocument(PathRef File, llvm::StringRef Contents,
   Inputs.Opts = std::move(Opts);
   Inputs.Index = Index;
   Inputs.ClangTidyProvider = ClangTidyProvider;
+  Inputs.FeatureModules = FeatureModules;
   bool NewFile = WorkScheduler->update(File, Inputs, WantDiags);
   // If we loaded Foo.h, we want to make sure Foo.cpp is indexed.
   if (NewFile && BackgroundIdx)

Diff for: clang-tools-extra/clangd/ClangdServer.h

@@ -38,6 +38,7 @@
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/StringRef.h"
 #include <functional>
+#include <memory>
 #include <string>
 #include <type_traits>
 #include <utility>

Diff for: clang-tools-extra/clangd/Compiler.h

@@ -15,13 +15,16 @@
 #ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_COMPILER_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_COMPILER_H
 
+#include "FeatureModule.h"
 #include "GlobalCompilationDatabase.h"
 #include "TidyProvider.h"
 #include "index/Index.h"
 #include "support/ThreadsafeFS.h"
 #include "clang/Frontend/CompilerInstance.h"
 #include "clang/Frontend/PrecompiledPreamble.h"
 #include "clang/Tooling/CompilationDatabase.h"
+#include <memory>
+#include <vector>
 
 namespace clang {
 namespace clangd {
@@ -54,6 +57,8 @@ struct ParseInputs {
   const SymbolIndex *Index = nullptr;
   ParseOptions Opts = ParseOptions();
   TidyProviderRef ClangTidyProvider = {};
+  // Used to acquire ASTListeners when parsing files.
+  FeatureModuleSet *FeatureModules = nullptr;
 };
 
 /// Builds compiler invocation that could be used to build AST or preamble.

Diff for: clang-tools-extra/clangd/Diagnostics.cpp

@@ -744,6 +744,8 @@ void StoreDiags::HandleDiagnostic(DiagnosticsEngine::Level DiagLevel,
       LastDiag->Fixes.insert(LastDiag->Fixes.end(), ExtraFixes.begin(),
                              ExtraFixes.end());
     }
+    if (DiagCB)
+      DiagCB(Info, *LastDiag);
   } else {
     // Handle a note to an existing diagnostic.
 

Diff for: clang-tools-extra/clangd/Diagnostics.h

@@ -22,7 +22,10 @@
 #include "llvm/ADT/StringSet.h"
 #include "llvm/Support/SourceMgr.h"
 #include <cassert>
+#include <functional>
+#include <memory>
 #include <string>
+#include <utility>
 
 namespace clang {
 namespace tidy {
@@ -136,18 +139,24 @@ class StoreDiags : public DiagnosticConsumer {
                                                    const clang::Diagnostic &)>;
   using LevelAdjuster = std::function<DiagnosticsEngine::Level(
       DiagnosticsEngine::Level, const clang::Diagnostic &)>;
+  using DiagCallback =
+      std::function<void(const clang::Diagnostic &, clangd::Diag &)>;
   /// If set, possibly adds fixes for diagnostics using \p Fixer.
   void contributeFixes(DiagFixer Fixer) { this->Fixer = Fixer; }
   /// If set, this allows the client of this class to adjust the level of
   /// diagnostics, such as promoting warnings to errors, or ignoring
   /// diagnostics.
   void setLevelAdjuster(LevelAdjuster Adjuster) { this->Adjuster = Adjuster; }
+  /// Invokes a callback every time a diagnostics is completely formed. Handler
+  /// of the callback can also mutate the diagnostic.
+  void setDiagCallback(DiagCallback CB) { DiagCB = std::move(CB); }
 
 private:
   void flushLastDiag();
 
   DiagFixer Fixer = nullptr;
   LevelAdjuster Adjuster = nullptr;
+  DiagCallback DiagCB = nullptr;
   std::vector<Diag> Output;
   llvm::Optional<LangOptions> LangOpts;
   llvm::Optional<Diag> LastDiag;

Diff for: clang-tools-extra/clangd/FeatureModule.h

@@ -11,6 +11,7 @@
 
 #include "support/Function.h"
 #include "support/Threading.h"
+#include "clang/Basic/Diagnostic.h"
 #include "llvm/ADT/FunctionExtras.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Compiler.h"
@@ -21,6 +22,7 @@
 
 namespace clang {
 namespace clangd {
+struct Diag;
 class LSPBinder;
 class SymbolIndex;
 class ThreadsafeFS;
@@ -96,6 +98,22 @@ class FeatureModule {
   /// enumerating or applying code actions.
   virtual void contributeTweaks(std::vector<std::unique_ptr<Tweak>> &Out) {}
 
+  /// Extension point that allows modules to observe and modify an AST build.
+  /// One instance is created each time clangd produces a ParsedAST or
+  /// PrecompiledPreamble. For a given instance, lifecycle methods are always
+  /// called on a single thread.
+  struct ASTListener {
+    /// Listeners are destroyed once the AST is built.
+    virtual ~ASTListener() = default;
+
+    /// Called everytime a diagnostic is encountered. Modules can use this
+    /// modify the final diagnostic, or store some information to surface code
+    /// actions later on.
+    virtual void sawDiagnostic(const clang::Diagnostic &, clangd::Diag &) {}
+  };
+  /// Can be called asynchronously before building an AST.
+  virtual std::unique_ptr<ASTListener> astListeners() { return nullptr; }
+
 protected:
   /// Accessors for modules to access shared server facilities they depend on.
   Facilities &facilities();

Diff for: clang-tools-extra/clangd/ParsedAST.cpp

@@ -14,6 +14,7 @@
 #include "Compiler.h"
 #include "Config.h"
 #include "Diagnostics.h"
+#include "FeatureModule.h"
 #include "Headers.h"
 #include "IncludeFixer.h"
 #include "Preamble.h"
@@ -25,6 +26,7 @@
 #include "support/Trace.h"
 #include "clang/AST/ASTContext.h"
 #include "clang/AST/Decl.h"
+#include "clang/Basic/Diagnostic.h"
 #include "clang/Basic/LangOptions.h"
 #include "clang/Basic/SourceLocation.h"
 #include "clang/Basic/SourceManager.h"
@@ -261,7 +263,19 @@ ParsedAST::build(llvm::StringRef Filename, const ParseInputs &Inputs,
   // breaks many features. Disable it for the main-file (not preamble).
   CI->getLangOpts()->DelayedTemplateParsing = false;
 
+  std::vector<std::unique_ptr<FeatureModule::ASTListener>> ASTListeners;
+  if (Inputs.FeatureModules) {
+    for (auto &M : *Inputs.FeatureModules) {
+      if (auto Listener = M.astListeners())
+        ASTListeners.emplace_back(std::move(Listener));
+    }
+  }
   StoreDiags ASTDiags;
+  ASTDiags.setDiagCallback(
+      [&ASTListeners](const clang::Diagnostic &D, clangd::Diag &Diag) {
+        llvm::for_each(ASTListeners,
+                       [&](const auto &L) { L->sawDiagnostic(D, Diag); });
+      });
 
   llvm::Optional<PreamblePatch> Patch;
   bool PreserveDiags = true;
@@ -318,7 +332,7 @@ ParsedAST::build(llvm::StringRef Filename, const ParseInputs &Inputs,
       Check->registerMatchers(&CTFinder);
     }
 
-    const Config& Cfg = Config::current();
+    const Config &Cfg = Config::current();
     ASTDiags.setLevelAdjuster([&](DiagnosticsEngine::Level DiagLevel,
                                   const clang::Diagnostic &Info) {
       if (Cfg.Diagnostics.SuppressAll ||

Diff for: clang-tools-extra/clangd/Preamble.cpp

@@ -325,7 +325,19 @@ buildPreamble(PathRef FileName, CompilerInvocation CI,
 
   trace::Span Tracer("BuildPreamble");
   SPAN_ATTACH(Tracer, "File", FileName);
+  std::vector<std::unique_ptr<FeatureModule::ASTListener>> ASTListeners;
+  if (Inputs.FeatureModules) {
+    for (auto &M : *Inputs.FeatureModules) {
+      if (auto Listener = M.astListeners())
+        ASTListeners.emplace_back(std::move(Listener));
+    }
+  }
   StoreDiags PreambleDiagnostics;
+  PreambleDiagnostics.setDiagCallback(
+      [&ASTListeners](const clang::Diagnostic &D, clangd::Diag &Diag) {
+        llvm::for_each(ASTListeners,
+                       [&](const auto &L) { L->sawDiagnostic(D, Diag); });
+      });
   llvm::IntrusiveRefCntPtr<DiagnosticsEngine> PreambleDiagsEngine =
       CompilerInstance::createDiagnostics(&CI.getDiagnosticOpts(),
                                           &PreambleDiagnostics, false);

Diff for: clang-tools-extra/clangd/unittests/ClangdLSPServerTests.cpp

@@ -365,6 +365,27 @@ TEST_F(LSPTest, FeatureModulesThreadingTest) {
   // And immediately shut down. FeatureModule destructor verifies we blocked.
 }
 
+TEST_F(LSPTest, DiagModuleTest) {
+  static constexpr llvm::StringLiteral DiagMsg = "DiagMsg";
+  class DiagModule final : public FeatureModule {
+    struct DiagHooks : public ASTListener {
+      void sawDiagnostic(const clang::Diagnostic &, clangd::Diag &D) override {
+        D.Message = DiagMsg.str();
+      }
+    };
+
+  public:
+    std::unique_ptr<ASTListener> astListeners() override {
+      return std::make_unique<DiagHooks>();
+    }
+  };
+  FeatureModules.add(std::make_unique<DiagModule>());
+
+  auto &Client = start();
+  Client.didOpen("foo.cpp", "test;");
+  EXPECT_THAT(Client.diagnostics("foo.cpp"),
+              llvm::ValueIs(testing::ElementsAre(DiagMessage(DiagMsg))));
+}
 } // namespace
 } // namespace clangd
 } // namespace clang

Diff for: clang-tools-extra/clangd/unittests/DiagnosticsTests.cpp

@@ -9,6 +9,7 @@
 #include "Annotations.h"
 #include "Config.h"
 #include "Diagnostics.h"
+#include "FeatureModule.h"
 #include "ParsedAST.h"
 #include "Protocol.h"
 #include "SourceCode.h"
@@ -26,6 +27,7 @@
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 #include <algorithm>
+#include <memory>
 
 namespace clang {
 namespace clangd {
@@ -1346,6 +1348,31 @@ TEST(ToLSPDiag, RangeIsInMain) {
              });
 }
 
+TEST(ParsedASTTest, ModuleSawDiag) {
+  static constexpr const llvm::StringLiteral KDiagMsg = "StampedDiag";
+  struct DiagModifierModule final : public FeatureModule {
+    struct Listener : public FeatureModule::ASTListener {
+      void sawDiagnostic(const clang::Diagnostic &Info,
+                         clangd::Diag &Diag) override {
+        Diag.Message = KDiagMsg.str();
+      }
+    };
+    std::unique_ptr<ASTListener> astListeners() override {
+      return std::make_unique<Listener>();
+    };
+  };
+  FeatureModuleSet FMS;
+  FMS.add(std::make_unique<DiagModifierModule>());
+
+  Annotations Code("[[test]]; /* error-ok */");
+  TestTU TU;
+  TU.Code = Code.code().str();
+  TU.FeatureModules = &FMS;
+
+  auto AST = TU.build();
+  EXPECT_THAT(*AST.getDiagnostics(),
+              testing::Contains(Diag(Code.range(), KDiagMsg.str())));
+}
 } // namespace
 } // namespace clangd
 } // namespace clang

Diff for: clang-tools-extra/clangd/unittests/TestTU.cpp

@@ -36,6 +36,7 @@ ParseInputs TestTU::inputs(MockFS &FS) const {
   FS.Files[ImportThunk] = ThunkContents;
 
   ParseInputs Inputs;
+  Inputs.FeatureModules = FeatureModules;
   auto &Argv = Inputs.CompileCommand.CommandLine;
   Argv = {"clang"};
   // FIXME: this shouldn't need to be conditional, but it breaks a

Diff for: clang-tools-extra/clangd/unittests/TestTU.h

@@ -19,12 +19,14 @@
 
 #include "../TidyProvider.h"
 #include "Compiler.h"
+#include "FeatureModule.h"
 #include "ParsedAST.h"
 #include "TestFS.h"
 #include "index/Index.h"
 #include "support/Path.h"
 #include "llvm/ADT/StringMap.h"
 #include "gtest/gtest.h"
+#include <memory>
 #include <string>
 #include <utility>
 #include <vector>
@@ -76,6 +78,8 @@ struct TestTU {
   // to eliminate this option some day.
   bool OverlayRealFileSystemForModules = false;
 
+  FeatureModuleSet *FeatureModules = nullptr;
+
   // By default, build() will report Error diagnostics as GTest errors.
   // Suppress this behavior by adding an 'error-ok' comment to the code.
   // The result will always have getDiagnostics() populated.