perf: add parser pooling and query caching for concurrent parsing

src/pkg/parser/parser_pool.go

@@ -0,0 +1,172 @@
+package parser
+
+import (
+	"context"
+	"fmt"
+	"sync"
+
+	sitter "github.com/smacker/go-tree-sitter"
+
+	"github.com/specvital/core/domain"
+)
+
+var (
+	goParserPool sync.Pool
+	jsParserPool sync.Pool
+	tsParserPool sync.Pool
+)
+
+func getParserPool(lang domain.Language) *sync.Pool {
+	switch lang {
+	case domain.LanguageGo:
+		return &goParserPool
+	case domain.LanguageJavaScript:
+		return &jsParserPool
+	default:
+		return &tsParserPool
+	}
+}
+
+func getPooledParser(lang domain.Language) *sitter.Parser {
+	pool := getParserPool(lang)
+
+	if p := pool.Get(); p != nil {
+		if parser, ok := p.(*sitter.Parser); ok {
+			return parser
+		}
+	}
+
+	initLanguages()
+	parser := sitter.NewParser()
+	parser.SetLanguage(getSitterLanguage(lang))
+	return parser
+}
+
+func putPooledParser(lang domain.Language, parser *sitter.Parser) {
+	if parser == nil {
+		return
+	}
+	pool := getParserPool(lang)
+	pool.Put(parser)
+}
+
+// ParseWithPool parses source using a pooled parser.
+// Caller must close the returned tree.
+func ParseWithPool(ctx context.Context, lang domain.Language, source []byte) (*sitter.Tree, error) {
+	parser := getPooledParser(lang)
+	defer putPooledParser(lang, parser)
+
+	tree, err := parser.ParseCtx(ctx, nil, source)
+	if err != nil {
+		return nil, fmt.Errorf("parse failed: %w", err)
+	}
+
+	return tree, nil
+}
+
+type queryCacheKey struct {
+	lang     domain.Language
+	queryStr string
+}
+
+type cachedQuery struct {
+	once  sync.Once
+	query *sitter.Query
+	err   error
+}
+
+var queryCache sync.Map
+
+// getCachedQuery returns a compiled query. The returned query must NOT be closed.
+func getCachedQuery(lang domain.Language, queryStr string) (*sitter.Query, error) {
+	key := queryCacheKey{
+		lang:     lang,
+		queryStr: queryStr,
+	}
+
+	if val, ok := queryCache.Load(key); ok {
+		cached, ok := val.(*cachedQuery)
+		if !ok {
+			return nil, fmt.Errorf("invalid cache entry type")
+		}
+		cached.once.Do(func() {})
+		return cached.query, cached.err
+	}
+
+	cached := &cachedQuery{}
+	actual, loaded := queryCache.LoadOrStore(key, cached)
+
+	if loaded {
+		var ok bool
+		cached, ok = actual.(*cachedQuery)
+		if !ok {
+			return nil, fmt.Errorf("invalid cache entry type")
+		}
+	}
+
+	initLanguages()
+
+	cached.once.Do(func() {
+		sitterLang := getSitterLanguage(lang)
+		cached.query, cached.err = sitter.NewQuery([]byte(queryStr), sitterLang)
+	})
+
+	return cached.query, cached.err
+}
+
+// QueryWithCache executes a query with cached compilation.
+func QueryWithCache(root *sitter.Node, source []byte, lang domain.Language, queryStr string) ([]QueryResult, error) {
+	query, err := getCachedQuery(lang, queryStr)
+	if err != nil {
+		return nil, fmt.Errorf("invalid query: %w", err)
+	}
+
+	cursor := sitter.NewQueryCursor()
+	defer cursor.Close()
+
+	cursor.Exec(query, root)
+
+	var results []QueryResult
+	for {
+		match, ok := cursor.NextMatch()
+		if !ok {
+			break
+		}
+
+		result := QueryResult{
+			Captures: make(map[string]*sitter.Node),
+		}
+
+		for _, capture := range match.Captures {
+			name := query.CaptureNameForId(capture.Index)
+			result.Captures[name] = capture.Node
+			if result.Node == nil {
+				result.Node = capture.Node
+			}
+		}
+
+		results = append(results, result)
+	}
+
+	return results, nil
+}
+
+// ClearQueryCache removes all cached queries. Only for testing.
+func ClearQueryCache() {
+	var toClose []*sitter.Query
+
+	queryCache.Range(func(key, value any) bool {
+		queryCache.Delete(key)
+		if cached, ok := value.(*cachedQuery); ok {
+			cached.once.Do(func() {})
+			if cached.query != nil && cached.err == nil {
+				toClose = append(toClose, cached.query)
+			}
+		}
+		return true
+	})
+
+	for _, q := range toClose {
+		q.Close()
+	}
+}