From efc51a596c6015b8b0dfd587122820412a5b97a8 Mon Sep 17 00:00:00 2001
From: Amaan Qureshi <amaanq12@gmail.com>
Date: Sat, 4 Jan 2025 01:28:46 -0500
Subject: [PATCH] fix(lib): don't consider unfinished captures definite when
 the following step is immediate

When collecting captures, we were treating unfinished ones as definite
even if they had pending immediate steps that weren't yet satisfied. Now
we only mark a capture as definite if the pattern is guaranteed and
there are no pending immediate steps to check.
---
 cli/src/tests/query_test.rs | 28 ++++++++++++++++++++++++++++
 lib/src/query.c             |  9 ++++++---
 2 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/cli/src/tests/query_test.rs b/cli/src/tests/query_test.rs
index b6b76cd1..132f7076 100644
--- a/cli/src/tests/query_test.rs
+++ b/cli/src/tests/query_test.rs
@@ -5550,3 +5550,31 @@ function foo() {
         &[(0, vec![("part", "function foo() {\n    \"bar\"\n}")])],
     );
 }
+
+#[test]
+fn test_unfinished_captures_are_not_definite_with_pending_anchors() {
+    let language = get_language("javascript");
+    let mut parser = Parser::new();
+    parser.set_language(&language).unwrap();
+
+    let source_code = "
+const foo = [
+  1, 2, 3
+]
+";
+
+    let tree = parser.parse(source_code, None).unwrap();
+    let query = Query::new(&language, r#"(array (_) @foo . "]")"#).unwrap();
+    let mut matches_cursor = QueryCursor::new();
+    let mut captures_cursor = QueryCursor::new();
+
+    let captures = captures_cursor.captures(&query, tree.root_node(), source_code.as_bytes());
+    let captures = collect_captures(captures, &query, source_code);
+
+    let matches = matches_cursor.matches(&query, tree.root_node(), source_code.as_bytes());
+    let matches = collect_matches(matches, &query, source_code);
+
+    assert_eq!(captures, vec![("foo", "3")]);
+    assert_eq!(matches.len(), 1);
+    assert_eq!(matches[0].1, captures);
+}
diff --git a/lib/src/query.c b/lib/src/query.c
index 12829a20..53864526 100644
--- a/lib/src/query.c
+++ b/lib/src/query.c
@@ -3206,7 +3206,7 @@ static bool ts_query_cursor__first_in_progress_capture(
   uint32_t *state_index,
   uint32_t *byte_offset,
   uint32_t *pattern_index,
-  bool *root_pattern_guaranteed
+  bool *is_definite
 ) {
   bool result = false;
   *state_index = UINT32_MAX;
@@ -3241,8 +3241,11 @@ static bool ts_query_cursor__first_in_progress_capture(
       (node_start_byte == *byte_offset && state->pattern_index < *pattern_index)
     ) {
       QueryStep *step = &self->query->steps.contents[state->step_index];
-      if (root_pattern_guaranteed) {
-        *root_pattern_guaranteed = step->root_pattern_guaranteed;
+      if (is_definite) {
+        // We're being a bit conservative here by asserting that the following step
+        // is not immediate, because this capture might end up being discarded if the
+        // following symbol in the tree isn't the required symbol for this step.
+        *is_definite = step->root_pattern_guaranteed && !step->is_immediate;
       } else if (step->root_pattern_guaranteed) {
         continue;
       }