feat: Add iterator methods to Query

binding.gyp

@@ -12,6 +12,7 @@
         "src/language.cc",
         "src/logger.cc",
         "src/lookaheaditerator.cc",
+        "src/query_iterator.cc",
         "src/node.cc",
         "src/parser.cc",
         "src/query.cc",

index.js

@@ -1,5 +1,5 @@
 const binding = require('node-gyp-build')(__dirname);
-const {Query, Parser, NodeMethods, Tree, TreeCursor, LookaheadIterator} = binding;
+const {Query, QueryIterator, Parser, NodeMethods, Tree, TreeCursor, LookaheadIterator} = binding;
 
 const util = require('util');
 
@@ -15,10 +15,10 @@ Object.defineProperty(Tree.prototype, 'rootNode', {
       Due to a race condition arising from Jest's worker pool, "this"
       has no knowledge of the native extension if the extension has not
       yet loaded when multiple Jest tests are being run simultaneously.
-      If the extension has correctly loaded, "this" should be an instance 
+      If the extension has correctly loaded, "this" should be an instance
       of the class whose prototype we are acting on (in this case, Tree).
-      Furthermore, the race condition sometimes results in the function in 
-      question being undefined even when the context is correct, so we also 
+      Furthermore, the race condition sometimes results in the function in
+      question being undefined even when the context is correct, so we also
       perform a null function check.
     */
     if (this instanceof Tree && rootNode) {
@@ -420,7 +420,7 @@ Object.defineProperties(TreeCursor.prototype, {
  * Query
  */
 
-const {_matches, _captures} = Query.prototype;
+const {_matches, _matchesIter, _captures, _capturesIter} = Query.prototype;
 
 const PREDICATE_STEP_TYPE = {
   DONE: 0,
@@ -660,6 +660,25 @@ Query.prototype.matches = function(
   return results;
 }
 
+Query.prototype.matchesIter = function(
+  node,
+  {
+    startPosition = ZERO_POINT,
+    endPosition = ZERO_POINT,
+    startIndex = 0,
+    endIndex = 0,
+    matchLimit = 0xFFFFFFFF,
+    maxStartDepth = 0xFFFFFFFF
+  } = {}
+) {
+  marshalNode(node);
+  return _matchesIter.call(this, node.tree,
+    startPosition.row, startPosition.column,
+    endPosition.row, endPosition.column,
+    startIndex, endIndex, matchLimit, maxStartDepth
+  );
+}
+
 Query.prototype.captures = function(
   node,
   {
@@ -710,6 +729,95 @@ Query.prototype.captures = function(
   return results;
 }
 
+Query.prototype.capturesIter = function(
+  node,
+  {
+    startPosition = ZERO_POINT,
+    endPosition = ZERO_POINT,
+    startIndex = 0,
+    endIndex = 0,
+    matchLimit = 0xFFFFFFFF,
+    maxStartDepth = 0xFFFFFFFF
+  } = {}
+) {
+  marshalNode(node);
+  return _capturesIter.call(this, node.tree,
+    startPosition.row, startPosition.column,
+    endPosition.row, endPosition.column,
+    startIndex, endIndex, matchLimit, maxStartDepth
+  );
+}
+
+/**
+ * QueryIterator
+ */
+
+const { _next: _matchesNext } = QueryIterator.prototype;
+
+QueryIterator.prototype.next = function () {
+  while (true) {
+    const { done, value } = _matchesNext.call(this);
+    if (done) {
+      return { done };
+    }
+    const [returnedMatches, returnedNodes] = value;
+    const nodes = unmarshalNodes(returnedNodes, this.tree);
+
+    let i = 0
+    let nodeIndex = 0;
+    if (this.captures) {
+      while (i < returnedMatches.length) {
+        const patternIndex = returnedMatches[i++];
+        const captureIndex = returnedMatches[i++];
+        const captures = [];
+
+        while (i < returnedMatches.length && typeof returnedMatches[i] === 'string') {
+          const captureName = returnedMatches[i++];
+          captures.push({
+            name: captureName,
+            node: nodes[nodeIndex++],
+          })
+        }
+
+        if (this.query.predicates[patternIndex].every(p => p(captures))) {
+          const result = captures[captureIndex];
+          const setProperties = this.query.setProperties[patternIndex];
+          const assertedProperties = this.query.assertedProperties[patternIndex];
+          const refutedProperties = this.query.refutedProperties[patternIndex];
+          if (setProperties) result.setProperties = setProperties;
+          if (assertedProperties) result.assertedProperties = assertedProperties;
+          if (refutedProperties) result.refutedProperties = refutedProperties;
+          return { value: result };
+        }
+      }
+    } else {
+      while (i < returnedMatches.length) {
+        const patternIndex = returnedMatches[i++];
+        const captures = [];
+
+        while (i < returnedMatches.length && typeof returnedMatches[i] === 'string') {
+          const captureName = returnedMatches[i++];
+          captures.push({
+            name: captureName,
+            node: nodes[nodeIndex++],
+          })
+        }
+
+        if (this.query.predicates[patternIndex].every(p => p(captures))) {
+          const result = {pattern: patternIndex, captures};
+          const setProperties = this.query.setProperties[patternIndex];
+          const assertedProperties = this.query.assertedProperties[patternIndex];
+          const refutedProperties = this.query.refutedProperties[patternIndex];
+          if (setProperties) result.setProperties = setProperties;
+          if (assertedProperties) result.assertedProperties = assertedProperties;
+          if (refutedProperties) result.refutedProperties = refutedProperties;
+          return { value: result };
+        }
+      }
+    }
+  }
+}
+
 /*
  * LookaheadIterator
  */

src/addon_data.h

@@ -40,6 +40,9 @@ class AddonData final {
 
   // lookaheaditerator
   Napi::FunctionReference lookahead_iterator_constructor;
+
+  // matches_iterator
+  Napi::FunctionReference query_iterator_constructor;
 };
 
 } // namespace node_tree_sitter

src/binding.cc

@@ -5,6 +5,7 @@
 #include "./node.h"
 #include "./parser.h"
 #include "./query.h"
+#include "./query_iterator.h"
 #include "./tree.h"
 #include "./tree_cursor.h"
 
@@ -24,6 +25,7 @@ Napi::Object InitAll(Napi::Env env, Napi::Object exports) {
   LookaheadIterator::Init(env, exports);
   Parser::Init(env, exports);
   Query::Init(env, exports);
+  QueryIterator::Init(env, exports);
   Tree::Init(env, exports);
   TreeCursor::Init(env, exports);
 

src/query.cc

@@ -39,7 +39,9 @@ void Query::Init(Napi::Env env, Napi::Object exports) {
     InstanceAccessor("matchLimit", &Query::MatchLimit, nullptr, napi_default_method),
 
     InstanceMethod("_matches", &Query::Matches, napi_default_method),
+    InstanceMethod("_matchesIter", &Query::MatchesIter, napi_default_method),
     InstanceMethod("_captures", &Query::Captures, napi_default_method),
+    InstanceMethod("_capturesIter", &Query::CapturesIter, napi_default_method),
     InstanceMethod("_getPredicates", &Query::GetPredicates, napi_default_method),
     InstanceMethod("disableCapture", &Query::DisableCapture, napi_default_method),
     InstanceMethod("disablePattern", &Query::DisablePattern, napi_default_method),
@@ -247,6 +249,18 @@ Napi::Value Query::Matches(const Napi::CallbackInfo &info) {
   return result;
 }
 
+Napi::Value Query::MatchesIter(const Napi::CallbackInfo &info) {
+  Napi::Env env = info.Env();
+  auto *data = info.Env().GetInstanceData<AddonData>();
+  size_t argc = info.Length();
+  std::vector<napi_value> args(argc + 2);
+  args[0] = Boolean::New(env, false);
+  args[1] = info.This();
+  napi_status status = napi_get_cb_info(env, static_cast<napi_callback_info>(info), &argc, args.data() + 2, nullptr, nullptr);
+  NAPI_THROW_IF_FAILED_VOID(env, status);
+  return data->query_iterator_constructor.New(args);
+}
+
 Napi::Value Query::Captures(const Napi::CallbackInfo &info) {
   Napi::Env env = info.Env();
   auto *data = env.GetInstanceData<AddonData>();
@@ -337,6 +351,18 @@ Napi::Value Query::Captures(const Napi::CallbackInfo &info) {
   return result;
 }
 
+Napi::Value Query::CapturesIter(const Napi::CallbackInfo &info) {
+  Napi::Env env = info.Env();
+  auto *data = info.Env().GetInstanceData<AddonData>();
+  size_t argc = info.Length();
+  std::vector<napi_value> args(argc + 2);
+  args[0] = Boolean::New(env, true);
+  args[1] = info.This();
+  napi_status status = napi_get_cb_info(env, static_cast<napi_callback_info>(info), &argc, args.data() + 2, nullptr, nullptr);
+  NAPI_THROW_IF_FAILED_VOID(env, status);
+  return data->query_iterator_constructor.New(args);
+}
+
 Napi::Value Query::DisableCapture(const Napi::CallbackInfo &info) {
   std::string string = info[0].As<String>().Utf8Value();
   const char *capture_name = string.c_str();

src/query.h

@@ -5,7 +5,6 @@
 #include "tree_sitter/api.h"
 
 #include <napi.h>
-#include <node_object_wrap.h>
 
 namespace node_tree_sitter {
 
@@ -17,12 +16,18 @@ class Query final : public Napi::ObjectWrap<Query> {
   explicit Query(const Napi::CallbackInfo &info);
   ~Query() final;
 
+  const TSQuery *Get() const {
+    return query_;
+  }
+
  private:
   TSQuery *query_;
 
   Napi::Value New(const Napi::CallbackInfo &);
   Napi::Value Matches(const Napi::CallbackInfo &);
   Napi::Value Captures(const Napi::CallbackInfo &);
+  Napi::Value MatchesIter(const Napi::CallbackInfo &);
+  Napi::Value CapturesIter(const Napi::CallbackInfo &);
   Napi::Value GetPredicates(const Napi::CallbackInfo &);
   Napi::Value DisableCapture(const Napi::CallbackInfo &);
   Napi::Value DisablePattern(const Napi::CallbackInfo &);