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

[rriddle]

Summary

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

[jpienaar]

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

161

Should this be debug?

[rriddle]

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

161

No, this matches what we already do with the other methods. It's also nice to have this logging be at a higher level, because it doesn't really clog up the terminal much and gives a high level indicator of what the server is doing.

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

void LSPServer::Impl::onShutdown(const NoParams &,
   reply(nullptr);
 }
 
 //===----------------------------------------------------------------------===//
 // Document Change
 
 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.
   publishDiagnostics(diagParams);
 }
 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) {
   // TODO: We currently only support full document updates, we should refactor
   // 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);
 }
 
 //===----------------------------------------------------------------------===//
 // Definitions and References
 

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

 /// 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.
   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,
                       std::vector<Location> &locations);
 
   /// Find all references of the object pointed at by the given position.
   void findReferencesOf(const URIForFile &uri, const Position &pos,
                         std::vector<Location> &references);

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

 //===----------------------------------------------------------------------===//
 // MLIRDocument
 //===----------------------------------------------------------------------===//
 
 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);
 
   //===--------------------------------------------------------------------===//
   // Definitions and References
   //===--------------------------------------------------------------------===//
 
   void getLocationsOf(const lsp::URIForFile &uri, const lsp::Position &defPos,
lsp::Hover buildHoverForOperationResult(llvm::SMRange hoverRange,
                                           unsigned resultEnd,
                                           llvm::SMLoc posLoc);
   lsp::Hover buildHoverForBlock(llvm::SMRange hoverRange,
                                 const AsmParserState::BlockDefinition &block);
   lsp::Hover
   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;
 
   /// 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.
   Block parsedIR;
 
   /// The source manager containing the contents of the input file.
   llvm::SourceMgr sourceMgr;
 };
 } // 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));
   });
 
   // Try to parsed the given IR string.
   auto memBuffer = llvm::MemoryBuffer::getMemBufferCopy(contents, uri.file());
   if (!memBuffer) {

 // MLIRServer
 //===----------------------------------------------------------------------===//
 
 lsp::MLIRServer::MLIRServer(DialectRegistry &registry)
     : impl(std::make_unique<Impl>(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,
                                      const Position &defPos,
                                      std::vector<Location> &locations) {
   auto fileIt = impl->documents.find(uri.file());
   if (fileIt != impl->documents.end())
     fileIt->second->getLocationsOf(uri, defPos, locations);

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

struct TextDocumentItem {
   /// The text document's URI.
   URIForFile uri;
 
   /// The text document's language identifier.
   std::string languageId;
 
   /// The content of the opened text document.
   std::string text;
+
+  /// The version number of this document.
+  int64_t version;
 };
 
 /// Add support for JSON serialization.
 bool fromJSON(const llvm::json::Value &value, TextDocumentItem &result,
               llvm::json::Path path);
 
 //===----------------------------------------------------------------------===//
 // TextDocumentIdentifier
 //===----------------------------------------------------------------------===//
 
 struct TextDocumentIdentifier {
   /// The text document's URI.
   URIForFile uri;
 };
 
 /// Add support for JSON serialization.
 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
 //===----------------------------------------------------------------------===//
 
 struct Position {
   /// Line position in a document (zero-based).
   int line = 0;
 
   /// Character offset on a line in a document (zero-based).

 };
 
 /// Add support for JSON serialization.
 bool fromJSON(const llvm::json::Value &value,
               TextDocumentContentChangeEvent &result, llvm::json::Path path);
 
 struct DidChangeTextDocumentParams {
   /// The document that changed.
-  TextDocumentIdentifier textDocument;
+  VersionedTextDocumentIdentifier textDocument;
 
   /// The actual content changes.
   std::vector<TextDocumentContentChangeEvent> contentChanges;
 };
 
 /// Add support for JSON serialization.
 bool fromJSON(const llvm::json::Value &value,
               DidChangeTextDocumentParams &result, llvm::json::Path path);

 /// Add support for JSON serialization.
 llvm::json::Value toJSON(const Diagnostic &diag);
 
 //===----------------------------------------------------------------------===//
 // PublishDiagnosticsParams
 //===----------------------------------------------------------------------===//
 
 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.
 llvm::json::Value toJSON(const PublishDiagnosticsParams &params);
 
 } // namespace lsp
 } // namespace mlir
 

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

 //===----------------------------------------------------------------------===//
 // TextDocumentItem
 //===----------------------------------------------------------------------===//
 
 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);
 }
 
 //===----------------------------------------------------------------------===//
 // TextDocumentIdentifier
 //===----------------------------------------------------------------------===//
 
 llvm::json::Value mlir::lsp::toJSON(const TextDocumentIdentifier &value) {
   return llvm::json::Object{{"uri", value.uri}};
 }
 
 bool mlir::lsp::fromJSON(const llvm::json::Value &value,
                          TextDocumentIdentifier &result,
                          llvm::json::Path path) {
   llvm::json::ObjectMapper o(value, path);
   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
 //===----------------------------------------------------------------------===//
 
 bool mlir::lsp::fromJSON(const llvm::json::Value &value, Position &result,
                          llvm::json::Path path) {
   llvm::json::ObjectMapper o(value, path);
   return o && o.map("line", result.line) &&
          o.map("character", result.character);

 //===----------------------------------------------------------------------===//
 // PublishDiagnosticsParams
 //===----------------------------------------------------------------------===//
 
 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/Transport.h

 // JSON-RPC over the stdin/stdout of a subprocess. This file contains a JSON
 // transport interface that handles this communication.
 //
 //===----------------------------------------------------------------------===//
 
 #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"
 #include "llvm/ADT/FunctionExtras.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/FormatAdapters.h"
 #include "llvm/Support/JSON.h"
notificationHandlers[method] = [method, handler,
       (thisPtr->*handler)(*param);
     };
   }
 
   /// Create an OutgoingNotification object used for the given method.
   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));
     };
   }
 
 private:
   template <typename HandlerT>
   using HandlerMap = llvm::StringMap<llvm::unique_function<HandlerT>>;
 

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

 
 void JSONTransport::sendMessage(llvm::json::Value msg) {
   outputBuffer.clear();
   llvm::raw_svector_ostream os(outputBuffer);
   os << llvm::formatv(prettyOutput ? "{0:2}\n" : "{0}", 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,
                                   MessageHandler &handler) {
   // Message must be an object with "jsonrpc":"2.0".
   llvm::json::Object *object = msg.getAsObject();
   if (!object ||
       object->getString("jsonrpc") != llvm::Optional<StringRef>("2.0"))

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

 // CHECK-NEXT:             "character": 6,
 // CHECK-NEXT:             "line": 0
 // CHECK-NEXT:           }
 // CHECK-NEXT:         },
 // CHECK-NEXT:         "severity": 1,
 // 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"}
 // -----
 {"jsonrpc":"2.0","method":"exit"}