Merge pull request #2714 from MrPrezident/test_point_range_captures_fix

fix "test_point_range_captures not working"

cli/src/tests/query_test.rs

@@ -2105,6 +2105,73 @@ fn test_query_captures_within_byte_range() {
     });
 }
 
+#[test]
+fn test_query_cursor_next_capture_with_byte_range() {
+    allocations::record(|| {
+        let language = get_language("python");
+        let query = Query::new(
+            language,
+            "(function_definition name: (identifier) @function)
+             (attribute attribute: (identifier) @property)
+             ((identifier) @variable)",
+        )
+        .unwrap();
+
+        let source = "def func():\n  foo.bar.baz()\n";
+        //            ^            ^    ^          ^
+        // byte_pos   0           12    17        27
+        // point_pos (0,0)      (1,0)  (1,5)    (1,15)
+
+        let mut parser = Parser::new();
+        parser.set_language(language).unwrap();
+        let tree = parser.parse(source, None).unwrap();
+
+        let mut cursor = QueryCursor::new();
+        let captures =
+            cursor
+                .set_byte_range(12..17)
+                .captures(&query, tree.root_node(), source.as_bytes());
+
+        assert_eq!(
+            collect_captures(captures, &query, source),
+            &[("variable", "foo"),]
+        );
+    });
+}
+
+#[test]
+fn test_query_cursor_next_capture_with_point_range() {
+    allocations::record(|| {
+        let language = get_language("python");
+        let query = Query::new(
+            language,
+            "(function_definition name: (identifier) @function)
+             (attribute attribute: (identifier) @property)
+             ((identifier) @variable)",
+        )
+        .unwrap();
+
+        let source = "def func():\n  foo.bar.baz()\n";
+        //            ^            ^    ^          ^
+        // byte_pos   0           12    17        27
+        // point_pos (0,0)      (1,0)  (1,5)    (1,15)
+
+        let mut parser = Parser::new();
+        parser.set_language(language).unwrap();
+        let tree = parser.parse(source, None).unwrap();
+
+        let mut cursor = QueryCursor::new();
+        let captures = cursor
+            .set_point_range(Point::new(1, 0)..Point::new(1, 5))
+            .captures(&query, tree.root_node(), source.as_bytes());
+
+        assert_eq!(
+            collect_captures(captures, &query, source),
+            &[("variable", "foo"),]
+        );
+    });
+}
+
 #[test]
 fn test_query_matches_with_unrooted_patterns_intersecting_byte_range() {
     allocations::record(|| {

lib/src/query.c

@@ -4048,9 +4048,20 @@ bool ts_query_cursor_next_capture(
         continue;
       }
 
-      // Skip captures that precede the cursor's start byte.
       TSNode node = captures->contents[state->consumed_capture_count].node;
-      if (ts_node_end_byte(node) <= self->start_byte) {
+
+      bool node_precedes_range = (
+        ts_node_end_byte(node) <= self->start_byte ||
+        point_lte(ts_node_end_point(node), self->start_point)
+      );
+      bool node_follows_range = (
+        ts_node_start_byte(node) >= self->end_byte ||
+        point_gte(ts_node_start_point(node), self->end_point)
+      );
+      bool node_outside_of_range = node_precedes_range || node_follows_range;
+
+      // Skip captures that are outside of the cursor's range.
+      if (node_outside_of_range) {
         state->consumed_capture_count++;
         continue;
       }