peeker 1.1.0

A CLI tool for extracting code structure using Tree-sitter
use anyhow::{Result, bail};
use tree_sitter::Language;

pub struct LanguageConfig {
    pub name: &'static str,
    pub language: Language,
    pub import_kinds: &'static [&'static str],
    pub struct_kinds: &'static [&'static str],
    pub function_kinds: &'static [&'static str],
    pub trait_kinds: &'static [&'static str],
    pub enum_kinds: &'static [&'static str],
    pub field_kinds: &'static [&'static str],
    pub param_kinds: &'static [&'static str],
    pub enum_variant_kinds: &'static [&'static str],
}

pub fn get_language_config(extension: &str) -> Result<LanguageConfig> {
    match extension {
        "rs" => Ok(rust_config()),
        "py" => Ok(python_config()),
        "ts" | "tsx" => Ok(typescript_config()),
        "js" | "jsx" => Ok(javascript_config()),
        "go" => Ok(go_config()),
        "java" => Ok(java_config()),
        "c" | "h" => Ok(c_config()),
        "cpp" | "cc" | "cxx" | "hpp" | "hxx" => Ok(cpp_config()),
        _ => bail!("Unsupported file extension: {}", extension),
    }
}

fn rust_config() -> LanguageConfig {
    LanguageConfig {
        name: "rust",
        language: tree_sitter_rust::LANGUAGE.into(),
        import_kinds: &["use_declaration"],
        struct_kinds: &["struct_item"],
        function_kinds: &["function_item"],
        trait_kinds: &["trait_item"],
        enum_kinds: &["enum_item"],
        field_kinds: &["field_declaration"],
        param_kinds: &["parameter", "self_parameter"],
        enum_variant_kinds: &["enum_variant"],
    }
}

fn python_config() -> LanguageConfig {
    LanguageConfig {
        name: "python",
        language: tree_sitter_python::LANGUAGE.into(),
        import_kinds: &["import_statement", "import_from_statement"],
        struct_kinds: &["class_definition"],
        function_kinds: &["function_definition"],
        trait_kinds: &[], // Python doesn't have traits
        enum_kinds: &[],  // Python enums are classes
        field_kinds: &["assignment"],
        param_kinds: &[
            "parameter",
            "default_parameter",
            "typed_parameter",
            "typed_default_parameter",
        ],
        enum_variant_kinds: &[],
    }
}

fn typescript_config() -> LanguageConfig {
    LanguageConfig {
        name: "typescript",
        language: tree_sitter_typescript::LANGUAGE_TYPESCRIPT.into(),
        import_kinds: &["import_statement"],
        struct_kinds: &[
            "class_declaration",
            "interface_declaration",
            "type_alias_declaration",
        ],
        function_kinds: &[
            "function_declaration",
            "method_definition",
            "arrow_function",
        ],
        trait_kinds: &["interface_declaration"],
        enum_kinds: &["enum_declaration"],
        field_kinds: &["public_field_definition", "property_signature"],
        param_kinds: &["required_parameter", "optional_parameter"],
        enum_variant_kinds: &["enum_assignment"],
    }
}

fn javascript_config() -> LanguageConfig {
    LanguageConfig {
        name: "javascript",
        language: tree_sitter_javascript::LANGUAGE.into(),
        import_kinds: &["import_statement"],
        struct_kinds: &["class_declaration"],
        function_kinds: &[
            "function_declaration",
            "method_definition",
            "arrow_function",
        ],
        trait_kinds: &[],
        enum_kinds: &[],
        field_kinds: &["field_definition"],
        param_kinds: &["formal_parameters"],
        enum_variant_kinds: &[],
    }
}

fn go_config() -> LanguageConfig {
    LanguageConfig {
        name: "go",
        language: tree_sitter_go::LANGUAGE.into(),
        import_kinds: &["import_declaration"],
        struct_kinds: &["type_declaration"],
        function_kinds: &["function_declaration", "method_declaration"],
        trait_kinds: &[], // Go interfaces handled via struct_kinds
        enum_kinds: &[],  // Go doesn't have native enums
        field_kinds: &["field_declaration"],
        param_kinds: &["parameter_declaration"],
        enum_variant_kinds: &[],
    }
}

fn java_config() -> LanguageConfig {
    LanguageConfig {
        name: "java",
        language: tree_sitter_java::LANGUAGE.into(),
        import_kinds: &["import_declaration"],
        struct_kinds: &["class_declaration", "interface_declaration"],
        function_kinds: &["method_declaration", "constructor_declaration"],
        trait_kinds: &["interface_declaration"],
        enum_kinds: &["enum_declaration"],
        field_kinds: &["field_declaration"],
        param_kinds: &["formal_parameter", "spread_parameter"],
        enum_variant_kinds: &["enum_constant"],
    }
}

fn c_config() -> LanguageConfig {
    LanguageConfig {
        name: "c",
        language: tree_sitter_c::LANGUAGE.into(),
        import_kinds: &["preproc_include"],
        struct_kinds: &["struct_specifier"],
        function_kinds: &["function_definition"],
        trait_kinds: &[],
        enum_kinds: &["enum_specifier"],
        field_kinds: &["field_declaration"],
        param_kinds: &["parameter_declaration"],
        enum_variant_kinds: &["enumerator"],
    }
}

fn cpp_config() -> LanguageConfig {
    LanguageConfig {
        name: "cpp",
        language: tree_sitter_cpp::LANGUAGE.into(),
        import_kinds: &["preproc_include"],
        struct_kinds: &["struct_specifier", "class_specifier"],
        function_kinds: &["function_definition"],
        trait_kinds: &[],
        enum_kinds: &["enum_specifier"],
        field_kinds: &["field_declaration"],
        param_kinds: &["parameter_declaration"],
        enum_variant_kinds: &["enumerator"],
    }
}