npm-snoop/internal/symbols/symbols.go

package symbols

import (
	"crypto/sha256"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"unsafe"

	tree_sitter "github.com/tree-sitter/go-tree-sitter"
	tree_sitter_javascript "github.com/tree-sitter/tree-sitter-javascript/bindings/go"
	tree_sitter_typescript "github.com/tree-sitter/tree-sitter-typescript/bindings/go"
)

// Symbol represents a named declaration in a JS/TS file.
type Symbol struct {
	Name string `json:"name"`
	Kind string `json:"kind"` // "function", "class", "variable", "method"
	File string `json:"file"`
}

// SymbolMap maps symbol keys (file:name:kind) to a hash of their source text.
type SymbolMap map[string]string

// Key returns the unique key for a symbol.
func (s Symbol) Key() string {
	return s.File + ":" + s.Name + ":" + s.Kind
}

// ExtractFromDir walks a directory and extracts symbols from all JS/TS files.
func ExtractFromDir(dir string) (map[string]Symbol, SymbolMap, error) {
	symbols := make(map[string]Symbol)
	hashes := make(SymbolMap)

	jsLang := tree_sitter.NewLanguage(tree_sitter_javascript.Language())
	tsLang := tree_sitter.NewLanguage(tree_sitter_typescript.LanguageTypescript())

	err := filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return nil // skip unreadable files
		}
		if info.IsDir() {
			base := filepath.Base(path)
			if base == "node_modules" || base == ".git" {
				return filepath.SkipDir
			}
			return nil
		}

		ext := filepath.Ext(path)
		var lang *tree_sitter.Language
		switch ext {
		case ".js", ".mjs", ".cjs":
			lang = jsLang
		case ".ts", ".mts", ".cts":
			lang = tsLang
		default:
			return nil
		}

		// Skip minified files.
		if strings.HasSuffix(strings.TrimSuffix(path, ext), ".min") {
			return nil
		}

		relPath, _ := filepath.Rel(dir, path)

		src, err := os.ReadFile(path)
		if err != nil {
			return nil
		}

		extracted, err := extractSymbols(src, lang, relPath)
		if err != nil {
			return nil // skip unparseable files
		}

		for _, sym := range extracted {
			symbols[sym.sym.Key()] = sym.sym
			hashes[sym.sym.Key()] = sym.hash
		}
		return nil
	})

	return symbols, hashes, err
}

type extractedSymbol struct {
	sym  Symbol
	hash string
}

func extractSymbols(src []byte, lang *tree_sitter.Language, file string) ([]extractedSymbol, error) {
	parser := tree_sitter.NewParser()
	defer parser.Close()

	if err := parser.SetLanguage(lang); err != nil {
		return nil, err
	}

	tree := parser.Parse(src, nil)
	defer tree.Close()

	root := tree.RootNode()
	var result []extractedSymbol

	walkTopLevel(root, src, file, &result)
	return result, nil
}

// walkTopLevel extracts declarations only at the module/program level.
// It recurses through structural nodes (program, export_statement, if_statement
// at top level, etc.) but does NOT descend into function/method bodies.
func walkTopLevel(node *tree_sitter.Node, src []byte, file string, result *[]extractedSymbol) {
	kind := node.Kind()

	switch kind {
	case "function_declaration", "generator_function_declaration":
		if name := childFieldText(node, "name", src); name != "" {
			*result = append(*result, makeSymbol(name, "function", file, node, src))
		}
		return // don't recurse into function body

	case "class_declaration":
		if name := childFieldText(node, "name", src); name != "" {
			*result = append(*result, makeSymbol(name, "class", file, node, src))
		}
		// Extract methods within the class.
		if body := node.ChildByFieldName("body"); body != nil {
			for i := uint(0); i < body.NamedChildCount(); i++ {
				child := body.NamedChild(i)
				if child.Kind() == "method_definition" {
					if mname := childFieldText(child, "name", src); mname != "" {
						qualName := childFieldText(node, "name", src) + "." + mname
						*result = append(*result, makeSymbol(qualName, "method", file, child, src))
					}
				}
			}
		}
		return

	case "lexical_declaration", "variable_declaration":
		for i := uint(0); i < node.NamedChildCount(); i++ {
			child := node.NamedChild(i)
			if child.Kind() == "variable_declarator" {
				nameNode := child.ChildByFieldName("name")
				if nameNode == nil {
					continue
				}
				switch nameNode.Kind() {
				case "identifier":
					name := nameNode.Utf8Text(src)
					declKind := "variable"
					if val := child.ChildByFieldName("value"); val != nil {
						switch val.Kind() {
						case "arrow_function", "function_expression", "generator_function":
							declKind = "function"
						}
					}
					*result = append(*result, makeSymbol(name, declKind, file, node, src))

				case "object_pattern":
					// Destructuring: const { a, b: c } = ...
					// Extract each bound identifier.
					for _, name := range extractDestructuredNames(nameNode, src) {
						*result = append(*result, makeSymbol(name, "variable", file, node, src))
					}

				case "array_pattern":
					for _, name := range extractDestructuredNames(nameNode, src) {
						*result = append(*result, makeSymbol(name, "variable", file, node, src))
					}
				}
			}
		}
		return

	case "export_statement":
		if decl := node.ChildByFieldName("declaration"); decl != nil {
			walkTopLevel(decl, src, file, result)
			return
		}

	case "expression_statement":
		// Handle module.exports = ... and module.exports.foo = ...
		extractModuleExports(node, src, file, result)
		return
	}

	// Only recurse into structural containers, not into function/class bodies.
	for i := uint(0); i < node.NamedChildCount(); i++ {
		child := node.NamedChild(i)
		switch child.Kind() {
		case "function_declaration", "generator_function_declaration",
			"class_declaration",
			"lexical_declaration", "variable_declaration",
			"export_statement",
			"expression_statement":
			walkTopLevel(child, src, file, result)
		}
	}
}

// extractDestructuredNames extracts bound identifiers from destructuring patterns.
// For { a, b: c, d } it returns ["a", "c", "d"].
// For [a, b, c] it returns ["a", "b", "c"].
func extractDestructuredNames(node *tree_sitter.Node, src []byte) []string {
	var names []string
	for i := uint(0); i < node.NamedChildCount(); i++ {
		child := node.NamedChild(i)
		switch child.Kind() {
		case "shorthand_property_identifier_pattern":
			// { foo } -> foo
			names = append(names, child.Utf8Text(src))
		case "pair_pattern":
			// { foo: bar } -> bar (the value side)
			if val := child.ChildByFieldName("value"); val != nil {
				if val.Kind() == "identifier" {
					names = append(names, val.Utf8Text(src))
				} else {
					// Nested destructuring
					names = append(names, extractDestructuredNames(val, src)...)
				}
			}
		case "identifier":
			// [a, b] -> a, b
			names = append(names, child.Utf8Text(src))
		case "rest_pattern":
			// { ...rest } or [...rest]
			for j := uint(0); j < child.NamedChildCount(); j++ {
				if id := child.NamedChild(j); id.Kind() == "identifier" {
					names = append(names, id.Utf8Text(src))
				}
			}
		case "object_pattern", "array_pattern":
			// Nested destructuring
			names = append(names, extractDestructuredNames(child, src)...)
		case "assignment_pattern", "object_assignment_pattern":
			// { foo = defaultVal } -> foo
			if left := child.ChildByFieldName("left"); left != nil {
				switch left.Kind() {
				case "identifier", "shorthand_property_identifier_pattern":
					names = append(names, left.Utf8Text(src))
				default:
					names = append(names, extractDestructuredNames(left, src)...)
				}
			}
		}
	}
	return names
}

// extractModuleExports handles CommonJS patterns like:
//   module.exports = function() { ... }
//   module.exports.foo = function() { ... }
//   exports.foo = function() { ... }
func extractModuleExports(node *tree_sitter.Node, src []byte, file string, result *[]extractedSymbol) {
	// Look for assignment_expression children.
	for i := uint(0); i < node.NamedChildCount(); i++ {
		child := node.NamedChild(i)
		if child.Kind() != "assignment_expression" {
			continue
		}
		left := child.ChildByFieldName("left")
		right := child.ChildByFieldName("right")
		if left == nil || right == nil {
			continue
		}
		leftText := left.Utf8Text(src)

		// module.exports.foo = ... or exports.foo = ...
		if strings.HasPrefix(leftText, "module.exports.") {
			name := strings.TrimPrefix(leftText, "module.exports.")
			kind := "variable"
			switch right.Kind() {
			case "arrow_function", "function_expression", "function_declaration":
				kind = "function"
			}
			*result = append(*result, makeSymbol(name, kind, file, child, src))
		} else if strings.HasPrefix(leftText, "exports.") {
			name := strings.TrimPrefix(leftText, "exports.")
			kind := "variable"
			switch right.Kind() {
			case "arrow_function", "function_expression", "function_declaration":
				kind = "function"
			}
			*result = append(*result, makeSymbol(name, kind, file, child, src))
		}
	}
}

func childFieldText(node *tree_sitter.Node, field string, src []byte) string {
	child := node.ChildByFieldName(field)
	if child == nil {
		return ""
	}
	return child.Utf8Text(src)
}

func makeSymbol(name, kind, file string, node *tree_sitter.Node, src []byte) extractedSymbol {
	text := nodeBytes(node, src)
	h := sha256.Sum256(text)
	return extractedSymbol{
		sym:  Symbol{Name: name, Kind: kind, File: file},
		hash: fmt.Sprintf("%x", h[:8]),
	}
}

func nodeBytes(node *tree_sitter.Node, src []byte) []byte {
	start := node.StartByte()
	end := node.EndByte()
	if int(end) > len(src) {
		end = uint(len(src))
	}
	return src[start:end]
}

// Ensure the unsafe import is used (needed by tree-sitter bindings).
var _ = unsafe.Pointer(nil)