feat: QueryCursor iterators

query.go

@@ -10,6 +10,7 @@ import "C"
 
 import (
 	"bytes"
+	"iter"
 	"fmt"
 	"math"
 	"regexp"
@@ -763,6 +764,30 @@ func (qc *QueryCursor) Matches(query *Query, node *Node, text []byte) QueryMatch
 	return qm
 }
 
+// Iterator yielding all of the matches in the order that they were found.
+//
+// Each match contains the index of the pattern that matched, and a list of
+// captures. Because multiple patterns can match the same set of nodes,
+// one match may contain captures that appear *before* some of the
+// captures from a previous match.
+//
+// Each match yielded by the iterator will overwrite the memory at the same location as prior matches, since the memory is reused. You can think of this as a stateful iterator.
+// If you need to keep the data of a prior match without it being overwritten, you should copy what you need before moving on to the next match
+func (qc *QueryCursor) AllMatches(query *Query, node *Node, text []byte) iter.Seq[*QueryMatch] {
+	qm := qc.Matches(query, node, text)
+	return func(yield func(*QueryMatch) bool) {
+		for {
+			c := qm.Next()
+			if c == nil {
+				break
+			}
+			if !yield(c) {
+				return
+			}
+		}
+	}
+}
+
 // This C function is passed to Tree-sitter as the progress callback.
 //
 //export queryProgressCallback
@@ -818,6 +843,29 @@ func (qc *QueryCursor) Captures(query *Query, node *Node, text []byte) QueryCapt
 	}
 }
 
+// Iterator yielding all of the individual captures in the order that they
+// appear.
+//
+// This is useful if you don't care about which pattern matched, and just
+// want a single, ordered sequence of captures.
+
+// Each capture yielded by the iterator will overwrite the memory at the same location as prior captures, since the memory is reused. You can think of this as a stateful iterator.
+// If you need to keep the data of a prior capture without it being overwritten, you should copy what you need before moving on to the next capture
+func (qc *QueryCursor) AllCaptures(query *Query, node *Node, text []byte) iter.Seq2[*QueryMatch, uint] {
+	qm := qc.Captures(query, node, text)
+	return func(yield func(*QueryMatch, uint) bool) {
+		for {
+			c, i := qm.Next()
+			if c == nil {
+				break
+			}
+			if !yield(c, i) {
+				return
+			}
+		}
+	}
+}
+
 // Set the range of bytes in which the query will be executed.
 //
 // The query cursor will return matches that intersect with the given point range.

query_test.go

@@ -3881,6 +3881,273 @@ func TestQueryCapturesAndMatchesIteratorsAreFused(t *testing.T) {
 	assert.Nil(t, matches.Next())
 }
 
+func TestAllMatchesBasic(t *testing.T) {
+	language := getLanguage("javascript")
+	query, err := NewQuery(
+		language,
+		`(comment) @comment`,
+	)
+	assert.Nil(t, err)
+	defer query.Close()
+
+	source := []byte(`
+		// one
+		// two
+		// three
+	`)
+
+	parser := NewParser()
+	defer parser.Close()
+	parser.SetLanguage(language)
+
+	tree := parser.Parse(source, nil)
+	defer tree.Close()
+
+	cursor := NewQueryCursor()
+	defer cursor.Close()
+
+	var count int
+	for match := range cursor.AllMatches(query, tree.RootNode(), source) {
+		assert.EqualValues(t, 0, match.PatternIndex)
+		assert.Len(t, match.Captures, 1)
+		assert.EqualValues(t, 0, match.Captures[0].Index)
+		count++
+	}
+	assert.Equal(t, 3, count)
+}
+
+func TestAllMatchesEarlyBreak(t *testing.T) {
+	language := getLanguage("javascript")
+	query, err := NewQuery(
+		language,
+		`(comment) @comment`,
+	)
+	assert.Nil(t, err)
+	defer query.Close()
+
+	source := []byte(`
+		// one
+		// two
+		// three
+	`)
+
+	parser := NewParser()
+	defer parser.Close()
+	parser.SetLanguage(language)
+
+	tree := parser.Parse(source, nil)
+	defer tree.Close()
+
+	cursor := NewQueryCursor()
+	defer cursor.Close()
+
+	var count int
+	for range cursor.AllMatches(query, tree.RootNode(), source) {
+		count++
+		if count == 2 {
+			break
+		}
+	}
+	assert.Equal(t, 2, count)
+}
+
+func TestAllMatchesEmpty(t *testing.T) {
+	language := getLanguage("javascript")
+	query, err := NewQuery(
+		language,
+		`(comment) @comment`,
+	)
+	assert.Nil(t, err)
+	defer query.Close()
+
+	source := []byte(`var x = 1;`)
+
+	parser := NewParser()
+	defer parser.Close()
+	parser.SetLanguage(language)
+
+	tree := parser.Parse(source, nil)
+	defer tree.Close()
+
+	cursor := NewQueryCursor()
+	defer cursor.Close()
+
+	var count int
+	for range cursor.AllMatches(query, tree.RootNode(), source) {
+		count++
+	}
+	assert.Equal(t, 0, count)
+}
+
+func TestAllMatchesMultiplePatterns(t *testing.T) {
+	language := getLanguage("javascript")
+	query, err := NewQuery(
+		language,
+		`
+		(comment) @comment
+		(variable_declarator name: (identifier) @varname)
+		`,
+	)
+	assert.Nil(t, err)
+	defer query.Close()
+
+	source := []byte(`
+		// a comment
+		var x = 1;
+	`)
+
+	parser := NewParser()
+	defer parser.Close()
+	parser.SetLanguage(language)
+
+	tree := parser.Parse(source, nil)
+	defer tree.Close()
+
+	cursor := NewQueryCursor()
+	defer cursor.Close()
+
+	patternsSeen := make(map[uint]int)
+	for match := range cursor.AllMatches(query, tree.RootNode(), source) {
+		patternsSeen[match.PatternIndex]++
+	}
+	assert.Equal(t, 1, patternsSeen[0], "expected 1 comment match")
+	assert.Equal(t, 1, patternsSeen[1], "expected 1 variable declarator match")
+}
+
+func TestAllCapturesBasic(t *testing.T) {
+	language := getLanguage("javascript")
+	query, err := NewQuery(
+		language,
+		`(comment) @comment`,
+	)
+	assert.Nil(t, err)
+	defer query.Close()
+
+	source := []byte(`
+		// one
+		// two
+		// three
+	`)
+
+	parser := NewParser()
+	defer parser.Close()
+	parser.SetLanguage(language)
+
+	tree := parser.Parse(source, nil)
+	defer tree.Close()
+
+	cursor := NewQueryCursor()
+	defer cursor.Close()
+
+	var matches []*QueryMatch
+	var indices []uint
+	for match, idx := range cursor.AllCaptures(query, tree.RootNode(), source) {
+		matches = append(matches, match)
+		indices = append(indices, idx)
+	}
+
+	assert.Len(t, matches, 3)
+	assert.EqualValues(t, []uint{0, 0, 0}, indices)
+}
+
+func TestAllCapturesEarlyBreak(t *testing.T) {
+	language := getLanguage("javascript")
+	query, err := NewQuery(
+		language,
+		`(comment) @comment`,
+	)
+	assert.Nil(t, err)
+	defer query.Close()
+
+	source := []byte(`
+		// one
+		// two
+		// three
+	`)
+
+	parser := NewParser()
+	defer parser.Close()
+	parser.SetLanguage(language)
+
+	tree := parser.Parse(source, nil)
+	defer tree.Close()
+
+	cursor := NewQueryCursor()
+	defer cursor.Close()
+
+	var count int
+	for range cursor.AllCaptures(query, tree.RootNode(), source) {
+		count++
+		if count == 1 {
+			break
+		}
+	}
+	assert.Equal(t, 1, count)
+}
+
+func TestAllCapturesEmpty(t *testing.T) {
+	language := getLanguage("javascript")
+	query, err := NewQuery(
+		language,
+		`(comment) @comment`,
+	)
+	assert.Nil(t, err)
+	defer query.Close()
+
+	source := []byte(`var x = 1;`)
+
+	parser := NewParser()
+	defer parser.Close()
+	parser.SetLanguage(language)
+
+	tree := parser.Parse(source, nil)
+	defer tree.Close()
+
+	cursor := NewQueryCursor()
+	defer cursor.Close()
+
+	var count int
+	for range cursor.AllCaptures(query, tree.RootNode(), source) {
+		count++
+	}
+	assert.Equal(t, 0, count)
+}
+
+func TestAllCapturesMultiplePatterns(t *testing.T) {
+	language := getLanguage("javascript")
+	query, err := NewQuery(
+		language,
+		`
+		(comment) @comment
+		(variable_declarator name: (identifier) @varname)
+		`,
+	)
+	assert.Nil(t, err)
+	defer query.Close()
+
+	source := []byte(`
+		// a comment
+		var x = 1;
+	`)
+
+	parser := NewParser()
+	defer parser.Close()
+	parser.SetLanguage(language)
+
+	tree := parser.Parse(source, nil)
+	defer tree.Close()
+
+	cursor := NewQueryCursor()
+	defer cursor.Close()
+
+	var indices []uint
+	for _, idx := range cursor.AllCaptures(query, tree.RootNode(), source) {
+		indices = append(indices, idx)
+	}
+	assert.Len(t, indices, 2)
+	assert.Contains(t, indices, uint(0))
+}
+
 func TestQueryStartEndByteForPattern(t *testing.T) {
 	language := getLanguage("javascript")