Add Support for the #set! directive and #is? predicate

tree-sitter/src/main/java/org/treesitter/TSQuery.java

@@ -265,6 +265,10 @@ private List<List<TSQueryPredicate>> parsePredicates() {
                         patternPredicates.add(handleMatch(name, steps, stepIndex, nargs));
                     } else if (TSQueryPredicate.TSQueryPredicateAnyOf.NAMES.contains(name)) {
                         patternPredicates.add(handleAnyOf(name, steps, stepIndex, nargs));
+                    } else if (TSQueryPredicate.TSQueryPredicateSet.NAMES.contains(name)) {
+                        patternPredicates.add(handleSet(name, steps, stepIndex, nargs));
+                    } else if (TSQueryPredicate.TSQueryPredicateIs.NAMES.contains(name)) {
+                        patternPredicates.add(handleIs(name, steps, stepIndex, nargs));
                     } else {
                         patternPredicates.add(new TSQueryPredicate.TSQueryPredicateGeneric(name));
                     }
@@ -335,6 +339,44 @@ private TSQueryPredicate handleAnyOf(String name, TSQueryPredicateStep[] steps,
         return new TSQueryPredicate.TSQueryPredicateAnyOf(name, captureId, values);
     }
 
+    private TSQueryPredicate handleSet(String name, TSQueryPredicateStep[] steps, int start, int nargs) {
+        if (nargs != 3) {
+            throw new TSQueryException(String.format("Predicate #%s expects 2 arguments, got %d", name, nargs - 1));
+        }
+        TSQueryPredicateStep arg1 = steps[start + 1];
+        if (arg1.getType() != TSQueryPredicateStepType.TSQueryPredicateStepTypeString) {
+            throw new TSQueryException(String.format("First argument to #%s must be a string literal (key)", name));
+        }
+        String key = getStringValueForId(arg1.getValueId());
+
+        TSQueryPredicateStep arg2 = steps[start + 2];
+        if (arg2.getType() != TSQueryPredicateStepType.TSQueryPredicateStepTypeString) {
+            throw new TSQueryException(String.format("Second argument to #%s must be a string literal (value)", name));
+        }
+        String value = getStringValueForId(arg2.getValueId());
+
+        return new TSQueryPredicate.TSQueryPredicateSet(name, key, value);
+    }
+
+    private TSQueryPredicate handleIs(String name, TSQueryPredicateStep[] steps, int start, int nargs) {
+        if (nargs != 3) {
+            throw new TSQueryException(String.format("Predicate #%s expects 2 arguments, got %d", name, nargs - 1));
+        }
+        TSQueryPredicateStep arg1 = steps[start + 1];
+        if (arg1.getType() != TSQueryPredicateStepType.TSQueryPredicateStepTypeString) {
+            throw new TSQueryException(String.format("First argument to #%s must be a string literal (key)", name));
+        }
+        String key = getStringValueForId(arg1.getValueId());
+
+        TSQueryPredicateStep arg2 = steps[start + 2];
+        if (arg2.getType() != TSQueryPredicateStepType.TSQueryPredicateStepTypeString) {
+            throw new TSQueryException(String.format("Second argument to #%s must be a string literal (value)", name));
+        }
+        String value = getStringValueForId(arg2.getValueId());
+
+        return new TSQueryPredicate.TSQueryPredicateIs(name, key, value);
+    }
+
     /**
      * Get the quantifier of the query's captures. Each capture is * associated
      * with a numeric id based on the order that it appeared in the query's source.
@@ -370,17 +412,11 @@ public TSQuantifier getCaptureQuantifierForId(int patternId, int captureId) {
      */
     public String getStringValueForId(int id) {
         ensureOpen();
-        int patternCount = getPatternCount();
-        for(int i = 0; i < patternCount; i++){
-            TSQueryPredicateStep[] predicates = getPredicateForPattern(i);
-            for(int j = 0; j < predicates.length; j++){
-                TSQueryPredicateStep predicate = predicates[j];
-                if(id == predicate.getValueId() && predicate.getType() == TSQueryPredicateStepType.TSQueryPredicateStepTypeString){
-                    return ts_query_string_value_for_id(ptr, predicate.getValueId());
-                }
-            }
+        int stringCount = getStringCount();
+        if (id < 0 || id >= stringCount) {
+            throw new TSException("Invalid string id: " + id);
         }
-        throw new TSException("Invalid string id: " + id);
+        return ts_query_string_value_for_id(ptr, id);
     }
 
     /**

tree-sitter/src/main/java/org/treesitter/TSQueryCursor.java

@@ -278,6 +278,7 @@ public boolean nextMatch(TSQueryMatch match){
         ensureOpen();
         assertExecuted();
         while (ts_query_cursor_next_match(ptr, match)) {
+            match.clearMetadata();
             addTsTreeRef(match);
             if (satisfiesPredicates(match)) {
                 return true;
@@ -311,6 +312,7 @@ public boolean nextCapture(TSQueryMatch match){
         ensureOpen();
         assertExecuted();
         while (ts_query_cursor_next_capture(ptr, match)) {
+            match.clearMetadata();
             addTsTreeRef(match);
             if (satisfiesPredicates(match)) {
                 return true;

tree-sitter/src/main/java/org/treesitter/TSQueryMatch.java

@@ -1,10 +1,20 @@
 package org.treesitter;
 
+import java.util.HashMap;
+import java.util.Map;
+
 public class TSQueryMatch {
     private int id;
     private int patternIndex;
     private int captureIndex;
     private TSQueryCapture[] captures;
+    /**
+     * Metadata associated with the match, typically populated by {@code #set!} predicates.
+     * <p>
+     * <b>Note:</b> {@code TSQueryMatch} objects are mutable and reused by {@link TSQueryCursor}.
+     * If you need to persist metadata across iterations, you must create a copy of the map.
+     */
+    private Map<String, String> metadata = null;
 
     public int getId() {
         return id;
@@ -21,4 +31,24 @@ public int getCaptureIndex() {
     public TSQueryCapture[] getCaptures() {
         return captures;
     }
+
+    /**
+     * Get the metadata for this match.
+     *
+     * @return A map of metadata. The map is lazily initialized and might be empty.
+     */
+    public Map<String, String> getMetadata() {
+        if (metadata == null) {
+            metadata = new HashMap<>(8);
+        }
+        return metadata;
+    }
+
+    /**
+     * Clear the metadata for this match.
+     * This allows for memory reclamation and avoids eager initialization during cursor iteration.
+     */
+    public void clearMetadata() {
+        metadata = null;
+    }
 }

tree-sitter/src/main/java/org/treesitter/TSQueryPredicate.java

@@ -241,6 +241,49 @@ public boolean test(TSQueryMatch match, byte[] sourceBytes) {
         }
     }
 
+    /**
+     * Handles {@code #set!}
+     */
+    public static final class TSQueryPredicateSet extends TSQueryPredicate {
+        private final String key;
+        private final String value;
+
+        public static final Set<String> NAMES = Set.of("set!");
+
+        public TSQueryPredicateSet(String name, String key, String value) {
+            super(name);
+            this.key = key;
+            this.value = value;
+        }
+
+        @Override
+        public boolean test(TSQueryMatch match, Function<TSNode, String> textProvider) {
+            match.getMetadata().put(key, value);
+            return true;
+        }
+    }
+
+    /**
+     * Handles {@code #is?}
+     */
+    public static final class TSQueryPredicateIs extends TSQueryPredicate {
+        private final String key;
+        private final String value;
+
+        public static final Set<String> NAMES = Set.of("is?");
+
+        public TSQueryPredicateIs(String name, String key, String value) {
+            super(name);
+            this.key = key;
+            this.value = value;
+        }
+
+        @Override
+        public boolean test(TSQueryMatch match, Function<TSNode, String> textProvider) {
+            return Objects.equals(match.getMetadata().get(key), value);
+        }
+    }
+
     /**
      * Handles unknown predicates or directives.
      */

tree-sitter/src/test/java/org/treesitter/TSQueryMetadataTest.java

@@ -0,0 +1,111 @@
+package org.treesitter;
+
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import java.util.Map;
+
+import static org.junit.jupiter.api.Assertions.*;
+
+class TSQueryMetadataTest {
+    public static final String JSON_SRC = "[1, 2]";
+    private TSTree tree;
+    private TSLanguage json;
+    private TSParser parser;
+    private TSQueryCursor cursor;
+    private TSNode rootNode;
+
+    @BeforeEach
+    void beforeEach() {
+        parser = new TSParser();
+        json = new TreeSitterJson();
+        parser.setLanguage(json);
+        tree = parser.parseString(null, JSON_SRC);
+        rootNode = tree.getRootNode();
+        cursor = new TSQueryCursor();
+    }
+
+    @Test
+    void testSetMetadataDirective() {
+        // ((number) @n (#set! role "foo"))
+        TSQuery query = new TSQuery(json, "((number) @n (#set! role \"foo\"))");
+        cursor.exec(query, rootNode);
+        TSQueryMatch match = new TSQueryMatch();
+
+        int count = 0;
+        while (cursor.nextMatch(match)) {
+            count++;
+            Map<String, String> metadata = match.getMetadata();
+            assertEquals("foo", metadata.get("role"), "Metadata 'role' should be 'foo'");
+        }
+        assertEquals(2, count, "Should have matched two numbers");
+    }
+
+    @Test
+    void testIsPredicateSuccess() {
+        // ((number) @n (#set! role "foo") (#is? role "foo"))
+        TSQuery query = new TSQuery(json, "((number) @n (#set! role \"foo\") (#is? role \"foo\"))");
+        cursor.exec(query, rootNode);
+        TSQueryMatch match = new TSQueryMatch();
+
+        int count = 0;
+        while (cursor.nextMatch(match)) {
+            count++;
+            assertEquals("foo", match.getMetadata().get("role"));
+        }
+        assertEquals(2, count, "Both matches should satisfy #is? role 'foo'");
+    }
+
+    @Test
+    void testIsPredicateFailure() {
+        // ((number) @n (#set! role "foo") (#is? role "bar"))
+        TSQuery query = new TSQuery(json, "((number) @n (#set! role \"foo\") (#is? role \"bar\"))");
+        cursor.exec(query, rootNode);
+        TSQueryMatch match = new TSQueryMatch();
+
+        int count = 0;
+        while (cursor.nextMatch(match)) {
+            count++;
+        }
+        assertEquals(0, count, "No matches should be returned because role is 'foo', not 'bar'");
+    }
+
+    @Test
+    void testMetadataWithNextCapture() {
+        // ((number) @n (#set! role "foo") (#is? role "foo"))
+        TSQuery query = new TSQuery(json, "((number) @n (#set! role \"foo\") (#is? role \"foo\"))");
+        cursor.exec(query, rootNode);
+        TSQueryMatch match = new TSQueryMatch();
+
+        int count = 0;
+        while (cursor.nextCapture(match)) {
+            count++;
+            assertEquals("foo", match.getMetadata().get("role"), "Metadata should be preserved in nextCapture");
+        }
+        assertEquals(2, count, "Both captures should satisfy #is? role 'foo'");
+    }
+
+    @Test
+    void testMetadataIsolationBetweenMatches() {
+        // We have [1, 2]. 
+        // Pattern 0 matches '1' and sets role=first
+        // Pattern 1 matches '2' and sets nothing
+        String queryString = 
+            "((number) @n1 (#eq? @n1 \"1\") (#set! role \"first\")) " +
+            "((number) @n2 (#eq? @n2 \"2\"))";
+
+        TSQuery query = new TSQuery(json, queryString);
+        cursor.exec(query, rootNode, JSON_SRC);
+        TSQueryMatch match = new TSQueryMatch();
+
+        // First match (the number 1)
+        assertTrue(cursor.nextMatch(match));
+        assertEquals("first", match.getMetadata().get("role"), "First match should have metadata");
+
+        // Second match (the number 2)
+        assertTrue(cursor.nextMatch(match));
+        // The metadata should have been cleared in TSQueryCursor.nextMatch() 
+        // before processing the second pattern.
+        assertNull(match.getMetadata().get("role"), "Second match metadata should be cleared/empty");
+    }
+}

tree-sitter/src/test/java/org/treesitter/TSQueryTest.java

@@ -3,6 +3,8 @@
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
+import java.util.List;
+
 import static org.junit.jupiter.api.Assertions.*;
 
 class TSQueryTest {
@@ -89,6 +91,14 @@ void getStringValueForId() {
         assertEquals("foo", query.getStringValueForId(1));
     }
 
+    @Test
+    void getPredicatesForPattern() {
+        List<TSQueryPredicate> predicates = query.getPredicatesForPattern(0);
+        assertNotNull(predicates);
+        assertFalse(predicates.isEmpty());
+        assertEquals("eq?", predicates.get(0).getName());
+    }
+
     @Test
     void disableCapture() {
         query.disableCapture("root");