From 60467ae701e37db8415d8bbdd1e3c67f1c51c958 Mon Sep 17 00:00:00 2001
From: Max Brunsfeld <maxbrunsfeld@gmail.com>
Date: Tue, 10 Sep 2019 22:30:27 -0700
Subject: [PATCH] Fix bugs in binary search used in tree queries

---
 cli/src/tests/query_test.rs        | 42 ++++++++++++++++++++++++++++++
 lib/binding_web/test/query-test.js | 11 ++++----
 lib/src/query.c                    |  6 ++---
 3 files changed, 51 insertions(+), 8 deletions(-)

diff --git a/cli/src/tests/query_test.rs b/cli/src/tests/query_test.rs
index 32adba62..28becda2 100644
--- a/cli/src/tests/query_test.rs
+++ b/cli/src/tests/query_test.rs
@@ -193,6 +193,48 @@ fn test_query_exec_with_multiple_matches_same_root() {
     allocations::stop_recording();
 }
 
+#[test]
+fn test_query_exec_multiple_patterns() {
+    allocations::start_recording();
+
+    let language = get_language("javascript");
+    let query = Query::new(
+        language,
+        "
+            (function_declaration name:(identifier) @fn-def)
+            (call_expression function:(identifier) @fn-ref)
+        ",
+    )
+    .unwrap();
+
+    let source = "
+        function f1() {
+            f2(f3());
+        }
+    ";
+
+    let mut parser = Parser::new();
+    parser.set_language(language).unwrap();
+    let tree = parser.parse(source, None).unwrap();
+    let context = query.context();
+    let matches = context.exec(tree.root_node());
+
+    assert_eq!(
+        collect_matches(matches, &query, source),
+        &[
+            (0, vec![("fn-def", "f1")]),
+            (1, vec![("fn-ref", "f2")]),
+            (1, vec![("fn-ref", "f3")]),
+        ],
+    );
+
+    drop(context);
+    drop(parser);
+    drop(query);
+    drop(tree);
+    allocations::stop_recording();
+}
+
 fn collect_matches<'a>(
     matches: impl Iterator<Item = QueryMatch<'a>>,
     query: &'a Query,
diff --git a/lib/binding_web/test/query-test.js b/lib/binding_web/test/query-test.js
index f02c5d86..38c3209e 100644
--- a/lib/binding_web/test/query-test.js
+++ b/lib/binding_web/test/query-test.js
@@ -32,8 +32,9 @@ describe("Query", () => {
 
   it('matches simple queries',  () => {
     tree = parser.parse("function one() { two(); function three() {} }");
-    const query = JavaScript.query(`
-      (function_declaration name:(identifier) @the-name)
+    query = JavaScript.query(`
+      (function_declaration name:(identifier) @fn-def)
+      (call_expression function:(identifier) @fn-ref)
     `);
     const matches = query.exec(tree.rootNode);
     assert.deepEqual(
@@ -42,9 +43,9 @@ describe("Query", () => {
         captures: captures.map(({name, node}) => ({name, text: node.text}))
       })),
       [
-        {pattern: 0, captures: [{name: 'the-name', text: 'one'}]},
-        // {pattern: 0, captures: [{name: 'the-function', text: 'two'}]},
-        {pattern: 0, captures: [{name: 'the-name', text: 'three'}]},
+        {pattern: 0, captures: [{name: 'fn-def', text: 'one'}]},
+        {pattern: 1, captures: [{name: 'fn-ref', text: 'two'}]},
+        {pattern: 0, captures: [{name: 'fn-def', text: 'three'}]},
       ]
     );
   });
diff --git a/lib/src/query.c b/lib/src/query.c
index f7836a86..167de1d7 100644
--- a/lib/src/query.c
+++ b/lib/src/query.c
@@ -252,20 +252,20 @@ static inline bool ts_query__pattern_map_search(
     TSSymbol mid_symbol = self->steps.contents[
       self->pattern_map.contents[mid_index].step_index
     ].symbol;
-    if (needle > mid_symbol) base_index = mid_index;
+    if (needle >= mid_symbol) base_index = mid_index;
     size -= half_size;
   }
   TSSymbol symbol = self->steps.contents[
     self->pattern_map.contents[base_index].step_index
   ].symbol;
   if (needle > symbol) {
-    *result = base_index;
+    *result = base_index + 1;
     return false;
   } else if (needle == symbol) {
     *result = base_index;
     return true;
   } else {
-    *result = base_index + 1;
+    *result = base_index;
     return false;
   }
 }