| 1 | //go:build treesitter && cgo |
| 2 | |
| 3 | package builtin |
| 4 | |
| 5 | import ( |
| 6 | "fmt" |
| 7 | "os" |
| 8 | "path/filepath" |
| 9 | "strings" |
| 10 | "unsafe" |
| 11 | |
| 12 | sitter "github.com/tree-sitter/go-tree-sitter" |
| 13 | tree_sitter_javascript "github.com/tree-sitter/tree-sitter-javascript/bindings/go" |
| 14 | tree_sitter_python "github.com/tree-sitter/tree-sitter-python/bindings/go" |
| 15 | tree_sitter_rust "github.com/tree-sitter/tree-sitter-rust/bindings/go" |
| 16 | tree_sitter_typescript "github.com/tree-sitter/tree-sitter-typescript/bindings/go" |
| 17 | ) |
| 18 | |
| 19 | type codeIndexTreeSitterSpec struct { |
| 20 | language func() unsafe.Pointer |
| 21 | query string |
| 22 | } |
| 23 | |
| 24 | func (c codeIndex) parseTreeSitter(path string) ([]codeSymbol, bool, error) { |
| 25 | spec, ok := codeIndexTreeSitterSpecForExt(filepath.Ext(path)) |
| 26 | if !ok { |
| 27 | return nil, false, nil |
| 28 | } |
| 29 | source, err := os.ReadFile(path) |
| 30 | if err != nil { |
| 31 | return nil, true, err |
| 32 | } |
| 33 | language := sitter.NewLanguage(spec.language()) |
| 34 | parser := sitter.NewParser() |
| 35 | defer parser.Close() |
| 36 | if err := parser.SetLanguage(language); err != nil { |
| 37 | return nil, true, err |
| 38 | } |
| 39 | tree := parser.Parse(source, nil) |
| 40 | if tree == nil { |
| 41 | return nil, true, fmt.Errorf("tree-sitter parse returned nil") |
| 42 | } |
| 43 | defer tree.Close() |
| 44 | query, qerr := sitter.NewQuery(language, spec.query) |
| 45 | if qerr != nil { |
| 46 | return nil, true, qerr |
| 47 | } |
| 48 | defer query.Close() |
| 49 | |
| 50 | cursor := sitter.NewQueryCursor() |
| 51 | defer cursor.Close() |
| 52 | cursor.SetTimeoutMicros(50_000) |
| 53 | |
| 54 | captureNames := query.CaptureNames() |
| 55 | matches := cursor.Matches(query, tree.RootNode(), source) |
| 56 | var out []codeSymbol |
| 57 | for match := matches.Next(); match != nil; match = matches.Next() { |
| 58 | var nameNode *sitter.Node |
| 59 | var symbolNode *sitter.Node |
| 60 | var captureKind string |
| 61 | for _, capture := range match.Captures { |
| 62 | captureName := captureNames[capture.Index] |
| 63 | kind, role, ok := splitCodeIndexTreeSitterCapture(captureName) |
| 64 | if !ok { |
| 65 | continue |
| 66 | } |
| 67 | switch role { |
| 68 | case "name": |
| 69 | captureKind = kind |
| 70 | node := capture.Node |
| 71 | nameNode = &node |
| 72 | case "symbol": |
| 73 | if captureKind == "" { |
| 74 | captureKind = kind |
| 75 | } |
| 76 | node := capture.Node |
| 77 | symbolNode = &node |
| 78 | } |
| 79 | } |
| 80 | if nameNode == nil || captureKind == "" { |
| 81 | continue |
| 82 | } |
| 83 | if symbolNode == nil { |
| 84 | symbolNode = nameNode |
| 85 | } |
| 86 | name := strings.TrimSpace(nameNode.Utf8Text(source)) |
| 87 | if name == "" { |
| 88 | continue |
| 89 | } |
| 90 | kind := normalizeCodeIndexKind(captureKind) |
| 91 | parent := treeSitterParentName(kind, symbolNode, source) |
| 92 | out = append(out, codeSymbol{ |
| 93 | Name: name, |
| 94 | Kind: kind, |
| 95 | File: c.displayPath(path), |
| 96 | Line: int(symbolNode.StartPosition().Row) + 1, |
| 97 | Parent: parent, |
| 98 | Signature: treeSitterLine(source, int(symbolNode.StartByte())), |
| 99 | }) |
| 100 | } |
| 101 | if cursor.DidExceedMatchLimit() { |
| 102 | return out, true, fmt.Errorf("tree-sitter query exceeded match limit") |
| 103 | } |
| 104 | return out, true, nil |
| 105 | } |
| 106 | |
| 107 | func codeIndexTreeSitterEnabled() bool { |
| 108 | return true |
| 109 | } |
| 110 | |
| 111 | func codeIndexTreeSitterSpecForExt(ext string) (codeIndexTreeSitterSpec, bool) { |
| 112 | switch ext { |
| 113 | case ".js", ".jsx": |
| 114 | return codeIndexTreeSitterSpec{language: tree_sitter_javascript.Language, query: treeSitterJavaScriptQuery}, true |
| 115 | case ".ts": |
| 116 | return codeIndexTreeSitterSpec{language: tree_sitter_typescript.LanguageTypescript, query: treeSitterTypeScriptQuery}, true |
| 117 | case ".tsx": |
| 118 | return codeIndexTreeSitterSpec{language: tree_sitter_typescript.LanguageTSX, query: treeSitterTypeScriptQuery}, true |
| 119 | case ".py": |
| 120 | return codeIndexTreeSitterSpec{language: tree_sitter_python.Language, query: treeSitterPythonQuery}, true |
| 121 | case ".rs": |
| 122 | return codeIndexTreeSitterSpec{language: tree_sitter_rust.Language, query: treeSitterRustQuery}, true |
| 123 | default: |
| 124 | return codeIndexTreeSitterSpec{}, false |
| 125 | } |
| 126 | } |
| 127 | |
| 128 | func splitCodeIndexTreeSitterCapture(name string) (kind, role string, ok bool) { |
| 129 | before, after, found := strings.Cut(name, ".") |
| 130 | if !found || before == "" || after == "" { |
| 131 | return "", "", false |
| 132 | } |
| 133 | if after != "name" && after != "symbol" { |
| 134 | return "", "", false |
| 135 | } |
| 136 | return before, after, true |
| 137 | } |
| 138 | |
| 139 | func treeSitterParentName(kind string, node *sitter.Node, source []byte) string { |
| 140 | if kind != "method" { |
| 141 | return "" |
| 142 | } |
| 143 | for parent := node.Parent(); parent != nil; parent = parent.Parent() { |
| 144 | switch parent.Kind() { |
| 145 | case "abstract_class_declaration", "class_declaration", "class": |
| 146 | if name := parent.ChildByFieldName("name"); name != nil { |
| 147 | return strings.TrimSpace(name.Utf8Text(source)) |
| 148 | } |
| 149 | } |
| 150 | } |
| 151 | return "" |
| 152 | } |
| 153 | |
| 154 | func treeSitterLine(source []byte, start int) string { |
| 155 | if start < 0 { |
| 156 | start = 0 |
| 157 | } |
| 158 | if start > len(source) { |
| 159 | start = len(source) |
| 160 | } |
| 161 | lineStart := start |
| 162 | for lineStart > 0 && source[lineStart-1] != '\n' && source[lineStart-1] != '\r' { |
| 163 | lineStart-- |
| 164 | } |
| 165 | lineEnd := start |
| 166 | for lineEnd < len(source) && source[lineEnd] != '\n' && source[lineEnd] != '\r' { |
| 167 | lineEnd++ |
| 168 | } |
| 169 | return strings.TrimSpace(string(source[lineStart:lineEnd])) |
| 170 | } |
| 171 | |
| 172 | const treeSitterJavaScriptQuery = ` |
| 173 | (function_declaration |
| 174 | name: (identifier) @func.name) @func.symbol |
| 175 | (generator_function_declaration |
| 176 | name: (identifier) @func.name) @func.symbol |
| 177 | (class_declaration |
| 178 | name: (identifier) @class.name) @class.symbol |
| 179 | (method_definition |
| 180 | name: [(property_identifier) (private_property_identifier)] @method.name) @method.symbol |
| 181 | (lexical_declaration |
| 182 | (variable_declarator |
| 183 | name: (identifier) @func.name |
| 184 | value: [(arrow_function) (function_expression)])) @func.symbol |
| 185 | (variable_declaration |
| 186 | (variable_declarator |
| 187 | name: (identifier) @func.name |
| 188 | value: [(arrow_function) (function_expression)])) @func.symbol |
| 189 | ` |
| 190 | |
| 191 | const treeSitterTypeScriptQuery = ` |
| 192 | (function_declaration |
| 193 | name: (identifier) @func.name) @func.symbol |
| 194 | (generator_function_declaration |
| 195 | name: (identifier) @func.name) @func.symbol |
| 196 | (class_declaration |
| 197 | name: (type_identifier) @class.name) @class.symbol |
| 198 | (abstract_class_declaration |
| 199 | name: (type_identifier) @class.name) @class.symbol |
| 200 | (method_definition |
| 201 | name: [(property_identifier) (private_property_identifier)] @method.name) @method.symbol |
| 202 | (interface_declaration |
| 203 | name: (type_identifier) @interface.name) @interface.symbol |
| 204 | (type_alias_declaration |
| 205 | name: (type_identifier) @type.name) @type.symbol |
| 206 | (enum_declaration |
| 207 | name: (identifier) @enum.name) @enum.symbol |
| 208 | (lexical_declaration |
| 209 | (variable_declarator |
| 210 | name: (identifier) @func.name |
| 211 | value: [(arrow_function) (function_expression)])) @func.symbol |
| 212 | (variable_declaration |
| 213 | (variable_declarator |
| 214 | name: (identifier) @func.name |
| 215 | value: [(arrow_function) (function_expression)])) @func.symbol |
| 216 | ` |
| 217 | |
| 218 | const treeSitterPythonQuery = ` |
| 219 | (function_definition |
| 220 | name: (identifier) @func.name) @func.symbol |
| 221 | (class_definition |
| 222 | name: (identifier) @class.name) @class.symbol |
| 223 | ` |
| 224 | |
| 225 | const treeSitterRustQuery = ` |
| 226 | (function_item |
| 227 | name: (identifier) @fn.name) @fn.symbol |
| 228 | (struct_item |
| 229 | name: (type_identifier) @struct.name) @struct.symbol |
| 230 | (enum_item |
| 231 | name: (type_identifier) @enum.name) @enum.symbol |
| 232 | (trait_item |
| 233 | name: (type_identifier) @trait.name) @trait.symbol |
| 234 | ` |
| 235 |