From 3340168097bd0375958ecdf2fce8ad694bfb69be Mon Sep 17 00:00:00 2001
From: Max Brunsfeld <maxbrunsfeld@gmail.com>
Date: Wed, 20 Mar 2019 13:14:02 -0700
Subject: [PATCH] Fix backwards logic for cancellation flag

---
 cli/src/tests/parser_test.rs              | 49 +++++++++++++++--------
 highlight/include/tree_sitter/highlight.h |  1 +
 highlight/src/c_lib.rs                    | 48 ++++++++++++----------
 lib/src/atomic.h                          |  2 +-
 lib/src/parser.c                          |  2 +-
 5 files changed, 61 insertions(+), 41 deletions(-)

diff --git a/cli/src/tests/parser_test.rs b/cli/src/tests/parser_test.rs
index c13bdd50..63a4f3e8 100644
--- a/cli/src/tests/parser_test.rs
+++ b/cli/src/tests/parser_test.rs
@@ -303,33 +303,48 @@ fn test_parsing_on_multiple_threads() {
 
 #[test]
 fn test_parsing_cancelled_by_another_thread() {
-    let cancellation_flag = AtomicU32::new(0);
+    let cancellation_flag = Box::new(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,
-        )
-    });
+    // Long input - parsing succeeds
+    let tree = parser.parse_with(
+        &mut |offset, _| {
+            if offset == 0 {
+                b" ["
+            } else if offset >= 20000 {
+                b""
+            } else {
+                b"0,"
+            }
+        },
+        None,
+    );
+    assert!(tree.is_some());
 
     let cancel_thread = thread::spawn(move || {
-        thread::sleep(time::Duration::from_millis(80));
-        cancellation_flag.store(1, Ordering::Relaxed);
+        thread::sleep(time::Duration::from_millis(100));
+        cancellation_flag.store(1, Ordering::SeqCst);
     });
 
+    // Infinite input
+    let tree = parser.parse_with(
+        &mut |offset, _| {
+            thread::yield_now();
+            thread::sleep(time::Duration::from_millis(10));
+            if offset == 0 {
+                b" ["
+            } else {
+                b"0,"
+            }
+        },
+        None,
+    );
+
+    // Parsing returns None because it was cancelled.
     cancel_thread.join().unwrap();
-    let tree = parse_thread.join().unwrap();
     assert!(tree.is_none());
 }
 
diff --git a/highlight/include/tree_sitter/highlight.h b/highlight/include/tree_sitter/highlight.h
index 458862b8..19f09fbd 100644
--- a/highlight/include/tree_sitter/highlight.h
+++ b/highlight/include/tree_sitter/highlight.h
@@ -10,6 +10,7 @@ extern "C" {
 typedef enum {
   TSHighlightOk,
   TSHighlightUnknownScope,
+  TSHighlightTimeout,
 } TSHighlightError;
 
 // The list of scopes which can be styled for syntax highlighting.
diff --git a/highlight/src/c_lib.rs b/highlight/src/c_lib.rs
index a283b0f6..38596007 100644
--- a/highlight/src/c_lib.rs
+++ b/highlight/src/c_lib.rs
@@ -29,6 +29,7 @@ pub struct TSHighlightBuffer {
 pub enum ErrorCode {
     Ok,
     UnknownScope,
+    Timeout,
 }
 
 #[no_mangle]
@@ -162,7 +163,7 @@ impl TSHighlighter {
         let configuration = configuration.unwrap();
         let languages = &self.languages;
 
-        let highlighter = unwrap(Highlighter::new(
+        let highlighter = Highlighter::new(
             source_code,
             configuration.language,
             &configuration.property_sheet,
@@ -178,29 +179,32 @@ impl TSHighlighter {
                 })
             },
             cancellation_flag,
-        ));
+        );
 
-        output.html.clear();
-        output.line_offsets.clear();
-        output.line_offsets.push(0);
-        let mut scopes = Vec::new();
-        for event in highlighter {
-            match event {
-                HighlightEvent::ScopeStart(s) => {
-                    scopes.push(s);
-                    output.start_scope(s, &self.attribute_strings);
-                }
-                HighlightEvent::ScopeEnd => {
-                    scopes.pop();
-                    output.end_scope();
-                }
-                HighlightEvent::Source(src) => {
-                    output.add_text(src, &scopes, &self.attribute_strings);
-                }
-            };
+        if let Ok(highlighter) = highlighter {
+            output.html.clear();
+            output.line_offsets.clear();
+            output.line_offsets.push(0);
+            let mut scopes = Vec::new();
+            for event in highlighter {
+                match event {
+                    HighlightEvent::ScopeStart(s) => {
+                        scopes.push(s);
+                        output.start_scope(s, &self.attribute_strings);
+                    }
+                    HighlightEvent::ScopeEnd => {
+                        scopes.pop();
+                        output.end_scope();
+                    }
+                    HighlightEvent::Source(src) => {
+                        output.add_text(src, &scopes, &self.attribute_strings);
+                    }
+                };
+            }
+            ErrorCode::Ok
+        } else {
+            ErrorCode::Timeout
         }
-
-        ErrorCode::Ok
     }
 }
 
diff --git a/lib/src/atomic.h b/lib/src/atomic.h
index 89f40e48..51abccf8 100644
--- a/lib/src/atomic.h
+++ b/lib/src/atomic.h
@@ -22,7 +22,7 @@ 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);
+  return __atomic_load_n(p, __ATOMIC_SEQ_CST);
 }
 
 static inline uint32_t atomic_inc(volatile uint32_t *p) {
diff --git a/lib/src/parser.c b/lib/src/parser.c
index 4bc455d4..39699866 100644
--- a/lib/src/parser.c
+++ b/lib/src/parser.c
@@ -1287,7 +1287,7 @@ static bool ts_parser__advance(
     if (++self->operation_count == OP_COUNT_PER_TIMEOUT_CHECK) {
       self->operation_count = 0;
       if (
-        (self->cancellation_flag && !atomic_load(self->cancellation_flag)) ||
+        (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);