ollie/tools/codeintel/internal/codeintel/query.go

package codeintel

import (
	"fmt"
	"sort"
	"strings"

	ts "github.com/tree-sitter/go-tree-sitter"
)

type Capture struct {
	Name, Text string
	Line       int
	Start, End uint
}

type RewriteMatch struct {
	Start, End uint
	Captures   map[string]string
}

func Query(path, language, source string) ([]Capture, error) {
	d, err := Parse(path, language)
	if err != nil {
		return nil, err
	}
	defer d.Close()
	q, qerr := ts.NewQuery(d.Root.Language(), source)
	if qerr != nil {
		return nil, fmt.Errorf("query: %s", qerr.Message)
	}
	defer q.Close()
	cursor := ts.NewQueryCursor()
	defer cursor.Close()
	var out []Capture
	matches := cursor.Captures(q, d.Root, d.Source)
	for {
		match, _ := matches.Next()
		if match == nil {
			break
		}
		for _, c := range match.Captures {
			name := q.CaptureNames()[c.Index]
			start, end := c.Node.ByteRange()
			out = append(out, Capture{Name: name, Text: Text(&c.Node, d.Source), Line: Line(&c.Node), Start: start, End: end})
		}
	}
	sort.Slice(out, func(i, j int) bool { return out[i].Start < out[j].Start })
	return out, nil
}

func rewriteMatches(d *Document, source string) ([]RewriteMatch, error) {
	q, qerr := ts.NewQuery(d.Root.Language(), source)
	if qerr != nil {
		return nil, fmt.Errorf("query: %s", qerr.Message)
	}
	defer q.Close()
	cursor := ts.NewQueryCursor()
	defer cursor.Close()
	var out []RewriteMatch
	matches := cursor.Matches(q, d.Root, d.Source)
	names := q.CaptureNames()
	for {
		match := matches.Next()
		if match == nil {
			break
		}
		if len(match.Captures) == 0 {
			continue
		}
		root := match.Captures[0].Node
		captures := make(map[string]string, len(match.Captures))
		for _, capture := range match.Captures {
			name := names[capture.Index]
			if name == "match" || name == "replace" {
				root = capture.Node
			}
			if _, exists := captures[name]; !exists {
				captures[name] = Text(&capture.Node, d.Source)
			}
		}
		start, end := root.ByteRange()
		out = append(out, RewriteMatch{Start: start, End: end, Captures: captures})
	}
	return out, nil
}

// Interpolate replaces $name and ${name} with named Tree-sitter captures.
func Interpolate(template string, captures map[string]string) string {
	var out strings.Builder
	for i := 0; i < len(template); {
		if template[i] != '$' {
			out.WriteByte(template[i])
			i++
			continue
		}
		if i+1 < len(template) && template[i+1] == '$' {
			out.WriteByte('$')
			i += 2
			continue
		}
		start := i + 1
		end := start
		nameStart := start
		braced := end < len(template) && template[end] == '{'
		if braced {
			nameStart++
			end++
		}
		for end < len(template) && ((template[end] >= 'a' && template[end] <= 'z') || (template[end] >= 'A' && template[end] <= 'Z') || (template[end] >= '0' && template[end] <= '9') || template[end] == '_') {
			end++
		}
		nameEnd := end
		if braced && end < len(template) && template[end] == '}' {
			nameEnd = end + 1
		}
		if end == nameStart || (braced && nameEnd == end) {
			out.WriteByte('$')
			i++
			continue
		}
		name := template[nameStart:end]
		if value, ok := captures[name]; ok {
			out.WriteString(value)
		} else {
			out.WriteString(template[i:nameEnd])
		}
		i = nameEnd
	}
	return out.String()
}

func ReplaceMatches(path, language, source, replacement string) ([]byte, error) {
	d, err := Parse(path, language)
	if err != nil {
		return nil, err
	}
	defer d.Close()
	matches, err := rewriteMatches(d, source)
	if err != nil {
		return nil, err
	}
	result := append([]byte(nil), d.Source...)
	for i := len(matches) - 1; i >= 0; i-- {
		m := matches[i]
		value := Interpolate(replacement, m.Captures)
		result = append(result[:m.Start], append([]byte(value), result[m.End:]...)...)
	}
	return result, nil
}

func Dependencies(path, language string) ([]string, error) {
	d, err := Parse(path, language)
	if err != nil {
		return nil, err
	}
	defer d.Close()
	var out []string
	Walk(d.Root, func(n *ts.Node) {
		k := n.Kind()
		if strings.Contains(k, "import") || strings.Contains(k, "include") || k == "use_declaration" || k == "namespace_use_declaration" {
			v := strings.TrimSpace(Text(n, d.Source))
			if v != "" {
				out = append(out, v)
			}
		}
	})
	return out, nil
}