highlight: Fix panic when cancelled before parsing a nested document

cli/src/highlight.rs

@@ -6,8 +6,10 @@ use serde::ser::SerializeMap;
 use serde::{Deserialize, Deserializer, Serialize, Serializer};
 use serde_json::{json, Value};
 use std::collections::HashMap;
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::Arc;
 use std::time::Instant;
-use std::{fmt, fs, io, path};
+use std::{fmt, fs, io, path, thread};
 use tree_sitter::{Language, PropertySheet};
 use tree_sitter_highlight::{highlight, highlight_html, Highlight, HighlightEvent, Properties};
 
@@ -252,6 +254,19 @@ fn color_to_css(color: Color) -> &'static str {
     }
 }
 
+fn cancel_on_stdin() -> Arc<AtomicUsize> {
+    let result = Arc::new(AtomicUsize::new(0));
+    thread::spawn({
+        let flag = result.clone();
+        move || {
+            let mut line = String::new();
+            io::stdin().read_line(&mut line).unwrap();
+            flag.store(1, Ordering::Relaxed);
+        }
+    });
+    result
+}
+
 pub fn ansi(
     loader: &Loader,
     theme: &Theme,
@@ -264,11 +279,19 @@ pub fn ansi(
     let stdout = io::stdout();
     let mut stdout = stdout.lock();
 
+    let cancellation_flag = cancel_on_stdin();
     let time = Instant::now();
     let mut highlight_stack = Vec::new();
-    for event in highlight(source, language, property_sheet, |s| {
-        language_for_injection_string(loader, s)
-    })? {
+    for event in highlight(
+        source,
+        language,
+        property_sheet,
+        Some(cancellation_flag.as_ref()),
+        |s| language_for_injection_string(loader, s),
+    )
+    .map_err(|e| e.to_string())?
+    {
+        let event = event.map_err(|e| e.to_string())?;
         match event {
             HighlightEvent::Source(s) => {
                 if let Some(style) = highlight_stack.last().and_then(|s| theme.ansi_style(*s)) {
@@ -332,10 +355,13 @@ pub fn html(
     let stdout = io::stdout();
     let mut stdout = stdout.lock();
     write!(&mut stdout, "<table>\n")?;
+
+    let cancellation_flag = cancel_on_stdin();
     let lines = highlight_html(
         source,
         language,
         property_sheet,
+        Some(cancellation_flag.as_ref()),
         |s| language_for_injection_string(loader, s),
         |highlight| {
             if let Some(css_style) = theme.css_style(highlight) {
@@ -344,7 +370,8 @@ pub fn html(
                 ""
             }
         },
-    )?;
+    )
+    .map_err(|e| e.to_string())?;
     for (i, line) in lines.into_iter().enumerate() {
         write!(
             &mut stdout,

cli/src/main.rs

@@ -14,6 +14,7 @@ const BUILD_SHA: Option<&'static str> = option_env!("BUILD_SHA");
 
 fn main() {
     if let Err(e) = run() {
+        println!("");
         eprintln!("{}", e.message());
         exit(1);
     }

cli/src/parse.rs

@@ -28,7 +28,7 @@ pub fn parse_file_at_path(
 ) -> Result<bool> {
     let mut _log_session = None;
     let mut parser = Parser::new();
-    parser.set_language(language)?;
+    parser.set_language(language).map_err(|e| e.to_string())?;
     let mut source_code = fs::read(path).map_err(Error::wrap(|| {
         format!("Error reading source file {:?}", path)
     }))?;

cli/src/test.rs

@@ -58,7 +58,7 @@ pub fn run_tests_at_path(
     let test_entry = parse_tests(path)?;
     let mut _log_session = None;
     let mut parser = Parser::new();
-    parser.set_language(language)?;
+    parser.set_language(language).map_err(|e| e.to_string())?;
 
     if debug_graph {
         _log_session = Some(util::log_graphs(&mut parser, "log.html")?);

cli/src/tests/highlight_test.rs

@@ -1,9 +1,13 @@
 use super::helpers::fixtures::{get_language, get_property_sheet, get_property_sheet_json};
 use lazy_static::lazy_static;
 use std::ffi::CString;
+
+use std::sync::atomic::{AtomicUsize, Ordering};
 use std::{ptr, slice, str};
 use tree_sitter::{Language, PropertySheet};
-use tree_sitter_highlight::{c, highlight, highlight_html, Highlight, HighlightEvent, Properties};
+use tree_sitter_highlight::{
+    c, highlight, highlight_html, Error, Highlight, HighlightEvent, Properties,
+};
 
 lazy_static! {
     static ref JS_SHEET: PropertySheet<Properties> =
@@ -209,9 +213,7 @@ fn test_highlighting_with_local_variable_tracking() {
                 (")", vec![Highlight::PunctuationBracket]),
                 (";", vec![Highlight::PunctuationDelimiter]),
             ],
-            vec![
-                ("}", vec![Highlight::PunctuationBracket])
-            ]
+            vec![("}", vec![Highlight::PunctuationBracket])]
         ],
     );
 }
@@ -307,6 +309,44 @@ fn test_highlighting_with_content_children_included() {
     );
 }
 
+#[test]
+fn test_highlighting_cancellation() {
+    // An HTML document with a large injected JavaScript document:
+    let mut source = "<script>\n".to_string();
+    for _ in 0..500 {
+        source += "function a() { console.log('hi'); }\n";
+    }
+    source += "</script>\n";
+
+    // Cancel the highlighting before parsing the injected document.
+    let cancellation_flag = AtomicUsize::new(0);
+    let injection_callback = |name: &str| {
+        cancellation_flag.store(1, Ordering::SeqCst);
+        test_language_for_injection_string(name)
+    };
+
+    // Constructing the highlighter, which eagerly parses the outer document,
+    // should not fail.
+    let highlighter = highlight(
+        source.as_bytes(),
+        get_language("html"),
+        &HTML_SHEET,
+        Some(&cancellation_flag),
+        injection_callback,
+    )
+    .unwrap();
+
+    // Iterating the scopes should not panic. It should return an error
+    // once the cancellation is detected.
+    for event in highlighter {
+        if let Err(e) = event {
+            assert_eq!(e, Error::Cancelled);
+            return;
+        }
+    }
+    panic!("Expected an error while iterating highlighter");
+}
+
 #[test]
 fn test_highlighting_via_c_api() {
     let js_lang = get_language("javascript");
@@ -410,11 +450,12 @@ fn to_html<'a>(
     src: &'a str,
     language: Language,
     property_sheet: &'a PropertySheet<Properties>,
-) -> Result<Vec<String>, String> {
+) -> Result<Vec<String>, Error> {
     highlight_html(
         src.as_bytes(),
         language,
         property_sheet,
+        None,
         &test_language_for_injection_string,
         &|highlight| SCOPE_CLASS_STRINGS[highlight as usize].as_str(),
     )
@@ -424,7 +465,7 @@ fn to_token_vector<'a>(
     src: &'a str,
     language: Language,
     property_sheet: &'a PropertySheet<Properties>,
-) -> Result<Vec<Vec<(&'a str, Vec<Highlight>)>>, String> {
+) -> Result<Vec<Vec<(&'a str, Vec<Highlight>)>>, Error> {
     let mut lines = Vec::new();
     let mut highlights = Vec::new();
     let mut line = Vec::new();
@@ -432,9 +473,10 @@ fn to_token_vector<'a>(
         src.as_bytes(),
         language,
         property_sheet,
+        None,
         &test_language_for_injection_string,
     )? {
-        match event {
+        match event? {
             HighlightEvent::HighlightStart(s) => highlights.push(s),
             HighlightEvent::HighlightEnd => {
                 highlights.pop();