From deeeb67a3b20043e05b7197022aa285fa6b1b58c Mon Sep 17 00:00:00 2001
From: Max Brunsfeld <maxbrunsfeld@gmail.com>
Date: Wed, 24 Jun 2020 14:20:56 -0700
Subject: [PATCH] query: Fix handling of alternations under field names (#661)

---
 cli/src/tests/query_test.rs | 35 +++++++++++++++++++++++++++++++++++
 lib/src/query.c             | 27 +++++++++++++++++++++------
 2 files changed, 56 insertions(+), 6 deletions(-)

diff --git a/cli/src/tests/query_test.rs b/cli/src/tests/query_test.rs
index 5499048e..d4f18c7d 100644
--- a/cli/src/tests/query_test.rs
+++ b/cli/src/tests/query_test.rs
@@ -1008,6 +1008,41 @@ fn test_query_matches_with_alternatives_at_root() {
     })
 }
 
+#[test]
+fn test_query_matches_with_alternatives_under_fields() {
+    allocations::record(|| {
+        let language = get_language("javascript");
+        let query = Query::new(
+            language,
+            r#"
+            (assignment_expression
+                left: [
+                    (identifier) @variable
+                    (member_expression property: (property_identifier) @variable)
+                ])
+            "#,
+        )
+        .unwrap();
+
+        assert_query_matches(
+            language,
+            &query,
+            "
+            a = b;
+            b = c.d;
+            e.f = g;
+            h.i = j.k;
+            ",
+            &[
+                (0, vec![("variable", "a")]),
+                (0, vec![("variable", "b")]),
+                (0, vec![("variable", "f")]),
+                (0, vec![("variable", "i")]),
+            ],
+        );
+    });
+}
+
 #[test]
 fn test_query_matches_in_language_with_simple_aliases() {
     allocations::record(|| {
diff --git a/lib/src/query.c b/lib/src/query.c
index 59902dee..ff243494 100644
--- a/lib/src/query.c
+++ b/lib/src/query.c
@@ -715,7 +715,7 @@ static TSQueryError ts_query__parse_pattern(
   uint32_t *capture_count,
   bool is_immediate
 ) {
-  uint32_t starting_step_index = self->steps.size;
+  const uint32_t starting_step_index = self->steps.size;
 
   if (stream->next == 0) return TSQueryErrorSyntax;
 
@@ -951,7 +951,6 @@ static TSQueryError ts_query__parse_pattern(
     stream_skip_whitespace(stream);
 
     // Parse the pattern
-    uint32_t step_index = self->steps.size;
     TSQueryError e = ts_query__parse_pattern(
       self,
       stream,
@@ -972,7 +971,22 @@ static TSQueryError ts_query__parse_pattern(
       stream->input = field_name;
       return TSQueryErrorField;
     }
-    self->steps.contents[step_index].field = field_id;
+
+    uint32_t step_index = starting_step_index;
+    QueryStep *step = &self->steps.contents[step_index];
+    for (;;) {
+      step->field = field_id;
+      if (
+        step->alternative_index != NONE &&
+        step->alternative_index > step_index &&
+        step->alternative_index < self->steps.size
+      ) {
+        step_index = step->alternative_index;
+        step = &self->steps.contents[step_index];
+      } else {
+        break;
+      }
+    }
   }
 
   else {
@@ -1041,15 +1055,16 @@ static TSQueryError ts_query__parse_pattern(
         length
       );
 
+      uint32_t step_index = starting_step_index;
       for (;;) {
         query_step__add_capture(step, capture_id);
         if (
           step->alternative_index != NONE &&
-          step->alternative_index > starting_step_index &&
+          step->alternative_index > step_index &&
           step->alternative_index < self->steps.size
         ) {
-          starting_step_index = step->alternative_index;
-          step = &self->steps.contents[starting_step_index];
+          step_index = step->alternative_index;
+          step = &self->steps.contents[step_index];
         } else {
           break;
         }