Updated TSQuery implementation

Author: itsaky

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

@@ -1,16 +1,74 @@
 package com.itsaky.androidide.treesitter;
 
-/**
- * TSQuery
- */
-public class TSQuery {
-    long pointer;
-
-    public TSQuery(TSLanguage language, String source) {
-        this.pointer = Native.newQuery(language.pointer, source);
-    }
-
-    private static class Native  {
-        public static native long newQuery(long pointer, String source);
-    }
+public class TSQuery implements AutoCloseable {
+  final long pointer;
+
+  private int errorOffset;
+  private int errorType;
+
+  /**
+   * Create a new query from a string containing one or more S-expression patterns. The query is
+   * associated with a particular language, and can only be run on syntax nodes parsed with that
+   * language.
+   *
+   * <p>If all of the given patterns are valid, this returns a `TSQuery`. If a pattern is invalid,
+   * this returns `NULL`, and provides two pieces of information about the problem: 1. The byte
+   * offset of the error is written to the `error_offset` parameter. 2. The type of error is written
+   * to the `error_type` parameter.
+   */
+  public TSQuery(TSLanguage language, String source) {
+    this.pointer = Native.newQuery(this, language.pointer, source);
+  }
+
+  public int getErrorOffset() {
+    return this.errorOffset;
+  }
+
+  public TSQueryError getErrorType() {
+    return TSQueryError.valueOf(this.errorType);
+  }
+
+  /**
+   * Get the number of captures in the query.
+   *
+   * @return The count.
+   */
+  public int getCaptureCount() {
+    return Native.captureCount(this.pointer);
+  }
+
+  /**
+   * Get the number of patterns in the query.
+   *
+   * @return The count.
+   */
+  public int getPatternCount() {
+    return Native.patternCount(this.pointer);
+  }
+
+  /**
+   * Get the number of string literals in the query.
+   *
+   * @return The count.
+   */
+  public int getStringCount() {
+    return Native.stringCount(this.pointer);
+  }
+
+  @Override
+  public void close() throws Exception {
+    Native.delete(this.pointer);
+  }
+
+  private static class Native {
+    public static native long newQuery(TSQuery query, long pointer, String source);
+
+    public static native void delete(long query);
+
+    public static native int captureCount(long query);
+
+    public static native int patternCount(long query);
+
+    public static native int stringCount(long query);
+  }
 }

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

@@ -0,0 +1,31 @@
+package com.itsaky.androidide.treesitter;
+
+/**
+ * The type of a {@link TSQuery} error.
+ *
+ * @author Akash Yadav
+ */
+public enum TSQueryError {
+  None(0),
+  Syntax(1),
+  NodeType(2),
+  Field(3),
+  Capture(4),
+  Structure(5),
+  Language(6);
+
+  private final int type;
+
+  TSQueryError(int type) {
+    this.type = type;
+  }
+
+  public static TSQueryError valueOf(int type) {
+    for (TSQueryError value : values()) {
+      if (value.type == type) {
+        return value;
+      }
+    }
+    return null;
+  }
+}

lib/com_itsaky_androidide_treesitter_TSQuery_Native.h

@@ -1,14 +1,71 @@
 #include "com_itsaky_androidide_treesitter_TSQuery_Native.h"
 #include "ts_utils.h"
 
-JNIEXPORT jlong JNICALL Java_com_itsaky_androidide_treesitter_TSQuery_00024Native_newQuery(
-  JNIEnv* env, jclass self, jlong language, jstring source) {
-
+JNIEXPORT jlong JNICALL
+Java_com_itsaky_androidide_treesitter_TSQuery_00024Native_newQuery(
+    JNIEnv* env, jclass self, jobject queryObject, jlong language,
+    jstring source) {
   const char* c_source;
   uint32_t source_length = env->GetStringLength(source);
   c_source = env->GetStringUTFChars(source, NULL);
   uint32_t* error_offset = new uint32_t;
   TSQueryError* error_type = new TSQueryError;
-  TSQuery* query = ts_query_new((TSLanguage*) language, c_source, source_length, error_offset, error_type);
-  return (jlong) query;
+  TSQuery* query = ts_query_new((TSLanguage*)language, c_source, source_length,
+                                error_offset, error_type);
+  fillQuery(env, queryObject, *error_offset, *error_type);
+  return (jlong)query;
+}
+
+JNIEXPORT void JNICALL
+Java_com_itsaky_androidide_treesitter_TSQuery_00024Native_delete(JNIEnv* env,
+                                                                 jclass self,
+                                                                 jlong query) {
+  ts_query_delete((TSQuery*)query);
+}
+
+JNIEXPORT jint JNICALL
+Java_com_itsaky_androidide_treesitter_TSQuery_00024Native_captureCount(
+    JNIEnv* env, jclass self, jlong query) {
+  return ts_query_capture_count((TSQuery*)query);
+}
+
+JNIEXPORT jint JNICALL
+Java_com_itsaky_androidide_treesitter_TSQuery_00024Native_patternCount(
+    JNIEnv* env, jclass self, jlong query) {
+  return ts_query_pattern_count((TSQuery*)query);
+}
+
+JNIEXPORT jint JNICALL
+Java_com_itsaky_androidide_treesitter_TSQuery_00024Native_stringCount(
+    JNIEnv* env, jclass self, jlong query) {
+  return ts_query_string_count((TSQuery*)query);
 }
+
+void fillQuery(JNIEnv* env, jobject query, uint32_t error_offset,
+               TSQueryError error_type) {
+  jclass klass = env->GetObjectClass(query);
+  jfieldID offsetField = env->GetFieldID(klass, "errorOffset", "I");
+  jfieldID errorTypeField = env->GetFieldID(klass, "errorType", "I");
+  env->SetIntField(query, offsetField, error_offset);
+  env->SetIntField(query, errorTypeField, getErrorType(error_type));
+}
+
+jint getErrorType(TSQueryError error) {
+  switch (error) {
+    default:
+    case TSQueryErrorNone:
+      return 0;
+    case TSQueryErrorSyntax:
+      return 1;
+    case TSQueryErrorNodeType:
+      return 2;
+    case TSQueryErrorField:
+      return 3;
+    case TSQueryErrorCapture:
+      return 4;
+    case TSQueryErrorStructure:
+      return 5;
+    case TSQueryErrorLanguage:
+      return 6;
+  }
+}