[clangd][WIP] Introduce DiagnosticAugmentationModule

[kadircet]

Summary

A module interface to enable custom contributions for diagnostic fixes.
Interface also allows contributors to return a Command which can be
computed outside of the latency sensitive diagnostic release cycle.

[kadircet]

Note that this doesn't handle the execution yet, it is mostly to demonstrate the idea and to get some early feedback.

[kadircet]

• Integrate with parseinputs

[kadircet]

see https://reviews.llvm.org/D97555 for revised version

clang-tools-extra/clangd/ClangdLSPServer.cpp

llvm::json::Object Result{
             {"selectionRangeProvider", true},
             {"documentSymbolProvider", true},
             {"workspaceSymbolProvider", true},
             {"referencesProvider", true},
             {"astProvider", true}, // clangd extension
             {"executeCommandProvider",
              llvm::json::Object{
                  {"commands",
-                  {ExecuteCommandParams::CLANGD_APPLY_FIX_COMMAND,
-                   ExecuteCommandParams::CLANGD_APPLY_TWEAK}},
+                  {
+                      ExecuteCommandParams::CLANGD_APPLY_FIX_COMMAND,
+                      ExecuteCommandParams::CLANGD_APPLY_TWEAK,
+                      ExecuteCommandParams::CLANGD_APPLY_AUGMENTED_FIX,
+                  }},
              }},
             {"typeHierarchyProvider", true},
             {"memoryUsageProvider", true}, // clangd extension
             {"compilationDatabase",        // clangd extension
              llvm::json::Object{{"automaticReload", true}}},
             {"callHierarchyProvider", true},
         }}}};
   if (Opts.Encoding)

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 = Version.str();
   Inputs.ForceRebuild = ForceRebuild;
   Inputs.Opts = std::move(Opts);
   Inputs.Index = Index;
   Inputs.ClangTidyProvider = ClangTidyProvider;
+  Inputs.Modules = Modules;
   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);
 }
 
 std::function<Context(PathRef)>
 ClangdServer::createConfiguredContextProvider(const config::Provider *Provider,

clang-tools-extra/clangd/Compiler.h

 // 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 "GlobalCompilationDatabase.h"
+#include "Module.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"
 
 namespace clang {
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 = {};
+  ModuleSet *Modules = 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

 //===--- Diagnostics.h -------------------------------------------*- 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_DIAGNOSTICS_H
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_DIAGNOSTICS_H
 
+#include "Module.h"
 #include "Protocol.h"
+#include "SourceCode.h"
 #include "support/Path.h"
 #include "clang/Basic/Diagnostic.h"
 #include "clang/Basic/LangOptions.h"
 #include "clang/Basic/SourceLocation.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 <string>
 
 namespace clang {
 namespace tidy {
 class ClangTidyContext;
 } // namespace tidy
 namespace clangd {
 
+/// Can be used to provide additional fixes to diagnostics.
+/// Returned codeaction might just contain a command instead to postpone
+/// computation of expensive edits or calculate them outside the critical path.
+/// This command will be invoked via executeCommand later on, if user requests.
+class DiagnosticsAugmentationModule : public Module {
+public:
+  /// Returns an action that can be invoked to fix the diagnostic.
+  virtual llvm::Optional<CodeAction>
+  augmentDiag(DiagnosticsEngine::Level L, const clang::Diagnostic &Diag) = 0;
+
+  // FIXME: Lift Tweak::Effect out of Tweak.h and use it here.
+  virtual llvm::Expected<FileEdits> executeCommand(llvm::json::Value Args) = 0;
+};
+
 struct ClangdDiagnosticOptions {
   /// If true, Clangd uses an LSP extension to embed the fixes with the
   /// diagnostics that are sent to the client.
   bool EmbedFixesInDiagnostics = false;
 
   /// If true, Clangd uses the relatedInformation field to include other
   /// locations (in particular attached notes).
   /// Otherwise, these are flattened into the diagnostic message.
public:
   using LevelAdjuster = std::function<DiagnosticsEngine::Level(
       DiagnosticsEngine::Level, const clang::Diagnostic &)>;
   /// 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; }
+  void setAugmenters(llvm::ArrayRef<DiagnosticsAugmentationModule *> Modules) {
+    this->Modules = Modules;
+  }
 
 private:
   void flushLastDiag();
 
   DiagFixer Fixer = nullptr;
   LevelAdjuster Adjuster = nullptr;
+  llvm::ArrayRef<DiagnosticsAugmentationModule *> Modules = {};
   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());
     }
+    for (auto *Module : Modules) {
+      if (auto Action = Module->augmentDiag(DiagLevel, Info))
+        LastDiag->Actions.emplace_back(std::move(*Action));
+    }
   } 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/ParsedAST.cpp

if (Inputs.Index && !BuildDir.getError()) {
     FixIncludes.emplace(Filename, Inserter, *Inputs.Index,
                         /*IndexRequestLimit=*/5);
     ASTDiags.contributeFixes([&FixIncludes](DiagnosticsEngine::Level DiagLevl,
                                             const clang::Diagnostic &Info) {
       return FixIncludes->fix(DiagLevl, Info);
     });
     Clang->setExternalSemaSource(FixIncludes->unresolvedNameRecorder());
   }
+  if (Inputs.Modules) {
+    ASTDiags.setAugmenters(
+        Inputs.Modules->getModules<DiagnosticsAugmentationModule>());
+  }
 
   IncludeStructure Includes;
   // If we are using a preamble, copy existing includes.
   if (Preamble) {
     Includes = Preamble->Includes;
     Includes.MainFileIncludes = Patch->preambleIncludes();
     // Replay the preamble includes so that clang-tidy checks can see them.
     ReplayPreamble::attach(Patch->preambleIncludes(), *Clang,

clang-tools-extra/clangd/Preamble.cpp

buildPreamble(PathRef FileName, CompilerInvocation CI,
   // 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);
   StoreDiags PreambleDiagnostics;
+  if (Inputs.Modules) {
+    PreambleDiagnostics.setAugmenters(
+        Inputs.Modules->getModules<DiagnosticsAugmentationModule>());
+  }
   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/Protocol.h

struct TweakArgs {
   /// A selection provided by the client on a textDocument/codeAction request.
   Range selection;
   /// ID of the tweak that should be executed. Corresponds to Tweak::id().
   std::string tweakID;
 };
 bool fromJSON(const llvm::json::Value &, TweakArgs &, llvm::json::Path);
 llvm::json::Value toJSON(const TweakArgs &A);
 
+/// Arguments for applyAugmentedFix command. These are send as code actions from
+/// the server in response to diagnostics.
+struct AugmentedDiagArgs {
+  /// ID of the module that generated the action.
+  std::string moduleID;
+  /// Augmenters arguments are opaque and should be parsed by them.
+  llvm::json::Value args;
+};
+
 /// Exact commands are not specified in the protocol so we define the
 /// ones supported by Clangd here. The protocol specifies the command arguments
 /// to be "any[]" but to make this safer and more manageable, each command we
 /// handle maps to a certain llvm::Optional of some struct to contain its
 /// arguments. Different commands could reuse the same llvm::Optional as
 /// arguments but a command that needs different arguments would simply add a
 /// new llvm::Optional and not use any other ones. In practice this means only
 /// one argument type will be parsed and set.
 struct ExecuteCommandParams {
   // Command to apply fix-its. Uses WorkspaceEdit as argument.
   const static llvm::StringLiteral CLANGD_APPLY_FIX_COMMAND;
   // Command to apply the code action. Uses TweakArgs as argument.
   const static llvm::StringLiteral CLANGD_APPLY_TWEAK;
+  // Command to apply a fix generated by a DiagnosticsAugmentationModule. Uses
+  // AugmentedDiagArgs as argument.
+  const static llvm::StringLiteral CLANGD_APPLY_AUGMENTED_FIX;
 
   /// The command identifier, e.g. CLANGD_APPLY_FIX_COMMAND
   std::string command;
 
   // Arguments
   llvm::Optional<WorkspaceEdit> workspaceEdit;
   llvm::Optional<TweakArgs> tweakArgs;
+  llvm::Optional<AugmentedDiagArgs> augmentDiagArgs;
 };
 bool fromJSON(const llvm::json::Value &, ExecuteCommandParams &,
               llvm::json::Path);
 
 struct Command : public ExecuteCommandParams {
   std::string title;
 };
 llvm::json::Value toJSON(const Command &C);

clang-tools-extra/clangd/Protocol.cpp

bool fromJSON(const llvm::json::Value &Params, WorkspaceEdit &R,
   llvm::json::ObjectMapper O(Params, P);
   return O && O.map("changes", R.changes);
 }
 
 const llvm::StringLiteral ExecuteCommandParams::CLANGD_APPLY_FIX_COMMAND =
     "clangd.applyFix";
 const llvm::StringLiteral ExecuteCommandParams::CLANGD_APPLY_TWEAK =
     "clangd.applyTweak";
+const llvm::StringLiteral ExecuteCommandParams::CLANGD_APPLY_AUGMENTED_FIX =
+    "clangd.applyAugmentedFix";
 
 bool fromJSON(const llvm::json::Value &Params, ExecuteCommandParams &R,
               llvm::json::Path P) {
   llvm::json::ObjectMapper O(Params, P);
   if (!O || !O.map("command", R.command))
     return false;
 
   const auto *Args = Params.getAsObject()->getArray("arguments");

clang-tools-extra/clangd/test/initialize-params.test

 # CHECK-NEXT:        "firstTriggerCharacter": "\n",
 # CHECK-NEXT:        "moreTriggerCharacter": []
 # CHECK-NEXT:      },
 # CHECK-NEXT:      "documentRangeFormattingProvider": true,
 # CHECK-NEXT:      "documentSymbolProvider": true,
 # CHECK-NEXT:      "executeCommandProvider": {
 # CHECK-NEXT:        "commands": [
 # CHECK-NEXT:          "clangd.applyFix",
-# CHECK-NEXT:          "clangd.applyTweak"
+# CHECK-NEXT:          "clangd.applyTweak",
+# CHECK-NEXT:          "clangd.applyAugmentedFix"
 # CHECK-NEXT:        ]
 # CHECK-NEXT:      },
 # CHECK-NEXT:      "hoverProvider": true,
 # CHECK-NEXT:      "implementationProvider": true,
 # CHECK-NEXT:      "memoryUsageProvider": true,
 # CHECK-NEXT:      "referencesProvider": true,
 # CHECK-NEXT:      "renameProvider": true,
 # CHECK-NEXT:      "selectionRangeProvider": true,