From 8941dc1dda8f418b974b9e4778a09c2d72f9c7ca Mon Sep 17 00:00:00 2001 From: Max Brunsfeld Date: Mon, 18 Mar 2019 09:52:02 -0700 Subject: [PATCH] Add cancellation flag parameter to highlight API --- cli/src/tests/highlight_test.rs | 1 + highlight/include/tree_sitter/highlight.h | 3 ++- highlight/src/c_lib.rs | 7 ++++- highlight/src/lib.rs | 31 ++++++++++++++++++----- 4 files changed, 34 insertions(+), 8 deletions(-) diff --git a/cli/src/tests/highlight_test.rs b/cli/src/tests/highlight_test.rs index 2847cb71..c4e86c6f 100644 --- a/cli/src/tests/highlight_test.rs +++ b/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); diff --git a/highlight/include/tree_sitter/highlight.h b/highlight/include/tree_sitter/highlight.h index 7b34aef9..458862b8 100644 --- a/highlight/include/tree_sitter/highlight.h +++ b/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 diff --git a/highlight/src/c_lib.rs b/highlight/src/c_lib.rs index ce9f3936..a283b0f6 100644 --- a/highlight/src/c_lib.rs +++ b/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(); diff --git a/highlight/src/lib.rs b/highlight/src/lib.rs index 7af0efb3..66c52c86 100644 --- a/highlight/src/lib.rs +++ b/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>, utf8_error_len: Option, + operation_count: usize, + cancellation_flag: Option<&'a AtomicU32>, } #[derive(Copy, Clone, Debug)] @@ -377,17 +382,22 @@ where language: Language, property_sheet: &'a PropertySheet, injection_callback: F, + cancellation_flag: Option<&'a AtomicU32>, ) -> Result { 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)>> Itera type Item = HighlightEvent<'a>; fn next(&mut self) -> Option { + 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)> + '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)>, 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 {