From f490befcde2b76df04b8ac335838bd2848deacb4 Mon Sep 17 00:00:00 2001
From: Max Brunsfeld <maxbrunsfeld@gmail.com>
Date: Thu, 26 Sep 2019 15:58:41 -0700
Subject: [PATCH] Add `ts_query_disable_capture` API

---
 lib/binding_rust/bindings.rs  | 10 ++++++++++
 lib/binding_rust/lib.rs       | 10 ++++++++++
 lib/include/tree_sitter/api.h |  7 +++++++
 lib/src/query.c               | 21 +++++++++++++++++++--
 4 files changed, 46 insertions(+), 2 deletions(-)

diff --git a/lib/binding_rust/bindings.rs b/lib/binding_rust/bindings.rs
index df1249a3..f84058b5 100644
--- a/lib/binding_rust/bindings.rs
+++ b/lib/binding_rust/bindings.rs
@@ -655,6 +655,16 @@ extern "C" {
         length: *mut u32,
     ) -> *const ::std::os::raw::c_char;
 }
+extern "C" {
+    #[doc = " Disable a certain capture within a query. This prevents the capture"]
+    #[doc = " from being returned in matches, and also avoids any resource usage"]
+    #[doc = " associated with recording the capture."]
+    pub fn ts_query_disable_capture(
+        arg1: *mut TSQuery,
+        arg2: *const ::std::os::raw::c_char,
+        arg3: u32,
+    );
+}
 extern "C" {
     #[doc = " Create a new cursor for executing a given query."]
     #[doc = ""]
diff --git a/lib/binding_rust/lib.rs b/lib/binding_rust/lib.rs
index d824f964..87771759 100644
--- a/lib/binding_rust/lib.rs
+++ b/lib/binding_rust/lib.rs
@@ -1202,6 +1202,16 @@ impl Query {
         &self.property_settings[index]
     }
 
+    pub fn disable_capture(&mut self, name: &str) {
+        unsafe {
+            ffi::ts_query_disable_capture(
+                self.ptr.as_ptr(),
+                name.as_bytes().as_ptr() as *const c_char,
+                name.len() as u32,
+            );
+        }
+    }
+
     fn parse_property(
         function_name: &str,
         capture_names: &[String],
diff --git a/lib/include/tree_sitter/api.h b/lib/include/tree_sitter/api.h
index b53174fa..04aed846 100644
--- a/lib/include/tree_sitter/api.h
+++ b/lib/include/tree_sitter/api.h
@@ -716,6 +716,13 @@ const char *ts_query_string_value_for_id(
   uint32_t *length
 );
 
+/**
+ * Disable a certain capture within a query. This prevents the capture
+ * from being returned in matches, and also avoids any resource usage
+ * associated with recording the capture.
+ */
+void ts_query_disable_capture(TSQuery *, const char *, uint32_t);
+
 /**
  * Create a new cursor for executing a given query.
  *
diff --git a/lib/src/query.c b/lib/src/query.c
index 10a436f4..c2ba3d30 100644
--- a/lib/src/query.c
+++ b/lib/src/query.c
@@ -889,6 +889,23 @@ uint32_t ts_query_start_byte_for_pattern(
   return self->start_bytes_by_pattern.contents[pattern_index];
 }
 
+void ts_query_disable_capture(
+  TSQuery *self,
+  const char *name,
+  uint32_t length
+) {
+  int id = symbol_table_id_for_name(&self->captures, name, length);
+  if (id != -1) {
+    for (unsigned i = 0; i < self->steps.size; i++) {
+      QueryStep *step = &self->steps.contents[i];
+      if (step->capture_id == id) {
+        step->capture_id = NONE;
+      }
+    }
+  }
+  ts_query__finalize_steps(self);
+}
+
 /***************
  * QueryCursor
  ***************/
@@ -1020,7 +1037,7 @@ static inline bool ts_query_cursor__advance(TSQueryCursor *self) {
       } else if (ts_tree_cursor_goto_parent(&self->cursor)) {
         self->depth--;
       } else {
-        return false;
+        return self->finished_states.size > 0;
       }
     } else {
       bool can_have_later_siblings;
@@ -1214,7 +1231,7 @@ static inline bool ts_query_cursor__advance(TSQueryCursor *self) {
         next_state->step_index++;
         QueryStep *next_step = step + 1;
         if (next_step->depth == PATTERN_DONE_MARKER) {
-          LOG("finish pattern %u\n", next_state->pattern_index);
+          LOG("  finish pattern %u\n", next_state->pattern_index);
 
           next_state->id = self->next_state_id++;
           array_push(&self->finished_states, *next_state);