cs/trace/
function_finder.rs

1use crate::error::{Result, SearchError};
2use crate::parse::Sitter; // Import Sitter
3use crate::search::TextSearcher;
4use regex::Regex;
5use std::collections::HashSet;
6use std::fs;
7use std::path::PathBuf;
8
9/// Represents a function definition found in code
10#[derive(Debug, Clone, PartialEq, Eq, Hash)]
11pub struct FunctionDef {
12    pub name: String,
13    pub file: PathBuf,
14    pub line: usize,
15    pub body: String,
16}
17
18/// Finds function definitions in code using Tree-sitter (primary) and pattern matching (fallback)
19pub struct FunctionFinder {
20    searcher: TextSearcher,
21    patterns: Vec<Regex>,
22    base_dir: PathBuf,
23    sitter: Sitter,
24}
25
26impl FunctionFinder {
27    /// Create a new FunctionFinder
28    ///
29    /// # Arguments
30    /// * `base_dir` - The base directory of the project to search in
31    pub fn new(base_dir: PathBuf) -> Self {
32        Self {
33            searcher: TextSearcher::new(base_dir.clone()),
34            patterns: Self::default_patterns(),
35            base_dir,
36            sitter: Sitter::new(),
37        }
38    }
39
40    /// Default patterns for finding function definitions across languages
41    fn default_patterns() -> Vec<Regex> {
42        vec![
43            // JavaScript/TypeScript - function declarations
44            Regex::new(r"function\s+(\w+)\s*\(").unwrap(),
45            // JavaScript/TypeScript - arrow functions
46            Regex::new(r"(?:const|let|var)\s+(\w+)\s*=\s*(?:async\s+)?\([^)]*\)\s*=>").unwrap(),
47            // JavaScript/TypeScript - method definitions
48            Regex::new(r"^\s*(?:async\s+)?(\w+)\s*\([^)]*\)\s*\{").unwrap(),
49            // JavaScript/TypeScript - export functions
50            Regex::new(r"export\s+function\s+(\w+)").unwrap(),
51            // JavaScript/TypeScript - class methods with modifiers
52            Regex::new(r"^\s*(?:public|private|protected|static|async)\s+(\w+)\s*\(").unwrap(),
53            // Ruby - method definitions
54            Regex::new(r"def\s+(\w+)").unwrap(),
55            // Ruby - class methods
56            Regex::new(r"def\s+self\.(\w+)").unwrap(),
57            // Python - function definitions
58            Regex::new(r"def\s+(\w+)\s*\(").unwrap(),
59            // Rust - function definitions
60            Regex::new(r"fn\s+(\w+)\s*[<(]").unwrap(),
61        ]
62    }
63
64    /// Generate case variants (omitted for brevity, same as before)
65    fn generate_case_variants(func_name: &str) -> Vec<String> {
66        let mut variants = HashSet::new();
67        variants.insert(func_name.to_string());
68        let snake_case = Self::to_snake_case(func_name);
69        variants.insert(snake_case.clone());
70        let camel_case = Self::to_camel_case(&snake_case);
71        variants.insert(camel_case.clone());
72        let pascal_case = Self::to_pascal_case(&snake_case);
73        variants.insert(pascal_case);
74        variants.into_iter().collect()
75    }
76
77    // ... helper methods (to_snake_case, etc.) unchanged ...
78    fn to_snake_case(input: &str) -> String {
79        let mut result = String::new();
80        for (i, ch) in input.chars().enumerate() {
81            if ch.is_uppercase() && i > 0 {
82                result.push('_');
83            }
84            result.push(ch.to_lowercase().next().unwrap());
85        }
86        result
87    }
88
89    fn to_camel_case(input: &str) -> String {
90        let parts: Vec<&str> = input.split('_').collect();
91        if parts.is_empty() {
92            return String::new();
93        }
94        let mut result = parts[0].to_lowercase();
95        for part in parts.iter().skip(1) {
96            if !part.is_empty() {
97                let mut chars = part.chars();
98                if let Some(first) = chars.next() {
99                    result.push(first.to_uppercase().next().unwrap());
100                    result.push_str(&chars.as_str().to_lowercase());
101                }
102            }
103        }
104        result
105    }
106
107    fn to_pascal_case(input: &str) -> String {
108        let parts: Vec<&str> = input.split('_').collect();
109        let mut result = String::new();
110        for part in parts {
111            if !part.is_empty() {
112                let mut chars = part.chars();
113                if let Some(first) = chars.next() {
114                    result.push(first.to_uppercase().next().unwrap());
115                    result.push_str(&chars.as_str().to_lowercase());
116                }
117            }
118        }
119        result
120    }
121
122    /// Find a single function definition, preferring exact matches
123    pub fn find_function(&mut self, func_name: &str) -> Option<FunctionDef> {
124        if let Ok(mut defs) = self.find_definition(func_name) {
125            if let Some(def) = defs.pop() {
126                return Some(def);
127            }
128        }
129        let variants = Self::generate_case_variants(func_name);
130        for variant in variants {
131            if variant != func_name {
132                if let Ok(mut defs) = self.find_definition(&variant) {
133                    if let Some(def) = defs.pop() {
134                        return Some(def);
135                    }
136                }
137            }
138        }
139        None
140    }
141
142    /// Find all definitions of a function by name
143    pub fn find_definition(&mut self, func_name: &str) -> Result<Vec<FunctionDef>> {
144        let mut results = Vec::new();
145
146        // 1. Search for files containing the function name
147        // We still use grep to find candidate files quickly
148        let matches = self.searcher.search(func_name)?;
149
150        // 2. Process each candidate file
151        for m in matches {
152            // Filter out tools/tests (same logic as before)
153            // Convert absolute path to relative path for filtering
154            let relative_path_buf = match m.file.strip_prefix(&self.base_dir) {
155                Ok(rel_path) => rel_path.to_path_buf(),
156                Err(_) => m.file.clone(),
157            };
158            let path_components: Vec<_> = relative_path_buf
159                .components()
160                .map(|c| c.as_os_str().to_string_lossy().to_lowercase())
161                .collect();
162            if !path_components.is_empty() {
163                if path_components[0] == "src" {
164                    continue;
165                }
166                if path_components[0] == "tests"
167                    && (path_components.len() < 2 || path_components[1] != "fixtures")
168                {
169                    continue;
170                }
171            }
172
173            let file_content = fs::read_to_string(&m.file)?;
174
175            // Try Tree-sitter parsing first
176            let is_supported_lang = self.sitter.is_supported(&m.file);
177
178            if is_supported_lang {
179                if let Ok(functions) = self.sitter.find_functions(&m.file, &file_content) {
180                    for func in functions {
181                        if func.name == func_name {
182                            // Simplify body extraction for now - just get from start line
183                            let body = file_content
184                                .lines()
185                                .skip(func.start_line - 1)
186                                .collect::<Vec<_>>()
187                                .join("\n");
188
189                            results.push(FunctionDef {
190                                name: func.name,
191                                file: m.file.clone(),
192                                line: func.start_line,
193                                body,
194                            });
195                        }
196                    }
197                }
198                // If it IS a supported language, we trust Sitter results (even if empty)
199                // and DO NOT fallback to regex, because regex gives false positives (like comments).
200            }
201
202            // Fallback to regex ONLY if language is not supported by Sitter
203            if !is_supported_lang {
204                // Skip ERB files as they are not reliable with regex
205                if m.file.extension().is_some_and(|ext| ext == "erb") {
206                    continue;
207                }
208
209                let content = &m.content;
210                for pattern in &self.patterns {
211                    if let Some(captures) = pattern.captures(content) {
212                        if let Some(name_match) = captures.get(1) {
213                            if name_match.as_str() == func_name {
214                                let body = file_content
215                                    .lines()
216                                    .skip(m.line - 1)
217                                    .collect::<Vec<_>>()
218                                    .join("\n");
219                                results.push(FunctionDef {
220                                    name: func_name.to_string(),
221                                    file: m.file.clone(),
222                                    line: m.line,
223                                    body,
224                                });
225                                break;
226                            }
227                        }
228                    }
229                }
230            }
231        }
232
233        if results.is_empty() {
234            Err(SearchError::Generic(format!(
235                "Function '{}' not found",
236                func_name
237            )))
238        } else {
239            results.sort_by(|a, b| a.file.cmp(&b.file).then(a.line.cmp(&b.line)));
240            Ok(results)
241        }
242    }
243}
244
245impl Default for FunctionFinder {
246    fn default() -> Self {
247        Self::new(std::env::current_dir().unwrap())
248    }
249}
250
251#[cfg(test)]
252mod tests {
253    use super::*;
254
255    #[test]
256    fn test_function_finder_creation() {
257        let finder = FunctionFinder::new(std::env::current_dir().unwrap());
258        assert!(!finder.patterns.is_empty());
259    }
260
261    #[test]
262    fn test_patterns_compile() {
263        let patterns = FunctionFinder::default_patterns();
264        assert_eq!(patterns.len(), 9);
265    }
266
267    #[test]
268    fn test_js_function_pattern() {
269        let patterns = FunctionFinder::default_patterns();
270        let js_pattern = &patterns[0];
271
272        assert!(js_pattern.is_match("function handleClick() {"));
273        assert!(js_pattern.is_match("function processData(x, y) {"));
274        assert!(!js_pattern.is_match("const x = function() {"));
275    }
276
277    #[test]
278    fn test_arrow_function_pattern() {
279        let patterns = FunctionFinder::default_patterns();
280        let arrow_pattern = &patterns[1];
281
282        assert!(arrow_pattern.is_match("const handleClick = () => {"));
283        assert!(arrow_pattern.is_match("let processData = async (x) => {"));
284        assert!(arrow_pattern.is_match("var foo = (a, b) => {"));
285    }
286
287    #[test]
288    fn test_ruby_pattern() {
289        let patterns = FunctionFinder::default_patterns();
290        let ruby_pattern = &patterns[5]; // Updated index for "def \w+" pattern
291
292        assert!(ruby_pattern.is_match("def process_order"));
293        assert!(ruby_pattern.is_match("  def calculate_total"));
294    }
295
296    #[test]
297    fn test_python_pattern() {
298        let patterns = FunctionFinder::default_patterns();
299        let python_pattern = &patterns[7]; // Updated index for Python pattern
300
301        assert!(python_pattern.is_match("def process_data(x):"));
302        assert!(python_pattern.is_match("    def helper():"));
303    }
304
305    #[test]
306    fn test_rust_pattern() {
307        let patterns = FunctionFinder::default_patterns();
308        let rust_pattern = &patterns[8]; // Updated index for Rust pattern
309
310        assert!(rust_pattern.is_match("fn main() {"));
311        assert!(rust_pattern.is_match("fn process<T>(x: T) {"));
312        assert!(rust_pattern.is_match("pub fn calculate("));
313    }
314
315    #[test]
316    fn test_javascript_export_patterns() {
317        let patterns = FunctionFinder::default_patterns();
318        let export_func_pattern = &patterns[3];
319
320        assert!(export_func_pattern.is_match("export function processData"));
321        assert!(export_func_pattern.is_match("export function calculate"));
322    }
323
324    #[test]
325    fn test_javascript_method_patterns() {
326        let patterns = FunctionFinder::default_patterns();
327        let method_pattern = &patterns[2];
328
329        assert!(method_pattern.is_match("  processData() {"));
330        assert!(method_pattern.is_match("    handleClick() {"));
331        assert!(method_pattern.is_match("  async methodName() {"));
332    }
333
334    #[test]
335    fn test_ruby_class_methods() {
336        let patterns = FunctionFinder::default_patterns();
337        let ruby_class_method_pattern = &patterns[6];
338
339        assert!(ruby_class_method_pattern.is_match("def self.create"));
340        assert!(ruby_class_method_pattern.is_match("  def self.find_by_name"));
341    }
342
343    #[test]
344    fn test_case_conversion() {
345        // Test snake_case conversion
346        assert_eq!(FunctionFinder::to_snake_case("createUser"), "create_user");
347        assert_eq!(
348            FunctionFinder::to_snake_case("validateEmailAddress"),
349            "validate_email_address"
350        );
351        assert_eq!(
352            FunctionFinder::to_snake_case("XMLHttpRequest"),
353            "x_m_l_http_request"
354        );
355        assert_eq!(
356            FunctionFinder::to_snake_case("already_snake"),
357            "already_snake"
358        );
359
360        // Test camelCase conversion
361        assert_eq!(FunctionFinder::to_camel_case("create_user"), "createUser");
362        assert_eq!(
363            FunctionFinder::to_camel_case("validate_email_address"),
364            "validateEmailAddress"
365        );
366        assert_eq!(FunctionFinder::to_camel_case("single"), "single");
367
368        // Test PascalCase conversion
369        assert_eq!(FunctionFinder::to_pascal_case("create_user"), "CreateUser");
370        assert_eq!(
371            FunctionFinder::to_pascal_case("validate_email_address"),
372            "ValidateEmailAddress"
373        );
374        assert_eq!(FunctionFinder::to_pascal_case("single"), "Single");
375    }
376
377    #[test]
378    fn test_generate_case_variants() {
379        // Test with camelCase input
380        let variants = FunctionFinder::generate_case_variants("createUser");
381        assert!(variants.contains(&"createUser".to_string()));
382        assert!(variants.contains(&"create_user".to_string()));
383        assert!(variants.contains(&"CreateUser".to_string()));
384
385        // Test with snake_case input
386        let variants = FunctionFinder::generate_case_variants("validate_email");
387        assert!(variants.contains(&"validate_email".to_string()));
388        assert!(variants.contains(&"validateEmail".to_string()));
389        assert!(variants.contains(&"ValidateEmail".to_string()));
390
391        // Test with PascalCase input
392        let variants = FunctionFinder::generate_case_variants("ProcessUserData");
393        assert!(variants.contains(&"ProcessUserData".to_string()));
394        assert!(variants.contains(&"process_user_data".to_string()));
395        assert!(variants.contains(&"processUserData".to_string()));
396    }
397}