Tweak QueryCursor to allow iterating either matches or captures

For syntax highlighting, we want to iterate over all of the captures in 
order, and don't care about grouping the captures by pattern.

lib/binding_web/binding.c

@@ -567,7 +567,7 @@ int ts_node_is_missing_wasm(const TSTree *tree) {
 /* Section - Query */
 /******************/
 
-void ts_query_exec_wasm(
+void ts_query_matches_wasm(
   const TSQuery *self,
   const TSTree *tree,
   uint32_t start_row,
@@ -580,24 +580,23 @@ void ts_query_exec_wasm(
   TSNode node = unmarshal_node(tree);
   TSPoint start_point = {start_row, code_unit_to_byte(start_column)};
   TSPoint end_point = {end_row, code_unit_to_byte(end_column)};
-
-  Array(const void *) result = array_new();
-
-  unsigned index = 0;
-  unsigned match_count = 0;
   ts_query_cursor_set_point_range(scratch_query_cursor, start_point, end_point);
   ts_query_cursor_exec(scratch_query_cursor, self, node);
-  while (ts_query_cursor_next(scratch_query_cursor)) {
+
+  uint32_t index = 0;
+  uint32_t match_count = 0;
+  Array(const void *) result = array_new();
+
+  uint32_t pattern_index, capture_count;
+  const TSQueryCapture *captures;
+  while (ts_query_cursor_next_match(
+    scratch_query_cursor,
+    &pattern_index,
+    &capture_count,
+    &captures
+  )) {
     match_count++;
-    uint32_t pattern_index = ts_query_cursor_matched_pattern_index(scratch_query_cursor);
-    uint32_t capture_count;
-    const TSQueryCapture *captures = ts_query_cursor_matched_captures(
-      scratch_query_cursor,
-      &capture_count
-    );
-
     array_grow_by(&result, 2 + 6 * capture_count);
-
     result.contents[index++] = (const void *)pattern_index;
     result.contents[index++] = (const void *)capture_count;
     for (unsigned i = 0; i < capture_count; i++) {
@@ -611,3 +610,37 @@ void ts_query_exec_wasm(
   TRANSFER_BUFFER[0] = (const void *)(match_count);
   TRANSFER_BUFFER[1] = result.contents;
 }
+
+void ts_query_captures_wasm(
+  const TSQuery *self,
+  const TSTree *tree,
+  uint32_t start_row,
+  uint32_t start_column,
+  uint32_t end_row,
+  uint32_t end_column
+) {
+  if (!scratch_query_cursor) scratch_query_cursor = ts_query_cursor_new();
+
+  TSNode node = unmarshal_node(tree);
+  TSPoint start_point = {start_row, code_unit_to_byte(start_column)};
+  TSPoint end_point = {end_row, code_unit_to_byte(end_column)};
+  ts_query_cursor_set_point_range(scratch_query_cursor, start_point, end_point);
+  ts_query_cursor_exec(scratch_query_cursor, self, node);
+
+  unsigned index = 0;
+  unsigned capture_count = 0;
+  Array(const void *) result = array_new();
+
+  TSQueryCapture capture;
+  while (ts_query_cursor_next_capture(scratch_query_cursor, &capture)) {
+    capture_count++;
+
+    array_grow_by(&result, 6);
+    result.contents[index++] = (const void *)capture.index;
+    marshal_node(result.contents + index, capture.node);
+    index += 5;
+  }
+
+  TRANSFER_BUFFER[0] = (const void *)(capture_count);
+  TRANSFER_BUFFER[1] = result.contents;
+}

lib/binding_web/binding.js

@@ -5,7 +5,7 @@ const SIZE_OF_NODE = 5 * SIZE_OF_INT;
 const SIZE_OF_POINT = 2 * SIZE_OF_INT;
 const SIZE_OF_RANGE = 2 * SIZE_OF_INT + 2 * SIZE_OF_POINT;
 const ZERO_POINT = {row: 0, column: 0};
-const QUERY_WORD_REGEX = /[\w-.]*/;
+const QUERY_WORD_REGEX = /[\w-.]*/g;
 
 var VERSION;
 var MIN_COMPATIBLE_VERSION;
@@ -694,7 +694,7 @@ class Language {
       const errorId = getValue(TRANSFER_BUFFER + SIZE_OF_INT, 'i32');
       const errorByte = getValue(TRANSFER_BUFFER, 'i32');
       const errorIndex = UTF8ToString(sourceAddress, errorByte).length;
-      const suffix = source.slice(errorIndex, 100);
+      const suffix = source.substr(errorIndex, 100);
       const word = suffix.match(QUERY_WORD_REGEX)[0];
       let error;
       switch (errorId) {
@@ -758,46 +758,75 @@ class Query {
     C._ts_query_delete(this[0]);
   }
 
-  exec(queryNode, startPosition, endPosition) {
+  matches(node, startPosition, endPosition) {
     if (!startPosition) startPosition = ZERO_POINT;
     if (!endPosition) endPosition = ZERO_POINT;
 
-    marshalNode(queryNode);
+    marshalNode(node);
 
-    C._ts_query_exec_wasm(
+    C._ts_query_matches_wasm(
       this[0],
-      queryNode.tree[0],
+      node.tree[0],
       startPosition.row,
       startPosition.column,
       endPosition.row,
       endPosition.column
     );
 
-    const matchCount = getValue(TRANSFER_BUFFER, 'i32');
-    const nodesAddress = getValue(TRANSFER_BUFFER + SIZE_OF_INT, 'i32');
-    const result = new Array(matchCount);
+    const count = getValue(TRANSFER_BUFFER, 'i32');
+    const startAddress = getValue(TRANSFER_BUFFER + SIZE_OF_INT, 'i32');
+    const result = new Array(count);
 
-    let address = nodesAddress;
-    for (let i = 0; i < matchCount; i++) {
+    let address = startAddress;
+    for (let i = 0; i < count; i++) {
       const pattern = getValue(address, 'i32');
       address += SIZE_OF_INT;
-      const captures = new Array(getValue(address, 'i32'));
+      const captureCount = getValue(address, 'i32');
       address += SIZE_OF_INT;
-      for (let j = 0, n = captures.length; j < n; j++) {
-        const captureIndex = getValue(address, 'i32');
-        address += SIZE_OF_INT;
-        const node = unmarshalNode(queryNode.tree, address);
-        address += SIZE_OF_NODE;
-        captures[j] = {name: this.captureNames[captureIndex], node};
-      }
+
+      const captures = new Array(captureCount);
+      address = unmarshalCaptures(this, node.tree, address, captures);
       result[i] = {pattern, captures};
     }
 
-    // Free the intermediate buffers
-    C._free(nodesAddress);
-
+    C._free(startAddress);
     return result;
   }
+
+  captures(node, startPosition, endPosition) {
+    if (!startPosition) startPosition = ZERO_POINT;
+    if (!endPosition) endPosition = ZERO_POINT;
+
+    marshalNode(node);
+
+    C._ts_query_captures_wasm(
+      this[0],
+      node.tree[0],
+      startPosition.row,
+      startPosition.column,
+      endPosition.row,
+      endPosition.column
+    );
+
+    const count = getValue(TRANSFER_BUFFER, 'i32');
+    const startAddress = getValue(TRANSFER_BUFFER + SIZE_OF_INT, 'i32');
+    const result = new Array(count);
+    unmarshalCaptures(this, node.tree, startAddress, result);
+
+    C._free(startAddress);
+    return result;
+  }
+}
+
+function unmarshalCaptures(query, tree, address, result) {
+  for (let i = 0, n = result.length; i < n; i++) {
+    const captureIndex = getValue(address, 'i32');
+    address += SIZE_OF_INT;
+    const node = unmarshalNode(tree, address);
+    address += SIZE_OF_NODE;
+    result[i] = {name: query.captureNames[captureIndex], node};
+  }
+  return address;
 }
 
 function assertInternal(x) {

lib/binding_web/exports.json

@@ -73,7 +73,8 @@
   "_ts_query_context_delete",
   "_ts_query_context_new",
   "_ts_query_delete",
-  "_ts_query_exec_wasm",
+  "_ts_query_matches_wasm",
+  "_ts_query_captures_wasm",
   "_ts_query_new",
   "_ts_tree_cursor_current_field_id_wasm",
   "_ts_tree_cursor_current_node_id_wasm",

lib/binding_web/test/query-test.js

@@ -18,64 +18,117 @@ describe("Query", () => {
     if (query) query.delete();
   });
 
-  it('throws an error on invalid syntax', () => {
-    assert.throws(() => {
-      JavaScript.query("(function_declaration wat)")
-    }, "Bad syntax at offset 22: \'wat)\'...");
-    assert.throws(() => {
-      JavaScript.query("(non_existent)")
-    }, "Bad node name 'non_existent'");
-    assert.throws(() => {
-      JavaScript.query("(a)")
-    }, "Bad node name 'a'");
-    assert.throws(() => {
-      JavaScript.query("(function_declaration non_existent:(identifier))")
-    }, "Bad field name 'non_existent'");
+  describe('construction', () => {
+    it('throws an error on invalid syntax', () => {
+      assert.throws(() => {
+        JavaScript.query("(function_declaration wat)")
+      }, "Bad syntax at offset 22: \'wat)\'...");
+      assert.throws(() => {
+        JavaScript.query("(non_existent)")
+      }, "Bad node name 'non_existent'");
+      assert.throws(() => {
+        JavaScript.query("(a)")
+      }, "Bad node name 'a'");
+      assert.throws(() => {
+        JavaScript.query("(function_declaration non_existent:(identifier))")
+      }, "Bad field name 'non_existent'");
+    });
   });
 
-  it('matches simple queries',  () => {
-    tree = parser.parse("function one() { two(); function three() {} }");
-    query = JavaScript.query(`
-      (function_declaration name:(identifier) @fn-def)
-      (call_expression function:(identifier) @fn-ref)
-    `);
-    const matches = query.exec(tree.rootNode);
-    assert.deepEqual(
-      formatMatches(matches),
-      [
-        {pattern: 0, captures: [{name: 'fn-def', text: 'one'}]},
-        {pattern: 1, captures: [{name: 'fn-ref', text: 'two'}]},
-        {pattern: 0, captures: [{name: 'fn-def', text: 'three'}]},
-      ]
-    );
+  describe('.matches', () => {
+    it('returns all of the matches for the given query',  () => {
+      tree = parser.parse("function one() { two(); function three() {} }");
+      query = JavaScript.query(`
+        (function_declaration name:(identifier) @fn-def)
+        (call_expression function:(identifier) @fn-ref)
+      `);
+      const matches = query.matches(tree.rootNode);
+      assert.deepEqual(
+        formatMatches(matches),
+        [
+          {pattern: 0, captures: [{name: 'fn-def', text: 'one'}]},
+          {pattern: 1, captures: [{name: 'fn-ref', text: 'two'}]},
+          {pattern: 0, captures: [{name: 'fn-def', text: 'three'}]},
+        ]
+      );
+    });
+
+    it('can search in a specified ranges',  () => {
+      tree = parser.parse("[a, b,\nc, d,\ne, f,\ng, h]");
+      query = JavaScript.query('(identifier) @element');
+      const matches = query.matches(
+        tree.rootNode,
+        {row: 1, column: 1},
+        {row: 3, column: 1}
+      );
+      assert.deepEqual(
+        formatMatches(matches),
+        [
+          {pattern: 0, captures: [{name: 'element', text: 'd'}]},
+          {pattern: 0, captures: [{name: 'element', text: 'e'}]},
+          {pattern: 0, captures: [{name: 'element', text: 'f'}]},
+          {pattern: 0, captures: [{name: 'element', text: 'g'}]},
+        ]
+      );
+    });
   });
 
-  it('matches queries in specified ranges',  () => {
-    tree = parser.parse("[a, b,\nc, d,\ne, f,\ng, h]");
-    query = JavaScript.query('(identifier) @element');
-    const matches = query.exec(
-      tree.rootNode,
-      {row: 1, column: 1},
-      {row: 3, column: 1}
-    );
-    assert.deepEqual(
-      formatMatches(matches),
-      [
-        {pattern: 0, captures: [{name: 'element', text: 'd'}]},
-        {pattern: 0, captures: [{name: 'element', text: 'e'}]},
-        {pattern: 0, captures: [{name: 'element', text: 'f'}]},
-        {pattern: 0, captures: [{name: 'element', text: 'g'}]},
-      ]
-    );
+  describe('.captures', () => {
+    it('returns all of the captures for the given query, in order', () => {
+      tree = parser.parse(`
+        a({
+          bc: function de() {
+            const fg = function hi() {}
+          },
+          jk: function lm() {
+            const no = function pq() {}
+          },
+        });
+      `);
+      query = JavaScript.query(`
+        (pair
+          key: * @method.def
+          (function
+            name: (identifier) @method.alias))
+
+        (variable_declarator
+          name: * @function.def
+          value: (function
+            name: (identifier) @function.alias))
+
+        ":" @delimiter
+        "=" @operator
+      `);
+
+      const captures = query.captures(tree.rootNode);
+      assert.deepEqual(
+        formatCaptures(captures),
+        [
+          {name: "method.def", text: "bc"},
+          {name: "delimiter", text: ":"},
+          {name: "method.alias", text: "de"},
+          {name: "function.def", text: "fg"},
+          {name: "operator", text: "="},
+          {name: "function.alias", text: "hi"},
+          {name: "method.def", text: "jk"},
+          {name: "delimiter", text: ":"},
+          {name: "method.alias", text: "lm"},
+          {name: "function.def", text: "no"},
+          {name: "operator", text: "="},
+          {name: "function.alias", text: "pq"},
+        ]
+      );
+    });
   });
 });
 
 function formatMatches(matches) {
   return matches.map(({pattern, captures}) => ({
     pattern,
-    captures: captures.map(({name, node}) => ({
-      name,
-      text: node.text
-    }))
+    captures: formatCaptures(captures)
   }))
 }
+
+function formatCaptures(captures) {
+  return captures.map(({name, node}) => ({ name, text: node.text }))
+}