From 6ba11d7778a7a7205a4b0a3c99677d831c261dd3 Mon Sep 17 00:00:00 2001
From: Max Brunsfeld <maxbrunsfeld@gmail.com>
Date: Mon, 8 Mar 2021 09:33:26 -0800
Subject: [PATCH] Fix query analysis with top-level non-terminal extras

Fixes https://github.com/tree-sitter/tree-sitter/issues/968#issuecomment-792314281
---
 cli/src/tests/query_test.rs | 14 ++++++++++++++
 lib/src/query.c             |  6 ++++--
 2 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/cli/src/tests/query_test.rs b/cli/src/tests/query_test.rs
index 3bfa4b68..30d8e9c9 100644
--- a/cli/src/tests/query_test.rs
+++ b/cli/src/tests/query_test.rs
@@ -3132,6 +3132,20 @@ fn test_query_step_is_definite() {
             "#,
             results_by_substring: &[("name:", true)],
         },
+        Row {
+            description: "top-level non-terminal extra nodes",
+            language: get_language("ruby"),
+            pattern: r#"
+            (heredoc_body
+                (interpolation)
+                (heredoc_end) @end)
+            "#,
+            results_by_substring: &[
+                ("(heredoc_body", false),
+                ("(interpolation)", false),
+                ("(heredoc_end)", true),
+            ],
+        },
     ];
 
     allocations::record(|| {
diff --git a/lib/src/query.c b/lib/src/query.c
index cc127079..eebae855 100644
--- a/lib/src/query.c
+++ b/lib/src/query.c
@@ -784,8 +784,10 @@ static bool ts_query__analyze_patterns(TSQuery *self, unsigned *error_offset) {
             state_predecessor_map_add(&predecessor_map, next_state, state);
           }
         }
-      } else if (lookahead_iterator.next_state != 0 && lookahead_iterator.next_state != state) {
-        state_predecessor_map_add(&predecessor_map, lookahead_iterator.next_state, state);
+      } else if (lookahead_iterator.next_state != 0) {
+        if (lookahead_iterator.next_state != state) {
+          state_predecessor_map_add(&predecessor_map, lookahead_iterator.next_state, state);
+        }
         const TSSymbol *aliases, *aliases_end;
         ts_language_aliases_for_symbol(
           self->language,