fix: cancellation flag is shared by all parsers

Author: itsaky

android-tree-sitter/src/main/cpp/ts_parser.cc

@@ -25,67 +25,153 @@
 #include "utils/ts_preconditions.h"
 #include "utils/ts_log.h"
 
-static std::mutex cancellation_flag_mutex;
-static std::atomic<size_t *> cancellation_flag(nullptr);
+/**
+ * `TSParserInternal` stores the actual tree sitter parser instance along
+ * with the cancellation flag and the cancellation flag mutex.
+ */
+class TSParserInternal {
+ public:
 
-static size_t *get_cancellation_flag() {
-  std::lock_guard<std::mutex> get_lock(cancellation_flag_mutex);
-  return cancellation_flag.load();
-}
+  TSParserInternal() {
+    cancellation_flag_mutex = new std::mutex();
+    cancellation_flag = new std::atomic<size_t *>(nullptr);
+    parser = ts_parser_new();
+  }
 
-static void set_cancellation_flag(size_t *flag) {
-  std::lock_guard<std::mutex> set_lock(cancellation_flag_mutex);
-  cancellation_flag.store(flag);
-}
+  ~TSParserInternal() {
+    delete cancellation_flag_mutex;
+    delete cancellation_flag;
+
+    ts_parser_delete(parser);
+
+    cancellation_flag_mutex = nullptr;
+    cancellation_flag = nullptr;
+    parser = nullptr;
+  }
+
+  TSParser *getParser(JNIEnv *env) {
+    if (check_destroyed(env)) {
+      return nullptr;
+    }
+
+    return this->parser;
+  }
+
+  bool begin_round(JNIEnv *env) {
+    auto flag = get_cancellation_flag(env);
+
+    if (flag) {
+      throw_illegal_state(env,
+                          "Parser is already parsing another syntax tree! You must cancel the current parse first!");
+      return false;
+    }
+
+    // allocate a new cancellation flag
+    flag = (size_t *) malloc(sizeof(int));
+    set_cancellation_flag(env, flag);
+
+    // set the cancellation flag to '0' to indicate that the parser should continue parsing
+    *flag = 0;
+    ts_parser_set_cancellation_flag(getParser(env), flag);
+
+    return true;
+  }
+
+  void end_round(JNIEnv *env) {
+
+    size_t *flag = get_cancellation_flag(env);
+
+    // release the cancellation flag
+    free((size_t *) flag);
+    set_cancellation_flag(env, nullptr);
+    ts_parser_set_cancellation_flag(getParser(env), nullptr);
+  }
+
+  size_t *get_cancellation_flag(JNIEnv *env) {
+    if (check_destroyed(env)) {
+      return nullptr;
+    }
+
+    std::lock_guard<std::mutex> get_lock(*cancellation_flag_mutex);
+    return cancellation_flag->load();
+  }
+
+  void set_cancellation_flag(JNIEnv *env, size_t *flag) {
+    if (check_destroyed(env)) {
+      return;
+    }
+
+    std::lock_guard<std::mutex> set_lock(*cancellation_flag_mutex);
+    cancellation_flag->store(flag);
+  }
+
+ private:
+  std::mutex *cancellation_flag_mutex;
+  std::atomic<size_t *> *cancellation_flag;
+
+  TSParser *parser;
+
+  bool check_destroyed(JNIEnv *env) {
+    if (cancellation_flag_mutex == nullptr || cancellation_flag == nullptr || parser == nullptr) {
+      throw_illegal_state(env, "TSParserInternal has already been destroyed");
+      return true;
+    }
+
+    return false;
+  }
+};
 
 extern "C" JNIEXPORT jlong JNICALL
 Java_com_itsaky_androidide_treesitter_TSParser_00024Native_newParser(
     JNIEnv *env, jclass self) {
-  return (jlong) ts_parser_new();
+  auto parser = new TSParserInternal;
+  return (jlong) parser;
 }
 
 extern "C" JNIEXPORT void JNICALL
 Java_com_itsaky_androidide_treesitter_TSParser_00024Native_delete(
-    JNIEnv *env, jclass self, jlong parser) {
-  req_nnp(env, parser);
-  ts_parser_delete((TSParser *) parser);
+    JNIEnv *env, jclass self, jlong parser_ptr) {
+  req_nnp(env, parser_ptr);
+
+  auto parser = (TSParserInternal *) parser_ptr;
+  delete parser;
 }
 
 extern "C" JNIEXPORT void JNICALL
 Java_com_itsaky_androidide_treesitter_TSParser_00024Native_setLanguage(
     JNIEnv *env, jclass self, jlong parser, jlong language) {
   req_nnp(env, parser, "parser");
   req_nnp(env, language, "language");
-  ts_parser_set_language((TSParser *) parser, (TSLanguage *) language);
+  ts_parser_set_language(((TSParserInternal *) parser)->getParser(env), (TSLanguage *) language);
 }
 
 extern "C" JNIEXPORT jlong JNICALL
 Java_com_itsaky_androidide_treesitter_TSParser_00024Native_getLanguage(
     JNIEnv *env, jclass self, jlong parser) {
   req_nnp(env, parser);
-  return (jlong) ts_parser_language((TSParser *) parser);
+  return (jlong) ts_parser_language(((TSParserInternal *) parser)->getParser(env));
 }
 
 extern "C" JNIEXPORT void JNICALL
 Java_com_itsaky_androidide_treesitter_TSParser_00024Native_reset(JNIEnv *env,
                                                                  jclass self,
                                                                  jlong parser) {
   req_nnp(env, parser);
-  ts_parser_reset((TSParser *) parser);
+  ts_parser_reset(((TSParserInternal *) parser)->getParser(env));
 }
 
 extern "C" JNIEXPORT void JNICALL
 Java_com_itsaky_androidide_treesitter_TSParser_00024Native_setTimeout(
     JNIEnv *env, jclass self, jlong parser, jlong macros) {
   req_nnp(env, parser);
-  ts_parser_set_timeout_micros((TSParser *) parser, macros);
+  ts_parser_set_timeout_micros(((TSParserInternal *) parser)->getParser(env), macros);
 }
 
 extern "C" JNIEXPORT jlong JNICALL
 Java_com_itsaky_androidide_treesitter_TSParser_00024Native_getTimeout(
     JNIEnv *env, jclass self, jlong parser) {
   req_nnp(env, parser);
-  return (jlong) ts_parser_timeout_micros((TSParser *) parser);
+  return (jlong) ts_parser_timeout_micros(((TSParserInternal *) parser)->getParser(env));
 }
 
 extern "C" JNIEXPORT jboolean JNICALL
@@ -102,7 +188,7 @@ Java_com_itsaky_androidide_treesitter_TSParser_00024Native_setIncludedRanges(
   }
 
   const TSRange *r = tsRanges;
-  return (jboolean) ts_parser_set_included_ranges((TSParser *) parser,
+  return (jboolean) ts_parser_set_included_ranges(((TSParserInternal *) parser)->getParser(env),
                                                   r,
                                                   count);
 }
@@ -112,7 +198,7 @@ Java_com_itsaky_androidide_treesitter_TSParser_00024Native_getIncludedRanges(
     JNIEnv *env, jclass self, jlong parser) {
   req_nnp(env, parser);
   jint count;
-  const TSRange *ranges = ts_parser_included_ranges((TSParser *) parser,
+  const TSRange *ranges = ts_parser_included_ranges(((TSParserInternal *) parser)->getParser(env),
                                                     reinterpret_cast<uint32_t *>(&count));
   jobjectArray result = createRangeArr(env, count);
   req_nnp(env, result, "TSRange[] from factory");
@@ -133,25 +219,12 @@ Java_com_itsaky_androidide_treesitter_TSParser_00024Native_parse(JNIEnv *env,
                                                                  jlong str_pointer) {
   req_nnp(env, parser);
   req_nnp(env, str_pointer, "string");
-  auto *ts_parser = (TSParser *) parser;
+  auto *ts_parser_internal = (TSParserInternal *) parser;
+  TSParser *ts_parser = ts_parser_internal->getParser(env);
   TSTree *old_tree = tree_pointer == 0 ? nullptr : (TSTree *) tree_pointer;
   auto *source = as_str(str_pointer);
 
-  auto flag = get_cancellation_flag();
-
-  if (flag) {
-    throw_illegal_state(env,
-                        "Parser is already parsing another syntax tree! You must cancel the current parse first!");
-    return 0;
-  }
-
-  // allocate a new cancellation flag
-  flag = (size_t *) malloc(sizeof(int));
-  set_cancellation_flag(flag);
-
-  // set the cancellation flag to '0' to indicate that the parser should continue parsing
-  *flag = 0;
-  ts_parser_set_cancellation_flag(ts_parser, flag);
+  ts_parser_internal->begin_round(env);
 
   // start parsing
   // if the user cancels the parse while this method is being executed
@@ -163,10 +236,7 @@ Java_com_itsaky_androidide_treesitter_TSParser_00024Native_parse(JNIEnv *env,
                                       source->byte_length(),
                                       TSInputEncodingUTF16);
 
-  // release the cancellation flag
-  free((size_t *) flag);
-  set_cancellation_flag(nullptr);
-  ts_parser_set_cancellation_flag(ts_parser, nullptr);
+  ts_parser_internal->end_round(env);
 
   return (jlong) tree;
 }
@@ -175,9 +245,11 @@ extern "C"
 JNIEXPORT jboolean JNICALL
 Java_com_itsaky_androidide_treesitter_TSParser_00024Native_requestCancellation(
     JNIEnv *env,
-    jclass clazz) {
+    jclass clazz,
+    jlong parser) {
 
-  auto flag = get_cancellation_flag();
+  auto *parserInternal = (TSParserInternal *) parser;
+  auto flag = parserInternal->get_cancellation_flag(env);
 
   // no parse is in progress
   if (flag == nullptr) {

android-tree-sitter/src/main/java/com/itsaky/androidide/treesitter/TSParser.java

@@ -255,7 +255,7 @@ protected boolean unsetParsingFlag() {
    * otherwise.
    */
   public boolean requestCancellationAsync() {
-    final var requested = Native.requestCancellation();
+    final var requested = Native.requestCancellation(getNativeObject());
     setCancellationRequested(requested);
     return requested;
   }
@@ -389,6 +389,6 @@ private static class Native {
 
     public static native long parse(long parser, long treePointer, long strPointer);
 
-    public static native boolean requestCancellation();
+    public static native boolean requestCancellation(long parser);
   }
 }

android-tree-sitter/src/test/java/com/itsaky/androidide/treesitter/ParserTest.java

@@ -446,6 +446,52 @@ public void testParserParseCallShouldSucceedIfAnotherParseIsInProgressAndCancell
     }
   }
 
+  @Test
+  public void testParserParseCallShouldNotFailIfWhenMultipleParsersAreParsing() {
+    try (final var parser1 = TSParser.create();
+      final var parser2 = TSParser.create();
+      final var mainParseContent = UTF16StringFactory.newString()
+    ) {
+      parser1.setLanguage(TSLanguageJava.getInstance());
+      parser2.setLanguage(TSLanguageJava.getInstance());
+
+      // Read the content before starting the threads
+      final var fileContent = readResource("View.java.txt");
+      mainParseContent.append(fileContent);
+      mainParseContent.append(fileContent);
+      mainParseContent.append(fileContent);
+
+      final var executor = Executors.newScheduledThreadPool(2);
+
+      final var parseFuture1 = executor.schedule(() -> {
+        assertThat(parser1.isParsing()).isFalse();
+        try (final var tree = parser1.parseString(mainParseContent)) {
+          assertThat(tree).isNotNull();
+        }
+      }, 0, TimeUnit.MICROSECONDS);
+
+      final var secondParseDelayMs = 1;
+      final var parseFuture2 = executor.schedule(() -> {
+        assertThat(parser2.isParsing()).isFalse();
+        try (var tree = parser2.parseString(fileContent)) {
+          assertThat(tree).isNotNull();
+        } catch (Throwable e) {
+          throw new RuntimeException(e);
+        }
+      }, secondParseDelayMs, TimeUnit.MILLISECONDS);
+
+      try {
+        // both parser instances are independent and hence, both of them must succeed
+        parseFuture1.get();
+        parseFuture2.get();
+      } catch (Throwable e) {
+        throw new RuntimeException(e);
+      } finally {
+        executor.shutdownNow();
+      }
+    }
+  }
+
   @Test
   public void testParserParseCallShouldFailIfAnotherParseIsInProgressAndCancellationWasNotRequested() {
     try (final var parser = TSParser.create(); final var mainParseContent = UTF16StringFactory.newString()) {