car-ast 0.15.2

use super::{extract_doc_comment, extract_signature, field_text, node_text};
use crate::types::*;

pub fn extract(tree: &tree_sitter::Tree, source: &[u8]) -> (Vec<Symbol>, Vec<Import>) {
    let root = tree.root_node();
    let mut symbols = Vec::new();
    let mut imports = Vec::new();
    let mut cursor = root.walk();

    for child in root.children(&mut cursor) {
        match child.kind() {
            "function_item" => {
                if let Some(sym) = extract_function(&child, source, None) {
                    symbols.push(sym);
                }
            }
            "struct_item" => {
                if let Some(sym) = extract_struct(&child, source) {
                    symbols.push(sym);
                }
            }
            "enum_item" => {
                if let Some(sym) = extract_enum(&child, source) {
                    symbols.push(sym);
                }
            }
            "trait_item" => {
                if let Some(sym) = extract_trait(&child, source) {
                    symbols.push(sym);
                }
            }
            "impl_item" => {
                symbols.push(extract_impl(&child, source));
            }
            "use_declaration" => {
                imports.push(Import {
                    path: node_text(&child, source)
                        .trim_start_matches("use ")
                        .trim_end_matches(';')
                        .to_string(),
                    alias: None,
                    span: Span::from_node(&child),
                });
            }
            "const_item" | "static_item" => {
                if let Some(name) = field_text(&child, "name", source) {
                    symbols.push(Symbol {
                        name: name.to_string(),
                        kind: SymbolKind::Const,
                        span: Span::from_node(&child),
                        signature: extract_signature(&child, "value", source),
                        doc_comment: extract_doc_comment(&child, source),
                        parent: None,
                        children: Vec::new(),
                    });
                }
            }
            "type_item" => {
                if let Some(name) = field_text(&child, "name", source) {
                    symbols.push(Symbol {
                        name: name.to_string(),
                        kind: SymbolKind::TypeAlias,
                        span: Span::from_node(&child),
                        signature: node_text(&child, source).to_string(),
                        doc_comment: extract_doc_comment(&child, source),
                        parent: None,
                        children: Vec::new(),
                    });
                }
            }
            "mod_item" => {
                if let Some(name) = field_text(&child, "name", source) {
                    symbols.push(Symbol {
                        name: name.to_string(),
                        kind: SymbolKind::Module,
                        span: Span::from_node(&child),
                        signature: format!("mod {}", name),
                        doc_comment: extract_doc_comment(&child, source),
                        parent: None,
                        children: Vec::new(),
                    });
                }
            }
            _ => {}
        }
    }

    (symbols, imports)
}

fn extract_function(
    node: &tree_sitter::Node,
    source: &[u8],
    parent: Option<&str>,
) -> Option<Symbol> {
    let name = field_text(node, "name", source)?;
    Some(Symbol {
        name: name.to_string(),
        kind: if parent.is_some() {
            SymbolKind::Method
        } else {
            SymbolKind::Function
        },
        span: Span::from_node(node),
        signature: extract_signature(node, "body", source),
        doc_comment: extract_doc_comment(node, source),
        parent: parent.map(|s| s.to_string()),
        children: Vec::new(),
    })
}

fn extract_struct(node: &tree_sitter::Node, source: &[u8]) -> Option<Symbol> {
    let name = field_text(node, "name", source)?;
    Some(Symbol {
        name: name.to_string(),
        kind: SymbolKind::Struct,
        span: Span::from_node(node),
        signature: extract_signature(node, "body", source),
        doc_comment: extract_doc_comment(node, source),
        parent: None,
        children: Vec::new(),
    })
}

fn extract_enum(node: &tree_sitter::Node, source: &[u8]) -> Option<Symbol> {
    let name = field_text(node, "name", source)?;
    Some(Symbol {
        name: name.to_string(),
        kind: SymbolKind::Enum,
        span: Span::from_node(node),
        signature: extract_signature(node, "body", source),
        doc_comment: extract_doc_comment(node, source),
        parent: None,
        children: Vec::new(),
    })
}

fn extract_trait(node: &tree_sitter::Node, source: &[u8]) -> Option<Symbol> {
    let name = field_text(node, "name", source)?;
    Some(Symbol {
        name: name.to_string(),
        kind: SymbolKind::Trait,
        span: Span::from_node(node),
        signature: extract_signature(node, "body", source),
        doc_comment: extract_doc_comment(node, source),
        parent: None,
        children: Vec::new(),
    })
}

fn extract_impl(node: &tree_sitter::Node, source: &[u8]) -> Symbol {
    // Get the type being implemented
    let type_name = field_text(node, "type", source)
        .or_else(|| {
            // For `impl Trait for Type`, the type is in a different position
            node.child_by_field_name("type")
                .map(|n| n.utf8_text(source).unwrap_or(""))
        })
        .unwrap_or("Unknown")
        .to_string();

    let mut methods = Vec::new();

    if let Some(body) = node.child_by_field_name("body") {
        let mut cursor = body.walk();
        for child in body.children(&mut cursor) {
            if child.kind() == "function_item" {
                if let Some(m) = extract_function(&child, source, Some(&type_name)) {
                    methods.push(m);
                }
            }
        }
    }

    Symbol {
        name: type_name,
        kind: SymbolKind::Impl,
        span: Span::from_node(node),
        signature: extract_signature(node, "body", source),
        doc_comment: extract_doc_comment(node, source),
        parent: None,
        children: methods,
    }
}