use crate::transform::structure::extract_markdown_headers_with_spans;
use crate::transform::truncate::NodeSpan;
use crate::transform::utils::{to_static_node_kind, FunctionNodeTypes};
use crate::{Language, Result, SkimError};
use tree_sitter::{Node, Tree};
const MAX_AST_DEPTH: usize = 500;
const MAX_SIGNATURES: usize = 10_000;
#[cfg(test)]
#[allow(dead_code)] pub(crate) fn transform_signatures(
source: &str,
tree: &Tree,
language: Language,
config: &crate::TransformConfig,
) -> Result<String> {
let (text, _spans) = transform_signatures_with_spans(source, tree, language, config)?;
Ok(text)
}
pub(crate) fn transform_signatures_with_spans(
source: &str,
tree: &Tree,
language: Language,
_config: &crate::TransformConfig,
) -> Result<(String, Vec<NodeSpan>)> {
if language == Language::Markdown {
return extract_markdown_headers_with_spans(source, tree, 1, 6);
}
let node_types = get_signature_node_types(language).ok_or_else(|| {
SkimError::ParseError(format!(
"Language {:?} does not support tree-sitter signature transformation",
language
))
})?;
let mut signatures: Vec<(String, &'static str)> = Vec::new();
collect_signatures_with_kinds(tree.root_node(), source, &node_types, &mut signatures, 0)?;
if signatures.len() > MAX_SIGNATURES {
return Err(SkimError::ParseError(format!(
"Too many signatures: {} (max: {}). Possible malicious input.",
signatures.len(),
MAX_SIGNATURES
)));
}
let mut spans = Vec::with_capacity(signatures.len());
let mut current_line = 0;
let texts: Vec<String> = signatures
.into_iter()
.map(|(sig, kind)| {
let line_count = sig.lines().count().max(1);
spans.push(NodeSpan::new(current_line..current_line + line_count, kind));
current_line += line_count;
sig
})
.collect();
Ok((texts.join("\n"), spans))
}
fn collect_signatures_with_kinds(
node: Node,
source: &str,
node_types: &SignatureNodeTypes,
signatures: &mut Vec<(String, &'static str)>,
depth: usize,
) -> Result<()> {
if depth > MAX_AST_DEPTH {
return Err(SkimError::ParseError(format!(
"Maximum AST depth exceeded: {} (possible malicious input)",
MAX_AST_DEPTH
)));
}
let kind = node.kind();
if is_signature_node(kind, node_types) {
if let Some(sig) = extract_signature(node, source, node_types)? {
let static_kind = to_static_node_kind(kind);
signatures.push((sig, static_kind));
}
}
let mut cursor = node.walk();
for child in node.children(&mut cursor) {
collect_signatures_with_kinds(child, source, node_types, signatures, depth + 1)?;
}
Ok(())
}
fn is_signature_node(kind: &str, node_types: &SignatureNodeTypes) -> bool {
kind == node_types.function
|| kind == node_types.method
|| kind == "arrow_function"
|| kind == "function_expression"
|| node_types.extra_function_kinds.contains(&kind)
}
fn extract_signature(
node: Node,
source: &str,
_node_types: &SignatureNodeTypes,
) -> Result<Option<String>> {
let body_node = find_body_for_signature(node);
let end_pos = if let Some(body) = body_node {
body.start_byte()
} else {
node.end_byte()
};
let start = node.start_byte();
if end_pos < start || end_pos > source.len() {
return Ok(None);
}
if !source.is_char_boundary(start) || !source.is_char_boundary(end_pos) {
return Err(SkimError::ParseError(format!(
"Invalid UTF-8 boundary at signature range [{}, {})",
start, end_pos
)));
}
let signature = source[start..end_pos].trim();
if signature.is_empty() {
return Ok(None);
}
Ok(Some(signature.to_string()))
}
fn find_body_for_signature(node: Node) -> Option<Node> {
crate::transform::utils::find_body_child(node)
}
type SignatureNodeTypes = FunctionNodeTypes;
fn get_signature_node_types(language: Language) -> Option<SignatureNodeTypes> {
match language {
Language::TypeScript | Language::JavaScript => Some(SignatureNodeTypes {
function: "function_declaration",
method: "method_definition",
extra_function_kinds: &[],
}),
Language::Python => Some(SignatureNodeTypes {
function: "function_definition",
method: "function_definition",
extra_function_kinds: &[],
}),
Language::Rust => Some(SignatureNodeTypes {
function: "function_item",
method: "function_item",
extra_function_kinds: &[],
}),
Language::Go => Some(SignatureNodeTypes {
function: "function_declaration",
method: "method_declaration",
extra_function_kinds: &[],
}),
Language::Java => Some(SignatureNodeTypes {
function: "method_declaration",
method: "method_declaration",
extra_function_kinds: &[],
}),
Language::Markdown => Some(SignatureNodeTypes {
function: "atx_heading",
method: "atx_heading",
extra_function_kinds: &[],
}),
Language::C | Language::Cpp => Some(SignatureNodeTypes {
function: "function_definition",
method: "function_definition",
extra_function_kinds: &[],
}),
Language::CSharp => Some(SignatureNodeTypes {
function: "method_declaration",
method: "constructor_declaration",
extra_function_kinds: &[],
}),
Language::Ruby => Some(SignatureNodeTypes {
function: "method",
method: "singleton_method",
extra_function_kinds: &[],
}),
Language::Sql => Some(SignatureNodeTypes {
function: "create_table",
method: "create_index",
extra_function_kinds: &[],
}),
Language::Kotlin => Some(SignatureNodeTypes {
function: "function_declaration",
method: "function_declaration",
extra_function_kinds: &["secondary_constructor"],
}),
Language::Swift => Some(SignatureNodeTypes {
function: "function_declaration",
method: "function_declaration",
extra_function_kinds: &["init_declaration"],
}),
Language::Json | Language::Yaml | Language::Toml => None,
}
}