[mlir:LSP] Add a quickfix code action for inserting expected-* diagnostic checks

[rriddle]

Summary

This allows for automatically inserting expected checks for parser and verifier
diagnostics, which simplifies the workflow when building new dialect
constructs or extending existing ones.

[jpienaar]

Nice!

mlir/lib/Tools/lsp-server-support/Protocol.h

1117

Map too heavy to create/not useful?

mlir/lib/Tools/mlir-lsp-server/MLIRServer.cpp

841

And the above flushes on destruction?

[rriddle]

mlir/lib/Tools/lsp-server-support/Protocol.h

1117

These classes generally map directly to how the LSP spec specifies the data, so I write them how they are described there. Outside of that though, this vector is generally so small (generally 1-3 elements), that using vector is going to be more optimal than map anyways.

mlir/lib/Tools/mlir-lsp-server/MLIRServer.cpp

841

raw_string_ostream is unbuffered and should always write directly to the string, e.g. .str() just returns the string directly.

https://github.com/llvm/llvm-project/blob/0c1b32717bcffcf8edf95294e98933bd4c1e76ed/llvm/include/llvm/Support/raw_ostream.h#L630

mlir/docs/Tools/MLIRLSP.md

 
 #### Diagnostics
 
 The language server actively runs verification on the IR as you type, showing
 any generated diagnostics in-place.
 
 ![IMG](/mlir-lsp-server/diagnostics.png)
 
+##### Automatically insert `expected-` diagnostic checks
+
+MLIR provides
+[infrastructure](https://mlir.llvm.org/docs/Diagnostics/#sourcemgr-diagnostic-verifier-handler)
+for checking expected diagnostics, which is heavily utilized when defining IR
+parsing and verification. The language server provides code actions for
+automatically inserting the checks for diagnostics it knows about.
+
+![IMG](/mlir-lsp-server/diagnostics_action.gif)
+
 #### Code completion
 
 The language server provides suggestions as you type, offering completions for
 dialect constructs (such as attributes, operations, and types), block names, SSA
 value names, keywords, and more.
 
 ![IMG](/mlir-lsp-server/code_complete.gif)
 

mlir/lib/Tools/lsp-server-support/Protocol.h

 //===----------------------------------------------------------------------===//
 // ClientCapabilities
 //===----------------------------------------------------------------------===//
 
 struct ClientCapabilities {
   /// Client supports hierarchical document symbols.
   /// textDocument.documentSymbol.hierarchicalDocumentSymbolSupport
   bool hierarchicalDocumentSymbol = false;
+
+  /// Client supports CodeAction return value for textDocument/codeAction.
+  /// textDocument.codeAction.codeActionLiteralSupport.
+  bool codeActionStructure = false;
 };
 
 /// Add support for JSON serialization.
 bool fromJSON(const llvm::json::Value &value, ClientCapabilities &result,
               llvm::json::Path path);
 
 //===----------------------------------------------------------------------===//
 // InitializeParams
struct Location {
   }
 
   friend bool operator<(const Location &lhs, const Location &rhs) {
     return std::tie(lhs.uri, lhs.range) < std::tie(rhs.uri, rhs.range);
   }
 };
 
 /// Add support for JSON serialization.
+bool fromJSON(const llvm::json::Value &value, Location &result,
+              llvm::json::Path path);
 llvm::json::Value toJSON(const Location &value);
 raw_ostream &operator<<(raw_ostream &os, const Location &value);
 
 //===----------------------------------------------------------------------===//
 // TextDocumentPositionParams
 //===----------------------------------------------------------------------===//
 
 struct TextDocumentPositionParams {

 //===----------------------------------------------------------------------===//
 // DiagnosticRelatedInformation
 //===----------------------------------------------------------------------===//
 
 /// Represents a related message and source code location for a diagnostic.
 /// This should be used to point to code locations that cause or related to a
 /// diagnostics, e.g. when duplicating a symbol in a scope.
 struct DiagnosticRelatedInformation {
+  DiagnosticRelatedInformation() = default;
   DiagnosticRelatedInformation(Location location, std::string message)
       : location(std::move(location)), message(std::move(message)) {}
 
   /// The location of this related diagnostic information.
   Location location;
   /// The message of this related diagnostic information.
   std::string message;
 };
 
 /// Add support for JSON serialization.
+bool fromJSON(const llvm::json::Value &value,
+              DiagnosticRelatedInformation &result, llvm::json::Path path);
 llvm::json::Value toJSON(const DiagnosticRelatedInformation &info);
 
 //===----------------------------------------------------------------------===//
 // Diagnostic
 //===----------------------------------------------------------------------===//
 
 enum class DiagnosticSeverity {
   /// It is up to the client to interpret diagnostics as error, warning, info or
struct Diagnostic {
   /// An LSP extension that's used to send the name of the category over to the
   /// client. The category typically describes the compilation stage during
   /// which the issue was produced, e.g. "Semantic Issue" or "Parse Issue".
   Optional<std::string> category;
 };
 
 /// Add support for JSON serialization.
 llvm::json::Value toJSON(const Diagnostic &diag);
+bool fromJSON(const llvm::json::Value &value, Diagnostic &result,
+              llvm::json::Path path);
 
 //===----------------------------------------------------------------------===//
 // PublishDiagnosticsParams
 //===----------------------------------------------------------------------===//
 
 struct PublishDiagnosticsParams {
   PublishDiagnosticsParams(URIForFile uri, int64_t version)
       : uri(std::move(uri)), version(version) {}

 };
 
 /// Add support for JSON serialization.
 llvm::json::Value toJSON(const InlayHint &);
 bool operator==(const InlayHint &lhs, const InlayHint &rhs);
 bool operator<(const InlayHint &lhs, const InlayHint &rhs);
 llvm::raw_ostream &operator<<(llvm::raw_ostream &os, InlayHintKind value);
 
+//===----------------------------------------------------------------------===//
+// CodeActionContext
+//===----------------------------------------------------------------------===//
+
+struct CodeActionContext {
+  /// An array of diagnostics known on the client side overlapping the range
+  /// provided to the `textDocument/codeAction` request. They are provided so
+  /// that the server knows which errors are currently presented to the user for
+  /// the given range. There is no guarantee that these accurately reflect the
+  /// error state of the resource. The primary parameter to compute code actions
+  /// is the provided range.
+  std::vector<Diagnostic> diagnostics;
+
+  /// Requested kind of actions to return.
+  ///
+  /// Actions not of this kind are filtered out by the client before being
+  /// shown. So servers can omit computing them.
+  std::vector<std::string> only;
+};
+
+/// Add support for JSON serialization.
+bool fromJSON(const llvm::json::Value &value, CodeActionContext &result,
+              llvm::json::Path path);
+
+//===----------------------------------------------------------------------===//
+// CodeActionParams
+//===----------------------------------------------------------------------===//
+
+struct CodeActionParams {
+  /// The document in which the command was invoked.
+  TextDocumentIdentifier textDocument;
+
+  /// The range for which the command was invoked.
+  Range range;
+
+  /// Context carrying additional information.
+  CodeActionContext context;
+};
+
+/// Add support for JSON serialization.
+bool fromJSON(const llvm::json::Value &value, CodeActionParams &result,
+              llvm::json::Path path);
+
+//===----------------------------------------------------------------------===//
+// WorkspaceEdit
+//===----------------------------------------------------------------------===//
+
+struct WorkspaceEdit {
+  /// Holds changes to existing resources.
+  std::map<std::string, std::vector<TextEdit>> changes;
+
+  /// Note: "documentChanges" is not currently used because currently there is
+  /// no support for versioned edits.
+};
+
+/// Add support for JSON serialization.
+bool fromJSON(const llvm::json::Value &value, WorkspaceEdit &result,
+              llvm::json::Path path);
+llvm::json::Value toJSON(const WorkspaceEdit &value);
+
+//===----------------------------------------------------------------------===//
+// CodeAction
+//===----------------------------------------------------------------------===//
+
+/// A code action represents a change that can be performed in code, e.g. to fix
+/// a problem or to refactor code.
+///
+/// A CodeAction must set either `edit` and/or a `command`. If both are
+/// supplied, the `edit` is applied first, then the `command` is executed.
+struct CodeAction {
+  /// A short, human-readable, title for this code action.
+  std::string title;
+
+  /// The kind of the code action.
+  /// Used to filter code actions.
+  Optional<std::string> kind;
+  const static llvm::StringLiteral kQuickFix;
+  const static llvm::StringLiteral kRefactor;
+  const static llvm::StringLiteral kInfo;
+
+  /// The diagnostics that this code action resolves.
+  Optional<std::vector<Diagnostic>> diagnostics;
+
+  /// Marks this as a preferred action. Preferred actions are used by the
+  /// `auto fix` command and can be targeted by keybindings.
+  /// A quick fix should be marked preferred if it properly addresses the
+  /// underlying error. A refactoring should be marked preferred if it is the
+  /// most reasonable choice of actions to take.
+  bool isPreferred = false;
+
+  /// The workspace edit this code action performs.
+  Optional<WorkspaceEdit> edit;
+};
+
+/// Add support for JSON serialization.
+llvm::json::Value toJSON(const CodeAction &);
+
 } // namespace lsp
 } // namespace mlir
 
 namespace llvm {
 template <>
 struct format_provider<mlir::lsp::Position> {
   static void format(const mlir::lsp::Position &pos, raw_ostream &os,
                      StringRef style) {
     assert(style.empty() && "style modifiers for this type are not supported");
     os << pos;
   }
 };
 } // namespace llvm
 
 #endif

mlir/lib/Tools/lsp-server-support/Protocol.cpp

bool mlir::lsp::fromJSON(const llvm::json::Value &value,
   }
   if (const llvm::json::Object *textDocument = o->getObject("textDocument")) {
     if (const llvm::json::Object *documentSymbol =
             textDocument->getObject("documentSymbol")) {
       if (Optional<bool> hierarchicalSupport =
               documentSymbol->getBoolean("hierarchicalDocumentSymbolSupport"))
         result.hierarchicalDocumentSymbol = *hierarchicalSupport;
     }
+    if (auto *codeAction = textDocument->getObject("codeAction")) {
+      if (codeAction->getObject("codeActionLiteralSupport"))
+        result.codeActionStructure = true;
+    }
   }
   return true;
 }
 
 //===----------------------------------------------------------------------===//
 // InitializeParams
 //===----------------------------------------------------------------------===//
 

 raw_ostream &mlir::lsp::operator<<(raw_ostream &os, const Range &value) {
   return os << value.start << '-' << value.end;
 }
 
 //===----------------------------------------------------------------------===//
 // Location
 //===----------------------------------------------------------------------===//
 
+bool mlir::lsp::fromJSON(const llvm::json::Value &value, Location &result,
+                         llvm::json::Path path) {
+  llvm::json::ObjectMapper o(value, path);
+  return o && o.map("uri", result.uri) && o.map("range", result.range);
+}
+
 llvm::json::Value mlir::lsp::toJSON(const Location &value) {
   return llvm::json::Object{
       {"uri", value.uri},
       {"range", value.range},
   };
 }
 
 raw_ostream &mlir::lsp::operator<<(raw_ostream &os, const Location &value) {
bool mlir::lsp::fromJSON(const llvm::json::Value &value,
   llvm::json::ObjectMapper o(value, path);
   return o && o.map("textDocument", result.textDocument);
 }
 
 //===----------------------------------------------------------------------===//
 // DiagnosticRelatedInformation
 //===----------------------------------------------------------------------===//
 
+bool mlir::lsp::fromJSON(const llvm::json::Value &value,
+                         DiagnosticRelatedInformation &result,
+                         llvm::json::Path path) {
+  llvm::json::ObjectMapper o(value, path);
+  return o && o.map("location", result.location) &&
+         o.map("message", result.message);
+}
+
 llvm::json::Value mlir::lsp::toJSON(const DiagnosticRelatedInformation &info) {
   return llvm::json::Object{
       {"location", info.location},
       {"message", info.message},
   };
 }
 
 //===----------------------------------------------------------------------===//
if (diag.category)
     result["category"] = *diag.category;
   if (!diag.source.empty())
     result["source"] = diag.source;
   if (diag.relatedInformation)
     result["relatedInformation"] = *diag.relatedInformation;
   return std::move(result);
 }
 
+bool mlir::lsp::fromJSON(const llvm::json::Value &value, Diagnostic &result,
+                         llvm::json::Path path) {
+  llvm::json::ObjectMapper o(value, path);
+  if (!o)
+    return false;
+  int severity = 0;
+  if (!mapOptOrNull(value, "severity", severity, path))
+    return false;
+  result.severity = (DiagnosticSeverity)severity;
+
+  return o.map("range", result.range) && o.map("message", result.message) &&
+         mapOptOrNull(value, "category", result.category, path) &&
+         mapOptOrNull(value, "source", result.source, path) &&
+         mapOptOrNull(value, "relatedInformation", result.relatedInformation,
+                      path);
+}
+
 //===----------------------------------------------------------------------===//
 // PublishDiagnosticsParams
 //===----------------------------------------------------------------------===//
 
 llvm::json::Value mlir::lsp::toJSON(const PublishDiagnosticsParams &params) {
   return llvm::json::Object{
       {"uri", params.uri},
       {"diagnostics", params.diagnostics},
llvm::raw_ostream &mlir::lsp::operator<<(llvm::raw_ostream &os,
   switch (value) {
   case InlayHintKind::Parameter:
     return os << "parameter";
   case InlayHintKind::Type:
     return os << "type";
   }
   llvm_unreachable("Unknown InlayHintKind");
 }
+
+//===----------------------------------------------------------------------===//
+// CodeActionContext
+//===----------------------------------------------------------------------===//
+
+bool mlir::lsp::fromJSON(const llvm::json::Value &value,
+                         CodeActionContext &result, llvm::json::Path path) {
+  llvm::json::ObjectMapper o(value, path);
+  if (!o || !o.map("diagnostics", result.diagnostics))
+    return false;
+  o.map("only", result.only);
+  return true;
+}
+
+//===----------------------------------------------------------------------===//
+// CodeActionParams
+//===----------------------------------------------------------------------===//
+
+bool mlir::lsp::fromJSON(const llvm::json::Value &value,
+                         CodeActionParams &result, llvm::json::Path path) {
+  llvm::json::ObjectMapper o(value, path);
+  return o && o.map("textDocument", result.textDocument) &&
+         o.map("range", result.range) && o.map("context", result.context);
+}
+
+//===----------------------------------------------------------------------===//
+// WorkspaceEdit
+//===----------------------------------------------------------------------===//
+
+bool mlir::lsp::fromJSON(const llvm::json::Value &value, WorkspaceEdit &result,
+                         llvm::json::Path path) {
+  llvm::json::ObjectMapper o(value, path);
+  return o && o.map("changes", result.changes);
+}
+
+llvm::json::Value mlir::lsp::toJSON(const WorkspaceEdit &value) {
+  llvm::json::Object fileChanges;
+  for (auto &change : value.changes)
+    fileChanges[change.first] = llvm::json::Array(change.second);
+  return llvm::json::Object{{"changes", std::move(fileChanges)}};
+}
+
+//===----------------------------------------------------------------------===//
+// CodeAction
+//===----------------------------------------------------------------------===//
+
+const llvm::StringLiteral CodeAction::kQuickFix = "quickfix";
+const llvm::StringLiteral CodeAction::kRefactor = "refactor";
+const llvm::StringLiteral CodeAction::kInfo = "info";
+
+llvm::json::Value mlir::lsp::toJSON(const CodeAction &value) {
+  llvm::json::Object codeAction{{"title", value.title}};
+  if (value.kind)
+    codeAction["kind"] = *value.kind;
+  if (value.diagnostics)
+    codeAction["diagnostics"] = llvm::json::Array(*value.diagnostics);
+  if (value.isPreferred)
+    codeAction["isPreferred"] = true;
+  if (value.edit)
+    codeAction["edit"] = *value.edit;
+  return std::move(codeAction);
+}

mlir/lib/Tools/mlir-lsp-server/LSPServer.cpp

struct LSPServer {
 
   //===--------------------------------------------------------------------===//
   // Code Completion
 
   void onCompletion(const CompletionParams &params,
                     Callback<CompletionList> reply);
 
   //===--------------------------------------------------------------------===//
+  // Code Action
+
+  void onCodeAction(const CodeActionParams &params,
+                    Callback<llvm::json::Value> reply);
+
+  //===--------------------------------------------------------------------===//
   // Fields
   //===--------------------------------------------------------------------===//
 
   MLIRServer &server;
 
   /// An outgoing notification used to send diagnostics to the client when they
   /// are ready to be processed.
   OutgoingNotification<PublishDiagnosticsParams> publishDiagnostics;
llvm::json::Object serverCaps{
       {"hoverProvider", true},
 
       // For now we only support documenting symbols when the client supports
       // hierarchical symbols.
       {"documentSymbolProvider",
        params.capabilities.hierarchicalDocumentSymbol},
   };
 
+  // Per LSP, codeActionProvider can be either boolean or CodeActionOptions.
+  // CodeActionOptions is only valid if the client supports action literal
+  // via textDocument.codeAction.codeActionLiteralSupport.
+  serverCaps["codeActionProvider"] =
+      params.capabilities.codeActionStructure
+          ? llvm::json::Object{{"codeActionKinds",
+                                {CodeAction::kQuickFix, CodeAction::kRefactor,
+                                 CodeAction::kInfo}}}
+          : llvm::json::Value(true);
+
   llvm::json::Object result{
       {{"serverInfo",
         llvm::json::Object{{"name", "mlir-lsp-server"}, {"version", "0.0.0"}}},
        {"capabilities", std::move(serverCaps)}}};
   reply(std::move(result));
 }
 void LSPServer::onInitialized(const InitializedParams &) {}
 void LSPServer::onShutdown(const NoParams &, Callback<std::nullptr_t> reply) {

 // Code Completion
 
 void LSPServer::onCompletion(const CompletionParams &params,
                              Callback<CompletionList> reply) {
   reply(server.getCodeCompletion(params.textDocument.uri, params.position));
 }
 
 //===----------------------------------------------------------------------===//
+// Code Action
+
+void LSPServer::onCodeAction(const CodeActionParams &params,
+                             Callback<llvm::json::Value> reply) {
+  URIForFile uri = params.textDocument.uri;
+
+  // Check whether a particular CodeActionKind is included in the response.
+  auto isKindAllowed = [only(params.context.only)](StringRef kind) {
+    if (only.empty())
+      return true;
+    return llvm::any_of(only, [&](StringRef base) {
+      return kind.consume_front(base) && (kind.empty() || kind.startswith("."));
+    });
+  };
+
+  // We provide a code action for fixes on the specified diagnostics.
+  std::vector<CodeAction> actions;
+  if (isKindAllowed(CodeAction::kQuickFix))
+    server.getCodeActions(uri, params.range.start, params.context, actions);
+  reply(std::move(actions));
+}
+
+//===----------------------------------------------------------------------===//
 // Entry point
 //===----------------------------------------------------------------------===//
 
 LogicalResult lsp::runMlirLSPServer(MLIRServer &server,
                                     JSONTransport &transport) {
   LSPServer lspServer(server);
   MessageHandler messageHandler(transport);
 
LogicalResult lsp::runMlirLSPServer(MLIRServer &server,
   // Document Symbols
   messageHandler.method("textDocument/documentSymbol", &lspServer,
                         &LSPServer::onDocumentSymbol);
 
   // Code Completion
   messageHandler.method("textDocument/completion", &lspServer,
                         &LSPServer::onCompletion);
 
+  // Code Action
+  messageHandler.method("textDocument/codeAction", &lspServer,
+                        &LSPServer::onCodeAction);
+
   // Diagnostics
   lspServer.publishDiagnostics =
       messageHandler.outgoingNotification<PublishDiagnosticsParams>(
           "textDocument/publishDiagnostics");
 
   // Run the main loop of the transport.
   LogicalResult result = success();
   if (llvm::Error error = transport.run(messageHandler)) {
     Logger::error("Transport error: {0}", error);
     llvm::consumeError(std::move(error));
     result = failure();
   } else {
     result = success(lspServer.shutdownRequestReceived);
   }
   return result;
 }

mlir/lib/Tools/mlir-lsp-server/MLIRServer.h

 
 #include "mlir/Support/LLVM.h"
 #include <memory>
 
 namespace mlir {
 class DialectRegistry;
 
 namespace lsp {
+struct CodeAction;
+struct CodeActionContext;
 struct CompletionList;
 struct Diagnostic;
 struct DocumentSymbol;
 struct Hover;
 struct Location;
 struct Position;
+struct Range;
 class URIForFile;
 
 /// This class implements all of the MLIR related functionality necessary for a
 /// language server. This class allows for keeping the MLIR specific logic
 /// separate from the logic that involves LSP server/client communication.
 class MLIRServer {
 public:
   /// Construct a new server with the given dialect regitstry.
public:
   /// Find all of the document symbols within the given file.
   void findDocumentSymbols(const URIForFile &uri,
                            std::vector<DocumentSymbol> &symbols);
 
   /// Get the code completion list for the position within the given file.
   CompletionList getCodeCompletion(const URIForFile &uri,
                                    const Position &completePos);
 
+  /// Get the set of code actions within the file.
+  void getCodeActions(const URIForFile &uri, const Range &pos,
+                      const CodeActionContext &context,
+                      std::vector<CodeAction> &actions);
+
 private:
   struct Impl;
 
   std::unique_ptr<Impl> impl;
 };
 
 } // namespace lsp
 } // namespace mlir
 
 #endif // LIB_MLIR_TOOLS_MLIRLSPSERVER_SERVER_H_

mlir/lib/Tools/mlir-lsp-server/MLIRServer.cpp

struct MLIRDocument {
   // Code Completion
   //===--------------------------------------------------------------------===//
 
   lsp::CompletionList getCodeCompletion(const lsp::URIForFile &uri,
                                         const lsp::Position &completePos,
                                         const DialectRegistry &registry);
 
   //===--------------------------------------------------------------------===//
+  // Code Action
+  //===--------------------------------------------------------------------===//
+
+  void getCodeActionForDiagnostic(const lsp::URIForFile &uri,
+                                  lsp::Position &pos, StringRef severity,
+                                  StringRef message,
+                                  std::vector<lsp::TextEdit> &edits);
+
+  //===--------------------------------------------------------------------===//
   // Fields
   //===--------------------------------------------------------------------===//
 
   /// The high level parser state used to find definitions and references within
   /// the source file.
   AsmParserState asmState;
 
   /// The container for the IR parsed from the input file.
MLIRDocument::getCodeCompletion(const lsp::URIForFile &uri,
   AsmParserState tmpState;
   (void)parseSourceFile(sourceMgr, &tmpIR, &tmpContext,
                         /*sourceFileLoc=*/nullptr, &tmpState,
                         &lspCompleteContext);
   return completionList;
 }
 
 //===----------------------------------------------------------------------===//
+// MLIRDocument: Code Action
+//===----------------------------------------------------------------------===//
+
+void MLIRDocument::getCodeActionForDiagnostic(
+    const lsp::URIForFile &uri, lsp::Position &pos, StringRef severity,
+    StringRef message, std::vector<lsp::TextEdit> &edits) {
+  // Ignore diagnostics that print the current operation. These are always
+  // enabled for the language server, but not generally during normal
+  // parsing/verification.
+  if (message.startswith("see current operation: "))
+    return;
+
+  // Get the start of the line containing the diagnostic.
+  const auto &buffer = sourceMgr.getBufferInfo(sourceMgr.getMainFileID());
+  const char *lineStart = buffer.getPointerForLineNumber(pos.line + 1);
+  if (!lineStart)
+    return;
+  StringRef line(lineStart, pos.character);
+
+  // Add a text edit for adding an expected-* diagnostic check for this
+  // diagnostic.
+  lsp::TextEdit edit;
+  edit.range = lsp::Range(lsp::Position(pos.line, 0));
+
+  // Use the indent of the current line for the expected-* diagnostic.
+  size_t indent = line.find_first_not_of(" ");
+  if (indent == StringRef::npos)
+    indent = line.size();
+
+  edit.newText.append(indent, ' ');
+  llvm::raw_string_ostream(edit.newText)
+      << "// expected-" << severity << " @below {{" << message << "}}\n";
+  edits.emplace_back(std::move(edit));
+}
+
+//===----------------------------------------------------------------------===//
 // MLIRTextFileChunk
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// This class represents a single chunk of an MLIR text file.
 struct MLIRTextFileChunk {
   MLIRTextFileChunk(MLIRContext &context, uint64_t lineOffset,
                     const lsp::URIForFile &uri, StringRef contents,
void getLocationsOf(const lsp::URIForFile &uri, lsp::Position defPos,
                       std::vector<lsp::Location> &locations);
   void findReferencesOf(const lsp::URIForFile &uri, lsp::Position pos,
                         std::vector<lsp::Location> &references);
   Optional<lsp::Hover> findHover(const lsp::URIForFile &uri,
                                  lsp::Position hoverPos);
   void findDocumentSymbols(std::vector<lsp::DocumentSymbol> &symbols);
   lsp::CompletionList getCodeCompletion(const lsp::URIForFile &uri,
                                         lsp::Position completePos);
+  void getCodeActions(const lsp::URIForFile &uri, const lsp::Range &pos,
+                      const lsp::CodeActionContext &context,
+                      std::vector<lsp::CodeAction> &actions);
 
 private:
   /// Find the MLIR document that contains the given position, and update the
   /// position to be anchored at the start of the found chunk instead of the
   /// beginning of the file.
   MLIRTextFileChunk &getChunkFor(lsp::Position &pos);
 
   /// The context used to hold the state contained by the parsed document.
for (lsp::CompletionItem &item : completionList.items) {
     if (item.textEdit)
       chunk.adjustLocForChunkOffset(item.textEdit->range);
     for (lsp::TextEdit &edit : item.additionalTextEdits)
       chunk.adjustLocForChunkOffset(edit.range);
   }
   return completionList;
 }
 
+void MLIRTextFile::getCodeActions(const lsp::URIForFile &uri,
+                                  const lsp::Range &pos,
+                                  const lsp::CodeActionContext &context,
+                                  std::vector<lsp::CodeAction> &actions) {
+  // Create actions for any diagnostics in this file.
+  for (auto &diag : context.diagnostics) {
+    if (diag.source != "mlir")
+      continue;
+    lsp::Position diagPos = diag.range.start;
+    MLIRTextFileChunk &chunk = getChunkFor(diagPos);
+
+    // Add a new code action that inserts a "expected" diagnostic check.
+    lsp::CodeAction action;
+    action.title = "Add expected-* diagnostic checks";
+    action.kind = lsp::CodeAction::kQuickFix.str();
+
+    StringRef severity;
+    switch (diag.severity) {
+    case lsp::DiagnosticSeverity::Error:
+      severity = "error";
+      break;
+    case lsp::DiagnosticSeverity::Warning:
+      severity = "warning";
+      break;
+    default:
+      continue;
+    }
+
+    // Get edits for the diagnostic.
+    std::vector<lsp::TextEdit> edits;
+    chunk.document.getCodeActionForDiagnostic(uri, diagPos, severity,
+                                              diag.message, edits);
+
+    // Walk the related diagnostics, this is how we encode notes.
+    if (diag.relatedInformation) {
+      for (auto &noteDiag : *diag.relatedInformation) {
+        if (noteDiag.location.uri != uri)
+          continue;
+        diagPos = noteDiag.location.range.start;
+        diagPos.line -= chunk.lineOffset;
+        chunk.document.getCodeActionForDiagnostic(uri, diagPos, "note",
+                                                  noteDiag.message, edits);
+      }
+    }
+    // Fixup the locations for any edits.
+    for (lsp::TextEdit &edit : edits)
+      chunk.adjustLocForChunkOffset(edit.range);
+
+    action.edit.emplace();
+    action.edit->changes[uri.uri().str()] = std::move(edits);
+    action.diagnostics = {diag};
+
+    actions.emplace_back(std::move(action));
+  }
+}
+
 MLIRTextFileChunk &MLIRTextFile::getChunkFor(lsp::Position &pos) {
   if (chunks.size() == 1)
     return *chunks.front();
 
   // Search for the first chunk with a greater line offset, the previous chunk
   // is the one that contains `pos`.
   auto it = llvm::upper_bound(
       chunks, pos, [](const lsp::Position &pos, const auto &chunk) {

 lsp::CompletionList
 lsp::MLIRServer::getCodeCompletion(const URIForFile &uri,
                                    const Position &completePos) {
   auto fileIt = impl->files.find(uri.file());
   if (fileIt != impl->files.end())
     return fileIt->second->getCodeCompletion(uri, completePos);
   return CompletionList();
 }
+
+void lsp::MLIRServer::getCodeActions(const URIForFile &uri, const Range &pos,
+                                     const CodeActionContext &context,
+                                     std::vector<CodeAction> &actions) {
+  auto fileIt = impl->files.find(uri.file());
+  if (fileIt != impl->files.end())
+    fileIt->second->getCodeActions(uri, pos, context, actions);
+}

mlir/test/mlir-lsp-server/code-action.test

+// RUN: mlir-lsp-server -lit-test < %s | FileCheck -strict-whitespace %s
+{"jsonrpc":"2.0","id":0,"method":"initialize","params":{"processId":123,"rootPath":"mlir","capabilities":{},"trace":"off"}}
+// -----
+{"jsonrpc":"2.0","method":"textDocument/didOpen","params":{"textDocument":{
+  "uri":"test:///foo.mlir",
+  "languageId":"mlir",
+  "version":1,
+  "text":"#attr = 42 : f32\n// -----\nfunc.func @foo(%arg: i32) -> i64 {\nreturn %arg : i64\n}\n"
+}}}
+// -----
+{"jsonrpc":"2.0","id":1,"method":"textDocument/codeAction","params":{
+  "textDocument":{
+    "uri":"file:///foo.mlir"
+  },
+  "range":{
+    "start":{"line":0,"character":8}, "end":{"line":0,"character":10}
+  },
+  "context":{
+    "diagnostics":[{
+      "range":{"start":{"line":0,"character":8}, "end":{"line":0,"character":10}},
+      "message":"unexpected decimal integer literal for a floating point value",
+      "severity":1,
+      "relatedInformation":[{
+        "message":"add a trailing dot to make the literal a float",
+        "location":{
+          "uri":"file:///foo.mlir",
+          "range":{"start":{"line":0,"character":8}, "end":{"line":0,"character":10}}
+        }
+      }],
+      "source":"mlir"
+    }],
+    "only":["quickfix"],
+    "triggerKind":1
+  }
+}}
+// CHECK-LABEL: "id": 1
+// CHECK-NEXT:  "jsonrpc": "2.0",
+// CHECK-NEXT:  "result": [
+// CHECK-NEXT:    {
+// CHECK-NEXT:      "diagnostics": [
+// CHECK-NEXT:        {
+// CHECK-NEXT:          "message": "unexpected decimal integer literal for a floating point value",
+// CHECK-NEXT:          "range": {
+// CHECK-NEXT:            "end": {
+// CHECK-NEXT:              "character": 10,
+// CHECK-NEXT:              "line": 0
+// CHECK-NEXT:            },
+// CHECK-NEXT:            "start": {
+// CHECK-NEXT:              "character": 8,
+// CHECK-NEXT:              "line": 0
+// CHECK-NEXT:            }
+// CHECK-NEXT:          },
+// CHECK-NEXT:          "relatedInformation": [
+// CHECK-NEXT:            {
+// CHECK-NEXT:              "location": {
+// CHECK-NEXT:                "range": {
+// CHECK-NEXT:                  "end": {
+// CHECK-NEXT:                    "character": 10,
+// CHECK-NEXT:                    "line": 0
+// CHECK-NEXT:                  },
+// CHECK-NEXT:                  "start": {
+// CHECK-NEXT:                    "character": 8,
+// CHECK-NEXT:                    "line": 0
+// CHECK-NEXT:                  }
+// CHECK-NEXT:                },
+// CHECK-NEXT:                "uri": "file:///foo.mlir"
+// CHECK-NEXT:              },
+// CHECK-NEXT:              "message": "add a trailing dot to make the literal a float"
+// CHECK-NEXT:            }
+// CHECK-NEXT:          ],
+// CHECK-NEXT:          "severity": 1,
+// CHECK-NEXT:          "source": "mlir"
+// CHECK-NEXT:        }
+// CHECK-NEXT:      ],
+// CHECK-NEXT:      "edit": {
+// CHECK-NEXT:        "changes": {
+// CHECK-NEXT:          "file:///foo.mlir": [
+// CHECK-NEXT:            {
+// CHECK-LITERAL:             "newText": "// expected-error @below {{unexpected decimal integer literal for a floating point value}}\n"
+// CHECK:                     "range": {
+// CHECK-NEXT:                "end": {
+// CHECK-NEXT:                  "character": 0,
+// CHECK-NEXT:                  "line": 0
+// CHECK-NEXT:                },
+// CHECK-NEXT:                "start": {
+// CHECK-NEXT:                  "character": 0,
+// CHECK-NEXT:                  "line": 0
+// CHECK-NEXT:                }
+// CHECK-NEXT:              }
+// CHECK-NEXT:            },
+// CHECK-NEXT:            {
+// CHECK-NEXT-LITERAL:      "newText": "// expected-note @below {{add a trailing dot to make the literal a float}}\n",
+// CHECK:                   "range": {
+// CHECK-NEXT:                "end": {
+// CHECK-NEXT:                  "character": 0,
+// CHECK-NEXT:                  "line": 0
+// CHECK-NEXT:                },
+// CHECK-NEXT:                "start": {
+// CHECK-NEXT:                  "character": 0,
+// CHECK-NEXT:                  "line": 0
+// CHECK-NEXT:                }
+// CHECK-NEXT:              }
+// CHECK-NEXT:            }
+// CHECK-NEXT:          ]
+// CHECK-NEXT:        }
+// CHECK-NEXT:      },
+// CHECK-NEXT:      "kind": "quickfix",
+// CHECK-NEXT:      "title": "Add expected-* diagnostic checks"
+// CHECK-NEXT:    }
+// CHECK-NEXT:  ]
+// -----
+{"jsonrpc":"2.0","id":2,"method":"textDocument/codeAction","params":{
+  "textDocument":{"uri":"file:///foo.mlir"},
+  "range":{"start":{"line":3,"character":9},"end":{"line":3,"character":13}},
+  "context":{
+    "diagnostics":[{
+      "range":{"start":{"line":3,"character":9},"end":{"line":3,"character":13}},
+      "message":"use of value '%arg' expects different type than prior uses: 'i64' vs 'i32'",
+      "severity":1,
+      "relatedInformation":[{
+        "message":"prior use here",
+        "location":{
+          "uri":"file:///foo.mlir",
+          "range":{"start":{"line":2,"character":15},"end":{"line":2,"character":19}}
+        }
+      }],
+      "source":"mlir"
+    }],
+    "only":["quickfix"],
+    "triggerKind":1
+  }
+}}
+// CHECK-LABEL: "id": 2
+// CHECK-NEXT:  "jsonrpc": "2.0",
+// CHECK-NEXT:  "result": [
+// CHECK-NEXT:     {
+// CHECK:            "edit": {
+// CHECK-NEXT:       "changes": {
+// CHECK-NEXT:         "file:///foo.mlir": [
+// CHECK-NEXT:           {
+// CHECK-NEXT-LITERAL:     "newText": "// expected-error @below {{use of value '%arg' expects different type than prior uses: 'i64' vs 'i32'}}\n",
+// CHECK:                  "range": {
+// CHECK-NEXT:               "end": {
+// CHECK-NEXT:                 "character": 0,
+// CHECK-NEXT:                 "line": 3
+// CHECK-NEXT:               },
+// CHECK-NEXT:               "start": {
+// CHECK-NEXT:                 "character": 0,
+// CHECK-NEXT:                 "line": 3
+// CHECK-NEXT:               }
+// CHECK-NEXT:             }
+// CHECK-NEXT:           },
+// CHECK-NEXT:           {
+// CHECK-NEXT-LITERAL:     "newText": "// expected-note @below {{prior use here}}\n",
+// CHECK:                  "range": {
+// CHECK-NEXT:               "end": {
+// CHECK-NEXT:                 "character": 0,
+// CHECK-NEXT:                 "line": 2
+// CHECK-NEXT:               },
+// CHECK-NEXT:               "start": {
+// CHECK-NEXT:                 "character": 0,
+// CHECK-NEXT:                 "line": 2
+// CHECK-NEXT:               }
+// CHECK-NEXT:             }
+// CHECK-NEXT:           }
+// CHECK-NEXT:         ]
+// CHECK-NEXT:       }
+// CHECK-NEXT:     },
+// CHECK-NEXT:     "kind": "quickfix",
+// CHECK-NEXT:     "title": "Add expected-* diagnostic checks"
+// CHECK-NEXT:   }
+// CHECK-NEXT: ]
+// -----
+{"jsonrpc":"2.0","id":10,"method":"shutdown"}
+// -----
+{"jsonrpc":"2.0","method":"exit"}

mlir/test/mlir-lsp-server/initialize-params.test

 // RUN: mlir-lsp-server -lit-test < %s | FileCheck -strict-whitespace %s
 // Test initialize request parameters with rootUri
 {"jsonrpc":"2.0","id":0,"method":"initialize","params":{"processId":123,"rootUri":"test:///workspace","capabilities":{},"trace":"off"}}
 //      CHECK:  "id": 0,
 // CHECK-NEXT:  "jsonrpc": "2.0",
 // CHECK-NEXT:  "result": {
 // CHECK-NEXT:    "capabilities": {
+// CHECK-NEXT:      "codeActionProvider": true,
 // CHECK-NEXT:      "completionProvider": {
 // CHECK-NEXT:        "allCommitCharacters": [
 // CHECK:             ],
 // CHECK-NEXT:        "resolveProvider": false,
 // CHECK-NEXT:        "triggerCharacters": [ 
 // CHECK:             ]
 // CHECK-NEXT:      },
 // CHECK-NEXT:      "definitionProvider": true,