Merge pull request #306 from tree-sitter/new-cancel-api

Rework the API for cancelling a parse from a different thread

.appveyor.yml

@@ -5,8 +5,8 @@ install:
 
   # Install rust
   - appveyor DownloadFile https://win.rustup.rs/ -FileName rustup-init.exe
-  - IF "%PLATFORM%" == "x86" rustup-init -y --default-toolchain stable --default-host i686-pc-windows-msvc
-  - IF "%PLATFORM%" == "x64" rustup-init -y --default-toolchain stable --default-host x86_64-pc-windows-msvc
+  - IF "%PLATFORM%" == "x86" rustup-init -y --default-toolchain beta --default-host i686-pc-windows-msvc
+  - IF "%PLATFORM%" == "x64" rustup-init -y --default-toolchain beta --default-host x86_64-pc-windows-msvc
   - set PATH=%PATH%;C:\Users\appveyor\.cargo\bin
   - rustc -vV
   - cargo -vV

.travis.yml

@@ -1,6 +1,6 @@
 language: rust
 rust:
-  - stable
+  - beta
 
 os:
   - linux

Cargo.lock

@@ -1,3 +1,5 @@
+# This file is automatically @generated by Cargo.
+# It is not intended for manual editing.
 [[package]]
 name = "aho-corasick"
 version = "0.6.9"

cli/src/tests/highlight_test.rs

@@ -229,6 +229,7 @@ fn test_highlighting_via_c_api() {
         source_code.as_ptr(),
         source_code.as_bytes().len() as u32,
         buffer,
+        ptr::null_mut(),
     );
 
     let output_bytes = c::ts_highlight_buffer_content(buffer);

cli/src/tests/parser_test.rs

@@ -1,6 +1,7 @@
 use super::helpers::edits::{perform_edit, Edit, ReadRecorder};
 use super::helpers::fixtures::{get_language, get_test_language};
 use crate::generate::generate_parser_for_grammar;
+use std::sync::atomic::{AtomicU32, Ordering};
 use std::{thread, time};
 use tree_sitter::{InputEdit, LogType, Parser, Point, Range};
 
@@ -81,7 +82,11 @@ fn test_parsing_with_debug_graph_enabled() {
         .lines()
         .map(|l| l.expect("Failed to read line from graph log"));
     for line in log_reader {
-        assert!(!has_zero_indexed_row(&line), "Graph log output includes zero-indexed row: {}", line);
+        assert!(
+            !has_zero_indexed_row(&line),
+            "Graph log output includes zero-indexed row: {}",
+            line
+        );
     }
 }
 
@@ -296,6 +301,38 @@ fn test_parsing_on_multiple_threads() {
     assert_eq!(child_count_differences, &[1, 2, 3, 4]);
 }
 
+#[test]
+fn test_parsing_cancelled_by_another_thread() {
+    let cancellation_flag = AtomicU32::new(0);
+
+    let mut parser = Parser::new();
+    parser.set_language(get_language("javascript")).unwrap();
+    unsafe { parser.set_cancellation_flag(Some(&cancellation_flag)) };
+
+    let parse_thread = thread::spawn(move || {
+        // Infinite input
+        parser.parse_with(
+            &mut |offset, _| {
+                if offset == 0 {
+                    b" ["
+                } else {
+                    b"0,"
+                }
+            },
+            None,
+        )
+    });
+
+    let cancel_thread = thread::spawn(move || {
+        thread::sleep(time::Duration::from_millis(80));
+        cancellation_flag.store(1, Ordering::Relaxed);
+    });
+
+    cancel_thread.join().unwrap();
+    let tree = parse_thread.join().unwrap();
+    assert!(tree.is_none());
+}
+
 // Timeouts
 
 #[test]

highlight/include/tree_sitter/highlight.h

@@ -79,7 +79,8 @@ int ts_highlighter_highlight(
   const char *scope_name,
   const char *source_code,
   uint32_t source_code_len,
-  TSHighlightBuffer *output
+  TSHighlightBuffer *output,
+  const uint32_t *cancellation_flag
 );
 
 // TSHighlightBuffer: This struct stores the HTML output of syntax

highlight/src/c_lib.rs

@@ -5,6 +5,7 @@ use std::ffi::CStr;
 use std::io::Write;
 use std::os::raw::c_char;
 use std::process::abort;
+use std::sync::atomic::AtomicU32;
 use std::{fmt, slice};
 use tree_sitter::{Language, PropertySheet};
 
@@ -135,13 +136,15 @@ pub extern "C" fn ts_highlighter_highlight(
     source_code: *const c_char,
     source_code_len: u32,
     output: *mut TSHighlightBuffer,
+    cancellation_flag: *const AtomicU32,
 ) -> ErrorCode {
     let this = unwrap_ptr(this);
     let output = unwrap_mut_ptr(output);
     let scope_name = unwrap(unsafe { CStr::from_ptr(scope_name).to_str() });
     let source_code =
         unsafe { slice::from_raw_parts(source_code as *const u8, source_code_len as usize) };
-    this.highlight(source_code, scope_name, output)
+    let cancellation_flag = unsafe { cancellation_flag.as_ref() };
+    this.highlight(source_code, scope_name, output, cancellation_flag)
 }
 
 impl TSHighlighter {
@@ -150,6 +153,7 @@ impl TSHighlighter {
         source_code: &[u8],
         scope_name: &str,
         output: &mut TSHighlightBuffer,
+        cancellation_flag: Option<&AtomicU32>,
     ) -> ErrorCode {
         let configuration = self.languages.get(scope_name);
         if configuration.is_none() {
@@ -173,6 +177,7 @@ impl TSHighlighter {
                     })
                 })
             },
+            cancellation_flag,
         ));
 
         output.html.clear();

highlight/src/lib.rs

@@ -6,9 +6,12 @@ use serde::{Deserialize, Deserializer, Serialize, Serializer};
 use serde_derive::*;
 use std::fmt::{self, Write};
 use std::mem::transmute;
+use std::sync::atomic::{AtomicU32, Ordering};
 use std::{cmp, str, usize};
 use tree_sitter::{Language, Node, Parser, Point, PropertySheet, Range, Tree, TreePropertyCursor};
 
+const CANCELLATION_CHECK_INTERVAL: usize = 100;
+
 #[derive(Debug)]
 enum TreeStep {
     Child {
@@ -91,6 +94,8 @@ where
     parser: Parser,
     layers: Vec<Layer<'a>>,
     utf8_error_len: Option<usize>,
+    operation_count: usize,
+    cancellation_flag: Option<&'a AtomicU32>,
 }
 
 #[derive(Copy, Clone, Debug)]
@@ -377,17 +382,22 @@ where
         language: Language,
         property_sheet: &'a PropertySheet<Properties>,
         injection_callback: F,
+        cancellation_flag: Option<&'a AtomicU32>,
     ) -> Result<Self, String> {
         let mut parser = Parser::new();
+        unsafe { parser.set_cancellation_flag(cancellation_flag.clone()) };
         parser.set_language(language)?;
         let tree = parser
             .parse(source, None)
             .ok_or_else(|| format!("Tree-sitter: failed to parse"))?;
         Ok(Self {
-            injection_callback,
-            source,
-            source_offset: 0,
             parser,
+            source,
+            cancellation_flag,
+            injection_callback,
+            source_offset: 0,
+            operation_count: 0,
+            utf8_error_len: None,
             layers: vec![Layer::new(
                 source,
                 tree,
@@ -400,7 +410,6 @@ where
                 }],
                 0,
             )],
-            utf8_error_len: None,
         })
     }
 
@@ -602,6 +611,16 @@ impl<'a, T: Fn(&str) -> Option<(Language, &'a PropertySheet<Properties>)>> Itera
     type Item = HighlightEvent<'a>;
 
     fn next(&mut self) -> Option<Self::Item> {
+        if let Some(cancellation_flag) = self.cancellation_flag {
+            self.operation_count += 1;
+            if self.operation_count >= CANCELLATION_CHECK_INTERVAL {
+                self.operation_count = 0;
+                if cancellation_flag.load(Ordering::Relaxed) != 0 {
+                    return None;
+                }
+            }
+        }
+
         if let Some(utf8_error_len) = self.utf8_error_len.take() {
             self.source_offset += utf8_error_len;
             return Some(HighlightEvent::Source("\u{FFFD}"));
@@ -824,7 +843,7 @@ pub fn highlight<'a, F>(
 where
     F: Fn(&str) -> Option<(Language, &'a PropertySheet<Properties>)> + 'a,
 {
-    Highlighter::new(source, language, property_sheet, injection_callback)
+    Highlighter::new(source, language, property_sheet, injection_callback, None)
 }
 
 pub fn highlight_html<'a, F1, F2>(
@@ -838,7 +857,7 @@ where
     F1: Fn(&str) -> Option<(Language, &'a PropertySheet<Properties>)>,
     F2: Fn(Scope) -> &'a str,
 {
-    let highlighter = Highlighter::new(source, language, property_sheet, injection_callback)?;
+    let highlighter = Highlighter::new(source, language, property_sheet, injection_callback, None)?;
     let mut renderer = HtmlRenderer::new(attribute_callback);
     let mut scopes = Vec::new();
     for event in highlighter {

lib/binding/bindings.rs

@@ -136,10 +136,10 @@ extern "C" {
     ) -> *mut TSTree;
 }
 extern "C" {
-    pub fn ts_parser_enabled(arg1: *const TSParser) -> bool;
+    pub fn ts_parser_cancellation_flag(arg1: *const TSParser) -> *const u32;
 }
 extern "C" {
-    pub fn ts_parser_set_enabled(arg1: *mut TSParser, arg2: bool);
+    pub fn ts_parser_set_cancellation_flag(arg1: *mut TSParser, arg2: *const u32);
 }
 extern "C" {
     pub fn ts_parser_timeout_micros(arg1: *const TSParser) -> u64;

lib/binding/lib.rs

@@ -13,13 +13,10 @@ use regex::Regex;
 use serde::de::DeserializeOwned;
 use std::collections::HashMap;
 use std::ffi::CStr;
-use std::fmt;
 use std::marker::PhantomData;
 use std::os::raw::{c_char, c_void};
-use std::ptr;
-use std::slice;
-use std::str;
-use std::u16;
+use std::sync::atomic::AtomicU32;
+use std::{fmt, ptr, slice, str, u16};
 
 pub const PARSER_HEADER: &'static str = include_str!("../include/tree_sitter/parser.h");
 
@@ -338,6 +335,18 @@ impl Parser {
             ffi::ts_parser_set_included_ranges(self.0, ts_ranges.as_ptr(), ts_ranges.len() as u32)
         };
     }
+
+    pub unsafe fn cancellation_flag(&self) -> Option<&AtomicU32> {
+        (ffi::ts_parser_cancellation_flag(self.0) as *const AtomicU32).as_ref()
+    }
+
+    pub unsafe fn set_cancellation_flag(&self, flag: Option<&AtomicU32>) {
+        if let Some(flag) = flag {
+            ffi::ts_parser_set_cancellation_flag(self.0, flag as *const AtomicU32 as *const u32);
+        } else {
+            ffi::ts_parser_set_cancellation_flag(self.0, ptr::null());
+        }
+    }
 }
 
 impl Drop for Parser {

lib/include/tree_sitter/api.h

@@ -88,8 +88,8 @@ void ts_parser_halt_on_error(TSParser *, bool);
 TSTree *ts_parser_parse(TSParser *, const TSTree *, TSInput);
 TSTree *ts_parser_parse_string(TSParser *, const TSTree *, const char *, uint32_t);
 TSTree *ts_parser_parse_string_encoding(TSParser *, const TSTree *, const char *, uint32_t, TSInputEncoding);
-bool ts_parser_enabled(const TSParser *);
-void ts_parser_set_enabled(TSParser *, bool);
+const uint32_t *ts_parser_cancellation_flag(const TSParser *);
+void ts_parser_set_cancellation_flag(TSParser *, const uint32_t *);
 uint64_t ts_parser_timeout_micros(const TSParser *);
 void ts_parser_set_timeout_micros(TSParser *, uint64_t);
 void ts_parser_reset(TSParser *);

lib/src/atomic.h

@@ -7,6 +7,10 @@
 
 #include <windows.h>
 
+static inline uint32_t atomic_load(const volatile uint32_t *p) {
+  return *p;
+}
+
 static inline uint32_t atomic_inc(volatile uint32_t *p) {
   return InterlockedIncrement(p);
 }
@@ -17,6 +21,10 @@ static inline uint32_t atomic_dec(volatile uint32_t *p) {
 
 #else
 
+static inline uint32_t atomic_load(const volatile uint32_t *p) {
+  return __atomic_load_n(p, __ATOMIC_RELAXED);
+}
+
 static inline uint32_t atomic_inc(volatile uint32_t *p) {
   return __sync_add_and_fetch(p, 1u);
 }

lib/src/parser.c

@@ -6,6 +6,7 @@
 #include "tree_sitter/api.h"
 #include "./alloc.h"
 #include "./array.h"
+#include "./atomic.h"
 #include "./clock.h"
 #include "./error_costs.h"
 #include "./get_changed_ranges.h"
@@ -69,7 +70,7 @@ struct TSParser {
   uint64_t clock_limit;
   uint64_t start_clock;
   unsigned operation_count;
-  volatile bool enabled;
+  const volatile uint32_t *cancellation_flag;
   bool halt_on_error;
   Subtree old_tree;
   TSRangeArray included_range_differences;
@@ -1283,9 +1284,12 @@ static bool ts_parser__advance(
   }
 
   for (;;) {
-    if (!self->enabled || ++self->operation_count == OP_COUNT_PER_TIMEOUT_CHECK) {
+    if (++self->operation_count == OP_COUNT_PER_TIMEOUT_CHECK) {
       self->operation_count = 0;
-      if ((uint64_t)(get_clock() - self->start_clock) > self->clock_limit) {
+      if (
+        (self->cancellation_flag && !atomic_load(self->cancellation_flag)) ||
+        (self->clock_limit && get_clock() - self->start_clock > self->clock_limit)
+      ) {
         ts_subtree_release(&self->tree_pool, lookahead);
         return false;
       }
@@ -1508,8 +1512,8 @@ TSParser *ts_parser_new() {
   self->reusable_node = reusable_node_new();
   self->dot_graph_file = NULL;
   self->halt_on_error = false;
-  self->enabled = true;
-  self->clock_limit = UINT64_MAX;
+  self->cancellation_flag = NULL;
+  self->clock_limit = 0;
   self->start_clock = 0;
   self->operation_count = 0;
   self->old_tree = NULL_SUBTREE;
@@ -1585,12 +1589,12 @@ void ts_parser_halt_on_error(TSParser *self, bool should_halt_on_error) {
   self->halt_on_error = should_halt_on_error;
 }
 
-bool ts_parser_enabled(const TSParser *self) {
-  return self->enabled;
+const uint32_t *ts_parser_cancellation_flag(const TSParser *self) {
+  return (const uint32_t *)self->cancellation_flag;
 }
 
-void ts_parser_set_enabled(TSParser *self, bool enabled) {
-  self->enabled = enabled;
+void ts_parser_set_cancellation_flag(TSParser *self, const uint32_t *flag) {
+  self->cancellation_flag = (const volatile uint32_t *)flag;
 }
 
 uint64_t ts_parser_timeout_micros(const TSParser *self) {
@@ -1599,7 +1603,6 @@ uint64_t ts_parser_timeout_micros(const TSParser *self) {
 
 void ts_parser_set_timeout_micros(TSParser *self, uint64_t timeout_micros) {
   self->clock_limit = timeout_micros * get_clocks_per_second() / 1000000;
-  if (self->clock_limit == 0) self->clock_limit = UINT64_MAX;
 }
 
 void ts_parser_set_included_ranges(TSParser *self, const TSRange *ranges, uint32_t count) {