amql_engine/resolver/

typescript.rs

//! TypeScript/JavaScript source file resolver using tree-sitter.
//!
//! Parses `.ts`, `.tsx`, `.js`, `.jsx`, `.mts`, `.mjs` files into `CodeElement`
//! trees, extracting functions, classes, methods, interfaces, type aliases,
//! enums, modules, and constants.

use super::{CodeElement, SourceLocation};
use crate::error::AqlError;
use crate::types::{AttrName, CodeElementName, RelativePath, TagName};
use rustc_hash::FxHashMap;
use std::cell::RefCell;
use std::path::Path;

/// TypeScript/JavaScript source file resolver using tree-sitter.
pub struct TypeScriptResolver;

impl super::CodeResolver for TypeScriptResolver {
    fn resolve(&self, file_path: &Path) -> Result<CodeElement, AqlError> {
        let source =
            std::fs::read_to_string(file_path).map_err(|e| format!("Failed to read file: {e}"))?;
        let is_tsx = file_path
            .extension()
            .is_some_and(|ext| ext == "tsx" || ext == "jsx");
        let root = parse_typescript_source(&source, file_path, is_tsx)?;
        Ok(root)
    }

    fn extensions(&self) -> &[&str] {
        &[".ts", ".tsx", ".js", ".jsx", ".mts", ".mjs"]
    }

    fn code_tags(&self) -> &[&str] {
        &[
            "function",
            "class",
            "method",
            "interface",
            "type",
            "enum",
            "module",
            "const",
        ]
    }
}

// Thread-local cached tree-sitter parsers to avoid re-creating on each call.
thread_local! {
    static TS_PARSER: RefCell<Option<tree_sitter::Parser>> = const { RefCell::new(None) };
    static TSX_PARSER: RefCell<Option<tree_sitter::Parser>> = const { RefCell::new(None) };
}

fn with_ts_parser<F, R>(is_tsx: bool, f: F) -> Result<R, String>
where
    F: FnOnce(&mut tree_sitter::Parser) -> Result<R, String>,
{
    if is_tsx {
        TSX_PARSER.with(|cell| {
            let mut opt = cell.borrow_mut();
            let parser = opt.get_or_insert_with(|| {
                let mut p = tree_sitter::Parser::new();
                p.set_language(&tree_sitter_typescript::LANGUAGE_TSX.into())
                    .expect("Failed to set TSX language for tree-sitter");
                p
            });
            f(parser)
        })
    } else {
        TS_PARSER.with(|cell| {
            let mut opt = cell.borrow_mut();
            let parser = opt.get_or_insert_with(|| {
                let mut p = tree_sitter::Parser::new();
                p.set_language(&tree_sitter_typescript::LANGUAGE_TYPESCRIPT.into())
                    .expect("Failed to set TypeScript language for tree-sitter");
                p
            });
            f(parser)
        })
    }
}

/// Parse a TypeScript/JavaScript source string into a CodeElement tree.
fn parse_typescript_source(
    source: &str,
    file_path: &Path,
    is_tsx: bool,
) -> Result<CodeElement, String> {
    let tree = with_ts_parser(is_tsx, |parser| {
        parser
            .parse(source, None)
            .ok_or_else(|| "Failed to parse source".to_string())
    })?;

    let root_node = tree.root_node();
    let src = source.as_bytes();
    let file_str = file_path.to_string_lossy().to_string();

    let mut children = Vec::new();
    let mut cursor = root_node.walk();
    for child in root_node.named_children(&mut cursor) {
        extract_elements(&child, src, &file_str, &mut children);
    }

    let filename = file_path
        .file_name()
        .map(|f| f.to_string_lossy().to_string())
        .unwrap_or_else(|| file_str.clone());

    Ok(CodeElement {
        tag: TagName::from("module"),
        name: CodeElementName::from(filename),
        attrs: FxHashMap::default(),
        children,
        source: SourceLocation {
            file: RelativePath::from(file_str),
            line: 1,
            column: 0,
            end_line: Some(root_node.end_position().row + 1),
            end_column: Some(root_node.end_position().column),
            start_byte: root_node.start_byte(),
            end_byte: root_node.end_byte(),
        },
    })
}

/// Extract CodeElement(s) from a tree-sitter node.
/// Some nodes (like export_statement) are wrappers that delegate to their child.
fn extract_elements(
    node: &tree_sitter::Node,
    src: &[u8],
    file: &str,
    result: &mut Vec<CodeElement>,
) {
    match node.kind() {
        "function_declaration" | "generator_function_declaration" => {
            result.push(extract_function(node, src, file));
        }
        "class_declaration" | "abstract_class_declaration" => {
            result.push(extract_class(node, src, file));
        }
        "interface_declaration" => {
            result.push(extract_interface(node, src, file));
        }
        "type_alias_declaration" => {
            result.push(extract_type_alias(node, src, file));
        }
        "enum_declaration" => {
            result.push(extract_enum(node, src, file));
        }
        "internal_module" | "module" => {
            result.push(extract_module(node, src, file));
        }
        "lexical_declaration" | "variable_declaration" => {
            extract_variable_declaration(node, src, file, result);
        }
        "export_statement" => {
            extract_export_statement(node, src, file, result);
        }
        _ => {}
    }
}

fn node_text<'a>(node: &tree_sitter::Node, src: &'a [u8]) -> &'a str {
    node.utf8_text(src).unwrap_or("")
}

fn get_name(node: &tree_sitter::Node, src: &[u8]) -> CodeElementName {
    CodeElementName::from(
        node.child_by_field_name("name")
            .map(|n| node_text(&n, src).to_string())
            .unwrap_or_default(),
    )
}

fn make_source_location(node: &tree_sitter::Node, file: &str) -> SourceLocation {
    let start = node.start_position();
    let end = node.end_position();
    SourceLocation {
        file: RelativePath::from(file),
        line: start.row + 1,
        column: start.column,
        end_line: Some(end.row + 1),
        end_column: Some(end.column),
        start_byte: node.start_byte(),
        end_byte: node.end_byte(),
    }
}

fn is_async(node: &tree_sitter::Node, src: &[u8]) -> bool {
    let mut cursor = node.walk();
    for child in node.children(&mut cursor) {
        if node_text(&child, src) == "async" {
            return true;
        }
    }
    false
}

fn has_keyword(node: &tree_sitter::Node, src: &[u8], keyword: &str) -> bool {
    let mut cursor = node.walk();
    let result = node
        .children(&mut cursor)
        .any(|c| node_text(&c, src) == keyword);
    result
}

fn is_generator(node: &tree_sitter::Node) -> bool {
    node.kind() == "generator_function_declaration" || node.kind() == "generator_function"
}

fn extract_function(node: &tree_sitter::Node, src: &[u8], file: &str) -> CodeElement {
    let name = get_name(node, src);
    let mut attrs = FxHashMap::default();
    attrs.insert(
        AttrName::from("name"),
        serde_json::Value::String(name.to_string()),
    );
    if is_async(node, src) {
        attrs.insert(AttrName::from("async"), serde_json::Value::Bool(true));
    }
    if is_generator(node) {
        attrs.insert(AttrName::from("generator"), serde_json::Value::Bool(true));
    }
    if let Some(ret) = node.child_by_field_name("return_type") {
        let text = node_text(&ret, src).trim_start_matches(':').trim();
        if !text.is_empty() {
            attrs.insert(
                AttrName::from("returnType"),
                serde_json::Value::String(text.to_string()),
            );
        }
    }
    CodeElement {
        tag: TagName::from("function"),
        name,
        attrs,
        children: vec![],
        source: make_source_location(node, file),
    }
}

fn extract_class(node: &tree_sitter::Node, src: &[u8], file: &str) -> CodeElement {
    let name = get_name(node, src);
    let mut attrs = FxHashMap::default();
    attrs.insert(
        AttrName::from("name"),
        serde_json::Value::String(name.to_string()),
    );
    if node.kind() == "abstract_class_declaration" {
        attrs.insert(AttrName::from("abstract"), serde_json::Value::Bool(true));
    }

    // Extract extends clause
    let mut cursor = node.walk();
    for child in node.named_children(&mut cursor) {
        if child.kind() == "class_heritage" {
            let heritage_text = node_text(&child, src);
            attrs.insert(
                AttrName::from("extends"),
                serde_json::Value::String(heritage_text.to_string()),
            );
        }
    }

    // Extract methods from class body
    let mut children = Vec::new();
    if let Some(body) = node.child_by_field_name("body") {
        let mut body_cursor = body.walk();
        for child in body.named_children(&mut body_cursor) {
            match child.kind() {
                "method_definition" => {
                    children.push(extract_method(&child, src, file));
                }
                "public_field_definition" | "property_definition" => {
                    // Check if it's an arrow function property
                    if let Some(value) = child.child_by_field_name("value") {
                        if value.kind() == "arrow_function" {
                            children.push(extract_method_from_property(&child, &value, src, file));
                        }
                    }
                }
                _ => {}
            }
        }
    }

    CodeElement {
        tag: TagName::from("class"),
        name,
        attrs,
        children,
        source: make_source_location(node, file),
    }
}

fn extract_method(node: &tree_sitter::Node, src: &[u8], file: &str) -> CodeElement {
    let name = get_name(node, src);
    let mut attrs = FxHashMap::default();
    attrs.insert(
        AttrName::from("name"),
        serde_json::Value::String(name.to_string()),
    );
    if is_async(node, src) {
        attrs.insert(AttrName::from("async"), serde_json::Value::Bool(true));
    }

    let mut cursor = node.walk();
    for child in node.children(&mut cursor) {
        let text = node_text(&child, src);
        match text {
            "static" => {
                attrs.insert(AttrName::from("static"), serde_json::Value::Bool(true));
            }
            "get" if child.kind() == "property_identifier" || !child.is_named() => {
                attrs.insert(AttrName::from("getter"), serde_json::Value::Bool(true));
            }
            "set" if child.kind() == "property_identifier" || !child.is_named() => {
                attrs.insert(AttrName::from("setter"), serde_json::Value::Bool(true));
            }
            _ => {}
        }
        if child.kind() == "accessibility_modifier" {
            attrs.insert(
                AttrName::from("visibility"),
                serde_json::Value::String(text.to_string()),
            );
        }
    }

    if let Some(ret) = node.child_by_field_name("return_type") {
        let text = node_text(&ret, src).trim_start_matches(':').trim();
        if !text.is_empty() {
            attrs.insert(
                AttrName::from("returnType"),
                serde_json::Value::String(text.to_string()),
            );
        }
    }

    CodeElement {
        tag: TagName::from("method"),
        name,
        attrs,
        children: vec![],
        source: make_source_location(node, file),
    }
}

fn extract_method_from_property(
    prop: &tree_sitter::Node,
    _arrow: &tree_sitter::Node,
    src: &[u8],
    file: &str,
) -> CodeElement {
    let name = get_name(prop, src);
    let mut attrs = FxHashMap::default();
    attrs.insert(
        AttrName::from("name"),
        serde_json::Value::String(name.to_string()),
    );
    attrs.insert(AttrName::from("arrow"), serde_json::Value::Bool(true));

    CodeElement {
        tag: TagName::from("method"),
        name,
        attrs,
        children: vec![],
        source: make_source_location(prop, file),
    }
}

fn extract_interface(node: &tree_sitter::Node, src: &[u8], file: &str) -> CodeElement {
    let name = get_name(node, src);
    let mut attrs = FxHashMap::default();
    attrs.insert(
        AttrName::from("name"),
        serde_json::Value::String(name.to_string()),
    );

    // Count members
    if let Some(body) = node.child_by_field_name("body") {
        let mut cursor = body.walk();
        let members = body.named_children(&mut cursor).count();
        attrs.insert(
            AttrName::from("members"),
            serde_json::Value::Number(members.into()),
        );
    }

    CodeElement {
        tag: TagName::from("interface"),
        name,
        attrs,
        children: vec![],
        source: make_source_location(node, file),
    }
}

fn extract_type_alias(node: &tree_sitter::Node, src: &[u8], file: &str) -> CodeElement {
    let name = get_name(node, src);
    let mut attrs = FxHashMap::default();
    attrs.insert(
        AttrName::from("name"),
        serde_json::Value::String(name.to_string()),
    );

    CodeElement {
        tag: TagName::from("type"),
        name,
        attrs,
        children: vec![],
        source: make_source_location(node, file),
    }
}

fn extract_enum(node: &tree_sitter::Node, src: &[u8], file: &str) -> CodeElement {
    let name = get_name(node, src);
    let mut attrs = FxHashMap::default();
    attrs.insert(
        AttrName::from("name"),
        serde_json::Value::String(name.to_string()),
    );

    // Count members
    if let Some(body) = node.child_by_field_name("body") {
        let mut cursor = body.walk();
        let members = body.named_children(&mut cursor).count();
        attrs.insert(
            AttrName::from("members"),
            serde_json::Value::Number(members.into()),
        );
    }

    CodeElement {
        tag: TagName::from("enum"),
        name,
        attrs,
        children: vec![],
        source: make_source_location(node, file),
    }
}

fn extract_module(node: &tree_sitter::Node, src: &[u8], file: &str) -> CodeElement {
    let name = get_name(node, src);
    let mut attrs = FxHashMap::default();
    attrs.insert(
        AttrName::from("name"),
        serde_json::Value::String(name.to_string()),
    );

    CodeElement {
        tag: TagName::from("module"),
        name,
        attrs,
        children: vec![],
        source: make_source_location(node, file),
    }
}

/// Extract from `const x = ...` or `let x = ...`.
/// Arrow functions and function expressions become function elements.
/// Plain constants become const elements.
fn extract_variable_declaration(
    node: &tree_sitter::Node,
    src: &[u8],
    file: &str,
    result: &mut Vec<CodeElement>,
) {
    let is_const = has_keyword(node, src, "const");

    let mut cursor = node.walk();
    for child in node.named_children(&mut cursor) {
        if child.kind() == "variable_declarator" {
            let name = get_name(&child, src);
            if name.is_empty() {
                continue;
            }

            if let Some(value) = child.child_by_field_name("value") {
                match value.kind() {
                    "arrow_function" | "function_expression" | "generator_function" => {
                        let mut attrs = FxHashMap::default();
                        attrs.insert(
                            AttrName::from("name"),
                            serde_json::Value::String(name.to_string()),
                        );
                        if value.kind() == "arrow_function" {
                            attrs.insert(AttrName::from("arrow"), serde_json::Value::Bool(true));
                        }
                        if is_async(&value, src) {
                            attrs.insert(AttrName::from("async"), serde_json::Value::Bool(true));
                        }
                        if value.kind() == "generator_function" {
                            attrs
                                .insert(AttrName::from("generator"), serde_json::Value::Bool(true));
                        }
                        result.push(CodeElement {
                            tag: TagName::from("function"),
                            name,
                            attrs,
                            children: vec![],
                            source: make_source_location(node, file),
                        });
                        return;
                    }
                    _ => {}
                }
            }

            if is_const {
                let mut attrs = FxHashMap::default();
                attrs.insert(
                    AttrName::from("name"),
                    serde_json::Value::String(name.to_string()),
                );
                result.push(CodeElement {
                    tag: TagName::from("const"),
                    name,
                    attrs,
                    children: vec![],
                    source: make_source_location(node, file),
                });
            }
        }
    }
}

/// Handle `export` wrapper: propagate `export` attr to the inner declaration.
fn extract_export_statement(
    node: &tree_sitter::Node,
    src: &[u8],
    file: &str,
    result: &mut Vec<CodeElement>,
) {
    let has_default = has_keyword(node, src, "default");

    let before = result.len();
    let mut cursor = node.walk();
    for child in node.named_children(&mut cursor) {
        extract_elements(&child, src, file, result);
    }

    // Tag all newly added elements with export (and default if applicable)
    for el in &mut result[before..] {
        el.attrs
            .insert(AttrName::from("export"), serde_json::Value::Bool(true));
        if has_default {
            el.attrs
                .insert(AttrName::from("default"), serde_json::Value::Bool(true));
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn parse_snippet(source: &str) -> CodeElement {
        parse_typescript_source(source, Path::new("test.ts"), false).unwrap()
    }

    fn parse_tsx_snippet(source: &str) -> CodeElement {
        parse_typescript_source(source, Path::new("test.tsx"), true).unwrap()
    }

    #[test]
    fn parses_function_declarations() {
        // Arrange
        let root = parse_snippet("function foo() {}");
        let async_root = parse_snippet("async function bar() {}");
        let gen_root = parse_snippet("function* gen() {}");

        // Act
        let foo = &root.children[0];
        let bar = &async_root.children[0];
        let gen = &gen_root.children[0];

        // Assert
        assert_eq!(foo.tag, "function", "should be function");
        assert_eq!(foo.name, "foo", "should be named foo");

        assert_eq!(bar.tag, "function", "should be function");
        assert_eq!(bar.name, "bar", "should be named bar");
        assert_eq!(
            bar.attrs.get("async"),
            Some(&serde_json::Value::Bool(true)),
            "should be async"
        );

        assert_eq!(gen.tag, "function", "should be function");
        assert_eq!(gen.name, "gen", "should be named gen");
        assert_eq!(
            gen.attrs.get("generator"),
            Some(&serde_json::Value::Bool(true)),
            "should be generator"
        );
    }

    #[test]
    fn parses_arrow_functions() {
        // Arrange
        let root = parse_snippet("const handler = async (req: Request) => { return 1; };");

        // Act
        let handler = &root.children[0];

        // Assert
        assert_eq!(handler.tag, "function", "arrow fn should be function");
        assert_eq!(handler.name, "handler", "should be named handler");
        assert_eq!(
            handler.attrs.get("arrow"),
            Some(&serde_json::Value::Bool(true)),
            "should be arrow"
        );
        assert_eq!(
            handler.attrs.get("async"),
            Some(&serde_json::Value::Bool(true)),
            "should be async"
        );
    }

    #[test]
    fn parses_classes_with_methods() {
        // Arrange
        let root = parse_snippet(
            r#"class UserService {
                async getById(id: string): Promise<User> { return null; }
                static create() { return new UserService(); }
            }"#,
        );

        // Act
        let cls = &root.children[0];

        // Assert
        assert_eq!(cls.tag, "class", "should be class");
        assert_eq!(cls.name, "UserService", "should be named UserService");
        assert_eq!(cls.children.len(), 2, "should have two methods");

        let get_by_id = &cls.children[0];
        assert_eq!(get_by_id.tag, "method", "should be method");
        assert_eq!(get_by_id.name, "getById", "should be named getById");
        assert_eq!(
            get_by_id.attrs.get("async"),
            Some(&serde_json::Value::Bool(true)),
            "should be async"
        );

        let create = &cls.children[1];
        assert_eq!(create.tag, "method", "should be method");
        assert_eq!(create.name, "create", "should be named create");
        assert_eq!(
            create.attrs.get("static"),
            Some(&serde_json::Value::Bool(true)),
            "should be static"
        );
    }

    #[test]
    fn parses_interfaces_and_types() {
        // Arrange
        let root = parse_snippet(
            r#"
            interface User { id: string; name: string; }
            type UserId = string;
        "#,
        );

        // Act
        let iface = &root.children[0];
        let alias = &root.children[1];

        // Assert
        assert_eq!(iface.tag, "interface", "should be interface");
        assert_eq!(iface.name, "User", "should be named User");
        assert_eq!(
            iface.attrs.get("members"),
            Some(&serde_json::Value::Number(2.into())),
            "should have 2 members"
        );

        assert_eq!(alias.tag, "type", "should be type");
        assert_eq!(alias.name, "UserId", "should be named UserId");
    }

    #[test]
    fn parses_enums() {
        // Arrange
        let root = parse_snippet("enum Role { Admin, User, Guest }");

        // Act
        let e = &root.children[0];

        // Assert
        assert_eq!(e.tag, "enum", "should be enum");
        assert_eq!(e.name, "Role", "should be named Role");
        assert_eq!(
            e.attrs.get("members"),
            Some(&serde_json::Value::Number(3.into())),
            "should have 3 members"
        );
    }

    #[test]
    fn parses_exports() {
        // Arrange
        let root = parse_snippet(
            r#"
            export function fetchUser() {}
            export default class App {}
            export const MAX = 3;
        "#,
        );

        // Act
        let func = &root.children[0];
        let cls = &root.children[1];
        let konst = &root.children[2];

        // Assert
        assert_eq!(func.tag, "function", "should be function");
        assert_eq!(
            func.attrs.get("export"),
            Some(&serde_json::Value::Bool(true)),
            "should be exported"
        );

        assert_eq!(cls.tag, "class", "should be class");
        assert_eq!(
            cls.attrs.get("export"),
            Some(&serde_json::Value::Bool(true)),
            "should be exported"
        );
        assert_eq!(
            cls.attrs.get("default"),
            Some(&serde_json::Value::Bool(true)),
            "should be default export"
        );

        assert_eq!(konst.tag, "const", "should be const");
        assert_eq!(konst.name, "MAX", "should be named MAX");
        assert_eq!(
            konst.attrs.get("export"),
            Some(&serde_json::Value::Bool(true)),
            "should be exported"
        );
    }

    #[test]
    fn parses_tsx_files() {
        // Arrange
        let root = parse_tsx_snippet(
            r#"
            export function App(): JSX.Element {
                return <div>Hello</div>;
            }
        "#,
        );

        // Act
        let func = &root.children[0];

        // Assert
        assert_eq!(func.tag, "function", "should be function");
        assert_eq!(func.name, "App", "should be named App");
        assert_eq!(
            func.attrs.get("export"),
            Some(&serde_json::Value::Bool(true)),
            "should be exported"
        );
    }

    #[test]
    fn parses_constants() {
        // Arrange
        let root = parse_snippet("const MAX_RETRIES = 3;");

        // Act
        let konst = &root.children[0];

        // Assert
        assert_eq!(konst.tag, "const", "should be const");
        assert_eq!(konst.name, "MAX_RETRIES", "should be named MAX_RETRIES");
    }
}