Updated native TSQueryCursor implementation

Author: itsaky

android-tree-sitter/src/main/java/com/itsaky/androidide/treesitter/TSQueryCursor.java

@@ -59,12 +59,7 @@ public void setPointRange(TSPoint start, TSPoint end) {
   }
 
   public TSQueryMatch nextMatch() {
-    TSQueryMatch match = new TSQueryMatch();
-    if (Native.nextMatch(this.pointer, match)) {
-      return match;
-    }
-
-    return null;
+    return Native.nextMatch(this.pointer);
   }
 
   public void removeMatch(int id) {
@@ -93,7 +88,7 @@ private static class Native {
 
     public static native void setPointRange(long cursor, TSPoint start, TSPoint end);
 
-    public static native boolean nextMatch(long cursor, TSQueryMatch match);
+    public static native TSQueryMatch nextMatch(long cursor);
 
     public static native void removeMatch(long cursor, int id);
   }

android-tree-sitter/src/main/jniLibs/arm64-v8a/libts.so

@@ -13,14 +13,20 @@ public class QueryTest extends TestBase {
   public void implementationTest() throws Exception {
     try (final var parser = new TSParser()) {
       parser.setLanguage(TSLanguages.java());
-      try (final var tree = parser.parseString("public class MyClass { int x = 0; }");
-          final var query = new TSQuery(tree.getLanguage(), "(class_body)");
-          final var cursor = new TSQueryCursor()) {
+      try (final var tree =
+          parser.parseString("public class MyClass { int x = 0; public void myFunc(){} }")) {
+        var query =
+            new TSQuery(
+                TSLanguages.java(), "(class_declaration name: (identifier) @MyClass)");
+        var cursor = new TSQueryCursor();
         cursor.exec(query, tree.getRootNode());
-        TSQueryMatch match;
-        while((match = cursor.nextMatch()) != null) {
-          assertThat(match.getCaptures()).isNotEmpty();
-        }
+        final var match = cursor.nextMatch();
+        assertThat(match).isNotNull();
+        assertThat(match.getCaptures()).hasLength(1);
+        query.close();
+        cursor.close();
+      } catch (Throwable err) {
+        throw new RuntimeException(err);
       }
     }
   }

android-tree-sitter/src/main/jniLibs/armeabi-v7a/libts.so

@@ -3,9 +3,6 @@
 
 #include <iostream>
 
-void fillMatch(JNIEnv *, TSQueryMatch *, jobject);
-jobject createCaptureObj(JNIEnv *, const TSQueryCapture *, int, jclass);
-
 JNIEXPORT jlong JNICALL
 Java_com_itsaky_androidide_treesitter_TSQueryCursor_00024Native_newCursor(
     JNIEnv *env, jclass self) {
@@ -58,59 +55,19 @@ Java_com_itsaky_androidide_treesitter_TSQueryCursor_00024Native_setPointRange(
                                   _unmarshalPoint(env, end));
 }
 
-JNIEXPORT jboolean JNICALL
+JNIEXPORT jobject JNICALL
 Java_com_itsaky_androidide_treesitter_TSQueryCursor_00024Native_nextMatch(
-    JNIEnv *env, jclass self, jlong cursor, jobject match) {
-  TSQueryMatch *m = new TSQueryMatch;
-  bool b = ts_query_cursor_next_match((TSQueryCursor *)cursor, m);
+    JNIEnv *env, jclass self, jlong cursor) {
+  TSQueryMatch m;
+  bool b = ts_query_cursor_next_match((TSQueryCursor *)cursor, &m);
   if (!b) {
-    return false;
+    return NULL;
   }
-  fillMatch(env, m, match);
-  return b;
+  return _marshalMatch(env, m);
 }
 
 JNIEXPORT void JNICALL
 Java_com_itsaky_androidide_treesitter_TSQueryCursor_00024Native_removeMatch(
     JNIEnv *env, jclass self, jlong cursor, jint id) {
   ts_query_cursor_remove_match((TSQueryCursor *)cursor, id);
-}
-
-void fillMatch(JNIEnv *env, TSQueryMatch *m, jobject match) {
-  jclass klass = env->GetObjectClass(match);
-  jfieldID id = env->GetFieldID(klass, "id", "I");
-  jfieldID pI = env->GetFieldID(klass, "patternIndex", "I");
-  jfieldID cs = env->GetFieldID(
-      klass, "captures", "[Lcom/itsaky/androidide/treesitter/TSQueryCapture;");
-  env->SetIntField(match, id, m->id);
-  env->SetIntField(match, pI, m->pattern_index);
-
-  uint16_t count = m->capture_count;
-  jclass qClass;
-  _loadClass(qClass, "com/itsaky/androidide/treesitter/TSQueryCapture");
-  jobjectArray arr = env->NewObjectArray(count, qClass, NULL);
-
-  std::cout << "Match ID:" << m->id << std::endl;
-  std::cout << "Match pattern index:" << m->pattern_index << std::endl;
-  std::cout << "Match capture count:" << m->capture_count << std::endl;
-  std::cout << "Number of captures in match:" << count << std::endl;
-  for (uint16_t i = 0; i < count; i++) {
-    env->SetObjectArrayElement(arr, i,
-                               createCaptureObj(env, m->captures, i, qClass));
-  }
-
-  env->SetObjectField(match, cs, arr);
-}
-
-jobject createCaptureObj(JNIEnv *env, TSQueryCapture *captures, int index,
-                         jclass klass) {
-  jobject obj = env->AllocObject(klass);
-  jfieldID nF = env->GetFieldID(klass, "node",
-                                "Lcom/itsaky/androidide/treesitter/TSNode;");
-  jfieldID iF = env->GetFieldID(klass, "index", "I");
-
-  TSQueryCapture *capture = (captures + index);
-  env->SetIntField(obj, iF, capture->index);
-  env->SetObjectField(obj, nF, _marshalNode(env, capture->node));
-  return obj;
 }

android-tree-sitter/src/main/jniLibs/x86/libts.so

@@ -34,4 +34,7 @@ jobject _marshalPoint(JNIEnv* env, TSPoint point);
 
 TSPoint _unmarshalPoint(JNIEnv* env, jobject javaObject);
 
-TSInputEdit _unmarshalInputEdit(JNIEnv* env, jobject inputEdit);
+TSInputEdit _unmarshalInputEdit(JNIEnv* env, jobject inputEdit);
+
+jobject _marshalMatch(JNIEnv *env, TSQueryMatch match);
+jobject _marshalCapture(JNIEnv *env, TSQueryCapture capture);

android-tree-sitter/src/main/jniLibs/x86_64/libts.so

@@ -33,9 +33,19 @@ static jfieldID _inputEditStartPointField;
 static jfieldID _inputEditOldEndPointField;
 static jfieldID _inputEditNewEndPointField;
 
+// TSQueryMatch
+static jclass _matchClass;
+static jfieldID _matchClassIdField;
+static jfieldID _matchClassPatternIndexField;
+static jfieldID _matchClassCapturesField;
+
+// TSQueryCapture
+static jclass _captureClass;
+static jfieldID _captureClassIndexField;
+static jfieldID _captureClassNodeField;
 
 void onLoad(JNIEnv* env) {
-    // Node
+  // Node
   _loadClass(_nodeClass, "com/itsaky/androidide/treesitter/TSNode");
   _loadField(_nodeContext0Field, _nodeClass, "context0", "I");
   _loadField(_nodeContext1Field, _nodeClass, "context1", "I");
@@ -45,7 +55,8 @@ void onLoad(JNIEnv* env) {
   _loadField(_nodeTreeField, _nodeClass, "tree", "J");
 
   // TreeCursorNode
-  _loadClass(_treeCursorNodeClass, "com/itsaky/androidide/treesitter/TSTreeCursorNode");
+  _loadClass(_treeCursorNodeClass,
+             "com/itsaky/androidide/treesitter/TSTreeCursorNode");
   _loadField(_treeCursorNodeTypeField, _treeCursorNodeClass, "type",
              "Ljava/lang/String;");
   _loadField(_treeCursorNodeNameField, _treeCursorNodeClass, "name",
@@ -64,14 +75,34 @@ void onLoad(JNIEnv* env) {
   _loadField(_inputEditStartByteField, _inputEditClass, "startByte", "I");
   _loadField(_inputEditOldEndByteField, _inputEditClass, "oldEndByte", "I");
   _loadField(_inputEditNewEndByteField, _inputEditClass, "newEndByte", "I");
-  _loadField(_inputEditStartPointField, _inputEditClass, "start_point", "Lcom/itsaky/androidide/treesitter/TSPoint;");
-  _loadField(_inputEditOldEndPointField, _inputEditClass, "old_end_point", "Lcom/itsaky/androidide/treesitter/TSPoint;");
-  _loadField(_inputEditNewEndPointField, _inputEditClass, "new_end_point", "Lcom/itsaky/androidide/treesitter/TSPoint;");
+  _loadField(_inputEditStartPointField, _inputEditClass, "start_point",
+             "Lcom/itsaky/androidide/treesitter/TSPoint;");
+  _loadField(_inputEditOldEndPointField, _inputEditClass, "old_end_point",
+             "Lcom/itsaky/androidide/treesitter/TSPoint;");
+  _loadField(_inputEditNewEndPointField, _inputEditClass, "new_end_point",
+             "Lcom/itsaky/androidide/treesitter/TSPoint;");
+
+  // TSQueryMatch
+  _loadClass(_matchClass, "com/itsaky/androidide/treesitter/TSQueryMatch");
+  _loadField(_matchClassIdField, _matchClass, "id", "I");
+  _loadField(_matchClassPatternIndexField, _matchClass, "patternIndex", "I");
+  _loadField(_matchClassCapturesField, _matchClass, "captures",
+             "[Lcom/itsaky/androidide/treesitter/TSQueryCapture;");
+
+  // TSQueryCapture
+  _loadClass(_captureClass, "com/itsaky/androidide/treesitter/TSQueryCapture");
+  _loadField(_captureClassIndexField, _captureClass, "index", "I");
+  _loadField(_captureClassNodeField, _captureClass, "node",
+             "Lcom/itsaky/androidide/treesitter/TSNode;");
 }
 
 void onUnload(JNIEnv* env) {
   env->DeleteGlobalRef(_nodeClass);
   env->DeleteGlobalRef(_treeCursorNodeClass);
+  env->DeleteGlobalRef(_pointClass);
+  env->DeleteGlobalRef(_inputEditClass);
+  env->DeleteGlobalRef(_matchClass);
+  env->DeleteGlobalRef(_captureClass);
 }
 
 // Node
@@ -87,16 +118,15 @@ jobject _marshalNode(JNIEnv* env, TSNode node) {
 }
 
 TSNode _unmarshalNode(JNIEnv* env, jobject javaObject) {
-  return (TSNode) {
-    {
-      (uint32_t)env->GetIntField(javaObject, _nodeContext0Field),
-      (uint32_t)env->GetIntField(javaObject, _nodeContext1Field),
-      (uint32_t)env->GetIntField(javaObject, _nodeContext2Field),
-      (uint32_t)env->GetIntField(javaObject, _nodeContext3Field),
-    },
-    (const void*)env->GetLongField(javaObject, _nodeIdField),
-    (const TSTree*)env->GetLongField(javaObject, _nodeTreeField)
-  };
+  return (TSNode){
+      {
+          (uint32_t)env->GetIntField(javaObject, _nodeContext0Field),
+          (uint32_t)env->GetIntField(javaObject, _nodeContext1Field),
+          (uint32_t)env->GetIntField(javaObject, _nodeContext2Field),
+          (uint32_t)env->GetIntField(javaObject, _nodeContext3Field),
+      },
+      (const void*)env->GetLongField(javaObject, _nodeIdField),
+      (const TSTree*)env->GetLongField(javaObject, _nodeTreeField)};
 }
 
 // TreeCursorNode
@@ -122,22 +152,45 @@ jobject _marshalPoint(JNIEnv* env, TSPoint point) {
 }
 
 TSPoint _unmarshalPoint(JNIEnv* env, jobject javaObject) {
-  return (TSPoint) {
-    (uint32_t)env->GetIntField(javaObject, _pointRowField),
-    (uint32_t)env->GetIntField(javaObject, _pointColumnField),
+  return (TSPoint){
+      (uint32_t)env->GetIntField(javaObject, _pointRowField),
+      (uint32_t)env->GetIntField(javaObject, _pointColumnField),
   };
 }
 
 // TSInputEdit
 TSInputEdit _unmarshalInputEdit(JNIEnv* env, jobject inputEdit) {
-  return (TSInputEdit) {
-    (uint32_t)env->GetIntField(inputEdit, _inputEditStartByteField),
-    (uint32_t)env->GetIntField(inputEdit, _inputEditOldEndByteField),
-    (uint32_t)env->GetIntField(inputEdit, _inputEditNewEndByteField),   
-    _unmarshalPoint(env, env->GetObjectField(inputEdit, _inputEditStartPointField)),
-    _unmarshalPoint(env, env->GetObjectField(inputEdit, _inputEditOldEndPointField)),
-    _unmarshalPoint(env, env->GetObjectField(inputEdit, _inputEditNewEndPointField)),
+  return (TSInputEdit){
+      (uint32_t)env->GetIntField(inputEdit, _inputEditStartByteField),
+      (uint32_t)env->GetIntField(inputEdit, _inputEditOldEndByteField),
+      (uint32_t)env->GetIntField(inputEdit, _inputEditNewEndByteField),
+      _unmarshalPoint(
+          env, env->GetObjectField(inputEdit, _inputEditStartPointField)),
+      _unmarshalPoint(
+          env, env->GetObjectField(inputEdit, _inputEditOldEndPointField)),
+      _unmarshalPoint(
+          env, env->GetObjectField(inputEdit, _inputEditNewEndPointField)),
   };
 }
 
+jobject _marshalMatch(JNIEnv* env, TSQueryMatch match) {
+  jobject obj = env->AllocObject(_matchClass);
+  env->SetIntField(obj, _matchClassIdField, match.id);
+  env->SetIntField(obj, _matchClassPatternIndexField, match.pattern_index);
+  jobjectArray captures = env->NewObjectArray(match.capture_count, _captureClass, NULL);
+  for (int i = 0; i < match.capture_count; i++) {
+    const TSQueryCapture* c = match.captures + i;
+    env->SetObjectArrayElement(captures, i, _marshalCapture(env, *c));
+  }
+  env->SetObjectField(obj, _matchClassCapturesField, captures);
+  return obj;
+}
+
+jobject _marshalCapture(JNIEnv* env, TSQueryCapture capture) {
+  jobject obj = env->AllocObject(_captureClass);
+  env->SetIntField(obj, _captureClassIndexField, capture.index);
+  env->SetObjectField(obj, _captureClassNodeField, _marshalNode(env, capture.node));
+  return obj;
+}
+
 #endif