refine Query / QueryCursor separation

Author: Emmeral

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

@@ -478,10 +478,21 @@ public Map<String, Optional<String>> getPatternAssertions(@Unsigned int index, b
     }
 
 
+    /**
+     * Execute the query on a given node.
+     * @param node The node that the query will run on.
+     * @return A cursor that can be used to iterate over the matches.
+     */
     public QueryCursor execute(Node node){
         return new QueryCursor(this, node, cursorOptions);
     }
 
+    /**
+     * Execute the query on a given node with the given options. The options override the default options set on the query.
+     * @param node The node that the query will run on.
+     * @param options The options that will be used for this query.
+     * @return A cursor that can be used to iterate over the matches.
+     */
     public QueryCursor execute(Node node, QueryCursorOptions options){
         return new QueryCursor(this, node, options);
     }
@@ -504,19 +515,20 @@ public Stream<QueryMatch> findMatches(Node node) {
      *
      * <h4 id="findMatches-example">Predicate Example</h4>
      * <p>
-     * {@snippet lang="java" :
+     * {@snippet lang = "java":
      * Stream<QueryMatch> matches = query.findMatches(tree.getRootNode(), (predicate, match) -> {
      *      if (!predicate.getName().equals("ieq?")) return true;
      *      List<QueryPredicateArg> args = predicate.getArgs();
      *      Node node = match.findNodes(args.getFirst().value()).getFirst();
      *      return args.getLast().value().equalsIgnoreCase(node.getText());
      *  });
-     * }
+     *}
      *
-     * @param node The node that the query will run on.
+     * @param node      The node that the query will run on.
      * @param predicate A function that handles custom predicates.
+     * @implNote The stream is not created lazily such that there is no open {@link QueryCursor} instance left behind. For creating a lazy stream use {@link #execute(Node)} and {@link QueryCursor#stream(BiPredicate)}.
      */
-    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);
     }
 
@@ -531,7 +543,9 @@ public Stream<QueryMatch> findMatches(Node node, @Nullable BiPredicate<QueryPred
      */
     public Stream<QueryMatch> findMatches(Node node, @Nullable BiPredicate<QueryPredicate, QueryMatch> predicate, SegmentAllocator allocator) {
         try(QueryCursor cursor = this.execute(node)){
-            return cursor.stream(allocator, predicate);
+            // make sure to load the entire stream into memory before closing the cursor.
+            // Otherwise, we call for nextMatch after closing the cursor which leads to undefined behavior.
+            return cursor.stream(allocator, predicate).toList().stream();
         }
     }
 

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

@@ -37,14 +37,14 @@ public class QueryCursor implements AutoCloseable{
     private final Arena arena;
 
     private final Query query;
-    private final Node cursorRootNode;
+    private final Tree tree;
 
 
     QueryCursor(Query query, Node cursorRootNode, @Nullable QueryCursorOptions options){
         arena = Arena.ofConfined();
         cursor = ts_query_cursor_new().reinterpret(arena, TreeSitter::ts_query_cursor_delete);
         this.query = query;
-        this.cursorRootNode = cursorRootNode;
+        this.tree = cursorRootNode.getTree();
         if(options != null){
             applyOptions(options);
         }
@@ -79,7 +79,6 @@ private void applyOptions(QueryCursorOptions options){
         if (options.getTimeoutMicros() >= 0) {
             ts_query_cursor_set_timeout_micros(cursor, options.getTimeoutMicros());
         }
-
     }
 
 
@@ -117,30 +116,61 @@ public void removeMatch(@Unsigned int matchId){
         ts_query_cursor_remove_match(cursor, matchId);
     }
 
+    /**
+     * Stream the matches produced by the query. The stream can not be consumed after the cursor is closed. The native
+     * nodes backing the matches are bound to the lifetime of the cursor.
+     * @return a stream of matches
+     */
     public Stream<QueryMatch> stream() {
         return stream(null);
     }
 
+    /**
+     * Like {@link #stream()} but allows for custom predicates to be applied to the matches.
+     * @param predicate a function to handle custom predicates.
+     * @return a stream of matches
+     */
     public Stream<QueryMatch> stream(@Nullable BiPredicate<QueryPredicate, QueryMatch> predicate) {
         return stream(arena , predicate);
     }
 
+    /**
+     * Like {@link #stream(BiPredicate)} but allows for a custom allocator to be used for allocating the native nodes.
+     * @param allocator allocator to use for allocating the native nodes backing the matches
+     * @param predicate a function to handle custom predicates.
+     * @return a stream of matches
+     */
     public Stream<QueryMatch> stream(SegmentAllocator allocator, @Nullable BiPredicate<QueryPredicate, QueryMatch> predicate) {
         return StreamSupport.stream(new MatchesIterator(this, allocator, predicate), false);
     }
 
 
+    /**
+     * Get the next match produced by the query. The native nodes backing the match are bound to the lifetime of the cursor.
+     * @return the next match, if available
+     */
     public Optional<QueryMatch> nextMatch() {
         return nextMatch(null);
     }
 
+    /**
+     * Like {@link #nextMatch()} but allows for custom predicates to be applied to the matches.
+     * @param predicate a function to handle custom predicates.
+     * @return the next match, if available
+     */
     public Optional<QueryMatch> nextMatch(@Nullable BiPredicate<QueryPredicate, QueryMatch> predicate) {
         return nextMatch(arena, predicate);
     }
 
+    /**
+     * Like {@link #nextMatch(BiPredicate)} but allows for a custom allocator to be used for allocating the native nodes.
+     * @param allocator allocator to use for allocating the native nodes backing the matches
+     * @param predicate a function to handle custom predicates.
+     * @return the next match, if available
+     */
     public Optional<QueryMatch> nextMatch(SegmentAllocator allocator, @Nullable BiPredicate<QueryPredicate, QueryMatch> predicate) {
 
-        var hasNoText = cursorRootNode.getTree().getText() == null;
+        var hasNoText = tree.getText() == null;
         MemorySegment match = arena.allocate(TSQueryMatch.layout());
         while (ts_query_cursor_next_match(cursor, match)) {
             var count = Short.toUnsignedInt(TSQueryMatch.capture_count(match));
@@ -150,11 +180,13 @@ public Optional<QueryMatch> nextMatch(SegmentAllocator allocator, @Nullable BiPr
                 var capture = TSQueryCapture.asSlice(matchCaptures, i);
                 var name = query.getCaptureNames().get(TSQueryCapture.index(capture));
                 var node = TSNode.allocate(allocator).copyFrom(TSQueryCapture.node(capture));
-                captureList.add(new QueryCapture(name, new Node(node, this.cursorRootNode.getTree())));
+                captureList.add(new QueryCapture(name, new Node(node, tree)));
             }
             var patternIndex = TSQueryMatch.pattern_index(match);
             var matchId = TSQueryMatch.id(match);
-            var result = new QueryMatch(matchId, patternIndex, captureList);
+            // we copy all the data. So we can directly remove the match from the cursor to free memory.
+            ts_query_cursor_remove_match(cursor, matchId);
+            var result = new QueryMatch(patternIndex, captureList);
             if (hasNoText || matches(predicate, result)) {
                 return Optional.of(result);
             }
@@ -190,6 +222,11 @@ public MatchesIterator(QueryCursor cursor, SegmentAllocator allocator, @Nullable
         }
         @Override
         public boolean tryAdvance(Consumer<? super QueryMatch> action) {
+
+            if(!cursor.arena.scope().isAlive()){
+                throw new IllegalStateException("Cursor is closed. Cannot produce more matches.");
+            }
+
             Optional<QueryMatch> queryMatch = cursor.nextMatch(allocator, predicate);
             queryMatch.ifPresent(action);
             return queryMatch.isPresent();

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

@@ -5,10 +5,10 @@
 
 /** A match that corresponds to a certain pattern in the query. */
 @NullMarked
-public record QueryMatch(int matchId, @Unsigned int patternIndex, List<QueryCapture> captures) {
+public record QueryMatch(@Unsigned int patternIndex, List<QueryCapture> captures) {
     /** Creates an instance of a QueryMatch record class. */
-    public QueryMatch(@Unsigned int matchId, @Unsigned int patternIndex, List<QueryCapture> captures) {
-        this.matchId = matchId;
+    public QueryMatch(@Unsigned int patternIndex, List<QueryCapture> captures) {
+
         this.patternIndex = patternIndex;
         this.captures = List.copyOf(captures);
     }
@@ -24,6 +24,6 @@ public List<Node> findNodes(String capture) {
     @Override
     public String toString() {
         return String.format(
-                "QueryMatch[matchId=%s,patternIndex=%s, captures=%s]", Integer.toUnsignedString(matchId), Integer.toUnsignedString(patternIndex), captures);
+                "QueryMatch[patternIndex=%s, captures=%s]", Integer.toUnsignedString(patternIndex), captures);
     }
 }