engine.go

package astkit

import (
	"context"
	"errors"
	"fmt"
	"sync"
	"time"

	sitter "github.com/smacker/go-tree-sitter"
	"github.com/smacker/go-tree-sitter/c"
	"github.com/smacker/go-tree-sitter/cpp"
	"github.com/smacker/go-tree-sitter/csharp"
	"github.com/smacker/go-tree-sitter/golang"
	"github.com/smacker/go-tree-sitter/java"
	"github.com/smacker/go-tree-sitter/javascript"
	"github.com/smacker/go-tree-sitter/kotlin"
	"github.com/smacker/go-tree-sitter/php"
	"github.com/smacker/go-tree-sitter/python"
	"github.com/smacker/go-tree-sitter/rust"
	"github.com/smacker/go-tree-sitter/swift"
	tstsx "github.com/smacker/go-tree-sitter/typescript/tsx"
	tstype "github.com/smacker/go-tree-sitter/typescript/typescript"

	"github.com/provasign/astkit/thirdparty/tsobjc"
)

// DefaultParseTimeout caps tree-sitter parse calls. Callers can override per
// Engine via WithTimeout.
const DefaultParseTimeout = 2 * time.Second

// Engine wraps tree-sitter parsing for every language astkit knows about.
// Engines are safe for concurrent use.
type Engine struct {
	timeout time.Duration
}

// NewEngine returns an Engine using DefaultParseTimeout.
func NewEngine() *Engine { return &Engine{timeout: DefaultParseTimeout} }

// WithTimeout returns a copy of e with the parse timeout overridden.
func (e *Engine) WithTimeout(d time.Duration) *Engine {
	cp := *e
	cp.timeout = d
	return &cp
}

// Parse parses src under lang and returns a tree-sitter tree. Caller is
// responsible for tree.Close().
//
// For non-AST languages (config formats, unknown) Parse returns (nil, nil)
// rather than an error — letting callers branch on `tree == nil` without
// inspecting error types.
func (e *Engine) Parse(ctx context.Context, lang LanguageKey, src []byte) (*sitter.Tree, error) {
	tsLang, ok := treeSitterLanguage(lang)
	if !ok {
		return nil, nil
	}
	p := sitter.NewParser()
	// Release the native parser eagerly: the returned tree does not depend on
	// it, and waiting for the finalizer lets large indexes accumulate
	// significant native memory before GC notices.
	defer p.Close()
	p.SetLanguage(tsLang)
	cctx, cancel := context.WithTimeout(ctx, e.timeout)
	defer cancel()
	var tree *sitter.Tree
	var err error
	if lang == LangKotlin {
		tree, err = parseKotlinWithoutComments(cctx, p, src)
	} else {
		tree, err = p.ParseCtx(cctx, nil, src)
	}
	if err != nil {
		return nil, fmt.Errorf("astkit: parse %s: %w", lang, err)
	}
	return tree, nil
}

// ErrSyntax is wrapped by the error Validate returns for source whose
// tree contains syntax errors (ERROR or MISSING nodes).
var ErrSyntax = errors.New("syntax error")

// Validate parses src and discards the tree. It returns the parse error, an
// error wrapping ErrSyntax (naming the first error's 1-based line and
// column) when the tree contains syntax errors, or nil. Languages without a
// grammar (config/data formats) parse to no tree and always validate.
func (e *Engine) Validate(ctx context.Context, lang LanguageKey, src []byte) error {
	tree, err := e.Parse(ctx, lang, src)
	if err != nil || tree == nil {
		return err
	}
	defer tree.Close()
	root := tree.RootNode()
	if !root.HasError() {
		return nil
	}
	p := firstSyntaxError(root).StartPoint()
	return fmt.Errorf("astkit: %s %w at %d:%d", lang, ErrSyntax, p.Row+1, p.Column+1)
}

// firstSyntaxError returns the first ERROR or MISSING node under n in
// source order, or n itself when none is found below it.
func firstSyntaxError(n *sitter.Node) *sitter.Node {
	for n != nil {
		if n.IsError() || n.IsMissing() {
			return n
		}
		var next *sitter.Node
		for i := 0; i < int(n.ChildCount()); i++ {
			if c := n.Child(i); c != nil && (c.HasError() || c.IsMissing()) {
				next = c
				break
			}
		}
		if next == nil {
			return n
		}
		n = next
	}
	return n
}

// treeSitterLanguage maps a LanguageKey to its tree-sitter grammar.
func treeSitterLanguage(lang LanguageKey) (*sitter.Language, bool) {
	switch lang {
	case LangGo:
		return golang.GetLanguage(), true
	case LangTypeScript:
		return tstype.GetLanguage(), true
	case LangTSX:
		return tstsx.GetLanguage(), true
	case LangJavaScript:
		return javascript.GetLanguage(), true
	case LangPython:
		return python.GetLanguage(), true
	case LangJava:
		return java.GetLanguage(), true
	case LangRust:
		return rust.GetLanguage(), true
	case LangC:
		return c.GetLanguage(), true
	case LangCPP:
		return cpp.GetLanguage(), true
	case LangCSharp:
		return csharp.GetLanguage(), true
	case LangPHP:
		return php.GetLanguage(), true
	case LangSwift:
		return swift.GetLanguage(), true
	case LangKotlin:
		return kotlin.GetLanguage(), true
	case LangObjC:
		return tsobjc.GetLanguage(), true
	default:
		return nil, false
	}
}

// Strategy is implemented by every supported language. A Strategy converts a
// parsed tree-sitter tree into the canonical astkit types.
type Strategy interface {
	Language() LanguageKey
	Extensions() []string
	Extract(tree *sitter.Tree, src []byte) ([]Symbol, error)
	ExtractImports(tree *sitter.Tree, src []byte) ([]ImportStatement, error)
}

// TextStrategy is an optional capability a Strategy may implement when it
// extracts from source text alone — no tree-sitter grammar exists for its
// language, and its Extract/ExtractImports accept a nil tree. Intended for
// line-structured artifact types (job control, data-division extracts,
// sort control cards); no current language implements it.
type TextStrategy interface {
	Strategy
	// ExtractsFromText reports that Extract tolerates a nil tree.
	ExtractsFromText() bool
}

// Registry is a thread-safe map of LanguageKey → Strategy.
type Registry struct {
	mu sync.RWMutex
	m  map[LanguageKey]Strategy
}

// NewRegistry returns an empty registry.
func NewRegistry() *Registry { return &Registry{m: make(map[LanguageKey]Strategy)} }

// Register adds s under its declared language.
func (r *Registry) Register(s Strategy) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.m[s.Language()] = s
}

// RegisterAs adds s under an additional key (used for JS/JSX, TS/TSX aliasing).
func (r *Registry) RegisterAs(key LanguageKey, s Strategy) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.m[key] = s
}

// Get returns the strategy for lang, or nil if none registered.
func (r *Registry) Get(lang LanguageKey) Strategy {
	r.mu.RLock()
	defer r.mu.RUnlock()
	return r.m[lang]
}

// Extract is a convenience that looks up the strategy and runs Extract.
// Returns (nil, nil) when no strategy is registered.
// TextCapable reports whether the strategy registered under lang extracts
// from source text alone (Extract accepts a nil tree). Callers use it to
// decide whether a language with no tree-sitter grammar still yields
// symbols; it is false for every current language.
func (r *Registry) TextCapable(lang LanguageKey) bool {
	ts, ok := r.Get(lang).(TextStrategy)
	return ok && ts.ExtractsFromText()
}

func (r *Registry) Extract(lang LanguageKey, tree *sitter.Tree, src []byte) ([]Symbol, error) {
	s := r.Get(lang)
	if s == nil {
		return nil, nil
	}
	if err := guardDepth(tree); err != nil {
		return nil, err
	}
	syms, err := s.Extract(tree, src)
	if err == nil {
		fillCallColumns(string(lang), syms, src)
	}
	return syms, err
}

// ExtractImports is a convenience analogous to Extract.
func (r *Registry) ExtractImports(lang LanguageKey, tree *sitter.Tree, src []byte) ([]ImportStatement, error) {
	s := r.Get(lang)
	if s == nil {
		return nil, nil
	}
	if err := guardDepth(tree); err != nil {
		return nil, err
	}
	return s.ExtractImports(tree, src)
}

// MaxASTDepth bounds tree nesting before extraction. Every per-language
// extractor walks the tree by direct Go recursion (one call frame per level);
// a crafted file nested a few hundred thousand levels deep ("((((((…))))))")
// parses fine but overflows the goroutine stack during extraction — a FATAL,
// non-recoverable runtime error that takes down the whole host process and
// every concurrent request Grove/Fuse is serving. A recover() cannot catch a
// stack overflow, so the depth must be rejected BEFORE the recursive walk.
// 4000 is far past any real source file (deeply-generated code rarely exceeds
// a few hundred) while leaving ample headroom below the stack limit.
const MaxASTDepth = 4000

// guardDepth rejects a tree whose nesting exceeds MaxASTDepth. It walks with a
// TreeCursor (iterative, O(1) stack — it cannot itself overflow) rather than
// recursion, so it is safe to run on the very input it is protecting against.
func guardDepth(tree *sitter.Tree) error {
	if tree == nil {
		return nil
	}
	root := tree.RootNode()
	if root == nil {
		return nil
	}
	c := sitter.NewTreeCursor(root)
	defer c.Close()
	depth := 0
	for {
		if c.GoToFirstChild() {
			depth++
			if depth > MaxASTDepth {
				return fmt.Errorf("astkit: AST nesting exceeds %d levels — refusing extraction (possible denial-of-service input)", MaxASTDepth)
			}
			continue
		}
		for !c.GoToNextSibling() {
			if !c.GoToParent() {
				return nil // walked the whole tree; within bounds
			}
			depth--
		}
	}
}