Implement cancellation for pending parser requests in the LSP server

Signed-off-by: Roberto Raggi <[email protected]>

Author: robertoraggi

src/lsp/cxx/lsp/cxx_document.cc

@@ -34,6 +34,10 @@
 #include <cxx/wasm32_wasi_toolchain.h>
 #include <cxx/windows_toolchain.h>
 
+#ifndef CXX_NO_THREADS
+#include <atomic>
+#endif
+
 namespace cxx::lsp {
 
 namespace {
@@ -65,12 +69,20 @@ struct Diagnostics final : cxx::DiagnosticsClient {
 
 struct CxxDocument::Private {
   const CLI& cli;
+  std::string fileName;
   long version;
   Diagnostics diagnosticsClient;
   TranslationUnit unit{&diagnosticsClient};
   std::shared_ptr<Toolchain> toolchain;
 
-  Private(const CLI& cli, long version) : cli(cli), version(version) {}
+#ifndef CXX_NO_THREADS
+  std::atomic<bool> cancelled{false};
+#else
+  bool cancelled{false};
+#endif
+
+  Private(const CLI& cli, std::string fileName, long version)
+      : cli(cli), fileName(std::move(fileName)), version(version) {}
 
   void configure();
 };
@@ -169,27 +181,53 @@ void CxxDocument::Private::configure() {
   }
 }
 
-CxxDocument::CxxDocument(const CLI& cli, long version)
-    : d(std::make_unique<Private>(cli, version)) {}
+CxxDocument::CxxDocument(const CLI& cli, std::string fileName, long version)
+    : d(std::make_unique<Private>(cli, std::move(fileName), version)) {}
 
-void CxxDocument::parse(std::string source, std::string fileName,
-                        std::function<bool()> stopParsingPredicate) {
+auto CxxDocument::isCancelled() const -> bool {
+#ifndef CXX_NO_THREADS
+  return d->cancelled.load();
+#else
+  return d->cancelled;
+#endif
+}
+
+void CxxDocument::cancel() {
+#ifndef CXX_NO_THREADS
+  d->cancelled.store(true);
+#else
+  d->cancelled = true;
+#endif
+}
+
+auto CxxDocument::fileName() const -> const std::string& { return d->fileName; }
+
+void CxxDocument::parse(std::string source) {
   d->configure();
 
   auto& unit = d->unit;
   auto& cli = d->cli;
 
-  unit.setSource(std::move(source), fileName);
-
   auto preprocessor = unit.preprocessor();
 
+  DefaultPreprocessorState state{*preprocessor};
+
+  unit.beginPreprocessing(std::move(source), d->fileName);
+
+  while (state) {
+    if (isCancelled()) break;
+    std::visit(state, unit.continuePreprocessing());
+  }
+
+  unit.endPreprocessing();
+
   unit.parse(ParserConfiguration{
       .checkTypes = cli.opt_fcheck,
       .fuzzyTemplateResolution = true,
       .staticAssert = cli.opt_fstatic_assert || cli.opt_fcheck,
       .reflect = !cli.opt_fno_reflect,
       .templates = cli.opt_ftemplates,
-      .stopParsingPredicate = std::move(stopParsingPredicate),
+      .stopParsingPredicate = [this] { return isCancelled(); },
   });
 }
 

src/lsp/cxx/lsp/cxx_document.h

@@ -31,11 +31,15 @@ namespace cxx::lsp {
 
 class CxxDocument {
  public:
-  explicit CxxDocument(const CLI& cli, long version);
+  explicit CxxDocument(const CLI& cli, std::string fileName, long version);
   ~CxxDocument();
 
-  void parse(std::string source, std::string fileName,
-             std::function<bool()> stopParsingPredicate = {});
+  [[nodiscard]] auto isCancelled() const -> bool;
+  void cancel();
+
+  [[nodiscard]] auto fileName() const -> const std::string&;
+
+  void parse(std::string source);
 
   [[nodiscard]] auto version() const -> long;
   [[nodiscard]] auto diagnostics() const -> Vector<Diagnostic>;

src/lsp/cxx/lsp/lsp_server.cc

@@ -286,21 +286,49 @@ void Server::run(std::function<void()> task) {
   task();
 }
 
+void Server::cancelPendingParserRequests(const std::string& fileName) {
+#ifndef CXX_NO_THREADS
+  auto lock = std::unique_lock(documentsMutex_);
+#endif
+
+  std::vector<std::shared_ptr<CxxDocument>> pendingParserRequests;
+  std::swap(pendingParserRequests_, pendingParserRequests);
+
+  for (const auto& doc : pendingParserRequests) {
+    if (doc->fileName() == fileName) {
+      doc->cancel();
+    } else {
+      pendingParserRequests_.push_back(doc);
+    }
+  }
+}
+
 void Server::parse(const std::string& uri) {
   const auto& doc = documentContents_.at(uri);
 
   auto text = doc.value;
   auto version = doc.version;
 
   run([text = std::move(text), uri = std::move(uri), version, this] {
-    auto doc = std::make_shared<CxxDocument>(cli, version);
-    doc->parse(std::move(text), pathFromUri(uri));
+    auto fileName = pathFromUri(uri);
+
+    cancelPendingParserRequests(fileName);
+
+    auto doc = std::make_shared<CxxDocument>(cli, std::move(fileName), version);
+    pendingParserRequests_.push_back(doc);
+
+    doc->parse(std::move(text));
 
     {
 #ifndef CXX_NO_THREADS
       auto locker = std::unique_lock(outputMutex_);
 #endif
 
+      if (auto it = std::ranges::find(pendingParserRequests_, doc);
+          it != pendingParserRequests_.end()) {
+        pendingParserRequests_.erase(it);
+      }
+
       if (documents_.contains(uri) && documents_.at(uri)->version() > version) {
         return;
       }

src/lsp/cxx/lsp/lsp_server.h

@@ -68,6 +68,8 @@ class Server {
   void startWorkersIfNeeded();
   void stopWorkersIfNeeded();
 
+  void cancelPendingParserRequests(const std::string& fileName);
+
   void run(std::function<void()> task);
 
   void parse(const std::string& uri);
@@ -105,6 +107,7 @@ class Server {
   std::ostream& log;
   std::unordered_map<std::string, std::shared_ptr<CxxDocument>> documents_;
   std::unordered_map<std::string, Text> documentContents_;
+  std::vector<std::shared_ptr<CxxDocument>> pendingParserRequests_;
 #ifndef CXX_NO_THREADS
   SyncQueue syncQueue_;
   std::vector<std::thread> workers_;