[clangd] Introduce ASTHooks to FeatureModules

[kadircet]

Summary

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.

[kadircet]

• Add tests

[sammccall]

I'm getting a little nervous about the amount of stuff we're packing into modules without in-tree examples.
I should split out some of the "standard" features into modules as that's possible already.

My model for modules using this diagnostic stuff (apart from the build-system stuff which sadly can't be meaningfully upstreamed) are IncludeFixer, ClangTidy, and IWYU - worth thinking about how we'd build those on top of this model. (Doesn't mean we need to add a hook to emit diagnostics, but it probably means we should know where it would go)

clang-tools-extra/clangd/FeatureModule.h

104

naming: giving this a plural name and having a collection of them is a little confusing (is Hooks the hooks from one module, or from all of them?). What about ASTListener?
(Listener suggests passive but listeners that "participate" is a common enough idiom I think)

112

This comment hints at what I *thought* was the idea: astHooks() is called each time we parse a file, the returned object has methods called on it while the file is being parsed, and is then destroyed.

But the code suggests we call once globally and it has effectively the same lifetime as the module.
This seems much less useful, e.g. if we want to observe several diagnostics, examine the preamble, and emit new diagnostics, then we have to plumb around some notion of "AST identity" rather than just tying it to the identity of the ParseASTHooks object itself.

(Lots of natural extensions here like providing ParseInputs to astHooks(), but YAGNI for now)

[kadircet]

In D98499#2626502, @sammccall wrote:

My model for modules using this diagnostic stuff (apart from the build-system stuff which sadly can't be meaningfully upstreamed) are IncludeFixer, ClangTidy, and IWYU - worth thinking about how we'd build those on top of this model. (Doesn't mean we need to add a hook to emit diagnostics, but it probably means we should know where it would go)

Agreed. I believe they all would need extra end points to ASTHooks though.

ClangTidy:

• needs extra hooks to register PP callbacks, and take in a diagnostics engine
• needs a new endpoint to traverse ast and *emit* diags
• CTContext needs to be alive until the parsing is done: so we can:
  • make ASTHooks own it and instantiate a new one on every parse (i think the cleanest and most explicit)
  • make the module own them and control the lifetime with entry/exit calls on the asthooks. (there's more burden on the modules now, and they'll need extra synchronisation on enter/exit calls)
  • return some other object on parsing entry hook that'll be kept alive by the caller (needs design around semantics of that object).

IncludeFixer:

• needs a new hook to act as an external sema source, so that it can fix unresolved names.
• similar to CTContext issue, it references per-tu state like HeaderSearchInfo, so will need to incorporate these into entry hook, and somehow disambiguate hooks for different TUs (all the options proposed for tidy).

IWYU:

• I'd expect this to make use of ~same API as IncludeFixer.

PreamblePatching:

• can mutate compiler instance via a hook
• drop diagnostics from ParsedAST on exit (requires some mutation on parsedast though)

clang-tools-extra/clangd/FeatureModule.h

112

This comment hints at what I *thought* was the idea: astHooks() is called each time we parse a file, the returned object has methods called on it while the file is being parsed, and is then destroyed.

This was the intent actually (to enable modularization of other features like clangtidy, includefixer, etc. as mentioned above), but looks like i got confused while integrating this into clangdserver :/

While trying to resolve my confusion i actually noticed that we cannot uphold the contract of "hooks being called synchronously", because we actually have both a preamblethread and astworker that can invoke hooks (embarrassing of me to forget that 😓).

So we can:

• Give up on that promise and make life slightly complicated for module implementers
• Don't invoke hooks from preamblethread, (at the cost of limiting functionality, we can still have downstream users that act on PPCallbacks, but no direct integration with compiler, and that's where layering violations are generated :/)
• Handle the synchronization ourselves, only complicates TUScheduler more, rather than all the module implementers.
• Propogate FeatureModuleSet into TUScheduler and create 2 set of hooks on each thread :)

I am leaning towards 4, but unsure (mostly hesitant about dependency schema, as featuremodules also depend on TUScheduler..). WDYT?

[kadircet]

• Rename ParsedASTHooks to Listeners.
• Generate list of hooks on each parse.

[sammccall]

In D98499#2644313, @kadircet wrote:

In D98499#2626502, @sammccall wrote:

My model for modules using this diagnostic stuff (apart from the build-system stuff which sadly can't be meaningfully upstreamed) are IncludeFixer, ClangTidy, and IWYU - worth thinking about how we'd build those on top of this model. (Doesn't mean we need to add a hook to emit diagnostics, but it probably means we should know where it would go)

Agreed. I believe they all would need extra end points to ASTHooks though.

Yup. Let's not bite off more for now.

• make ASTHooks own it and instantiate a new one on every parse (i think the cleanest and most explicit)

Agreed. (And this only one that overlaps this patch a lot)

clang-tools-extra/clangd/FeatureModule.h

112

we actually have both a preamblethread and astworker that can invoke hooks (embarrassing of me to forget that 😓)

Ha! And yes, in our motivating case of fixing build system rules, the diagnostics are mostly going to be in the preamble.

The options that seem most tempting to me are:

• don't attempt to create/run astHooks while building preambles, but *do* feed the preamble's Diags into the main-file's AST hooks every time it's used. (you won't have a clang::Diagnostic, so that param would have to be gone/optional, and we'd scrape the messages instead of extracting args). This is kind of thematic, remember how we replay PPCallbacks for clang-tidy :-). This is the smallest tweak to your #2 that actually works for us, I think.
• create one astHooks for the preamble, and another for the AST build, and try to make the interface suitable for both. This is cute but there may be too much tension between the two cases. (Is this what you mean by #4?)
• or have separate ASTHooks & PreambleHooks interfaces and support all this crap for both. (Or is *this* what you mean by #4?) Hybrid is also possible, give the interfaces an inheritance relationship, or have one interface but pass boolean parameters to indicate which version we're doing, or...
• preambles and ASTs are a tree, so... give modules a PreambleHooks factory, and the factory function for ASTHooks is PreambleHooks::astHooks(). Holy overengineering batman...

These are roughly in order of complexity so we should probably start toward the top of the list somewhere. Up to you.

(I don't like the #1 or #3 in your list above much at all.)

[kadircet]

clang-tools-extra/clangd/FeatureModule.h

112

create one astHooks for the preamble, and another for the AST build, and try to make the interface suitable for both. This is cute but there may be too much tension between the two cases. (Is this what you mean by #4?)

yes this is what i meant by #4.

These are roughly in order of complexity so we should probably start toward the top of the list somewhere. Up to you.

I'd lean towards the second option(creating separate ASTHooks with the same interface for preamble and astworker). As discussed before first one (i.e. scraping diag message) is a limited solution that's likely to get us into a corner in the future. let me know if you have any concerns around this approach.

[kadircet]

• Have 2 separate hooks for preamble and mainfile ASTs, as they are produced async.

[sammccall]

clang-tools-extra/clangd/ClangdServer.cpp

237

Now we're creating two ASTListeners for each version of the file seen, but most won't be parsed twice (and some won't even be parsed once).

This seems like it's going to be a weird wart for e.g. hooks that expect to do something interesting when a file finishes, or those that allocate something heavyweight.

I'd suggest sinking the ModuleSet into ParsedAST::build etc, or at least moving the astHooks() calls into TUScheduler. (And documenting that astHooks() should be threadsafe)

Otherwise, need to document that sometimes hooks are created and then never used.

clang-tools-extra/clangd/Compiler.h

60

I think having these in ParseInputs is conceptually confusing and shared_ptr is risky.
For example we copy inputs into ASTWorker::FileInputs, which extends the lifetime of the listeners in a way that seems undesirable.

I think these listeners should be clearly either:

• bound to the lifetime of the ParsedAST (and therefore owned by it)
• bound to the *creation* of the ParsedAST (and therefore strictly scoped within ParsedAST::build)

You could consider the ModuleSet to be part of ParseInputs, though...

clang-tools-extra/clangd/Diagnostics.h

150

ArrayRef<shared_ptr<ASTListener>> seems like we're leaking a lot. And even mentioning AST hooks seems dubious layering-wise.
Can we use a std::function like LevelAdjuster?

clang-tools-extra/clangd/FeatureModule.h

101

nit: this describes the interaction more but not what it's for.
Maybe

/// 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.

113

the "hooks" name is still used here and throughout

[kadircet]

• Pass FeatureModuleSet rather than astListeners in ParseInputs.
• Create listeners only when building ASTs.
• Use a callback function in StoreDiags to notify about diagnostics.

[sammccall]

clang-tools-extra/clangd/FeatureModule.h

106

comment: listeners are destroyed once the AST is built.

182

this seems like a slightly weird thing to have in FeatureModule.h, I'd be tempted to inline it in the callsites instead, but up to you.

[kadircet]

• Inline helper to call-sites
• Add comments about destruction of ast listeners.

clang-tools-extra/clangd/ClangdServer.h

 #include "support/ThreadsafeFS.h"
 #include "clang/Tooling/CompilationDatabase.h"
 #include "clang/Tooling/Core/Replacement.h"
 #include "llvm/ADT/FunctionExtras.h"
 #include "llvm/ADT/IntrusiveRefCntPtr.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/StringRef.h"
 #include <functional>
+#include <memory>
 #include <string>
 #include <type_traits>
 #include <utility>
 #include <vector>
 
 namespace clang {
 namespace clangd {
 /// Manages a collection of source files and derived data (ASTs, indexes),

clang-tools-extra/clangd/ClangdServer.cpp

void ClangdServer::addDocument(PathRef File, llvm::StringRef Contents,
   ParseInputs Inputs;
   Inputs.TFS = &TFS;
   Inputs.Contents = std::string(Contents);
   Inputs.Version = std::move(ActualVersion);
   Inputs.ForceRebuild = ForceRebuild;
   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)
     BackgroundIdx->boostRelated(File);
 }
 
 void ClangdServer::reparseOpenFilesIfNeeded(
     llvm::function_ref<bool(llvm::StringRef File)> Filter) {
   // Reparse only opened files that were modified.
   for (const Path &FilePath : DraftMgr.getActiveFiles())
     if (Filter(FilePath))

clang-tools-extra/clangd/Compiler.h

 // Most callers will use this through Preamble/ParsedAST, but some features like
 // CodeComplete run their own compile actions that share these low-level pieces.
 //
 //===----------------------------------------------------------------------===//
 
 #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 {
 
 class IgnoreDiagnostics : public DiagnosticConsumer {
 public:
   static void log(DiagnosticsEngine::Level DiagLevel,
                   const clang::Diagnostic &Info);
struct ParseInputs {
   std::string Version = "null";
   // Prevent reuse of the cached preamble/AST. Slow! Useful to workaround
   // clangd's assumption that missing header files will stay missing.
   bool ForceRebuild = false;
   // Used to recover from diagnostics (e.g. find missing includes for symbol).
   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.
 std::unique_ptr<CompilerInvocation>
 buildCompilerInvocation(const ParseInputs &Inputs, clang::DiagnosticConsumer &D,
                         std::vector<std::string> *CC1Args = nullptr);
 
 /// Creates a compiler instance, configured so that:

clang-tools-extra/clangd/Diagnostics.h

 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/None.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/StringSet.h"
 #include "llvm/Support/SourceMgr.h"
 #include <cassert>
+#include <functional>
+#include <memory>
 #include <string>
+#include <utility>
 
 namespace clang {
 namespace tidy {
 class ClangTidyContext;
 } // namespace tidy
 namespace clangd {
 
 struct ClangdDiagnosticOptions {
public:
   void EndSourceFile() override;
   void HandleDiagnostic(DiagnosticsEngine::Level DiagLevel,
                         const clang::Diagnostic &Info) override;
 
   using DiagFixer = std::function<std::vector<Fix>(DiagnosticsEngine::Level,
                                                    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;
   llvm::Optional<FullSourceLoc> LastDiagLoc; // Valid only when LastDiag is set.
   bool LastDiagOriginallyError = false;      // Valid only when LastDiag is set.
   SourceManager *OrigSrcMgr = nullptr;
 
   llvm::DenseSet<std::pair<unsigned, unsigned>> IncludedErrorLocations;

clang-tools-extra/clangd/Diagnostics.cpp

if (!isNote(DiagLevel)) {
 
     if (!Info.getFixItHints().empty())
       AddFix(true /* try to invent a message instead of repeating the diag */);
     if (Fixer) {
       auto ExtraFixes = Fixer(DiagLevel, Info);
       LastDiag->Fixes.insert(LastDiag->Fixes.end(), ExtraFixes.begin(),
                              ExtraFixes.end());
     }
+    if (DiagCB)
+      DiagCB(Info, *LastDiag);
   } else {
     // Handle a note to an existing diagnostic.
 
     // If a diagnostic was suppressed due to the suppression filter,
     // also suppress notes associated with it.
     if (LastPrimaryDiagnosticWasSuppressed) {
       return;
     }

clang-tools-extra/clangd/FeatureModule.h

 //===--- FeatureModule.h - Plugging features into clangd ----------*-C++-*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_CLANG_TOOLS_EXTRA_CLANGD_FEATUREMODULE_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_FEATUREMODULE_H
 
 #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"
 #include "llvm/Support/JSON.h"
 #include <memory>
 #include <type_traits>
 #include <vector>
 
 namespace clang {
 namespace clangd {
+struct Diag;
 class LSPBinder;
 class SymbolIndex;
 class ThreadsafeFS;
 class TUScheduler;
 class Tweak;
 
 /// A FeatureModule contributes a vertical feature to clangd.
 ///
public:
   /// FeatureModules should go idle quickly if stop() has been called.
   /// Called by the server when shutting down, and also by tests.
   virtual bool blockUntilIdle(Deadline) { return true; }
 
   /// Tweaks implemented by this module. Can be called asynchronously when
   /// 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();
   /// The scheduler is used to run tasks on worker threads and access ASTs.
   TUScheduler &scheduler() { return facilities().Scheduler; }
   /// The index is used to get information about the whole codebase.
   const SymbolIndex *index() { return facilities().Index; }
   /// The filesystem is used to read source files on disk.
public:
   }
   template <typename Mod> const Mod *get() const {
     return const_cast<FeatureModuleSet *>(this)->get<Mod>();
   }
 };
 
 template <typename Mod> int FeatureModuleSet::ID<Mod>::Key;
 
 } // namespace clangd
 } // namespace clang
 #endif

clang-tools-extra/clangd/ParsedAST.cpp

 //===--- ParsedAST.cpp -------------------------------------------*- C++-*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "ParsedAST.h"
 #include "../clang-tidy/ClangTidyCheck.h"
 #include "../clang-tidy/ClangTidyDiagnosticConsumer.h"
 #include "../clang-tidy/ClangTidyModuleRegistry.h"
 #include "AST.h"
 #include "Compiler.h"
 #include "Config.h"
 #include "Diagnostics.h"
+#include "FeatureModule.h"
 #include "Headers.h"
 #include "IncludeFixer.h"
 #include "Preamble.h"
 #include "SourceCode.h"
 #include "TidyProvider.h"
 #include "index/CanonicalIncludes.h"
 #include "index/Index.h"
 #include "support/Logger.h"
 #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"
 #include "clang/Basic/TokenKinds.h"
 #include "clang/Frontend/CompilerInstance.h"
 #include "clang/Frontend/CompilerInvocation.h"
 #include "clang/Frontend/FrontendActions.h"
 #include "clang/Frontend/Utils.h"
ParsedAST::build(llvm::StringRef Filename, const ParseInputs &Inputs,
   CI->getFrontendOpts().DisableFree = false;
   const PrecompiledPreamble *PreamblePCH =
       Preamble ? &Preamble->Preamble : nullptr;
 
   // This is on-by-default in windows to allow parsing SDK headers, but it
   // 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;
   if (Preamble) {
     Patch = PreamblePatch::create(Filename, Inputs, *Preamble);
     Patch->apply(*CI);
     PreserveDiags = Patch->preserveDiagnostics();
   }
llvm::erase_if(CTChecks, [&](const auto &Check) {
       return !Check->isLanguageVersionSupported(CTContext->getLangOpts());
     });
     Preprocessor *PP = &Clang->getPreprocessor();
     for (const auto &Check : CTChecks) {
       Check->registerPPCallbacks(Clang->getSourceManager(), PP, PP);
       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 ||
           isBuiltinDiagnosticSuppressed(Info.getID(), Cfg.Diagnostics.Suppress))
         return DiagnosticsEngine::Ignored;
       if (!CTChecks.empty()) {
         std::string CheckName = CTContext->getCheckName(Info.getID());
         bool IsClangTidyDiag = !CheckName.empty();

clang-tools-extra/clangd/Preamble.cpp

buildPreamble(PathRef FileName, CompilerInvocation CI,
   // Note that we don't need to copy the input contents, preamble can live
   // without those.
   auto ContentsBuffer =
       llvm::MemoryBuffer::getMemBuffer(Inputs.Contents, FileName);
   auto Bounds = ComputePreambleBounds(*CI.getLangOpts(), *ContentsBuffer, 0);
 
   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);
 
   // Skip function bodies when building the preamble to speed up building
   // the preamble and make it smaller.
   assert(!CI.getFrontendOpts().SkipFunctionBodies);
   CI.getFrontendOpts().SkipFunctionBodies = true;

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

TEST_F(LSPTest, FeatureModulesThreadingTest) {
   EXPECT_EQ(3, FeatureModules.get<AsyncCounter>()->getSync());
   // Throw some work on the queue to make sure shutdown blocks on it.
   Client.notify("increment", nullptr);
   Client.notify("increment", nullptr);
   Client.notify("increment", nullptr);
   // 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

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

 //===--- DiagnosticsTests.cpp ------------------------------------*- C++-*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "Annotations.h"
 #include "Config.h"
 #include "Diagnostics.h"
+#include "FeatureModule.h"
 #include "ParsedAST.h"
 #include "Protocol.h"
 #include "SourceCode.h"
 #include "TestFS.h"
 #include "TestIndex.h"
 #include "TestTU.h"
 #include "TidyProvider.h"
 #include "index/MemIndex.h"
 #include "support/Context.h"
 #include "support/Path.h"
 #include "clang/Basic/Diagnostic.h"
 #include "clang/Basic/DiagnosticSema.h"
 #include "llvm/Support/ScopedPrinter.h"
 #include "llvm/Support/TargetSelect.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 #include <algorithm>
+#include <memory>
 
 namespace clang {
 namespace clangd {
 namespace {
 
 using ::testing::_;
 using ::testing::AllOf;
 using ::testing::Contains;
TEST(ToLSPDiag, RangeIsInMain) {
   D.InsideMainFile = false;
   N.InsideMainFile = true;
   toLSPDiags(D, {}, Opts,
              [&](clangd::Diagnostic LSPDiag, ArrayRef<clangd::Fix>) {
                EXPECT_EQ(LSPDiag.range, N.Range);
              });
 }
 
+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

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

 //
 //===---------------------------------------------------------------------===//
 
 #ifndef LLVM_CLANG_TOOLS_EXTRA_UNITTESTS_CLANGD_TESTTU_H
 #define LLVM_CLANG_TOOLS_EXTRA_UNITTESTS_CLANGD_TESTTU_H
 
 #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>
 
 namespace clang {
 namespace clangd {
 
 struct TestTU {
struct TestTU {
   // disk and later read back.
   // FIXME: Change the way reading/writing modules work to allow us to keep them
   // in memory across multiple clang invocations, at least in tests, to
   // eliminate the need for real file system here.
   // Please avoid using this for things other than implicit modules. The plan is
   // 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.
   ParsedAST build() const;
   std::shared_ptr<const PreambleData>
   preamble(PreambleParsedCallback PreambleCallback = nullptr) const;
   ParseInputs inputs(MockFS &FS) const;
   SymbolSlab headerSymbols() const;

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

ParseInputs TestTU::inputs(MockFS &FS) const {
   std::string ThunkContents = "#import \"" + FullHeaderName + "\"\n";
 
   FS.Files = AdditionalFiles;
   FS.Files[FullFilename] = Code;
   FS.Files[FullHeaderName] = HeaderCode;
   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
   // GoToDefinition test for some reason (getMacroArgExpandedLocation fails).
   if (!HeaderCode.empty()) {
     Argv.push_back("-include");
     Argv.push_back(ImplicitHeaderGuard ? ImportThunk : FullHeaderName);
     // ms-compatibility changes the meaning of #import.