Add cancellation support for parser requests

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,25 +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) {
+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 = [this] { return isCancelled(); },
   });
 }
 

src/lsp/cxx/lsp/cxx_document.h

@@ -31,10 +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);
+  [[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_;

src/parser/cxx/parser.cc

@@ -869,6 +869,12 @@ auto Parser::expect(TokenKind tk, SourceLocation& location) -> bool {
 
 void Parser::operator()(UnitAST*& ast) { parse(ast); }
 
+auto Parser::config() const -> const ParserConfiguration& { return config_; }
+
+void Parser::setConfig(ParserConfiguration config) {
+  config_ = std::move(config);
+}
+
 void Parser::parse(UnitAST*& ast) { parse_translation_unit(ast); }
 
 void Parser::parse_warn(std::string message) {
@@ -1211,6 +1217,8 @@ void Parser::parse_top_level_declaration_seq(UnitAST*& yyast) {
   LoopParser loop(this);
 
   while (LA()) {
+    if (shouldStopParsing()) break;
+
     loop.start();
 
     DeclarationAST* declaration = nullptr;
@@ -1233,6 +1241,8 @@ void Parser::parse_declaration_seq(List<DeclarationAST*>*& yyast) {
   LoopParser loop(this);
 
   while (LA()) {
+    if (shouldStopParsing()) break;
+
     if (lookat(TokenKind::T_RBRACE)) break;
 
     if (parse_maybe_module()) break;
@@ -8111,6 +8121,8 @@ void Parser::parse_namespace_body(NamespaceDefinitionAST* yyast) {
   LoopParser loop{this};
 
   while (LA()) {
+    if (shouldStopParsing()) break;
+
     if (lookat(TokenKind::T_RBRACE)) break;
 
     loop.start();
@@ -9160,6 +9172,8 @@ void Parser::parse_class_body(List<DeclarationAST*>*& yyast) {
   LoopParser loop{this};
 
   while (LA()) {
+    if (shouldStopParsing()) break;
+
     if (lookat(TokenKind::T_RBRACE)) break;
 
     loop.start();

src/parser/cxx/parser.h

@@ -29,6 +29,7 @@
 
 #include <algorithm>
 #include <deque>
+#include <functional>
 #include <optional>
 #include <unordered_map>
 #include <unordered_set>
@@ -54,26 +55,6 @@ class Parser final {
 
   [[nodiscard]] auto control() const -> Control* { return control_; }
 
-  [[nodiscard]] auto config() const -> const ParserConfiguration& {
-    return config_;
-  }
-
-  void setConfig(const ParserConfiguration& config) { config_ = config; }
-
-  /**
-   * Whether to enable fuzzy template resolution.
-   */
-  [[nodiscard]] auto fuzzyTemplateResolution() const -> bool;
-
-  /**
-   * Sets whether to enable fuzzy template resolution.
-   *
-   * When enabled, the parser will try to resolve template names
-   *
-   * @param fuzzyTemplateResolution whether to enable fuzzy template resolution
-   */
-  void setFuzzyTemplateResolution(bool fuzzyTemplateResolution);
-
   /**
    * Parse the given unit.
    */
@@ -84,6 +65,9 @@ class Parser final {
    */
   void operator()(UnitAST*& ast);
 
+  [[nodiscard]] auto config() const -> const ParserConfiguration&;
+  void setConfig(ParserConfiguration config);
+
  private:
   struct DeclSpecs;
   struct Decl;
@@ -123,6 +107,11 @@ class Parser final {
 
   static auto prec(TokenKind tk) -> Prec;
 
+  [[nodiscard]] auto shouldStopParsing() const -> bool {
+    if (config_.stopParsingPredicate) return config_.stopParsingPredicate();
+    return false;
+  }
+
   [[nodiscard]] auto context_allows_function_definition(
       BindingContext ctx) const -> bool {
     if (ctx == BindingContext::kBlock) return false;

src/parser/cxx/parser_fwd.h

@@ -21,6 +21,8 @@
 
 #pragma once
 
+#include <functional>
+
 namespace cxx {
 
 class Parser;
@@ -31,6 +33,7 @@ struct ParserConfiguration {
   bool staticAssert = false;
   bool reflect = true;
   bool templates = false;
+  std::function<bool()> stopParsingPredicate;
 };
 
 }  // namespace cxx

src/parser/cxx/translation_unit.cc

@@ -96,6 +96,27 @@ void TranslationUnit::endPreprocessing() {
   preprocessor_->squeeze();
 }
 
+auto TranslationUnit::fatalErrors() const -> bool {
+  return diagnosticsClient_->fatalErrors();
+}
+
+void TranslationUnit::setFatalErrors(bool fatalErrors) {
+  diagnosticsClient_->setFatalErrors(fatalErrors);
+}
+
+auto TranslationUnit::blockErrors(bool blockErrors) -> bool {
+  return diagnosticsClient_->blockErrors(blockErrors);
+}
+
+void TranslationUnit::error(SourceLocation loc, std::string message) const {
+  diagnosticsClient_->report(tokenAt(loc), Severity::Error, std::move(message));
+}
+
+void TranslationUnit::warning(SourceLocation loc, std::string message) const {
+  TranslationUnit::diagnosticsClient_->report(tokenAt(loc), Severity::Warning,
+                                              std::move(message));
+}
+
 auto TranslationUnit::tokenLength(SourceLocation loc) const -> int {
   const auto& tk = tokenAt(loc);
   if (tk.kind() == TokenKind::T_IDENTIFIER) {
@@ -143,13 +164,13 @@ auto TranslationUnit::tokenEndPosition(SourceLocation loc) const
   return preprocessor_->tokenEndPosition(tokenAt(loc));
 }
 
-void TranslationUnit::parse(const ParserConfiguration& config) {
+void TranslationUnit::parse(ParserConfiguration config) {
   if (ast_) {
     cxx_runtime_error("translation unit already parsed");
   }
 
   Parser parse(this);
-  parse.setConfig(config);
+  parse.setConfig(std::move(config));
   parse(ast_);
 }
 
@@ -158,6 +179,10 @@ auto TranslationUnit::globalScope() const -> Scope* {
   return globalNamespace_->scope();
 }
 
+auto TranslationUnit::fileName() const -> const std::string& {
+  return fileName_;
+}
+
 auto TranslationUnit::load(std::span<const std::uint8_t> data) -> bool {
 #ifndef CXX_NO_FLATBUFFERS
   ASTDecoder decode{this};