| 1 | |
|
| 2 | |
|
| 3 | |
|
| 4 | |
|
| 5 | |
|
| 6 | |
|
| 7 | | #[derive(Clone, Debug, PartialEq, Eq)] |
| 8 | | pub struct Function { |
| 9 | | |
| 10 | | pub name: String, |
| 11 | | pub start: usize, |
| 12 | | pub end: usize, |
| 13 | | } |
| 14 | |
|
| 15 | |
|
| 16 | |
|
| 17 | 14× | fn grammar(path: &str) -> Option<(tree_sitter::Language, &'static [&'static str])> { |
| 18 | 14× | let ext = path.rsplit_once('.').map(|(_, e)| e.to_ascii_lowercase())?; |
| 19 | 13× | Some(match ext.as_str() { |
| 20 | 13× | "rs" => (tree_sitter_rust::LANGUAGE.into(), &["function_item"][..]), |
| 21 | 6× | "ts" | "mts" | "cts" => ( |
| 22 | 1× | tree_sitter_typescript::LANGUAGE_TYPESCRIPT.into(), |
| 23 | 1× | &["function_declaration", "method_definition", "generator_function_declaration"][..], |
| 24 | 1× | ), |
| 25 | 5× | "tsx" => ( |
| 26 | 0× | tree_sitter_typescript::LANGUAGE_TSX.into(), |
| 27 | 0× | &["function_declaration", "method_definition", "generator_function_declaration"][..], |
| 28 | 0× | ), |
| 29 | 5× | "py" => (tree_sitter_python::LANGUAGE.into(), &["function_definition"][..]), |
| 30 | 4× | "sh" | "bash" => (tree_sitter_bash::LANGUAGE.into(), &["function_definition"][..]), |
| 31 | 3× | _ => return None, |
| 32 | | }) |
| 33 | 14× | } |
| 34 | |
|
| 35 | |
|
| 36 | | const CONTAINERS: &[&str] = &["impl_item", "trait_item", "class_declaration", "class_definition", "mod_item"]; |
| 37 | |
|
| 38 | |
|
| 39 | |
|
| 40 | 14× | pub fn functions(path: &str, text: &str) -> Vec<Function> { |
| 41 | 14× | let Some((language, kinds)) = grammar(path) else { return vec![] }; |
| 42 | 10× | let mut parser = tree_sitter::Parser::new(); |
| 43 | 10× | if parser.set_language(&language).is_err() { |
| 44 | 0× | return vec![]; |
| 45 | 10× | } |
| 46 | 10× | let Some(tree) = parser.parse(text, None) else { return vec![] }; |
| 47 | 10× | let mut out = Vec::new(); |
| 48 | 10× | walk(tree.root_node(), text.as_bytes(), kinds, &mut Vec::new(), &mut out); |
| 49 | 10× | out |
| 50 | 14× | } |
| 51 | |
|
| 52 | 22× | fn name_of(node: tree_sitter::Node, src: &[u8]) -> Option<String> { |
| 53 | 22× | let n = node.child_by_field_name("name").or_else(|| node.child_by_field_name("type"))?; |
| 54 | 22× | n.utf8_text(src).ok().map(|s| s.split_whitespace().collect::<Vec<_>>().join(" ")) |
| 55 | 22× | } |
| 56 | |
|
| 57 | 335× | fn walk(node: tree_sitter::Node, src: &[u8], kinds: &[&str], scope: &mut Vec<String>, out: &mut Vec<Function>) { |
| 58 | 335× | let kind = node.kind(); |
| 59 | 335× | let mut pushed = false; |
| 60 | 335× | if kinds.contains(&kind) { |
| 61 | 16× | if let Some(name) = name_of(node, src) { |
| 62 | 16× | let mut full = scope.clone(); |
| 63 | 16× | full.push(name); |
| 64 | 16× | out.push(Function { name: full.join("::"), start: node.start_position().row + 1, end: node.end_position().row + 1 }); |
| 65 | 16× | } |
| 66 | 319× | } else if CONTAINERS.contains(&kind) { |
| 67 | 6× | if let Some(name) = name_of(node, src) { |
| 68 | 6× | scope.push(name); |
| 69 | 6× | pushed = true; |
| 70 | 6× | } |
| 71 | 313× | } |
| 72 | 335× | let mut cursor = node.walk(); |
| 73 | 335× | for child in node.children(&mut cursor) { |
| 74 | 325× | walk(child, src, kinds, scope, out); |
| 75 | 325× | } |
| 76 | 335× | if pushed { |
| 77 | 6× | scope.pop(); |
| 78 | 329× | } |
| 79 | 335× | } |
| 80 | |
|
| 81 | | #[cfg(test)] |
| 82 | | mod tests { |
| 83 | | use super::*; |
| 84 | |
|
| 85 | 5× | fn names(path: &str, text: &str) -> Vec<(String, usize, usize)> { |
| 86 | 8× | functions(path, text).into_iter().map(|f| (f.name, f.start, f.end)).collect() |
| 87 | 5× | } |
| 88 | |
|
| 89 | | #[test] |
| 90 | 1× | fn rust_functions_and_methods() { |
| 91 | 1× | let src = "use x;\n\nfn free() {\n a();\n}\n\nimpl Repo {\n pub fn log(&self) -> u8 {\n 1\n }\n}\n\ntrait T {\n fn t(&self);\n}\n"; |
| 92 | 1× | assert_eq!(names("src/lib.rs", src), [("free".to_string(), 3, 5), ("Repo::log".to_string(), 8, 10)]); |
| 93 | 1× | let generic = "impl<T> Wrap<T> {\n fn get(&self) {}\n}\n"; |
| 94 | 1× | assert_eq!(names("a.rs", generic), [("Wrap<T>::get".to_string(), 2, 2)]); |
| 95 | 1× | } |
| 96 | |
|
| 97 | | #[test] |
| 98 | 1× | fn typescript_python_and_shell() { |
| 99 | 1× | let ts = "export function activate(ctx: X): void {\n go();\n}\nclass Repo {\n refresh(): void {\n x();\n }\n}\n"; |
| 100 | 1× | assert_eq!(names("src/extension.ts", ts), [("activate".to_string(), 1, 3), ("Repo::refresh".to_string(), 5, 7)]); |
| 101 | 1× | let py = "def top():\n pass\n\nclass Recipe:\n def main(self):\n return 1\n"; |
| 102 | 1× | assert_eq!(names("images/rust/recipe.py", py), [("top".to_string(), 1, 2), ("Recipe::main".to_string(), 5, 6)]); |
| 103 | 1× | let sh = "#!/bin/sh\nbuild() {\n cargo build\n}\n"; |
| 104 | 1× | assert_eq!(names("build.sh", sh), [("build".to_string(), 2, 4)]); |
| 105 | 1× | } |
| 106 | |
|
| 107 | | #[test] |
| 108 | 1× | fn other_files_have_none() { |
| 109 | 1× | assert!(functions("README.md", "# fn main() {}\n").is_empty()); |
| 110 | 1× | assert!(functions("Makefile", "all:\n").is_empty()); |
| 111 | 1× | } |
| 112 | | } |