Skip to main content

project_map_cli_rust/core/
parser.rs

1use tree_sitter::{Parser, Query, QueryCursor};
2use streaming_iterator::StreamingIterator;
3use std::fs;
4use std::path::Path;
5use crate::error::{AppError, Result};
6use serde::{Serialize, Deserialize};
7
8#[derive(Debug, Serialize, Deserialize, Clone)]
9pub struct Symbol {
10    pub name: String,
11    pub kind: String,
12    pub line: usize,
13    pub start_byte: usize,
14    pub end_byte: usize,
15}
16
17#[derive(Debug, Serialize, Deserialize)]
18pub struct FileOutline {
19    pub path: String,
20    pub language: String,
21    pub symbols: Vec<Symbol>,
22    pub imports: Vec<String>,
23}
24
25pub struct CodeParser {
26    parser: Parser,
27}
28
29impl CodeParser {
30    pub fn new() -> Self {
31        Self {
32            parser: Parser::new(),
33        }
34    }
35
36    pub fn parse_file(&mut self, path: &Path) -> Result<FileOutline> {
37        let extension = path.extension()
38            .and_then(|s| s.to_str())
39            .unwrap_or("");
40
41        let (language, ts_language) = match extension {
42            "py" => ("python", tree_sitter_python::LANGUAGE.into()),
43            "rs" => ("rust", tree_sitter_rust::LANGUAGE.into()),
44            "ts" | "tsx" => ("typescript", tree_sitter_typescript::LANGUAGE_TYPESCRIPT.into()),
45            "kt" => ("kotlin", tree_sitter_kotlin_ng::LANGUAGE.into()),
46            "sql" => ("sql", tree_sitter_sequel::LANGUAGE.into()),
47            "vue" => ("vue", tree_sitter_vue_updated::language().into()),
48            _ => return Err(AppError::Parser(format!("Unsupported extension: {}", extension))),
49        };
50
51        self.parser.set_language(&ts_language)
52            .map_err(|e| AppError::Parser(format!("Failed to set language: {}", e)))?;
53
54        let content = fs::read_to_string(path)?;
55        let tree = self.parser.parse(&content, None)
56            .ok_or_else(|| AppError::Parser("Failed to parse file".to_string()))?;
57
58        let query_str = match language {
59            "python" => "((class_definition name: (identifier) @name) @class)
60                         ((function_definition name: (identifier) @name) @function)
61                         (import_statement (dotted_name) @import)
62                         (import_from_statement module_name: (dotted_name) @import)",
63            "rust" => "((struct_item name: (type_identifier) @name) @struct)
64                       ((enum_item name: (type_identifier) @name) @enum)
65                       ((function_item name: (identifier) @name) @function)
66                       ((trait_item name: (type_identifier) @name) @trait)
67                       ((impl_item type: (_) @name) @impl)",
68            "typescript" => "((class_declaration name: (type_identifier) @name) @class)
69                             ((function_declaration name: (identifier) @name) @function)
70                             ((interface_declaration name: (type_identifier) @name) @interface)
71                             ((type_alias_declaration name: (type_identifier) @name) @type)
72                             ((method_definition name: (property_identifier) @name) @function)
73                             (import_statement source: (string (string_fragment) @import))",
74            "kotlin" => "((class_declaration) @name) @class
75                         ((function_declaration) @name) @function
76                         ((_) @import)",
77            "sql" => "((identifier) @name) @symbol",
78            "vue" => "((tag_name) @name) @component",
79            _ => unreachable!(),
80        };
81
82        let query = Query::new(&ts_language, query_str)
83            .map_err(|e| AppError::Parser(format!("Failed to create query: {}", e)))?;
84        
85        let mut cursor = QueryCursor::new();
86        let mut matches = cursor.matches(&query, tree.root_node(), content.as_bytes());
87
88        let mut symbols = Vec::new();
89        let mut imports = Vec::new();
90        while let Some(m) = matches.next() {
91            let mut name = String::new();
92            let mut kind = String::new();
93            let mut line = 0;
94            let mut start_byte = 0;
95            let mut end_byte = 0;
96            let mut is_import = false;
97
98            for capture in m.captures {
99                let capture_name = query.capture_names()[capture.index as usize].to_string();
100                if capture_name == "import" {
101                    let imp = capture.node.utf8_text(content.as_bytes())
102                        .unwrap_or("")
103                        .to_string();
104                    if !imp.is_empty() {
105                        imports.push(imp);
106                    }
107                    is_import = true;
108                    break;
109                } else if capture_name == "name" {
110                    name = capture.node.utf8_text(content.as_bytes())
111                        .unwrap_or("unknown")
112                        .to_string();
113                } else {
114                    kind = capture_name;
115                    line = capture.node.start_position().row + 1;
116                    start_byte = capture.node.start_byte();
117                    end_byte = capture.node.end_byte();
118                }
119            }
120            
121            if !is_import && !name.is_empty() && !kind.is_empty() {
122                // Clean up name: remove excessive whitespace and truncate
123                let mut clean_name = name.replace("\n", " ")
124                    .split_whitespace()
125                    .collect::<Vec<_>>()
126                    .join(" ");
127                
128                if clean_name.chars().count() > 100 {
129                    clean_name = format!("{}...", clean_name.chars().take(97).collect::<String>());
130                }
131
132                symbols.push(Symbol {
133                    name: clean_name,
134                    kind,
135                    line,
136                    start_byte,
137                    end_byte,
138                });
139            }
140        }
141
142        // For Vue, always add a component symbol based on the filename
143        if language == "vue" {
144            let file_name = path.file_name().and_then(|s| s.to_str()).unwrap_or("Component");
145            symbols.push(Symbol {
146                name: file_name.trim_end_matches(".vue").to_string(),
147                kind: "component".to_string(),
148                line: 1,
149                start_byte: 0,
150                end_byte: content.len(),
151            });
152        }
153
154        Ok(FileOutline {
155            path: path.to_string_lossy().to_string(),
156            language: language.to_string(),
157            symbols,
158            imports,
159        })
160    }
161}