tags: Handle cancellation

cli/src/highlight.rs

@@ -1,3 +1,4 @@
+use super::util;
 use crate::error::Result;
 use crate::loader::Loader;
 use ansi_term::Color;
@@ -6,10 +7,8 @@ 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::{fs, io, path, str, thread, usize};
+use std::{fs, io, path, str, usize};
 use tree_sitter_highlight::{HighlightConfiguration, HighlightEvent, Highlighter, HtmlRenderer};
 
 pub const HTML_HEADER: &'static str = "
@@ -273,19 +272,6 @@ 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,
@@ -296,7 +282,7 @@ pub fn ansi(
     let stdout = io::stdout();
     let mut stdout = stdout.lock();
     let time = Instant::now();
-    let cancellation_flag = cancel_on_stdin();
+    let cancellation_flag = util::cancel_on_stdin();
     let mut highlighter = Highlighter::new();
 
     let events = highlighter.highlight(config, source, Some(&cancellation_flag), |string| {
@@ -341,7 +327,7 @@ pub fn html(
     let stdout = io::stdout();
     let mut stdout = stdout.lock();
     let time = Instant::now();
-    let cancellation_flag = cancel_on_stdin();
+    let cancellation_flag = util::cancel_on_stdin();
     let mut highlighter = Highlighter::new();
 
     let events = highlighter.highlight(config, source, Some(&cancellation_flag), |string| {

cli/src/tags.rs

@@ -1,4 +1,5 @@
 use super::loader::Loader;
+use super::util;
 use crate::error::{Error, Result};
 use std::io::{self, Write};
 use std::path::Path;
@@ -15,6 +16,7 @@ pub fn generate_tags(loader: &Loader, scope: Option<&str>, paths: &[String]) ->
     }
 
     let mut context = TagsContext::new();
+    let cancellation_flag = util::cancel_on_stdin();
     let stdout = io::stdout();
     let mut stdout = stdout.lock();
 
@@ -36,7 +38,8 @@ pub fn generate_tags(loader: &Loader, scope: Option<&str>, paths: &[String]) ->
             writeln!(&mut stdout, "{}", &path_str[1..path_str.len() - 1])?;
 
             let source = fs::read(path)?;
-            for tag in context.generate_tags(tags_config, &source) {
+            for tag in context.generate_tags(tags_config, &source, Some(&cancellation_flag))? {
+                let tag = tag?;
                 write!(
                     &mut stdout,
                     "  {:<8} {:<40}\t{:>9}-{:<9}",

cli/src/tests/tags_test.rs

@@ -3,7 +3,7 @@ use super::helpers::fixtures::{get_language, get_language_queries_path};
 use std::ffi::CString;
 use std::{fs, ptr, slice, str};
 use tree_sitter_tags::c_lib as c;
-use tree_sitter_tags::{TagKind, TagsConfiguration, TagsContext};
+use tree_sitter_tags::{Error, TagKind, TagsConfiguration, TagsContext};
 
 const PYTHON_TAG_QUERY: &'static str = r#"
 ((function_definition
@@ -79,8 +79,10 @@ fn test_tags_python() {
     "#;
 
     let tags = tag_context
-        .generate_tags(&tags_config, source)
-        .collect::<Vec<_>>();
+        .generate_tags(&tags_config, source, None)
+        .unwrap()
+        .collect::<Result<Vec<_>, _>>()
+        .unwrap();
 
     assert_eq!(
         tags.iter()
@@ -128,8 +130,10 @@ fn test_tags_javascript() {
 
     let mut tag_context = TagsContext::new();
     let tags = tag_context
-        .generate_tags(&tags_config, source)
-        .collect::<Vec<_>>();
+        .generate_tags(&tags_config, source, None)
+        .unwrap()
+        .collect::<Result<Vec<_>, _>>()
+        .unwrap();
 
     assert_eq!(
         tags.iter()
@@ -178,8 +182,10 @@ fn test_tags_ruby() {
 
     let mut tag_context = TagsContext::new();
     let tags = tag_context
-        .generate_tags(&tags_config, source.as_bytes())
-        .collect::<Vec<_>>();
+        .generate_tags(&tags_config, source.as_bytes(), None)
+        .unwrap()
+        .collect::<Result<Vec<_>, _>>()
+        .unwrap();
 
     assert_eq!(
         tags.iter()
@@ -201,6 +207,39 @@ fn test_tags_ruby() {
     );
 }
 
+#[test]
+fn test_tags_cancellation() {
+    use std::sync::atomic::{AtomicUsize, Ordering};
+
+    allocations::record(|| {
+        // Large javascript document
+        let source = (0..500)
+            .map(|_| "/* hi */ class A { /* ok */ b() {} }\n")
+            .collect::<String>();
+
+        let cancellation_flag = AtomicUsize::new(0);
+        let language = get_language("javascript");
+        let tags_config = TagsConfiguration::new(language, JS_TAG_QUERY, "").unwrap();
+
+        let mut tag_context = TagsContext::new();
+        let tags = tag_context
+            .generate_tags(&tags_config, source.as_bytes(), Some(&cancellation_flag))
+            .unwrap();
+
+        for (i, tag) in tags.enumerate() {
+            if i == 150 {
+                cancellation_flag.store(1, Ordering::SeqCst);
+            }
+            if let Err(e) = tag {
+                assert_eq!(e, Error::Cancelled);
+                return;
+            }
+        }
+
+        panic!("Expected to halt tagging with an error");
+    });
+}
+
 #[test]
 fn test_tags_via_c_api() {
     allocations::record(|| {

cli/src/util.rs

@@ -1,12 +1,29 @@
+use std::io;
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::Arc;
+use std::thread;
+use tree_sitter::Parser;
+
 #[cfg(unix)]
 use std::path::PathBuf;
 #[cfg(unix)]
 use std::process::{Child, ChildStdin, Command, Stdio};
-use tree_sitter::Parser;
 
 #[cfg(unix)]
 const HTML_HEADER: &[u8] = b"<!DOCTYPE html>\n<style>svg { width: 100%; }</style>\n\n";
 
+pub 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
+}
 #[cfg(windows)]
 pub struct LogSession();