Merge pull request tree-sitter#1127 from dcreager/query-mempool

query: Allow unlimited pending matches

Author: maxbrunsfeld

cli/src/tests/query_test.rs

@@ -1645,6 +1645,7 @@ fn test_query_matches_with_too_many_permutations_to_track() {
         parser.set_language(language).unwrap();
         let tree = parser.parse(&source, None).unwrap();
         let mut cursor = QueryCursor::new();
+        cursor.set_match_limit(32);
         let matches = cursor.matches(&query, tree.root_node(), to_callback(&source));
 
         // For this pathological query, some match permutations will be dropped.
@@ -1686,6 +1687,7 @@ fn test_query_matches_with_alternatives_and_too_many_permutations_to_track() {
         parser.set_language(language).unwrap();
         let tree = parser.parse(&source, None).unwrap();
         let mut cursor = QueryCursor::new();
+        cursor.set_match_limit(32);
         let matches = cursor.matches(&query, tree.root_node(), to_callback(&source));
 
         assert_eq!(
@@ -2765,6 +2767,7 @@ fn test_query_captures_with_too_many_nested_results() {
         parser.set_language(language).unwrap();
         let tree = parser.parse(&source, None).unwrap();
         let mut cursor = QueryCursor::new();
+        cursor.set_match_limit(32);
         let captures = cursor.captures(&query, tree.root_node(), to_callback(&source));
         let captures = collect_captures(captures, &query, &source);
 

lib/binding_rust/bindings.rs

@@ -726,15 +726,27 @@ extern "C" {
     pub fn ts_query_cursor_exec(arg1: *mut TSQueryCursor, arg2: *const TSQuery, arg3: TSNode);
 }
 extern "C" {
-    #[doc = " Check if this cursor has exceeded its maximum number of in-progress"]
-    #[doc = " matches."]
+    #[doc = " Manage the maximum number of in-progress matches allowed by this query"]
+    #[doc = " cursor."]
     #[doc = ""]
-    #[doc = " Currently, query cursors have a fixed capacity for storing lists"]
-    #[doc = " of in-progress captures. If this capacity is exceeded, then the"]
-    #[doc = " earliest-starting match will silently be dropped to make room for"]
-    #[doc = " further matches."]
+    #[doc = " Query cursors have a maximum capacity for storing lists of in-progress"]
+    #[doc = " captures. If this capacity is exceeded, then the earliest-starting match will"]
+    #[doc = " silently be dropped to make room for further matches."]
+    #[doc = ""]
+    #[doc = " By default, this limit is 65,536 pending matches, which is effectively"]
+    #[doc = " unlimited for most queries and syntax trees. You can optionally set this to a"]
+    #[doc = " lower number if you want to have (and check) a tighter bound on query"]
+    #[doc = " complexity."]
+    #[doc = ""]
+    #[doc = " If you update the match limit, it must be > 0 and <= 65536."]
     pub fn ts_query_cursor_did_exceed_match_limit(arg1: *const TSQueryCursor) -> bool;
 }
+extern "C" {
+    pub fn ts_query_cursor_match_limit(arg1: *const TSQueryCursor) -> u32;
+}
+extern "C" {
+    pub fn ts_query_cursor_set_match_limit(arg1: *mut TSQueryCursor, arg2: u32);
+}
 extern "C" {
     #[doc = " Set the range of bytes or (row, column) positions in which the query"]
     #[doc = " will be executed."]

lib/binding_rust/lib.rs

@@ -1598,6 +1598,19 @@ impl<'a> QueryCursor {
         QueryCursor(unsafe { NonNull::new_unchecked(ffi::ts_query_cursor_new()) })
     }
 
+    /// Return the maximum number of in-progress matches for this cursor.
+    pub fn match_limit(&self) -> u32 {
+        unsafe { ffi::ts_query_cursor_match_limit(self.0.as_ptr()) }
+    }
+
+    /// Set the maximum number of in-progress matches for this cursor.  The limit must be > 0 and
+    /// <= 65536.
+    pub fn set_match_limit(&mut self, limit: u32) {
+        unsafe {
+            ffi::ts_query_cursor_set_match_limit(self.0.as_ptr(), limit);
+        }
+    }
+
     /// Check if, on its last execution, this cursor exceeded its maximum number of
     /// in-progress matches.
     pub fn did_exceed_match_limit(&self) -> bool {

lib/binding_web/binding.c

@@ -594,9 +594,15 @@ void ts_query_matches_wasm(
   uint32_t start_row,
   uint32_t start_column,
   uint32_t end_row,
-  uint32_t end_column
+  uint32_t end_column,
+  uint32_t match_limit
 ) {
   if (!scratch_query_cursor) scratch_query_cursor = ts_query_cursor_new();
+  if (match_limit == 0) {
+    ts_query_cursor_set_match_limit(scratch_query_cursor, UINT32_MAX);
+  } else {
+    ts_query_cursor_set_match_limit(scratch_query_cursor, match_limit);
+  }
 
   TSNode node = unmarshal_node(tree);
   TSPoint start_point = {start_row, code_unit_to_byte(start_column)};
@@ -635,9 +641,15 @@ void ts_query_captures_wasm(
   uint32_t start_row,
   uint32_t start_column,
   uint32_t end_row,
-  uint32_t end_column
+  uint32_t end_column,
+  uint32_t match_limit
 ) {
   if (!scratch_query_cursor) scratch_query_cursor = ts_query_cursor_new();
+  if (match_limit == 0) {
+    ts_query_cursor_set_match_limit(scratch_query_cursor, UINT32_MAX);
+  } else {
+    ts_query_cursor_set_match_limit(scratch_query_cursor, match_limit);
+  }
 
   TSNode node = unmarshal_node(tree);
   TSPoint start_point = {start_row, code_unit_to_byte(start_column)};

lib/binding_web/binding.js

@@ -953,9 +953,17 @@ class Query {
     this[0] = 0;
   }
 
-  matches(node, startPosition, endPosition) {
+  matches(node, startPosition, endPosition, options) {
     if (!startPosition) startPosition = ZERO_POINT;
     if (!endPosition) endPosition = ZERO_POINT;
+    if (!options) options = {};
+
+    let matchLimit = options.matchLimit;
+    if (typeof matchLimit === 'undefined') {
+      matchLimit = 0;
+    } else if (typeof matchLimit !== 'number') {
+      throw new Error('Arguments must be numbers');
+    }
 
     marshalNode(node);
 
@@ -965,7 +973,8 @@ class Query {
       startPosition.row,
       startPosition.column,
       endPosition.row,
-      endPosition.column
+      endPosition.column,
+      matchLimit
     );
 
     const rawCount = getValue(TRANSFER_BUFFER, 'i32');
@@ -1000,9 +1009,17 @@ class Query {
     return result;
   }
 
-  captures(node, startPosition, endPosition) {
+  captures(node, startPosition, endPosition, options) {
     if (!startPosition) startPosition = ZERO_POINT;
     if (!endPosition) endPosition = ZERO_POINT;
+    if (!options) options = {};
+
+    let matchLimit = options.matchLimit;
+    if (typeof matchLimit === 'undefined') {
+      matchLimit = 0;
+    } else if (typeof matchLimit !== 'number') {
+      throw new Error('Arguments must be numbers');
+    }
 
     marshalNode(node);
 
@@ -1012,7 +1029,8 @@ class Query {
       startPosition.row,
       startPosition.column,
       endPosition.row,
-      endPosition.column
+      endPosition.column,
+      matchLimit
     );
 
     const count = getValue(TRANSFER_BUFFER, 'i32');

lib/binding_web/test/query-test.js

@@ -256,7 +256,7 @@ describe("Query", () => {
         (array (identifier) @pre (identifier) @post)
       `);
 
-      const captures = query.captures(tree.rootNode);
+      const captures = query.captures(tree.rootNode, null, null, {matchLimit: 32});
       assert.ok(query.didExceedMatchLimit());
     });
   });

lib/include/tree_sitter/api.h

@@ -798,15 +798,19 @@ void ts_query_cursor_delete(TSQueryCursor *);
 void ts_query_cursor_exec(TSQueryCursor *, const TSQuery *, TSNode);
 
 /**
- * Check if this cursor has exceeded its maximum number of in-progress
- * matches.
+ * Manage the maximum number of in-progress matches allowed by this query
+ * cursor.
  *
- * Currently, query cursors have a fixed capacity for storing lists
- * of in-progress captures. If this capacity is exceeded, then the
- * earliest-starting match will silently be dropped to make room for
- * further matches.
+ * Query cursors have an optional maximum capacity for storing lists of
+ * in-progress captures. If this capacity is exceeded, then the
+ * earliest-starting match will silently be dropped to make room for further
+ * matches. This maximum capacity is optional — by default, query cursors allow
+ * any number of pending matches, dynamically allocating new space for them as
+ * needed as the query is executed.
  */
 bool ts_query_cursor_did_exceed_match_limit(const TSQueryCursor *);
+uint32_t ts_query_cursor_match_limit(const TSQueryCursor *);
+void ts_query_cursor_set_match_limit(TSQueryCursor *, uint32_t);
 
 /**
  * Set the range of bytes or (row, column) positions in which the query