Possibility to use custom allocators for nodes retrieved from cursors and queries

src/main/java/io/github/treesitter/jtreesitter/Query.java

@@ -5,6 +5,7 @@
 import io.github.treesitter.jtreesitter.internal.*;
 import java.lang.foreign.Arena;
 import java.lang.foreign.MemorySegment;
+import java.lang.foreign.SegmentAllocator;
 import java.util.*;
 import java.util.function.BiPredicate;
 import java.util.function.Consumer;
@@ -478,16 +479,19 @@ public Map<String, Optional<String>> getPatternAssertions(@Unsigned int index, b
     }
 
     /**
-     * Iterate over all the matches in the order that they were found.
+     * Iterate over all the matches in the order that they were found. The lifetime of the native memory of the returned
+     * matches is bound to the lifetime of this query object.
      *
      * @param node The node that the query will run on.
      */
     public Stream<QueryMatch> findMatches(Node node) {
-        return findMatches(node, null);
+        return findMatches(node, null, arena);
     }
 
+
     /**
-     * Iterate over all the matches in the order that they were found.
+     * Iterate over all the matches in the order that they were found. The lifetime of the native memory of the returned
+     * matches is bound to the lifetime of this query object.
      *
      * <h4 id="findMatches-example">Predicate Example</h4>
      * <p>
@@ -503,11 +507,24 @@ public Stream<QueryMatch> findMatches(Node node) {
      * @param node The node that the query will run on.
      * @param predicate A function that handles custom predicates.
      */
-    public Stream<QueryMatch> findMatches(Node node, @Nullable BiPredicate<QueryPredicate, QueryMatch> predicate) {
+    public Stream<QueryMatch> findMatches(Node node, @Nullable BiPredicate<QueryPredicate, QueryMatch> predicate){
+        return findMatches(node, predicate, arena);
+    }
+
+
+    /**
+     * Like {@link #findMatches(Node, BiPredicate)} but the native memory of the returned matches is created using the
+     * given allocator.
+     *
+     * @param node The node that the query will run on.
+     * @param predicate A function that handles custom predicates.
+     * @param allocator The allocator that is used to allocate the native memory of the returned matches.
+     */
+    public Stream<QueryMatch> findMatches(Node node, @Nullable BiPredicate<QueryPredicate, QueryMatch> predicate, SegmentAllocator allocator) {
         try (var alloc = Arena.ofConfined()) {
             ts_query_cursor_exec(cursor, query, node.copy(alloc));
         }
-        return StreamSupport.stream(new MatchesIterator(node.getTree(), predicate), false);
+        return StreamSupport.stream(new MatchesIterator(node.getTree(), predicate, allocator), false);
     }
 
     @Override
@@ -537,11 +554,13 @@ private void checkIndex(@Unsigned int index) throws IndexOutOfBoundsException {
     private final class MatchesIterator extends Spliterators.AbstractSpliterator<QueryMatch> {
         private final @Nullable BiPredicate<QueryPredicate, QueryMatch> predicate;
         private final Tree tree;
+        private final SegmentAllocator allocator;
 
-        public MatchesIterator(Tree tree, @Nullable BiPredicate<QueryPredicate, QueryMatch> predicate) {
+        public MatchesIterator(Tree tree, @Nullable BiPredicate<QueryPredicate, QueryMatch> predicate, SegmentAllocator allocator) {
             super(Long.MAX_VALUE, Spliterator.IMMUTABLE | Spliterator.NONNULL);
             this.predicate = predicate;
             this.tree = tree;
+            this.allocator = allocator;
         }
 
         @Override
@@ -555,7 +574,7 @@ public boolean tryAdvance(Consumer<? super QueryMatch> action) {
                 for (int i = 0; i < count; ++i) {
                     var capture = TSQueryCapture.asSlice(matchCaptures, i);
                     var name = captureNames.get(TSQueryCapture.index(capture));
-                    var node = TSNode.allocate(arena).copyFrom(TSQueryCapture.node(capture));
+                    var node = TSNode.allocate(allocator).copyFrom(TSQueryCapture.node(capture));
                     captureList.add(new QueryCapture(name, new Node(node, tree)));
                 }
                 var patternIndex = TSQueryMatch.pattern_index(match);

src/main/java/io/github/treesitter/jtreesitter/TreeCursor.java

@@ -4,6 +4,7 @@
 
 import java.lang.foreign.Arena;
 import java.lang.foreign.MemorySegment;
+import java.lang.foreign.SegmentAllocator;
 import java.util.OptionalInt;
 import org.jspecify.annotations.NullMarked;
 import org.jspecify.annotations.Nullable;
@@ -14,7 +15,6 @@ public final class TreeCursor implements AutoCloseable, Cloneable {
     private final MemorySegment self;
     private final Arena arena;
     private final Tree tree;
-    private @Nullable Node node;
 
     TreeCursor(Node node, Tree tree) {
         arena = Arena.ofShared();
@@ -33,16 +33,25 @@ private TreeCursor(TreeCursor cursor) {
         arena = Arena.ofShared();
         self = ts_tree_cursor_copy(arena, cursor.self);
         tree = cursor.tree.clone();
-        node = cursor.node;
     }
 
-    /** Get the current node of the cursor. */
+    /**
+     * Get the current node of the cursor. Its native memory will be managed by the cursor. It will become invalid
+     * once the cursor is closed by calling {@link #close()}.
+     * @return the current node
+     */
     public Node getCurrentNode() {
-        if (this.node == null) {
-            var node = ts_tree_cursor_current_node(arena, self);
-            this.node = new Node(node, tree);
-        }
-        return this.node;
+        return getCurrentNode(arena);
+    }
+
+    /**
+     * Get the current node of the cursor. Its native memory will be managed by the given allocator.
+     * @param allocator the allocator to use for managing the native memory of the node
+     * @return the current node
+     */
+    public Node getCurrentNode(SegmentAllocator allocator) {
+        var node = ts_tree_cursor_current_node(allocator, self);
+        return new Node(node, tree);
     }
 
     /**
@@ -88,9 +97,7 @@ public Node getCurrentNode() {
      *         {@code false} if there were no children.
      */
     public boolean gotoFirstChild() {
-        var result = ts_tree_cursor_goto_first_child(self);
-        if (result) node = null;
-        return result;
+        return ts_tree_cursor_goto_first_child(self);
     }
 
     /**
@@ -100,9 +107,7 @@ public boolean gotoFirstChild() {
      *         {@code false} if there were no children.
      */
     public boolean gotoLastChild() {
-        var result = ts_tree_cursor_goto_last_child(self);
-        if (result) node = null;
-        return result;
+        return ts_tree_cursor_goto_last_child(self);
     }
 
     /**
@@ -112,9 +117,7 @@ public boolean gotoLastChild() {
      *         {@code false} if there was no parent node.
      */
     public boolean gotoParent() {
-        var result = ts_tree_cursor_goto_parent(self);
-        if (result) node = null;
-        return result;
+        return ts_tree_cursor_goto_parent(self);
     }
 
     /**
@@ -124,9 +127,7 @@ public boolean gotoParent() {
      *         {@code false} if there was no next sibling node.
      */
     public boolean gotoNextSibling() {
-        var result = ts_tree_cursor_goto_next_sibling(self);
-        if (result) node = null;
-        return result;
+        return ts_tree_cursor_goto_next_sibling(self);
     }
 
     /**
@@ -136,9 +137,7 @@ public boolean gotoNextSibling() {
      *         {@code false} if there was no previous sibling node.
      */
     public boolean gotoPreviousSibling() {
-        var result = ts_tree_cursor_goto_previous_sibling(self);
-        if (result) node = null;
-        return result;
+        return ts_tree_cursor_goto_previous_sibling(self);
     }
 
     /**
@@ -149,7 +148,7 @@ public boolean gotoPreviousSibling() {
      */
     public void gotoDescendant(@Unsigned int index) {
         ts_tree_cursor_goto_descendant(self, index);
-        node = null;
+
     }
 
     /**
@@ -161,7 +160,6 @@ public void gotoDescendant(@Unsigned int index) {
     public @Unsigned OptionalInt gotoFirstChildForByte(@Unsigned int offset) {
         var index = ts_tree_cursor_goto_first_child_for_byte(self, offset);
         if (index == -1L) return OptionalInt.empty();
-        node = null;
         return OptionalInt.of((int) index);
     }
 
@@ -176,7 +174,6 @@ public void gotoDescendant(@Unsigned int index) {
             var goal = point.into(arena);
             var index = ts_tree_cursor_goto_first_child_for_point(self, goal);
             if (index == -1L) return OptionalInt.empty();
-            node = null;
             return OptionalInt.of((int) index);
         }
     }
@@ -185,15 +182,12 @@ public void gotoDescendant(@Unsigned int index) {
     public void reset(Node node) {
         try (var arena = Arena.ofConfined()) {
             ts_tree_cursor_reset(self, node.copy(arena));
-        } finally {
-            this.node = null;
         }
     }
 
     /** Reset the cursor to start at the same position as another cursor. */
     public void reset(TreeCursor cursor) {
         ts_tree_cursor_reset_to(self, cursor.self);
-        this.node = null;
     }
 
     /** Create a shallow copy of the tree cursor. */

src/test/java/io/github/treesitter/jtreesitter/TreeCursorTest.java

@@ -5,6 +5,8 @@
 import io.github.treesitter.jtreesitter.languages.TreeSitterJava;
 import org.junit.jupiter.api.*;
 
+import java.lang.foreign.Arena;
+
 class TreeCursorTest {
     private static Tree tree;
     private TreeCursor cursor;
@@ -29,14 +31,29 @@ void setUp() {
 
     @AfterEach
     void tearDown() {
-        cursor.close();
+        if(cursor != null){
+            cursor.close();
+        }
     }
 
     @Test
     void getCurrentNode() {
         var node = cursor.getCurrentNode();
         assertEquals(tree.getRootNode(), node);
-        assertSame(node, cursor.getCurrentNode());
+    }
+
+    @Test
+    void getCurrentNodeWithCustomAllocator() {
+
+        try(var arena = Arena.ofConfined()){
+            var node = cursor.getCurrentNode(arena);
+            assertEquals(tree.getRootNode(), node);
+            cursor.close();
+            cursor = null; // avoid double close
+            // can still access node after cursor was closed
+            assertEquals(tree.getRootNode(), node);
+        }
+
     }
 
     @Test