Add cancellation flag parameter to highlight API

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);

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 {