[mlir:LSP] Switch document sync mode to Incremental

[rriddle]

Summary

This is much more efficient over the full mode, as it only requires sending
smalls chunks of files. It also works around a weird command ordering
issue (full document updates are being sent after other commands like
code completion) in newer versions of vscode.

[jpienaar]

Nice, injecting errors for testing :-)

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

476

And these can't get out of sync? E.g., would there be a failure rather than the state we believe internally here doesn't match what the user sees?

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

330

Why not getAsSMRange?

[rriddle]

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

476

Given that we communicate LSP commands via stdin/stdout it would only really get out-of-sync if the client is sending commands in an invalid order, but in that case given that we only get ranges and positions there isn't much we could do in terms of fully guaranteeing the text file is what it is supposed to be. If we do detect an invalid range though, we essentially forget the file (no longer serve e.g. diagnostics to the user). This is what clangd does in this situation as well.

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

 #include "llvm/Support/raw_ostream.h"
 #include <bitset>
 #include <memory>
 #include <string>
 #include <utility>
 #include <vector>
 
 namespace mlir {
+struct LogicalResult;
+
 namespace lsp {
 
 enum class ErrorCode {
   // Defined by JSON RPC.
   ParseError = -32700,
   InvalidRequest = -32600,
   MethodNotFound = -32601,
   InvalidParams = -32602,
struct Range {
   friend bool operator<(const Range &lhs, const Range &rhs) {
     return std::tie(lhs.start, lhs.end) < std::tie(rhs.start, rhs.end);
   }
 
   bool contains(Position pos) const { return start <= pos && pos < end; }
   bool contains(Range range) const {
     return start <= range.start && range.end <= end;
   }
+
+  /// Convert this range into a source range in the main file of the given
+  /// source manager.
+  SMRange getAsSMLoc(llvm::SourceMgr &mgr) const {
+    SMLoc startLoc = start.getAsSMLoc(mgr);
+    SMLoc endLoc = end.getAsSMLoc(mgr);
+    // Check that the start and end locations are valid.
+    if (!startLoc.isValid() || !endLoc.isValid() ||
+        startLoc.getPointer() > endLoc.getPointer())
+      return SMRange();
+    return SMRange(startLoc, endLoc);
+  }
 };
 
 /// Add support for JSON serialization.
 bool fromJSON(const llvm::json::Value &value, Range &result,
               llvm::json::Path path);
 llvm::json::Value toJSON(const Range &value);
 raw_ostream &operator<<(raw_ostream &os, const Range &value);
 

 bool fromJSON(const llvm::json::Value &value,
               DidCloseTextDocumentParams &result, llvm::json::Path path);
 
 //===----------------------------------------------------------------------===//
 // DidChangeTextDocumentParams
 //===----------------------------------------------------------------------===//
 
 struct TextDocumentContentChangeEvent {
+  /// Try to apply this change to the given contents string.
+  LogicalResult applyTo(std::string &contents) const;
+  /// Try to apply a set of changes to the given contents string.
+  static LogicalResult applyTo(ArrayRef<TextDocumentContentChangeEvent> changes,
+                               std::string &contents);
+
   /// The range of the document that changed.
   Optional<Range> range;
 
   /// The length of the range that got replaced.
   Optional<int> rangeLength;
 
   /// The new text of the range/document.
   std::string text;

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

 //===--- Protocol.cpp - Language Server Protocol Implementation -----------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file contains the serialization code for the LSP structs.
 //
 //===----------------------------------------------------------------------===//
 
 #include "Protocol.h"
+#include "Logging.h"
+#include "mlir/Support/LogicalResult.h"
 #include "llvm/ADT/Hashing.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/StringSwitch.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/Format.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/JSON.h"
+#include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/raw_ostream.h"
 
 using namespace mlir;
 using namespace mlir::lsp;
 
 // Helper that doesn't treat `null` and absent fields as failures.
 template <typename T>
bool mlir::lsp::fromJSON(const llvm::json::Value &value,
   llvm::json::ObjectMapper o(value, path);
   return o && o.map("textDocument", result.textDocument);
 }
 
 //===----------------------------------------------------------------------===//
 // DidChangeTextDocumentParams
 //===----------------------------------------------------------------------===//
 
+LogicalResult
+TextDocumentContentChangeEvent::applyTo(std::string &contents) const {
+  // If there is no range, the full document changed.
+  if (!range) {
+    contents = text;
+    return success();
+  }
+
+  // Try to map the replacement range to the content.
+  llvm::SourceMgr tmpScrMgr;
+  tmpScrMgr.AddNewSourceBuffer(llvm::MemoryBuffer::getMemBuffer(contents),
+                               SMLoc());
+  SMRange rangeLoc = range->getAsSMLoc(tmpScrMgr);
+  if (!rangeLoc.isValid())
+    return failure();
+
+  contents.replace(rangeLoc.Start.getPointer() - contents.data(),
+                   rangeLoc.End.getPointer() - rangeLoc.Start.getPointer(),
+                   text);
+  return success();
+}
+
+LogicalResult TextDocumentContentChangeEvent::applyTo(
+    ArrayRef<TextDocumentContentChangeEvent> changes, std::string &contents) {
+  for (const auto &change : changes)
+    if (failed(change.applyTo(contents)))
+      return failure();
+  return success();
+}
+
 bool mlir::lsp::fromJSON(const llvm::json::Value &value,
                          TextDocumentContentChangeEvent &result,
                          llvm::json::Path path) {
   llvm::json::ObjectMapper o(value, path);
   return o && o.map("range", result.range) &&
          o.map("rangeLength", result.rangeLength) && o.map("text", result.text);
 }
 

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

 
 void LSPServer::onInitialize(const InitializeParams &params,
                              Callback<llvm::json::Value> reply) {
   // Send a response with the capabilities of this server.
   llvm::json::Object serverCaps{
       {"textDocumentSync",
        llvm::json::Object{
            {"openClose", true},
-           {"change", (int)TextDocumentSyncKind::Full},
+           {"change", (int)TextDocumentSyncKind::Incremental},
            {"save", true},
        }},
       {"completionProvider",
        llvm::json::Object{
            {"allCommitCharacters",
             {"\t", "(", ")", "[", "]", "{",  "}", "<", ">",
              ":",  ";", ",", "+", "-", "/",  "*", "%", "^",
              "&",  "#", "?", ".", "=", "\"", "'", "|"}},

 }
 
 //===----------------------------------------------------------------------===//
 // Document Change
 
 void LSPServer::onDocumentDidOpen(const DidOpenTextDocumentParams &params) {
   PublishDiagnosticsParams diagParams(params.textDocument.uri,
                                       params.textDocument.version);
-  server.addOrUpdateDocument(params.textDocument.uri, params.textDocument.text,
-                             params.textDocument.version,
-                             diagParams.diagnostics);
+  server.addDocument(params.textDocument.uri, params.textDocument.text,
+                     params.textDocument.version, diagParams.diagnostics);
 
   // Publish any recorded diagnostics.
   publishDiagnostics(diagParams);
 }
 void LSPServer::onDocumentDidClose(const DidCloseTextDocumentParams &params) {
   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, *version));
 }
 void LSPServer::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,
                                       params.textDocument.version);
-  server.addOrUpdateDocument(
-      params.textDocument.uri, params.contentChanges.front().text,
+  server.updateDocument(params.textDocument.uri, params.contentChanges,
                         params.textDocument.version, diagParams.diagnostics);
 
   // Publish any recorded diagnostics.
   publishDiagnostics(diagParams);
 }
 
 //===----------------------------------------------------------------------===//
 // Definitions and References
 

mlir/lib/Tools/mlir-pdll-lsp-server/PDLLServer.h

 enum class PDLLViewOutputKind;
 struct CompletionList;
 struct DocumentLink;
 struct DocumentSymbol;
 struct Hover;
 struct Location;
 struct Position;
 struct SignatureHelp;
+struct TextDocumentContentChangeEvent;
 class URIForFile;
 
 /// This class implements all of the PDLL related functionality necessary for a
 /// language server. This class allows for keeping the PDLL specific logic
 /// separate from the logic that involves LSP server/client communication.
 class PDLLServer {
 public:
   struct Options {
     Options(const std::vector<std::string> &compilationDatabases,
             const std::vector<std::string> &extraDirs)
         : compilationDatabases(compilationDatabases), extraDirs(extraDirs) {}
 
     /// The filenames for databases containing compilation commands for PDLL
     /// files passed to the server.
     const std::vector<std::string> &compilationDatabases;
 
     /// Additional list of include directories to search.
     const std::vector<std::string> &extraDirs;
   };
 
   PDLLServer(const Options &options);
   ~PDLLServer();
 
-  /// 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);
+  /// Add the document, with the provided `version`, at the given URI. Any
+  /// diagnostics emitted for this document should be added to `diagnostics`.
+  void addDocument(const URIForFile &uri, StringRef contents, int64_t version,
 
+  /// Update the document, with the provided `version`, at the given URI. Any
+  /// diagnostics emitted for this document should be added to `diagnostics`.
+  void updateDocument(const URIForFile &uri,
+                      ArrayRef<TextDocumentContentChangeEvent> changes,
+                      int64_t version, std::vector<Diagnostic> &diagnostics);
+
   /// 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);

mlir/lib/Tools/mlir-pdll-lsp-server/PDLLServer.cpp

 public:
   PDLTextFile(const lsp::URIForFile &uri, StringRef fileContents,
               int64_t version, const std::vector<std::string> &extraDirs,
               std::vector<lsp::Diagnostic> &diagnostics);
 
   /// Return the current version of this text file.
   int64_t getVersion() const { return version; }
 
+  /// Update the file to the new version using the provided set of content
+  /// changes. Returns failure if the update was unsuccessful.
+  LogicalResult update(const lsp::URIForFile &uri, int64_t newVersion,
+                       ArrayRef<lsp::TextDocumentContentChangeEvent> changes,
+                       std::vector<lsp::Diagnostic> &diagnostics);
+
   //===--------------------------------------------------------------------===//
   // LSP Queries
   //===--------------------------------------------------------------------===//
 
   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);
   void getDocumentLinks(const lsp::URIForFile &uri,
                         std::vector<lsp::DocumentLink> &links);
   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);
   lsp::SignatureHelp getSignatureHelp(const lsp::URIForFile &uri,
                                       lsp::Position helpPos);
   lsp::PDLLViewOutputResult getPDLLViewOutput(lsp::PDLLViewOutputKind kind);
 
 private:
+  /// Initialize the text file from the given file contents.
+  void initialize(const lsp::URIForFile &uri, int64_t newVersion,
+                  std::vector<lsp::Diagnostic> &diagnostics);
+
   /// Find the PDL 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.
   PDLTextFileChunk &getChunkFor(lsp::Position &pos);
 
   /// The full string contents of the file.
   std::string contents;
 
   /// The version of this file.
-  int64_t version;
+  int64_t version = 0;
 
   /// The number of lines in the file.
   int64_t totalNumLines = 0;
 
   /// The chunks of this file. The order of these chunks is the order in which
   /// they appear in the text file.
   std::vector<std::unique_ptr<PDLTextFileChunk>> chunks;
+
+  /// The extra set of include directories for this file.
+  std::vector<std::string> extraIncludeDirs;
 };
 } // namespace
 
 PDLTextFile::PDLTextFile(const lsp::URIForFile &uri, StringRef fileContents,
                          int64_t version,
                          const std::vector<std::string> &extraDirs,
                          std::vector<lsp::Diagnostic> &diagnostics)
-    : contents(fileContents.str()), version(version) {
-  // Split the file into separate PDL documents.
-  // TODO: Find a way to share the split file marker with other tools. We don't
-  // want to use `splitAndProcessBuffer` here, but we do want to make sure this
-  // marker doesn't go out of sync.
-  SmallVector<StringRef, 8> subContents;
-  StringRef(contents).split(subContents, "// -----");
-  chunks.emplace_back(std::make_unique<PDLTextFileChunk>(
-      /*lineOffset=*/0, uri, subContents.front(), extraDirs, diagnostics));
-
-  uint64_t lineOffset = subContents.front().count('\n');
-  for (StringRef docContents : llvm::drop_begin(subContents)) {
-    unsigned currentNumDiags = diagnostics.size();
-    auto chunk = std::make_unique<PDLTextFileChunk>(
-        lineOffset, uri, docContents, extraDirs, diagnostics);
-    lineOffset += docContents.count('\n');
-
-    // Adjust locations used in diagnostics to account for the offset from the
-    // beginning of the file.
-    for (lsp::Diagnostic &diag :
-         llvm::drop_begin(diagnostics, currentNumDiags)) {
-      chunk->adjustLocForChunkOffset(diag.range);
-
-      if (!diag.relatedInformation)
-        continue;
-      for (auto &it : *diag.relatedInformation)
-        if (it.location.uri == uri)
-          chunk->adjustLocForChunkOffset(it.location.range);
-    }
-    chunks.emplace_back(std::move(chunk));
-  }
-  totalNumLines = lineOffset;
+    : contents(fileContents.str()), extraIncludeDirs(extraDirs) {
+  initialize(uri, version, diagnostics);
+
+LogicalResult
+PDLTextFile::update(const lsp::URIForFile &uri, int64_t newVersion,
+                    ArrayRef<lsp::TextDocumentContentChangeEvent> changes,
+                    std::vector<lsp::Diagnostic> &diagnostics) {
+  if (failed(lsp::TextDocumentContentChangeEvent::applyTo(changes, contents))) {
+    lsp::Logger::error("Failed to update contents of {0}", uri.file());
+    return failure();
+  }
+
+  // If the file contents were properly changed, reinitialize the text file.
+  initialize(uri, newVersion, diagnostics);
+  return success();
 }
 
 void PDLTextFile::getLocationsOf(const lsp::URIForFile &uri,
                                  lsp::Position defPos,
                                  std::vector<lsp::Location> &locations) {
   PDLTextFileChunk &chunk = getChunkFor(defPos);
   chunk.document.getLocationsOf(uri, defPos, locations);
 
llvm::interleave(
         [&](PDLTextFileChunk &chunk) {
           chunk.document.getPDLLViewOutput(outputOS, kind);
         },
         [&] { outputOS << "\n// -----\n\n"; });
   }
   return result;
 }
 
+void PDLTextFile::initialize(const lsp::URIForFile &uri, int64_t newVersion,
+                             std::vector<lsp::Diagnostic> &diagnostics) {
+  version = newVersion;
+  chunks.clear();
+
+  // Split the file into separate PDL documents.
+  // TODO: Find a way to share the split file marker with other tools. We don't
+  // want to use `splitAndProcessBuffer` here, but we do want to make sure this
+  // marker doesn't go out of sync.
+  SmallVector<StringRef, 8> subContents;
+  StringRef(contents).split(subContents, "// -----");
+  chunks.emplace_back(std::make_unique<PDLTextFileChunk>(
+      /*lineOffset=*/0, uri, subContents.front(), extraIncludeDirs,
+      diagnostics));
+
+  uint64_t lineOffset = subContents.front().count('\n');
+  for (StringRef docContents : llvm::drop_begin(subContents)) {
+    unsigned currentNumDiags = diagnostics.size();
+    auto chunk = std::make_unique<PDLTextFileChunk>(
+        lineOffset, uri, docContents, extraIncludeDirs, diagnostics);
+    lineOffset += docContents.count('\n');
+
+    // Adjust locations used in diagnostics to account for the offset from the
+    // beginning of the file.
+    for (lsp::Diagnostic &diag :
+         llvm::drop_begin(diagnostics, currentNumDiags)) {
+      chunk->adjustLocForChunkOffset(diag.range);
+
+      if (!diag.relatedInformation)
+        continue;
+      for (auto &it : *diag.relatedInformation)
+        if (it.location.uri == uri)
+          chunk->adjustLocForChunkOffset(it.location.range);
+    }
+    chunks.emplace_back(std::move(chunk));
+  }
+  totalNumLines = lineOffset;
+}
+
 PDLTextFileChunk &PDLTextFile::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) {

 //===----------------------------------------------------------------------===//
 // PDLLServer
 //===----------------------------------------------------------------------===//
 
 lsp::PDLLServer::PDLLServer(const Options &options)
     : impl(std::make_unique<Impl>(options)) {}
 lsp::PDLLServer::~PDLLServer() = default;
 
-void lsp::PDLLServer::addOrUpdateDocument(
-    const URIForFile &uri, StringRef contents, int64_t version,
+void lsp::PDLLServer::addDocument(const URIForFile &uri, StringRef contents,
+                                  int64_t version,
                                   std::vector<Diagnostic> &diagnostics) {
   std::vector<std::string> additionalIncludeDirs = impl->options.extraDirs;
   if (auto *fileInfo = impl->compilationDatabase.getFileInfo(uri.file()))
     llvm::append_range(additionalIncludeDirs, fileInfo->includeDirs);
 
   impl->files[uri.file()] = std::make_unique<PDLTextFile>(
       uri, contents, version, additionalIncludeDirs, diagnostics);
 }
 
+void lsp::PDLLServer::updateDocument(
+    const URIForFile &uri, ArrayRef<TextDocumentContentChangeEvent> changes,
+    int64_t version, std::vector<Diagnostic> &diagnostics) {
+  // Check that we actually have a document for this uri.
+  auto it = impl->files.find(uri.file());
+  if (it == impl->files.end())
+    return;
+
+  // Try to update the document. If we fail, erase the file from the server. A
+  // failed updated generally means we've fallen out of sync somewhere.
+  if (failed(it->second->update(uri, version, changes, diagnostics)))
+    impl->files.erase(it);
+}
+
 Optional<int64_t> lsp::PDLLServer::removeDocument(const URIForFile &uri) {
   auto it = impl->files.find(uri.file());
   if (it == impl->files.end())
     return llvm::None;
 
   int64_t version = it->second->getVersion();
   impl->files.erase(it);
   return version;

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

 
 void LSPServer::onInitialize(const InitializeParams &params,
                              Callback<llvm::json::Value> reply) {
   // Send a response with the capabilities of this server.
   llvm::json::Object serverCaps{
       {"textDocumentSync",
        llvm::json::Object{
            {"openClose", true},
-           {"change", (int)TextDocumentSyncKind::Full},
+           {"change", (int)TextDocumentSyncKind::Incremental},
            {"save", true},
        }},
   };
 
   llvm::json::Object result{
       {{"serverInfo", llvm::json::Object{{"name", "tblgen-lsp-server"},
                                          {"version", "0.0.1"}}},
        {"capabilities", std::move(serverCaps)}}};
   reply(std::move(result));
 }
 void LSPServer::onInitialized(const InitializedParams &) {}
 void LSPServer::onShutdown(const NoParams &, Callback<std::nullptr_t> reply) {
   shutdownRequestReceived = true;
   reply(nullptr);
 }
 
 //===----------------------------------------------------------------------===//
 // Document Change
 
 void LSPServer::onDocumentDidOpen(const DidOpenTextDocumentParams &params) {
   PublishDiagnosticsParams diagParams(params.textDocument.uri,
                                       params.textDocument.version);
-  server.addOrUpdateDocument(params.textDocument.uri, params.textDocument.text,
-                             params.textDocument.version,
-                             diagParams.diagnostics);
+  server.addDocument(params.textDocument.uri, params.textDocument.text,
+                     params.textDocument.version, diagParams.diagnostics);
 
   // Publish any recorded diagnostics.
   publishDiagnostics(diagParams);
 }
 void LSPServer::onDocumentDidClose(const DidCloseTextDocumentParams &params) {
   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, *version));
 }
 void LSPServer::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,
                                       params.textDocument.version);
-  server.addOrUpdateDocument(
-      params.textDocument.uri, params.contentChanges.front().text,
+  server.updateDocument(params.textDocument.uri, params.contentChanges,
                         params.textDocument.version, diagParams.diagnostics);
 
   // Publish any recorded diagnostics.
   publishDiagnostics(diagParams);
 }
 
 //===----------------------------------------------------------------------===//
 // Entry Point
 //===----------------------------------------------------------------------===//

mlir/lib/Tools/tblgen-lsp-server/TableGenServer.h

 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/StringRef.h"
 #include <memory>
 #include <string>
 
 namespace mlir {
 namespace lsp {
 struct Diagnostic;
+struct TextDocumentContentChangeEvent;
 class URIForFile;
 
 /// This class implements all of the TableGen related functionality necessary
 /// for a language server. This class allows for keeping the TableGen specific
 /// logic separate from the logic that involves LSP server/client communication.
 class TableGenServer {
 public:
   TableGenServer();
   ~TableGenServer();
 
-  /// 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);
+  /// Add the document, with the provided `version`, at the given URI. Any
+  /// diagnostics emitted for this document should be added to `diagnostics`.
+  void addDocument(const URIForFile &uri, StringRef contents, int64_t version,
 
+  /// Update the document, with the provided `version`, at the given URI. Any
+  /// diagnostics emitted for this document should be added to `diagnostics`.
+  void updateDocument(const URIForFile &uri,
+                      ArrayRef<TextDocumentContentChangeEvent> changes,
+                      int64_t version, std::vector<Diagnostic> &diagnostics);
+
   /// 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);
 
 private:
   struct Impl;
   std::unique_ptr<Impl> impl;
 };
 
 } // namespace lsp
 } // namespace mlir
 
 #endif // LIB_MLIR_TOOLS_TBLGENLSPSERVER_TABLEGENSERVER_H_

mlir/lib/Tools/tblgen-lsp-server/TableGenServer.cpp

 //===- TableGenServer.cpp - TableGen Language Server ----------------------===//
 //
 // 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 "TableGenServer.h"
 
 #include "../lsp-server-support/Logging.h"
 #include "../lsp-server-support/Protocol.h"
 #include "../lsp-server-support/SourceMgrUtils.h"
+#include "mlir/Support/LogicalResult.h"
 #include "llvm/ADT/IntervalMap.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringSet.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/Path.h"
 #include "llvm/TableGen/Parser.h"
 #include "llvm/TableGen/Record.h"

 class TableGenTextFile {
 public:
   TableGenTextFile(const lsp::URIForFile &uri, StringRef fileContents,
                    int64_t version, std::vector<lsp::Diagnostic> &diagnostics);
 
   /// Return the current version of this text file.
   int64_t getVersion() const { return version; }
 
+  /// Update the file to the new version using the provided set of content
+  /// changes. Returns failure if the update was unsuccessful.
+  LogicalResult update(const lsp::URIForFile &uri, int64_t newVersion,
+                       ArrayRef<lsp::TextDocumentContentChangeEvent> changes,
+                       std::vector<lsp::Diagnostic> &diagnostics);
+
 private:
+  /// Initialize the text file from the given file contents.
+  void initialize(const lsp::URIForFile &uri, int64_t newVersion,
+                  std::vector<lsp::Diagnostic> &diagnostics);
+
   /// The full string contents of the file.
   std::string contents;
 
   /// The version of this file.
   int64_t version;
 
   /// The include directories for this file.
   std::vector<std::string> includeDirs;
 
   /// The source manager containing the contents of the input file.
   llvm::SourceMgr sourceMgr;
 
   /// The record keeper containing the parsed tablegen constructs.
-  llvm::RecordKeeper recordKeeper;
+  std::unique_ptr<llvm::RecordKeeper> recordKeeper;
 };
 } // namespace
 
 TableGenTextFile::TableGenTextFile(const lsp::URIForFile &uri,
                                    StringRef fileContents, int64_t version,
                                    std::vector<lsp::Diagnostic> &diagnostics)
     : contents(fileContents.str()), version(version) {
-  auto memBuffer = llvm::MemoryBuffer::getMemBufferCopy(contents, uri.file());
-  if (!memBuffer) {
-    lsp::Logger::error("Failed to create memory buffer for file", uri.file());
-    return;
-  }
-
   // Build the set of include directories for this file.
   // TODO: Setup external include directories.
   llvm::SmallString<32> uriDirectory(uri.file());
   llvm::sys::path::remove_filename(uriDirectory);
   includeDirs.push_back(uriDirectory.str().str());
 
+  // Initialize the file.
+  initialize(uri, version, diagnostics);
+}
+
+LogicalResult
+TableGenTextFile::update(const lsp::URIForFile &uri, int64_t newVersion,
+                         ArrayRef<lsp::TextDocumentContentChangeEvent> changes,
+                         std::vector<lsp::Diagnostic> &diagnostics) {
+  if (failed(lsp::TextDocumentContentChangeEvent::applyTo(changes, contents))) {
+    lsp::Logger::error("Failed to update contents of {0}", uri.file());
+    return failure();
+  }
+
+  // If the file contents were properly changed, reinitialize the text file.
+  initialize(uri, newVersion, diagnostics);
+  return success();
+}
+
+void TableGenTextFile::initialize(const lsp::URIForFile &uri,
+                                  int64_t newVersion,
+                                  std::vector<lsp::Diagnostic> &diagnostics) {
+  version = newVersion;
+  sourceMgr = llvm::SourceMgr();
+  recordKeeper = std::make_unique<llvm::RecordKeeper>();
+
+  // Build a buffer for this file.
+  auto memBuffer = llvm::MemoryBuffer::getMemBuffer(contents, uri.file());
+  if (!memBuffer) {
+    lsp::Logger::error("Failed to create memory buffer for file", uri.file());
+    return;
+  }
   sourceMgr.setIncludeDirs(includeDirs);
   sourceMgr.AddNewSourceBuffer(std::move(memBuffer), SMLoc());
 
   // This class provides a context argument for the llvm::SourceMgr diagnostic
   // handler.
   struct DiagHandlerContext {
     std::vector<lsp::Diagnostic> &diagnostics;
     const lsp::URIForFile &uri;
   } handlerContext{diagnostics, uri};
 
   // Set the diagnostic handler for the tablegen source manager.
   sourceMgr.setDiagHandler(
       [](const llvm::SMDiagnostic &diag, void *rawHandlerContext) {
         auto *ctx = reinterpret_cast<DiagHandlerContext *>(rawHandlerContext);
         if (auto lspDiag = getLspDiagnoticFromDiag(diag, ctx->uri))
           ctx->diagnostics.push_back(*lspDiag);
       },
       &handlerContext);
-  if (llvm::TableGenParseFile(sourceMgr, recordKeeper))
+  if (llvm::TableGenParseFile(sourceMgr, *recordKeeper))
     return;
 }
 
 //===----------------------------------------------------------------------===//
 // TableGenServer::Impl
 //===----------------------------------------------------------------------===//
 
 struct lsp::TableGenServer::Impl {
   /// The files held by the server, mapped by their URI file name.
   llvm::StringMap<std::unique_ptr<TableGenTextFile>> files;
 };
 
 //===----------------------------------------------------------------------===//
 // TableGenServer
 //===----------------------------------------------------------------------===//
 
 lsp::TableGenServer::TableGenServer() : impl(std::make_unique<Impl>()) {}
 lsp::TableGenServer::~TableGenServer() = default;
 
-void lsp::TableGenServer::addOrUpdateDocument(
-    const URIForFile &uri, StringRef contents, int64_t version,
+void lsp::TableGenServer::addDocument(const URIForFile &uri, StringRef contents,
+                                      int64_t version,
                                       std::vector<Diagnostic> &diagnostics) {
   impl->files[uri.file()] =
       std::make_unique<TableGenTextFile>(uri, contents, version, diagnostics);
 }
 
+void lsp::TableGenServer::updateDocument(
+    const URIForFile &uri, ArrayRef<TextDocumentContentChangeEvent> changes,
+    int64_t version, std::vector<Diagnostic> &diagnostics) {
+  // Check that we actually have a document for this uri.
+  auto it = impl->files.find(uri.file());
+  if (it == impl->files.end())
+    return;
+
+  // Try to update the document. If we fail, erase the file from the server. A
+  // failed updated generally means we've fallen out of sync somewhere.
+  if (failed(it->second->update(uri, version, changes, diagnostics)))
+    impl->files.erase(it);
+}
+
 Optional<int64_t> lsp::TableGenServer::removeDocument(const URIForFile &uri) {
   auto it = impl->files.find(uri.file());
   if (it == impl->files.end())
     return llvm::None;
 
   int64_t version = it->second->getVersion();
   impl->files.erase(it);
   return version;
 }

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

 // CHECK-NEXT:      "referencesProvider": true,
 // CHECK-NEXT:      "signatureHelpProvider": {
 // CHECK-NEXT:        "triggerCharacters": [
 // CHECK-NEXT:          "(",
 // CHECK-NEXT:          ","
 // CHECK-NEXT:        ]
 // CHECK-NEXT:      },
 // CHECK-NEXT:      "textDocumentSync": {
-// CHECK-NEXT:        "change": 1,
+// CHECK-NEXT:        "change": 2,
 // CHECK-NEXT:        "openClose": true,
 // CHECK-NEXT:        "save": true
 // CHECK-NEXT:      }
 // CHECK-NEXT:    },
 // CHECK-NEXT:    "serverInfo": {
 // CHECK-NEXT:      "name": "mlir-pdll-lsp-server",
 // CHECK-NEXT:      "version": "{{.*}}"
 // CHECK-NEXT:    }
 // CHECK-NEXT:  }
 // -----
 {"jsonrpc":"2.0","id":3,"method":"shutdown"}
 //      CHECK:  "id": 3,
 // CHECK-NEXT:  "jsonrpc": "2.0",
 // CHECK-NEXT:  "result": null
 // -----
 {"jsonrpc":"2.0","method":"exit"}

mlir/test/mlir-pdll-lsp-server/textdocument-didchange.test

+// RUN: mlir-pdll-lsp-server -lit-test < %s | FileCheck -strict-whitespace %s
+{"jsonrpc":"2.0","id":0,"method":"initialize","params":{"processId":123,"rootPath":"pdll","capabilities":{},"trace":"off"}}
+// -----
+{"jsonrpc":"2.0","method":"textDocument/didOpen","params":{"textDocument":{
+  "uri":"test:///foo.pdll",
+  "languageId":"pdll",
+  "version":1,
+  "text":"Pattern => replace with ;"
+}}}
+// CHECK: "method": "textDocument/publishDiagnostics",
+// CHECK-NEXT: "params": {
+// CHECK-NEXT:     "diagnostics": [
+// CHECK-NEXT:       {
+// CHECK-NEXT:         "category": "Parse Error",
+// CHECK-NEXT:         "message": "expected expression",
+// CHECK-NEXT:         "range": {
+// CHECK-NEXT:           "end": {
+// CHECK-NEXT:             "character": 23,
+// CHECK-NEXT:             "line": 0
+// CHECK-NEXT:           },
+// CHECK-NEXT:           "start": {
+// CHECK-NEXT:             "character": 19,
+// CHECK-NEXT:             "line": 0
+// CHECK-NEXT:           }
+// CHECK-NEXT:         },
+// CHECK-NEXT:         "severity": 1,
+// CHECK-NEXT:         "source": "pdll"
+// CHECK-NEXT:       }
+// CHECK-NEXT:     ],
+// CHECK-NEXT:     "uri": "test:///foo.pdll",
+// CHECK-NEXT:     "version": 1
+// CHECK-NEXT:   }
+// -----
+{"jsonrpc":"2.0","method":"textDocument/didChange","params":{"textDocument":{
+  "uri":"test:///foo.pdll",
+  "version":2
+}, "contentChanges": [{
+  "range":{
+    "start":{"line":0,"character":18},
+    "end":{"line":0,"character":18}
+  },
+  "text": " op<test.op>"
+}]}}
+// CHECK: "method": "textDocument/publishDiagnostics",
+// CHECK-NEXT: "params": {
+// CHECK-NEXT:     "diagnostics": [
+// CHECK-NEXT:       {
+// CHECK-NEXT:         "category": "Parse Error",
+// CHECK-NEXT:         "message": "expected expression",
+// CHECK-NEXT:         "range": {
+// CHECK-NEXT:           "end": {
+// CHECK-NEXT:             "character": 37,
+// CHECK-NEXT:             "line": 0
+// CHECK-NEXT:           },
+// CHECK-NEXT:           "start": {
+// CHECK-NEXT:             "character": 36,
+// CHECK-NEXT:             "line": 0
+// CHECK-NEXT:           }
+// CHECK-NEXT:         },
+// CHECK-NEXT:         "severity": 1,
+// CHECK-NEXT:         "source": "pdll"
+// CHECK-NEXT:       }
+// CHECK-NEXT:     ],
+// CHECK-NEXT:     "uri": "test:///foo.pdll",
+// CHECK-NEXT:     "version": 2
+// CHECK-NEXT:   }
+// -----
+{"jsonrpc":"2.0","method":"textDocument/didChange","params":{"textDocument":{
+  "uri":"test:///foo.pdll",
+  "version":3
+}, "contentChanges": [
+  {
+    "range":{
+      "start":{"line":0,"character":30},
+      "end":{"line":0,"character":30}
+    },
+    "text": "(values: ValueRange)"
+  },
+  {
+    "range":{
+      "start":{"line":0,"character":56},
+      "end":{"line":0,"character":57}
+    },
+    "text": "values;"
+  }
+]}}
+// CHECK: "method": "textDocument/publishDiagnostics",
+// CHECK-NEXT: "params": {
+// CHECK-NEXT:     "diagnostics": [],
+// CHECK-NEXT:     "uri": "test:///foo.pdll",
+// CHECK-NEXT:     "version": 3
+// CHECK-NEXT:   }
+// -----
+{"jsonrpc":"2.0","id":3,"method":"shutdown"}
+// -----
+{"jsonrpc":"2.0","method":"exit"}

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

 // RUN: tblgen-lsp-server -lit-test < %s | FileCheck %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:      "textDocumentSync": {
-// CHECK-NEXT:        "change": 1,
+// CHECK-NEXT:        "change": 2,
 // CHECK-NEXT:        "openClose": true,
 // CHECK-NEXT:        "save": true
 // CHECK-NEXT:      }
 // CHECK-NEXT:    },
 // CHECK-NEXT:    "serverInfo": {
 // CHECK-NEXT:      "name": "tblgen-lsp-server",
 // CHECK-NEXT:      "version": "{{.*}}"
 // CHECK-NEXT:    }
 // CHECK-NEXT:  }
 // -----
 {"jsonrpc":"2.0","id":3,"method":"shutdown"}
 //      CHECK:  "id": 3,
 // CHECK-NEXT:  "jsonrpc": "2.0",
 // CHECK-NEXT:  "result": null
 // -----
 {"jsonrpc":"2.0","method":"exit"}

mlir/test/tblgen-lsp-server/textdocument-didchange.test

+// RUN: tblgen-lsp-server -lit-test < %s | FileCheck -strict-whitespace %s
+{"jsonrpc":"2.0","id":0,"method":"initialize","params":{"processId":123,"rootPath":"tablegen","capabilities":{},"trace":"off"}}
+// -----
+{"jsonrpc":"2.0","method":"textDocument/didOpen","params":{"textDocument":{
+  "uri":"test:///foo.td",
+  "languageId":"tablegen",
+  "version":1,
+  "text":"class Foo<>;"
+}}}
+// CHECK: "method": "textDocument/publishDiagnostics",
+// CHECK-NEXT: "params": {
+// CHECK-NEXT:     "diagnostics": [
+// CHECK-NEXT:       {
+// CHECK-NEXT:         "category": "Parse Error",
+// CHECK-NEXT:         "message": "Unknown token when expecting a type",
+// CHECK-NEXT:         "range": {
+// CHECK-NEXT:           "end": {
+// CHECK-NEXT:             "character": 11,
+// CHECK-NEXT:             "line": 0
+// CHECK-NEXT:           },
+// CHECK-NEXT:           "start": {
+// CHECK-NEXT:             "character": 10,
+// CHECK-NEXT:             "line": 0
+// CHECK-NEXT:           }
+// CHECK-NEXT:         },
+// CHECK-NEXT:         "severity": 1,
+// CHECK-NEXT:         "source": "tablegen"
+// CHECK-NEXT:       }
+// CHECK-NEXT:     ],
+// CHECK-NEXT:     "uri": "test:///foo.td",
+// CHECK-NEXT:     "version": 1
+// CHECK-NEXT:   }
+// -----
+{"jsonrpc":"2.0","method":"textDocument/didChange","params":{"textDocument":{
+  "uri":"test:///foo.td",
+  "version":2
+}, "contentChanges": [{
+  "range":{
+    "start":{"line":0,"character":10},
+    "end":{"line":0,"character":10}
+  },
+  "text": "int"
+}]}}
+// CHECK: "method": "textDocument/publishDiagnostics",
+// CHECK-NEXT: "params": {
+// CHECK-NEXT:     "diagnostics": [
+// CHECK-NEXT:       {
+// CHECK-NEXT:         "category": "Parse Error",
+// CHECK-NEXT:         "message": "Expected identifier in declaration",
+// CHECK-NEXT:         "range": {
+// CHECK-NEXT:           "end": {
+// CHECK-NEXT:             "character": 14,
+// CHECK-NEXT:             "line": 0
+// CHECK-NEXT:           },
+// CHECK-NEXT:           "start": {
+// CHECK-NEXT:             "character": 13,
+// CHECK-NEXT:             "line": 0
+// CHECK-NEXT:           }
+// CHECK-NEXT:         },
+// CHECK-NEXT:         "severity": 1,
+// CHECK-NEXT:         "source": "tablegen"
+// CHECK-NEXT:       }
+// CHECK-NEXT:     ],
+// CHECK-NEXT:     "uri": "test:///foo.td",
+// CHECK-NEXT:     "version": 2
+// CHECK-NEXT:   }
+// -----
+{"jsonrpc":"2.0","method":"textDocument/didChange","params":{"textDocument":{
+  "uri":"test:///foo.td",
+  "version":3
+}, "contentChanges": [
+  {
+    "range":{
+      "start":{"line":0,"character":13},
+      "end":{"line":0,"character":13}
+    },
+    "text": " i"
+  },
+  {
+    "range":{
+      "start":{"line":0,"character":15},
+      "end":{"line":0,"character":17}
+    },
+    "text": "> { int x = i; }"
+  }
+]}}
+// CHECK: "method": "textDocument/publishDiagnostics",
+// CHECK-NEXT: "params": {
+// CHECK-NEXT:     "diagnostics": [],
+// CHECK-NEXT:     "uri": "test:///foo.td",
+// CHECK-NEXT:     "version": 3
+// CHECK-NEXT:   }
+// -----
+{"jsonrpc":"2.0","id":3,"method":"shutdown"}
+// -----
+{"jsonrpc":"2.0","method":"exit"}