From f18c36ca62954fcdb038195f6fafcb78092bf347 Mon Sep 17 00:00:00 2001
From: Max Brunsfeld <maxbrunsfeld@gmail.com>
Date: Sat, 6 Mar 2021 14:35:41 -0800
Subject: [PATCH] Fix handling of extra nodes in query analysis

Fixes #967
---
 cli/src/tests/query_test.rs | 43 +++++++++++++++++++++++++++++++++++++
 lib/src/query.c             | 12 ++++++++---
 2 files changed, 52 insertions(+), 3 deletions(-)

diff --git a/cli/src/tests/query_test.rs b/cli/src/tests/query_test.rs
index f41cdf3f..f1139c6c 100644
--- a/cli/src/tests/query_test.rs
+++ b/cli/src/tests/query_test.rs
@@ -642,6 +642,49 @@ fn test_query_matches_capturing_error_nodes() {
     });
 }
 
+#[test]
+fn test_query_matches_with_extra_children() {
+    allocations::record(|| {
+        let language = get_language("ruby");
+        let query = Query::new(
+            language,
+            "
+            (program(comment) @top_level_comment)
+            (argument_list (heredoc_body) @heredoc_in_args)
+            ",
+        )
+        .unwrap();
+
+        assert_query_matches(
+            language,
+            &query,
+            "
+            # top-level
+            puts(
+                # not-top-level
+                <<-IN_ARGS, bar.baz
+                HELLO
+                IN_ARGS
+            )
+
+            puts <<-NOT_IN_ARGS
+            NO
+            NOT_IN_ARGS
+            ",
+            &[
+                (0, vec![("top_level_comment", "# top-level")]),
+                (
+                    1,
+                    vec![(
+                        "heredoc_in_args",
+                        "\n                HELLO\n                IN_ARGS",
+                    )],
+                ),
+            ],
+        );
+    });
+}
+
 #[test]
 fn test_query_matches_with_named_wildcard() {
     allocations::record(|| {
diff --git a/lib/src/query.c b/lib/src/query.c
index 5a20603d..c3497d8d 100644
--- a/lib/src/query.c
+++ b/lib/src/query.c
@@ -959,6 +959,10 @@ static bool ts_query__analyze_patterns(TSQuery *self, unsigned *error_offset) {
         }
       #endif
 
+      // If no further progress can be made within the current recursion depth limit, then
+      // bump the depth limit by one, and continue to process the states the exceeded the
+      // limit. But only allow this if progress has been made since the last time the depth
+      // limit was increased.
       if (states.size == 0) {
         if (deeper_states.size > 0 && final_step_indices.size > prev_final_step_count) {
           #ifdef DEBUG_ANALYZE_QUERY
@@ -1019,12 +1023,12 @@ static bool ts_query__analyze_patterns(TSQuery *self, unsigned *error_offset) {
           TSStateId next_parse_state;
           if (lookahead_iterator.action_count) {
             const TSParseAction *action = &lookahead_iterator.actions[lookahead_iterator.action_count - 1];
-            if (action->type == TSParseActionTypeShift && !action->shift.extra) {
-              next_parse_state = action->shift.state;
+            if (action->type == TSParseActionTypeShift) {
+              next_parse_state = action->shift.extra ? parse_state : action->shift.state;
             } else {
               continue;
             }
-          } else if (lookahead_iterator.next_state != 0 && lookahead_iterator.next_state != parse_state) {
+          } else if (lookahead_iterator.next_state != 0) {
             next_parse_state = lookahead_iterator.next_state;
           } else {
             continue;
@@ -1127,6 +1131,8 @@ static bool ts_query__analyze_patterns(TSQuery *self, unsigned *error_offset) {
                   next_step->depth <= parent_depth + 1
                 ) break;
               }
+            } else if (next_parse_state == parse_state) {
+              continue;
             }
 
             for (;;) {