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

package codeintel

import (
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strings"

	yaml "github.com/tree-sitter-grammars/tree-sitter-yaml/bindings/go"
	ts "github.com/tree-sitter/go-tree-sitter"
	c "github.com/tree-sitter/tree-sitter-c/bindings/go"
	cpp "github.com/tree-sitter/tree-sitter-cpp/bindings/go"
	golang "github.com/tree-sitter/tree-sitter-go/bindings/go"
	javascript "github.com/tree-sitter/tree-sitter-javascript/bindings/go"
	jsonlang "github.com/tree-sitter/tree-sitter-json/bindings/go"
	markdown "github.com/tree-sitter/tree-sitter-markdown/bindings/go"
	php "github.com/tree-sitter/tree-sitter-php/bindings/go"
	python "github.com/tree-sitter/tree-sitter-python/bindings/go"
	rust "github.com/tree-sitter/tree-sitter-rust/bindings/go"
	typescript "github.com/tree-sitter/tree-sitter-typescript/bindings/go"
)

type Language struct {
	Name       string
	Extensions []string
	Grammar    *ts.Language
}

var languages = []Language{
	{"go", []string{".go"}, ts.NewLanguage(golang.Language())},
	{"javascript", []string{".js", ".jsx", ".mjs", ".cjs"}, ts.NewLanguage(javascript.Language())},
	{"typescript", []string{".ts"}, ts.NewLanguage(typescript.LanguageTypescript())},
	{"tsx", []string{".tsx"}, ts.NewLanguage(typescript.LanguageTSX())},
	{"python", []string{".py"}, ts.NewLanguage(python.Language())},
	{"rust", []string{".rs"}, ts.NewLanguage(rust.Language())},
	{"c", []string{".c", ".h"}, ts.NewLanguage(c.Language())},
	{"cpp", []string{".cc", ".cpp", ".cxx", ".hpp", ".hxx"}, ts.NewLanguage(cpp.Language())},
	{"json", []string{".json"}, ts.NewLanguage(jsonlang.Language())},
	{"yaml", []string{".yaml", ".yml"}, ts.NewLanguage(yaml.Language())},
	{"markdown", []string{".md", ".markdown"}, ts.NewLanguage(markdown.Language())},
	{"php", []string{".php"}, ts.NewLanguage(php.LanguagePHP())},
}

func LanguageFor(path, name string) (Language, bool) {
	if name != "" {
		for _, l := range languages {
			if l.Name == name {
				return l, true
			}
		}
		return Language{}, false
	}
	ext := strings.ToLower(filepath.Ext(path))
	for _, l := range languages {
		for _, e := range l.Extensions {
			if e == ext {
				return l, true
			}
		}
	}
	return Language{}, false
}
func Languages() []string {
	out := make([]string, len(languages))
	for i, l := range languages {
		out[i] = l.Name
	}
	return out
}

type Document struct {
	Path, Language string
	Source         []byte
	Tree           *ts.Tree
	Root           *ts.Node
}

func Parse(path, language string) (*Document, error) {
	l, ok := LanguageFor(path, language)
	if !ok {
		return nil, fmt.Errorf("unsupported language or extension: %s", path)
	}
	src, err := os.ReadFile(path)
	if err != nil {
		return nil, err
	}
	p := ts.NewParser()
	defer p.Close()
	if err := p.SetLanguage(l.Grammar); err != nil {
		return nil, err
	}
	tree := p.Parse(src, nil)
	if tree == nil {
		return nil, fmt.Errorf("parse failed: %s", path)
	}
	return &Document{Path: path, Language: l.Name, Source: src, Tree: tree, Root: tree.RootNode()}, nil
}
func (d *Document) Close() { d.Tree.Close() }

type File struct{ Path, Language string }

func Files(root string) ([]File, error) {
	var out []File
	err := filepath.Walk(root, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		if info.IsDir() {
			if path != root && (info.Name() == ".git" || info.Name() == "vendor" || info.Name() == "node_modules" || info.Name() == ".kiro") {
				return filepath.SkipDir
			}
			return nil
		}
		if l, ok := LanguageFor(path, ""); ok {
			out = append(out, File{path, l.Name})
		}
		return nil
	})
	sort.Slice(out, func(i, j int) bool { return out[i].Path < out[j].Path })
	return out, err
}
func Text(n *ts.Node, source []byte) string { return n.Utf8Text(source) }
func Line(n *ts.Node) int                   { return int(n.StartPosition().Row) + 1 }
func SymbolKinds() map[string]bool {
	return map[string]bool{"function_declaration": true, "method_declaration": true, "func_literal": true, "type_declaration": true, "type_spec": true, "class_declaration": true, "class_definition": true, "struct_item": true, "enum_item": true, "interface_declaration": true, "namespace_definition": true, "function_definition": true, "function_item": true, "const_declaration": true, "var_declaration": true, "lexical_declaration": true, "variable_declaration": true, "import_declaration": true, "use_declaration": true, "callable_declaration": true, "property_declaration": true, "atx_heading": true, "setext_heading": true}
}
func Walk(n *ts.Node, fn func(*ts.Node)) {
	fn(n)
	for i := uint(0); i < n.NamedChildCount(); i++ {
		Walk(n.NamedChild(i), fn)
	}
}