amql_engine/extractor/
ts_structure.rs

1//! General-purpose TypeScript/JavaScript structure extractor.
2//!
3//! Extracts functions, classes, methods, interfaces, type aliases, enums,
4//! and constants as annotations. Mirrors the TypeScript resolver logic but
5//! produces `Annotation` metadata instead of `CodeElement` trees, discarding
6//! implementation bodies entirely.
7
8use super::BuiltinExtractor;
9use crate::store::Annotation;
10use crate::types::{AttrName, Binding, RelativePath, TagName};
11use rustc_hash::FxHashMap;
12use serde_json::Value as JsonValue;
13use std::cell::RefCell;
14
15/// General-purpose TypeScript/JavaScript structure extractor.
16pub struct TypeScriptStructureExtractor;
17
18impl BuiltinExtractor for TypeScriptStructureExtractor {
19    fn name(&self) -> &str {
20        "structure"
21    }
22
23    fn extensions(&self) -> &[&str] {
24        &[".ts", ".tsx", ".js", ".jsx", ".mts", ".mjs"]
25    }
26
27    fn extract(&self, source: &str, file: &RelativePath) -> Vec<Annotation> {
28        let tree = match parse_ts(source, file) {
29            Some(t) => t,
30            None => return vec![],
31        };
32        let mut annotations = Vec::new();
33        let root = tree.root_node();
34        let src = source.as_bytes();
35        let mut cursor = root.walk();
36        for child in root.named_children(&mut cursor) {
37            extract_elements(&child, src, file, &mut annotations, false, false);
38        }
39        annotations
40    }
41}
42
43// ---------------------------------------------------------------------------
44// Tree-sitter helpers
45// ---------------------------------------------------------------------------
46
47fn node_text<'a>(node: &tree_sitter::Node, src: &'a [u8]) -> &'a str {
48    node.utf8_text(src).unwrap_or("")
49}
50
51fn get_name(node: &tree_sitter::Node, src: &[u8]) -> String {
52    node.child_by_field_name("name")
53        .map(|n| node_text(&n, src).to_string())
54        .unwrap_or_default()
55}
56
57fn is_async(node: &tree_sitter::Node, src: &[u8]) -> bool {
58    let mut cursor = node.walk();
59    let result = node
60        .children(&mut cursor)
61        .any(|c| node_text(&c, src) == "async");
62    result
63}
64
65fn has_keyword(node: &tree_sitter::Node, src: &[u8], keyword: &str) -> bool {
66    let mut cursor = node.walk();
67    let result = node
68        .children(&mut cursor)
69        .any(|c| node_text(&c, src) == keyword);
70    result
71}
72
73fn is_generator(node: &tree_sitter::Node) -> bool {
74    node.kind() == "generator_function_declaration" || node.kind() == "generator_function"
75}
76
77fn make_annotation(
78    tag: &str,
79    binding: String,
80    attrs: FxHashMap<AttrName, JsonValue>,
81    file: &RelativePath,
82    children: Vec<Annotation>,
83) -> Annotation {
84    Annotation {
85        tag: TagName::from(tag),
86        attrs,
87        binding: Binding::from(binding),
88        file: file.clone(),
89        children,
90    }
91}
92
93// ---------------------------------------------------------------------------
94// Element extraction
95// ---------------------------------------------------------------------------
96
97/// Extract elements from a tree-sitter node, propagating export/default flags.
98fn extract_elements(
99    node: &tree_sitter::Node,
100    src: &[u8],
101    file: &RelativePath,
102    result: &mut Vec<Annotation>,
103    is_export: bool,
104    is_default: bool,
105) {
106    match node.kind() {
107        "function_declaration" | "generator_function_declaration" => {
108            if let Some(ann) = extract_function(node, src, file, is_export, is_default) {
109                result.push(ann);
110            }
111        }
112        "class_declaration" | "abstract_class_declaration" => {
113            if let Some(ann) = extract_class(node, src, file, is_export, is_default) {
114                result.push(ann);
115            }
116        }
117        "interface_declaration" => {
118            if let Some(ann) = extract_interface(node, src, file, is_export, is_default) {
119                result.push(ann);
120            }
121        }
122        "type_alias_declaration" => {
123            if let Some(ann) = extract_type_alias(node, src, file, is_export, is_default) {
124                result.push(ann);
125            }
126        }
127        "enum_declaration" => {
128            if let Some(ann) = extract_enum(node, src, file, is_export, is_default) {
129                result.push(ann);
130            }
131        }
132        "lexical_declaration" | "variable_declaration" => {
133            extract_variable_declaration(node, src, file, result, is_export, is_default);
134        }
135        "export_statement" => {
136            let has_default = has_keyword(node, src, "default");
137            let mut cursor = node.walk();
138            for child in node.named_children(&mut cursor) {
139                extract_elements(&child, src, file, result, true, has_default);
140            }
141        }
142        _ => {}
143    }
144}
145
146fn extract_function(
147    node: &tree_sitter::Node,
148    src: &[u8],
149    file: &RelativePath,
150    is_export: bool,
151    is_default: bool,
152) -> Option<Annotation> {
153    let name = get_name(node, src);
154    if name.is_empty() {
155        return None;
156    }
157    let mut attrs = FxHashMap::default();
158    attrs.insert(AttrName::from("name"), JsonValue::String(name.clone()));
159    if is_async(node, src) {
160        attrs.insert(AttrName::from("async"), JsonValue::Bool(true));
161    }
162    if is_generator(node) {
163        attrs.insert(AttrName::from("generator"), JsonValue::Bool(true));
164    }
165    if is_export {
166        attrs.insert(AttrName::from("export"), JsonValue::Bool(true));
167    }
168    if is_default {
169        attrs.insert(AttrName::from("default"), JsonValue::Bool(true));
170    }
171    Some(make_annotation("function", name, attrs, file, vec![]))
172}
173
174fn extract_class(
175    node: &tree_sitter::Node,
176    src: &[u8],
177    file: &RelativePath,
178    is_export: bool,
179    is_default: bool,
180) -> Option<Annotation> {
181    let name = get_name(node, src);
182    if name.is_empty() {
183        return None;
184    }
185    let mut attrs = FxHashMap::default();
186    attrs.insert(AttrName::from("name"), JsonValue::String(name.clone()));
187    if node.kind() == "abstract_class_declaration" {
188        attrs.insert(AttrName::from("abstract"), JsonValue::Bool(true));
189    }
190    if is_export {
191        attrs.insert(AttrName::from("export"), JsonValue::Bool(true));
192    }
193    if is_default {
194        attrs.insert(AttrName::from("default"), JsonValue::Bool(true));
195    }
196
197    // Extract extends clause
198    let mut cursor = node.walk();
199    for child in node.named_children(&mut cursor) {
200        if child.kind() == "class_heritage" {
201            let heritage_text = node_text(&child, src);
202            attrs.insert(
203                AttrName::from("extends"),
204                JsonValue::String(heritage_text.to_string()),
205            );
206        }
207    }
208
209    // Extract methods from class body
210    let mut children = Vec::new();
211    if let Some(body) = node.child_by_field_name("body") {
212        let mut body_cursor = body.walk();
213        for child in body.named_children(&mut body_cursor) {
214            match child.kind() {
215                "method_definition" => {
216                    if let Some(ann) = extract_method(&child, src, file) {
217                        children.push(ann);
218                    }
219                }
220                "public_field_definition" | "property_definition" => {
221                    if let Some(value) = child.child_by_field_name("value") {
222                        if value.kind() == "arrow_function" {
223                            if let Some(ann) = extract_method_from_property(&child, src, file) {
224                                children.push(ann);
225                            }
226                        }
227                    }
228                }
229                _ => {}
230            }
231        }
232    }
233
234    Some(make_annotation("class", name, attrs, file, children))
235}
236
237fn extract_method(node: &tree_sitter::Node, src: &[u8], file: &RelativePath) -> Option<Annotation> {
238    let name = get_name(node, src);
239    if name.is_empty() {
240        return None;
241    }
242    let mut attrs = FxHashMap::default();
243    attrs.insert(AttrName::from("name"), JsonValue::String(name.clone()));
244    if is_async(node, src) {
245        attrs.insert(AttrName::from("async"), JsonValue::Bool(true));
246    }
247
248    let mut cursor = node.walk();
249    for child in node.children(&mut cursor) {
250        let text = node_text(&child, src);
251        match text {
252            "static" => {
253                attrs.insert(AttrName::from("static"), JsonValue::Bool(true));
254            }
255            "get" if child.kind() == "property_identifier" || !child.is_named() => {
256                attrs.insert(AttrName::from("getter"), JsonValue::Bool(true));
257            }
258            "set" if child.kind() == "property_identifier" || !child.is_named() => {
259                attrs.insert(AttrName::from("setter"), JsonValue::Bool(true));
260            }
261            _ => {}
262        }
263        if child.kind() == "accessibility_modifier" {
264            attrs.insert(
265                AttrName::from("visibility"),
266                JsonValue::String(text.to_string()),
267            );
268        }
269    }
270
271    Some(make_annotation("method", name, attrs, file, vec![]))
272}
273
274fn extract_method_from_property(
275    prop: &tree_sitter::Node,
276    src: &[u8],
277    file: &RelativePath,
278) -> Option<Annotation> {
279    let name = get_name(prop, src);
280    if name.is_empty() {
281        return None;
282    }
283    let mut attrs = FxHashMap::default();
284    attrs.insert(AttrName::from("name"), JsonValue::String(name.clone()));
285    attrs.insert(AttrName::from("arrow"), JsonValue::Bool(true));
286    Some(make_annotation("method", name, attrs, file, vec![]))
287}
288
289fn extract_interface(
290    node: &tree_sitter::Node,
291    src: &[u8],
292    file: &RelativePath,
293    is_export: bool,
294    is_default: bool,
295) -> Option<Annotation> {
296    let name = get_name(node, src);
297    if name.is_empty() {
298        return None;
299    }
300    let mut attrs = FxHashMap::default();
301    attrs.insert(AttrName::from("name"), JsonValue::String(name.clone()));
302    if is_export {
303        attrs.insert(AttrName::from("export"), JsonValue::Bool(true));
304    }
305    if is_default {
306        attrs.insert(AttrName::from("default"), JsonValue::Bool(true));
307    }
308    Some(make_annotation("interface", name, attrs, file, vec![]))
309}
310
311fn extract_type_alias(
312    node: &tree_sitter::Node,
313    src: &[u8],
314    file: &RelativePath,
315    is_export: bool,
316    is_default: bool,
317) -> Option<Annotation> {
318    let name = get_name(node, src);
319    if name.is_empty() {
320        return None;
321    }
322    let mut attrs = FxHashMap::default();
323    attrs.insert(AttrName::from("name"), JsonValue::String(name.clone()));
324    if is_export {
325        attrs.insert(AttrName::from("export"), JsonValue::Bool(true));
326    }
327    if is_default {
328        attrs.insert(AttrName::from("default"), JsonValue::Bool(true));
329    }
330    Some(make_annotation("type", name, attrs, file, vec![]))
331}
332
333fn extract_enum(
334    node: &tree_sitter::Node,
335    src: &[u8],
336    file: &RelativePath,
337    is_export: bool,
338    is_default: bool,
339) -> Option<Annotation> {
340    let name = get_name(node, src);
341    if name.is_empty() {
342        return None;
343    }
344    let mut attrs = FxHashMap::default();
345    attrs.insert(AttrName::from("name"), JsonValue::String(name.clone()));
346    if is_export {
347        attrs.insert(AttrName::from("export"), JsonValue::Bool(true));
348    }
349    if is_default {
350        attrs.insert(AttrName::from("default"), JsonValue::Bool(true));
351    }
352
353    // Extract members as children
354    let mut children = Vec::new();
355    if let Some(body) = node.child_by_field_name("body") {
356        let mut cursor = body.walk();
357        for child in body.named_children(&mut cursor) {
358            // property_identifier = bare member (e.g. Admin)
359            // enum_assignment = member with value (e.g. Admin = 0); first named child is the name
360            if child.kind() == "enum_assignment" || child.kind() == "property_identifier" {
361                let member_name = if child.kind() == "enum_assignment" {
362                    // No "name" field label — the identifier is the first named child
363                    child
364                        .named_child(0)
365                        .map(|n| node_text(&n, src).to_string())
366                        .unwrap_or_default()
367                } else {
368                    node_text(&child, src).to_string()
369                };
370                if !member_name.is_empty() {
371                    let mut member_attrs = FxHashMap::default();
372                    member_attrs.insert(
373                        AttrName::from("name"),
374                        JsonValue::String(member_name.clone()),
375                    );
376                    children.push(make_annotation(
377                        "member",
378                        member_name,
379                        member_attrs,
380                        file,
381                        vec![],
382                    ));
383                }
384            }
385        }
386    }
387
388    Some(make_annotation("enum", name, attrs, file, children))
389}
390
391/// Extract from `const x = ...` or `let x = ...`.
392/// Arrow functions and function expressions become function annotations.
393/// Plain constants become const annotations.
394fn extract_variable_declaration(
395    node: &tree_sitter::Node,
396    src: &[u8],
397    file: &RelativePath,
398    result: &mut Vec<Annotation>,
399    is_export: bool,
400    is_default: bool,
401) {
402    let is_const = has_keyword(node, src, "const");
403
404    let mut cursor = node.walk();
405    for child in node.named_children(&mut cursor) {
406        if child.kind() == "variable_declarator" {
407            let name = get_name(&child, src);
408            if name.is_empty() {
409                continue;
410            }
411
412            if let Some(value) = child.child_by_field_name("value") {
413                match value.kind() {
414                    "arrow_function" | "function_expression" | "generator_function" => {
415                        let mut attrs = FxHashMap::default();
416                        attrs.insert(AttrName::from("name"), JsonValue::String(name.clone()));
417                        if value.kind() == "arrow_function" {
418                            attrs.insert(AttrName::from("arrow"), JsonValue::Bool(true));
419                        }
420                        if is_async(&value, src) {
421                            attrs.insert(AttrName::from("async"), JsonValue::Bool(true));
422                        }
423                        if value.kind() == "generator_function" {
424                            attrs.insert(AttrName::from("generator"), JsonValue::Bool(true));
425                        }
426                        if is_const {
427                            attrs.insert(AttrName::from("const"), JsonValue::Bool(true));
428                        }
429                        if is_export {
430                            attrs.insert(AttrName::from("export"), JsonValue::Bool(true));
431                        }
432                        if is_default {
433                            attrs.insert(AttrName::from("default"), JsonValue::Bool(true));
434                        }
435                        result.push(make_annotation("function", name, attrs, file, vec![]));
436                        return;
437                    }
438                    _ => {}
439                }
440            }
441
442            if is_const {
443                let mut attrs = FxHashMap::default();
444                attrs.insert(AttrName::from("name"), JsonValue::String(name.clone()));
445                if is_export {
446                    attrs.insert(AttrName::from("export"), JsonValue::Bool(true));
447                }
448                if is_default {
449                    attrs.insert(AttrName::from("default"), JsonValue::Bool(true));
450                }
451                result.push(make_annotation("const", name, attrs, file, vec![]));
452            }
453        }
454    }
455}
456
457// ---------------------------------------------------------------------------
458// Parser cache (thread-local)
459// ---------------------------------------------------------------------------
460
461thread_local! {
462    static TS_PARSER: RefCell<Option<tree_sitter::Parser>> = const { RefCell::new(None) };
463    static TSX_PARSER: RefCell<Option<tree_sitter::Parser>> = const { RefCell::new(None) };
464}
465
466fn parse_ts(source: &str, file: &RelativePath) -> Option<tree_sitter::Tree> {
467    let path: &str = file.as_ref();
468    let is_tsx = path.ends_with(".tsx") || path.ends_with(".jsx");
469
470    if is_tsx {
471        TSX_PARSER.with(|cell| {
472            let mut opt = cell.borrow_mut();
473            let parser = opt.get_or_insert_with(|| {
474                let mut p = tree_sitter::Parser::new();
475                p.set_language(&tree_sitter_typescript::LANGUAGE_TSX.into())
476                    .expect("Failed to set TSX language");
477                p
478            });
479            parser.parse(source, None)
480        })
481    } else {
482        TS_PARSER.with(|cell| {
483            let mut opt = cell.borrow_mut();
484            let parser = opt.get_or_insert_with(|| {
485                let mut p = tree_sitter::Parser::new();
486                p.set_language(&tree_sitter_typescript::LANGUAGE_TYPESCRIPT.into())
487                    .expect("Failed to set TypeScript language");
488                p
489            });
490            parser.parse(source, None)
491        })
492    }
493}
494
495#[cfg(test)]
496mod tests {
497    use super::*;
498
499    fn run(source: &str) -> Vec<Annotation> {
500        let file = RelativePath::from("src/app.ts");
501        TypeScriptStructureExtractor.extract(source, &file)
502    }
503
504    fn run_tsx(source: &str) -> Vec<Annotation> {
505        let file = RelativePath::from("src/app.tsx");
506        TypeScriptStructureExtractor.extract(source, &file)
507    }
508
509    #[test]
510    fn extracts_function_declarations() {
511        // Arrange
512        let source = "function foo() {}\nasync function bar() {}\nfunction* gen() {}";
513
514        // Act
515        let anns = run(source);
516
517        // Assert
518        assert_eq!(anns.len(), 3, "should find 3 functions");
519        assert_eq!(anns[0].binding.as_ref(), "foo", "first function name");
520        assert_eq!(anns[0].tag.as_ref(), "function", "first function tag");
521        assert_eq!(
522            anns[1].attrs.get(&AttrName::from("async")),
523            Some(&JsonValue::Bool(true)),
524            "bar should be async"
525        );
526        assert_eq!(
527            anns[2].attrs.get(&AttrName::from("generator")),
528            Some(&JsonValue::Bool(true)),
529            "gen should be generator"
530        );
531    }
532
533    #[test]
534    fn extracts_arrow_functions() {
535        // Arrange
536        let source = "const handler = async (req: Request) => { return 1; };";
537
538        // Act
539        let anns = run(source);
540
541        // Assert
542        assert_eq!(anns.len(), 1, "should find 1 function");
543        assert_eq!(anns[0].tag.as_ref(), "function", "should be function tag");
544        assert_eq!(
545            anns[0].binding.as_ref(),
546            "handler",
547            "should be named handler"
548        );
549        assert_eq!(
550            anns[0].attrs.get(&AttrName::from("arrow")),
551            Some(&JsonValue::Bool(true)),
552            "should be arrow"
553        );
554        assert_eq!(
555            anns[0].attrs.get(&AttrName::from("async")),
556            Some(&JsonValue::Bool(true)),
557            "should be async"
558        );
559    }
560
561    #[test]
562    fn extracts_classes_with_methods() {
563        // Arrange
564        let source = r#"class UserService {
565            async getById(id: string): Promise<User> { return null; }
566            static create() { return new UserService(); }
567        }"#;
568
569        // Act
570        let anns = run(source);
571
572        // Assert
573        assert_eq!(anns.len(), 1, "should find 1 class");
574        assert_eq!(anns[0].tag.as_ref(), "class", "should be class");
575        assert_eq!(anns[0].binding.as_ref(), "UserService", "class name");
576        assert_eq!(anns[0].children.len(), 2, "should have 2 methods");
577        assert_eq!(
578            anns[0].children[0].binding.as_ref(),
579            "getById",
580            "method name"
581        );
582        assert_eq!(
583            anns[0].children[0].attrs.get(&AttrName::from("async")),
584            Some(&JsonValue::Bool(true)),
585            "getById should be async"
586        );
587        assert_eq!(
588            anns[0].children[1].attrs.get(&AttrName::from("static")),
589            Some(&JsonValue::Bool(true)),
590            "create should be static"
591        );
592    }
593
594    #[test]
595    fn extracts_interfaces_and_types() {
596        // Arrange
597        let source = "interface User { id: string; name: string; }\ntype UserId = string;";
598
599        // Act
600        let anns = run(source);
601
602        // Assert
603        assert_eq!(anns.len(), 2, "should find interface + type");
604        assert_eq!(
605            anns[0].tag.as_ref(),
606            "interface",
607            "first should be interface"
608        );
609        assert_eq!(anns[0].binding.as_ref(), "User", "interface name");
610        assert_eq!(anns[1].tag.as_ref(), "type", "second should be type");
611        assert_eq!(anns[1].binding.as_ref(), "UserId", "type name");
612    }
613
614    #[test]
615    fn extracts_enums_with_members() {
616        // Arrange
617        let source = "enum Role { Admin, User, Guest }";
618
619        // Act
620        let anns = run(source);
621
622        // Assert
623        assert_eq!(anns.len(), 1, "should find 1 enum");
624        assert_eq!(anns[0].tag.as_ref(), "enum", "should be enum");
625        assert_eq!(anns[0].binding.as_ref(), "Role", "enum name");
626        assert_eq!(anns[0].children.len(), 3, "should have 3 members");
627    }
628
629    #[test]
630    fn extracts_enums_with_assigned_members() {
631        // Arrange — enum members with explicit value assignments
632        let source = "enum Status { Active = 0, Inactive = 1, Pending = 2 }";
633
634        // Act
635        let anns = run(source);
636
637        // Assert
638        assert_eq!(anns.len(), 1, "should find 1 enum");
639        assert_eq!(anns[0].children.len(), 3, "should have 3 assigned members");
640        assert_eq!(
641            anns[0].children[0].binding.as_ref(),
642            "Active",
643            "first member name"
644        );
645        assert_eq!(
646            anns[0].children[1].binding.as_ref(),
647            "Inactive",
648            "second member name"
649        );
650        assert_eq!(
651            anns[0].children[2].binding.as_ref(),
652            "Pending",
653            "third member name"
654        );
655    }
656
657    #[test]
658    fn extracts_exports() {
659        // Arrange
660        let source =
661            "export function fetchUser() {}\nexport default class App {}\nexport const MAX = 3;";
662
663        // Act
664        let anns = run(source);
665
666        // Assert
667        assert_eq!(anns.len(), 3, "should find 3 exports");
668        assert_eq!(
669            anns[0].attrs.get(&AttrName::from("export")),
670            Some(&JsonValue::Bool(true)),
671            "fetchUser should be exported"
672        );
673        assert_eq!(
674            anns[1].attrs.get(&AttrName::from("default")),
675            Some(&JsonValue::Bool(true)),
676            "App should be default export"
677        );
678        assert_eq!(
679            anns[2].attrs.get(&AttrName::from("export")),
680            Some(&JsonValue::Bool(true)),
681            "MAX should be exported"
682        );
683    }
684
685    #[test]
686    fn extracts_constants() {
687        // Arrange
688        let source = "const MAX_RETRIES = 3;";
689
690        // Act
691        let anns = run(source);
692
693        // Assert
694        assert_eq!(anns.len(), 1, "should find 1 const");
695        assert_eq!(anns[0].tag.as_ref(), "const", "should be const");
696        assert_eq!(anns[0].binding.as_ref(), "MAX_RETRIES", "const name");
697    }
698
699    #[test]
700    fn extracts_tsx() {
701        // Arrange
702        let source = r#"export function App(): JSX.Element { return <div>Hello</div>; }"#;
703
704        // Act
705        let anns = run_tsx(source);
706
707        // Assert
708        assert_eq!(anns.len(), 1, "should find 1 function");
709        assert_eq!(anns[0].binding.as_ref(), "App", "function name");
710        assert_eq!(
711            anns[0].attrs.get(&AttrName::from("export")),
712            Some(&JsonValue::Bool(true)),
713            "should be exported"
714        );
715    }
716}
amql_engine/extractor/ts_structure.rs

amql_engine/extractor/
ts_structure.rs
