[mlir-lsp-server] Add support for recording text document versions

The version is used by LSP clients to ignore stale diagnostics, and can be used in a followup to help verify incremental changes.

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

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

@@ -103,8 +103,10 @@ void LSPServer::Impl::onShutdown(const NoParams &,
 
 void LSPServer::Impl::onDocumentDidOpen(
     const DidOpenTextDocumentParams &params) {
-  PublishDiagnosticsParams diagParams(params.textDocument.uri);
+  PublishDiagnosticsParams diagParams(params.textDocument.uri,
+                                      params.textDocument.version);
   server.addOrUpdateDocument(params.textDocument.uri, params.textDocument.text,
+                             params.textDocument.version,
                              diagParams.diagnostics);
 
   // Publish any recorded diagnostics.
@@ -112,12 +114,15 @@ void LSPServer::Impl::onDocumentDidOpen(
 }
 void LSPServer::Impl::onDocumentDidClose(
     const DidCloseTextDocumentParams &params) {
-  server.removeDocument(params.textDocument.uri);
+  Optional<int64_t> version = server.removeDocument(params.textDocument.uri);
+  if (!version)
+    return;
 
   // Empty out the diagnostics shown for this document. This will clear out
   // anything currently displayed by the client for this document (e.g. in the
   // "Problems" pane of VSCode).
-  publishDiagnostics(PublishDiagnosticsParams(params.textDocument.uri));
+  publishDiagnostics(
+      PublishDiagnosticsParams(params.textDocument.uri, *version));
 }
 void LSPServer::Impl::onDocumentDidChange(
     const DidChangeTextDocumentParams &params) {
@@ -125,10 +130,11 @@ void LSPServer::Impl::onDocumentDidChange(
   // to avoid this.
   if (params.contentChanges.size() != 1)
     return;
-  PublishDiagnosticsParams diagParams(params.textDocument.uri);
-  server.addOrUpdateDocument(params.textDocument.uri,
-                             params.contentChanges.front().text,
-                             diagParams.diagnostics);
+  PublishDiagnosticsParams diagParams(params.textDocument.uri,
+                                      params.textDocument.version);
+  server.addOrUpdateDocument(
+      params.textDocument.uri, params.contentChanges.front().text,
+      params.textDocument.version, diagParams.diagnostics);
 
   // Publish any recorded diagnostics.
   publishDiagnostics(diagParams);

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

@@ -252,7 +252,7 @@ namespace {
 /// This class represents all of the information pertaining to a specific MLIR
 /// document.
 struct MLIRDocument {
-  MLIRDocument(const lsp::URIForFile &uri, StringRef contents,
+  MLIRDocument(const lsp::URIForFile &uri, StringRef contents, int64_t version,
                DialectRegistry &registry,
                std::vector<lsp::Diagnostic> &diagnostics);
 
@@ -283,6 +283,9 @@ struct MLIRDocument {
   buildHoverForBlockArgument(llvm::SMRange hoverRange, BlockArgument arg,
                              const AsmParserState::BlockDefinition &block);
 
+  /// The version of this document.
+  int64_t version;
+
   /// The context used to hold the state contained by the parsed document.
   MLIRContext context;
 
@@ -299,9 +302,9 @@ struct MLIRDocument {
 } // namespace
 
 MLIRDocument::MLIRDocument(const lsp::URIForFile &uri, StringRef contents,
-                           DialectRegistry &registry,
+                           int64_t version, DialectRegistry &registry,
                            std::vector<lsp::Diagnostic> &diagnostics)
-    : context(registry) {
+    : version(version), context(registry) {
   context.allowUnregisteredDialects();
   ScopedDiagnosticHandler handler(&context, [&](Diagnostic &diag) {
     diagnostics.push_back(getLspDiagnoticFromDiag(diag, uri));
@@ -569,14 +572,20 @@ lsp::MLIRServer::MLIRServer(DialectRegistry &registry)
 lsp::MLIRServer::~MLIRServer() {}
 
 void lsp::MLIRServer::addOrUpdateDocument(
-    const URIForFile &uri, StringRef contents,
+    const URIForFile &uri, StringRef contents, int64_t version,
     std::vector<Diagnostic> &diagnostics) {
   impl->documents[uri.file()] = std::make_unique<MLIRDocument>(
-      uri, contents, impl->registry, diagnostics);
+      uri, contents, version, impl->registry, diagnostics);
 }
 
-void lsp::MLIRServer::removeDocument(const URIForFile &uri) {
-  impl->documents.erase(uri.file());
+Optional<int64_t> lsp::MLIRServer::removeDocument(const URIForFile &uri) {
+  auto it = impl->documents.find(uri.file());
+  if (it == impl->documents.end())
+    return llvm::None;
+
+  int64_t version = it->second->version;
+  impl->documents.erase(it);
+  return version;
 }
 
 void lsp::MLIRServer::getLocationsOf(const URIForFile &uri,

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

@@ -31,13 +31,17 @@ public:
   MLIRServer(DialectRegistry &registry);
   ~MLIRServer();
 
-  /// Add or update the document at the given URI. Any diagnostics emitted for
-  /// this document should be added to `diagnostics`
+  /// Add or update the document, with the provided `version`, at the given URI.
+  /// Any diagnostics emitted for this document should be added to
+  /// `diagnostics`.
   void addOrUpdateDocument(const URIForFile &uri, StringRef contents,
+                           int64_t version,
                            std::vector<Diagnostic> &diagnostics);
 
-  /// Remove the document with the given uri.
-  void removeDocument(const URIForFile &uri);
+  /// Remove the document with the given uri. Returns the version of the removed
+  /// document, or None if the uri did not have a corresponding document within
+  /// the server.
+  Optional<int64_t> removeDocument(const URIForFile &uri);
 
   /// Return the locations of the object pointed at by the given position.
   void getLocationsOf(const URIForFile &uri, const Position &defPos,

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

@@ -287,7 +287,8 @@ bool mlir::lsp::fromJSON(const llvm::json::Value &value,
                          TextDocumentItem &result, llvm::json::Path path) {
   llvm::json::ObjectMapper o(value, path);
   return o && o.map("uri", result.uri) &&
-         o.map("languageId", result.languageId) && o.map("text", result.text);
+         o.map("languageId", result.languageId) && o.map("text", result.text) &&
+         o.map("version", result.version);
 }
 
 //===----------------------------------------------------------------------===//
@@ -305,6 +306,25 @@ bool mlir::lsp::fromJSON(const llvm::json::Value &value,
   return o && o.map("uri", result.uri);
 }
 
+//===----------------------------------------------------------------------===//
+// VersionedTextDocumentIdentifier
+//===----------------------------------------------------------------------===//
+
+llvm::json::Value
+mlir::lsp::toJSON(const VersionedTextDocumentIdentifier &value) {
+  return llvm::json::Object{
+      {"uri", value.uri},
+      {"version", value.version},
+  };
+}
+
+bool mlir::lsp::fromJSON(const llvm::json::Value &value,
+                         VersionedTextDocumentIdentifier &result,
+                         llvm::json::Path path) {
+  llvm::json::ObjectMapper o(value, path);
+  return o && o.map("uri", result.uri) && o.map("version", result.version);
+}
+
 //===----------------------------------------------------------------------===//
 // Position
 //===----------------------------------------------------------------------===//
@@ -512,5 +532,6 @@ llvm::json::Value mlir::lsp::toJSON(const PublishDiagnosticsParams &params) {
   return llvm::json::Object{
       {"uri", params.uri},
       {"diagnostics", params.diagnostics},
+      {"version", params.version},
   };
 }

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

@@ -180,6 +180,9 @@ struct TextDocumentItem {
 
   /// The content of the opened text document.
   std::string text;
+
+  /// The version number of this document.
+  int64_t version;
 };
 
 /// Add support for JSON serialization.
@@ -200,6 +203,22 @@ llvm::json::Value toJSON(const TextDocumentIdentifier &value);
 bool fromJSON(const llvm::json::Value &value, TextDocumentIdentifier &result,
               llvm::json::Path path);
 
+//===----------------------------------------------------------------------===//
+// VersionedTextDocumentIdentifier
+//===----------------------------------------------------------------------===//
+
+struct VersionedTextDocumentIdentifier {
+  /// The text document's URI.
+  URIForFile uri;
+  /// The version number of this document.
+  int64_t version;
+};
+
+/// Add support for JSON serialization.
+llvm::json::Value toJSON(const VersionedTextDocumentIdentifier &value);
+bool fromJSON(const llvm::json::Value &value,
+              VersionedTextDocumentIdentifier &result, llvm::json::Path path);
+
 //===----------------------------------------------------------------------===//
 // Position
 //===----------------------------------------------------------------------===//
@@ -381,7 +400,7 @@ bool fromJSON(const llvm::json::Value &value,
 
 struct DidChangeTextDocumentParams {
   /// The document that changed.
-  TextDocumentIdentifier textDocument;
+  VersionedTextDocumentIdentifier textDocument;
 
   /// The actual content changes.
   std::vector<TextDocumentContentChangeEvent> contentChanges;
@@ -498,12 +517,15 @@ llvm::json::Value toJSON(const Diagnostic &diag);
 //===----------------------------------------------------------------------===//
 
 struct PublishDiagnosticsParams {
-  PublishDiagnosticsParams(URIForFile uri) : uri(uri) {}
+  PublishDiagnosticsParams(URIForFile uri, int64_t version)
+      : uri(uri), version(version) {}
 
   /// The URI for which diagnostic information is reported.
   URIForFile uri;
   /// The list of reported diagnostics.
   std::vector<Diagnostic> diagnostics;
+  /// The version number of the document the diagnostics are published for.
+  int64_t version;
 };
 
 /// Add support for JSON serialization.

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

@@ -217,6 +217,7 @@ void JSONTransport::sendMessage(llvm::json::Value msg) {
   out << "Content-Length: " << outputBuffer.size() << "\r\n\r\n"
       << outputBuffer;
   out.flush();
+  Logger::debug(">>> {0}\n", outputBuffer);
 }
 
 bool JSONTransport::handleMessage(llvm::json::Value msg,

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

@@ -15,6 +15,7 @@
 #ifndef LIB_MLIR_TOOLS_MLIRLSPSERVER_LSP_TRANSPORT_H_
 #define LIB_MLIR_TOOLS_MLIRLSPSERVER_LSP_TRANSPORT_H_
 
+#include "Logging.h"
 #include "Protocol.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Support/LogicalResult.h"
@@ -157,6 +158,7 @@ public:
   template <typename T>
   OutgoingNotification<T> outgoingNotification(llvm::StringLiteral method) {
     return [&, method](const T &params) {
+      Logger::info("--> {0}", method);
       transport.notify(method, llvm::json::Value(params));
     };
   }

mlir/test/mlir-lsp-server/diagnostics.test

@@ -27,7 +27,8 @@
 // CHECK-NEXT:         "source": "mlir"
 // CHECK-NEXT:       }
 // CHECK-NEXT:     ],
-// CHECK-NEXT:     "uri": "test:///foo.mlir"
+// CHECK-NEXT:     "uri": "test:///foo.mlir",
+// CHECK-NEXT:     "version": 1
 // CHECK-NEXT:   }
 // -----
 {"jsonrpc":"2.0","id":3,"method":"shutdown"}