use tree_sitter::{Node, Tree};
pub(super) fn for_each_function(
tree: &Tree,
content: &str,
language: &str,
callback: &mut dyn FnMut(Node<'_>, &str, bool),
) {
visit(tree.root_node(), content, language, callback);
}
fn visit(
node: Node<'_>,
content: &str,
language: &str,
callback: &mut dyn FnMut(Node<'_>, &str, bool),
) {
if let Some((name, is_anonymous)) = function_like(node, language, content) {
callback(node, &name, is_anonymous);
}
let mut cursor = node.walk();
for child in node.children(&mut cursor) {
visit(child, content, language, callback);
}
}
pub(super) fn is_function_node(node: Node<'_>, language: &str) -> bool {
match language {
"Rust" => node.kind() == "function_item",
"Python" => node.kind() == "function_definition",
"Go" => matches!(node.kind(), "function_declaration" | "method_declaration"),
"Java" => matches!(
node.kind(),
"method_declaration" | "constructor_declaration"
),
"CSharp" | "C#" => matches!(
node.kind(),
"method_declaration" | "constructor_declaration" | "local_function_statement"
),
"Kotlin" => node.kind() == "function_declaration",
_ => matches!(
node.kind(),
"function_declaration"
| "generator_function_declaration"
| "method_definition"
| "function_expression"
| "generator_function"
| "arrow_function"
),
}
}
fn function_like(node: Node<'_>, language: &str, content: &str) -> Option<(String, bool)> {
if !is_function_node(node, language) {
return None;
}
match node.kind() {
"function_expression" | "generator_function" => match field_name(node, content) {
Some(name) => Some((name, false)),
None => Some((String::new(), true)),
},
"arrow_function" => match arrow_function_name(node, content) {
Some(name) => Some((name, false)),
None => Some((String::new(), true)),
},
_ => Some((field_name(node, content).unwrap_or_default(), false)),
}
}
fn field_name(node: Node<'_>, content: &str) -> Option<String> {
let name = node.child_by_field_name("name")?;
name.utf8_text(content.as_bytes()).ok().map(str::to_string)
}
fn arrow_function_name(node: Node<'_>, content: &str) -> Option<String> {
let parent = node.parent()?;
let field = match parent.kind() {
"variable_declarator" | "public_field_definition" | "field_definition" => "name",
"pair" => "key",
"assignment_expression" => "left",
_ => return None,
};
let name = parent.child_by_field_name(field)?;
name.utf8_text(content.as_bytes()).ok().map(str::to_string)
}