use crate::transform::compute_line_starts;
use crate::transform::minimal::{MAX_AST_DEPTH, MAX_AST_NODES};
use crate::transform::truncate::NodeSpan;
use crate::transform::utils::{FunctionNodeTypes, to_static_node_kind};
use crate::{Language, Result, SkimError, TransformConfig};
use std::collections::HashMap;
use tree_sitter::{Node, Tree};
const MAX_MARKDOWN_HEADERS: usize = 10_000;
#[cfg(test)]
#[allow(dead_code)] pub(crate) fn transform_structure(
source: &str,
tree: &Tree,
language: Language,
config: &TransformConfig,
) -> Result<String> {
let (text, _spans) = transform_structure_with_spans(source, tree, language, config)?;
Ok(text)
}
pub(crate) fn transform_structure_with_spans(
source: &str,
tree: &Tree,
language: Language,
config: &TransformConfig,
) -> Result<(String, Vec<NodeSpan>)> {
let (text, spans, _line_map) =
transform_structure_with_spans_and_line_map(source, tree, language, config)?;
Ok((text, spans))
}
pub(crate) fn transform_structure_with_spans_and_line_map(
source: &str,
tree: &Tree,
language: Language,
_config: &TransformConfig,
) -> Result<(String, Vec<NodeSpan>, Vec<usize>)> {
if language == Language::Markdown {
let (text, spans, line_map) = extract_markdown_headers_with_spans(source, tree, 1, 3)?;
return Ok((text, spans, line_map));
}
let node_types = get_node_types_for_language(language).ok_or_else(|| {
SkimError::ParseError(format!(
"Language {:?} does not support tree-sitter structure transformation",
language
))
})?;
let mut replacements: HashMap<(usize, usize), &'static str> = HashMap::new();
collect_body_replacements(tree.root_node(), &node_types, &mut replacements, 0)?;
if replacements.len() > MAX_AST_NODES {
return Err(SkimError::ParseError(format!(
"Too many AST nodes: {} (max: {}). Possible malicious input.",
replacements.len(),
MAX_AST_NODES
)));
}
let estimated_capacity = source.len() + (replacements.len() * 20);
let mut result = String::with_capacity(estimated_capacity);
let mut last_pos = 0;
let mut sorted_replacements: Vec<_> = replacements.into_iter().collect();
sorted_replacements.sort_unstable_by_key(|(range, _)| range.0);
let mut offset_delta: i64 = 0;
let mut offset_map: Vec<(usize, i64)> = Vec::new();
for ((start, end), replacement) in sorted_replacements {
if end < start {
return Err(SkimError::ParseError(format!(
"Invalid AST range: start={} end={}",
start, end
)));
}
if end > source.len() {
return Err(SkimError::ParseError(format!(
"AST range exceeds source length: end={} len={}",
end,
source.len()
)));
}
if start < last_pos {
continue;
}
if !source.is_char_boundary(start) || !source.is_char_boundary(end) {
return Err(SkimError::ParseError(format!(
"Invalid UTF-8 boundary at range [{}, {})",
start, end
)));
}
result.push_str(&source[last_pos..start]);
result.push_str(replacement);
let replaced_len = end - start;
let replacement_len = replacement.len();
offset_delta += replacement_len as i64 - replaced_len as i64;
offset_map.push((end, offset_delta));
last_pos = end;
}
if !source.is_char_boundary(last_pos) {
return Err(SkimError::ParseError(format!(
"Invalid UTF-8 boundary at position {}",
last_pos
)));
}
result.push_str(&source[last_pos..]);
let spans = build_spans_from_top_level_nodes(tree, &result, &offset_map);
let source_line_map = compute_source_line_map_from_offset_map(source, &result, &offset_map);
Ok((result, spans, source_line_map))
}
pub(crate) fn compute_source_line_map_from_offset_map(
source: &str,
output: &str,
offset_map: &[(usize, i64)],
) -> Vec<usize> {
if output.is_empty() {
return Vec::new();
}
let source_line_starts: Vec<usize> = compute_line_starts(source.as_bytes());
let output_line_starts: Vec<usize> = compute_line_starts(output.as_bytes());
let output_lines = if output.ends_with('\n') {
output_line_starts.len().saturating_sub(1)
} else {
output_line_starts.len()
};
let mut cursor_idx = 0usize; let mut applicable_delta = 0i64;
output_line_starts
.iter()
.take(output_lines)
.map(|&output_byte| {
while cursor_idx < offset_map.len() {
let (src_end, d) = offset_map[cursor_idx];
let output_end = src_end as i64 + d;
if output_end < 0 {
cursor_idx += 1;
continue;
}
if output_end as usize <= output_byte {
applicable_delta = d;
cursor_idx += 1;
} else {
break;
}
}
let source_byte =
((output_byte as i64 - applicable_delta).max(0) as usize).min(source.len());
match source_line_starts.binary_search(&source_byte) {
Ok(idx) => idx + 1, Err(idx) => idx.max(1), }
})
.collect()
}
fn collect_body_replacements(
node: Node,
node_types: &NodeTypes,
replacements: &mut HashMap<(usize, usize), &'static str>,
depth: usize,
) -> Result<()> {
if depth > MAX_AST_DEPTH {
return Err(SkimError::ParseError(format!(
"Maximum AST depth exceeded: {} (possible malicious input with deeply nested functions)",
MAX_AST_DEPTH
)));
}
let kind = node.kind();
if matches_function_node(kind, node_types)
&& let Some(body) = find_body_node(node)
{
let start = body.start_byte();
let end = body.end_byte();
replacements.insert((start, end), " {...}");
}
let mut cursor = node.walk();
for child in node.children(&mut cursor) {
collect_body_replacements(child, node_types, replacements, depth + 1)?;
}
Ok(())
}
fn matches_function_node(kind: &str, node_types: &NodeTypes) -> bool {
kind == node_types.function
|| kind == node_types.method
|| kind == "arrow_function"
|| kind == "function_expression"
|| node_types.extra_function_kinds.contains(&kind)
}
fn find_body_node(node: Node) -> Option<Node> {
crate::transform::utils::find_body_child(node)
}
type NodeTypes = FunctionNodeTypes;
fn get_node_types_for_language(language: Language) -> Option<NodeTypes> {
match language {
Language::TypeScript | Language::JavaScript => Some(NodeTypes {
function: "function_declaration",
method: "method_definition",
extra_function_kinds: &[],
}),
Language::Python => Some(NodeTypes {
function: "function_definition",
method: "function_definition",
extra_function_kinds: &[],
}),
Language::Rust => Some(NodeTypes {
function: "function_item",
method: "function_item",
extra_function_kinds: &[],
}),
Language::Go => Some(NodeTypes {
function: "function_declaration",
method: "method_declaration",
extra_function_kinds: &[],
}),
Language::Java => Some(NodeTypes {
function: "method_declaration",
method: "method_declaration",
extra_function_kinds: &[],
}),
Language::Markdown => Some(NodeTypes {
function: "atx_heading",
method: "atx_heading",
extra_function_kinds: &[],
}),
Language::C | Language::Cpp => Some(NodeTypes {
function: "function_definition",
method: "function_definition",
extra_function_kinds: &[],
}),
Language::CSharp => Some(NodeTypes {
function: "method_declaration",
method: "constructor_declaration",
extra_function_kinds: &[],
}),
Language::Ruby => Some(NodeTypes {
function: "method",
method: "singleton_method",
extra_function_kinds: &[],
}),
Language::Sql => Some(NodeTypes {
function: "statement",
method: "statement",
extra_function_kinds: &[],
}),
Language::Kotlin => Some(NodeTypes {
function: "function_declaration",
method: "function_declaration", extra_function_kinds: &["secondary_constructor", "anonymous_initializer"],
}),
Language::Swift => Some(NodeTypes {
function: "function_declaration",
method: "function_declaration", extra_function_kinds: &["init_declaration", "deinit_declaration"],
}),
Language::Json | Language::Yaml | Language::Toml => None,
}
}
fn build_spans_from_top_level_nodes(
tree: &Tree,
output: &str,
offset_map: &[(usize, i64)],
) -> Vec<NodeSpan> {
let root = tree.root_node();
let mut cursor = root.walk();
let mut spans = Vec::new();
let line_starts = crate::transform::compute_line_starts(output.as_bytes());
let byte_to_line = |byte_pos: usize| -> usize {
match line_starts.binary_search(&byte_pos) {
Ok(idx) => idx,
Err(idx) => idx.saturating_sub(1),
}
};
let source_to_output_byte = |source_byte: usize| -> usize {
let delta = match offset_map.binary_search_by_key(&source_byte, |&(pos, _)| pos) {
Ok(idx) => offset_map[idx].1,
Err(0) => 0,
Err(idx) => offset_map[idx - 1].1,
};
(source_byte as i64 + delta).max(0) as usize
};
for child in root.children(&mut cursor) {
let kind = child.kind();
let source_start = child.start_byte();
let source_end = child.end_byte();
let output_start = source_to_output_byte(source_start).min(output.len());
let output_end = source_to_output_byte(source_end).min(output.len());
let start_line = byte_to_line(output_start);
let end_line = byte_to_line(output_end.saturating_sub(1)) + 1;
let static_kind = to_static_node_kind(kind);
if start_line < end_line {
spans.push(NodeSpan::new(start_line..end_line, static_kind));
}
}
spans
}
#[cfg(test)]
#[allow(dead_code)] pub(crate) fn extract_markdown_headers(
source: &str,
tree: &Tree,
min_level: u32,
max_level: u32,
) -> Result<String> {
let (text, _spans, _line_map) =
extract_markdown_headers_with_spans(source, tree, min_level, max_level)?;
Ok(text)
}
pub(crate) fn extract_markdown_headers_with_spans(
source: &str,
tree: &Tree,
min_level: u32,
max_level: u32,
) -> Result<(String, Vec<NodeSpan>, Vec<usize>)> {
let mut headers: Vec<(String, &'static str, usize)> = Vec::new();
let root = tree.root_node();
let mut visit_stack = vec![(0_usize, root)];
while let Some((depth, node)) = visit_stack.pop() {
if depth > MAX_AST_DEPTH {
return Err(SkimError::ParseError(format!(
"Maximum markdown depth exceeded: {} (possible malicious input)",
MAX_AST_DEPTH
)));
}
if headers.len() > MAX_MARKDOWN_HEADERS {
return Err(SkimError::ParseError(format!(
"Too many markdown headers: {} (max: {}). Possible malicious input.",
headers.len(),
MAX_MARKDOWN_HEADERS
)));
}
let node_type = node.kind();
if node_type == "atx_heading" {
let mut cursor = node.walk();
let marker = node.children(&mut cursor).find(|child| {
child.kind().starts_with("atx_h") && child.kind().ends_with("_marker")
});
if let Some(marker) = marker {
let marker_kind = marker.kind();
let level = marker_kind
.chars()
.find(|c| c.is_ascii_digit())
.and_then(|c| c.to_digit(10))
.unwrap_or(1);
if level >= min_level && level <= max_level {
let header_text = node.utf8_text(source.as_bytes()).map_err(|e| {
SkimError::ParseError(format!("UTF-8 error in header: {}", e))
})?;
let source_start_line = node.start_position().row + 1;
headers.push((header_text.to_string(), "atx_heading", source_start_line));
}
}
} else if node_type == "setext_heading" {
let mut cursor = node.walk();
let underline = node.children(&mut cursor).find(|child| {
let kind = child.kind();
kind == "setext_h1_underline" || kind == "setext_h2_underline"
});
let level = if let Some(underline_node) = underline {
if underline_node.kind() == "setext_h1_underline" {
1
} else {
2
}
} else {
1
};
if level >= min_level && level <= max_level {
let header_text = node.utf8_text(source.as_bytes()).map_err(|e| {
SkimError::ParseError(format!("UTF-8 error in setext header: {}", e))
})?;
let source_start_line = node.start_position().row + 1;
headers.push((header_text.to_string(), "setext_heading", source_start_line));
}
}
let mut child_cursor = node.walk();
for child in node.children(&mut child_cursor) {
visit_stack.push((depth + 1, child));
}
}
let mut spans = Vec::with_capacity(headers.len());
let mut source_line_map: Vec<usize> = Vec::new();
let mut current_output_line = 0;
let texts: Vec<String> = headers
.into_iter()
.map(|(text, kind, source_start_line)| {
let line_count = text.lines().count().max(1);
spans.push(NodeSpan::new(
current_output_line..current_output_line + line_count,
kind,
));
for i in 0..line_count {
source_line_map.push(source_start_line + i);
}
current_output_line += line_count;
text
})
.collect();
Ok((texts.join("\n"), spans, source_line_map))
}
#[cfg(test)]
#[allow(clippy::unwrap_used)] mod offset_map_tests {
use super::compute_source_line_map_from_offset_map;
fn newlines_before(s: &str, end: usize) -> usize {
s[..end].bytes().filter(|&b| b == b'\n').count()
}
#[test]
fn test_no_replacements_identity() {
let source = "line one\nline two\nline three\n";
let output = "line one\nline two\nline three\n";
let offset_map: Vec<(usize, i64)> = vec![];
let map = compute_source_line_map_from_offset_map(source, output, &offset_map);
assert_eq!(map.len(), 3, "expected 3 output lines");
assert_eq!(map[0], 1, "output line 1 -> source line 1");
assert_eq!(map[1], 2, "output line 2 -> source line 2");
assert_eq!(map[2], 3, "output line 3 -> source line 3");
}
#[test]
fn test_single_replacement_shrinks_output() {
let source = "function foo() {\n return 42;\n}\n// end\n";
let body_start: usize = 15; let body_end: usize = 31;
assert_eq!(
&source[body_start..body_end],
"{\n return 42;\n}",
"sanity-check body slice"
);
let repl = " {...}"; let delta: i64 = repl.len() as i64 - (body_end - body_start) as i64;
let output = format!("{}{}{}", &source[..body_start], repl, &source[body_end..]);
assert_eq!(
output.lines().count(),
2,
"output should have 2 lines, got: {:?}",
output
);
let offset_map = vec![(body_end, delta)];
let map = compute_source_line_map_from_offset_map(source, &output, &offset_map);
assert_eq!(map.len(), 2, "expected 2 output lines, got {}", map.len());
assert_eq!(map[0], 1, "output line 1 should map to source line 1");
assert_eq!(
map[1], 4,
"output line 2 ('// end') should map to source line 4"
);
}
#[test]
fn test_multiple_replacements() {
let source = "function a() {\n return 1;\n}\nfunction b() {\n return 2;\n}\n// done\n";
let repl = " {...}";
let a_body_start = "function a() ".len(); let a_body_end = a_body_start + "{\n return 1;\n}".len(); assert_eq!(
&source[a_body_start..a_body_end],
"{\n return 1;\n}",
"sanity-check a body"
);
let delta1: i64 = repl.len() as i64 - (a_body_end - a_body_start) as i64;
let b_sig_start = a_body_end + 1; let b_body_start = b_sig_start + "function b() ".len(); let b_body_end = b_body_start + "{\n return 2;\n}".len(); assert_eq!(
&source[b_body_start..b_body_end],
"{\n return 2;\n}",
"sanity-check b body"
);
let delta2: i64 = delta1 + repl.len() as i64 - (b_body_end - b_body_start) as i64;
let output = format!(
"{}{}{}{}{}",
&source[..a_body_start],
repl,
&source[a_body_end..b_body_start],
repl,
&source[b_body_end..]
);
assert_eq!(
output.lines().count(),
3,
"output should have 3 lines, got: {:?}",
output
);
let offset_map = vec![(a_body_end, delta1), (b_body_end, delta2)];
let map = compute_source_line_map_from_offset_map(source, &output, &offset_map);
assert_eq!(map.len(), 3, "expected 3 output lines, got {}", map.len());
assert_eq!(map[0], 1, "first function -> source line 1");
assert_eq!(map[1], 4, "second function -> source line 4");
let done_src_line = newlines_before(source, source.find("// done").unwrap()) + 1;
assert_eq!(
done_src_line, 7,
"sanity: '// done' should be source line 7"
);
assert_eq!(
map[2], done_src_line,
"'// done' should map to source line {}, got {}",
done_src_line, map[2]
);
}
#[test]
fn test_replacement_at_file_start() {
let source = "/**\n * docs\n */\nexport const X = 1;\n";
let block_end: usize = "/**\n * docs\n */".len(); assert_eq!(&source[..block_end], "/**\n * docs\n */");
let repl = "/* ... */"; let delta: i64 = repl.len() as i64 - block_end as i64;
let output = format!("{}{}", repl, &source[block_end..]);
assert_eq!(output.lines().count(), 2, "output should have 2 lines");
let offset_map = vec![(block_end, delta)];
let map = compute_source_line_map_from_offset_map(source, &output, &offset_map);
assert_eq!(map.len(), 2, "expected 2 output lines, got {}", map.len());
assert_eq!(map[0], 1, "replacement-at-start -> source line 1");
let export_src_line = newlines_before(source, source.find("export").unwrap()) + 1;
assert_eq!(export_src_line, 4, "sanity: export is source line 4");
assert_eq!(
map[1], export_src_line,
"'export const X = 1;' should map to source line {}, got {}",
export_src_line, map[1]
);
}
#[test]
fn test_multiline_signature_line_numbers() {
let source =
"function complex(\n a: number,\n b: string\n) {\n return a;\n}\nconst x = 1;\n";
let repl = " {...}";
let prefix = "function complex(\n a: number,\n b: string\n) ";
let body_start = prefix.len();
let body_len = "{\n return a;\n}".len(); let body_end = body_start + body_len;
assert_eq!(
&source[body_start..body_end],
"{\n return a;\n}",
"sanity-check body slice"
);
let delta: i64 = repl.len() as i64 - body_len as i64;
let output = format!("{}{}{}", &source[..body_start], repl, &source[body_end..]);
assert_eq!(output.lines().count(), 5, "output should have 5 lines");
let offset_map = vec![(body_end, delta)];
let map = compute_source_line_map_from_offset_map(source, &output, &offset_map);
assert_eq!(map.len(), 5, "expected 5 output lines, got {}", map.len());
assert_eq!(map[0], 1, "output line 1 -> source line 1");
assert_eq!(map[1], 2, "output line 2 -> source line 2");
assert_eq!(map[2], 3, "output line 3 -> source line 3");
assert_eq!(map[3], 4, "output line 4 (collapsed body) -> source line 4");
let const_src_line = newlines_before(source, source.find("const x").unwrap()) + 1;
assert_eq!(const_src_line, 7, "sanity: 'const x' is source line 7");
assert_eq!(
map[4], const_src_line,
"'const x = 1;' should map to source line {}, got {}",
const_src_line, map[4]
);
}
}
#[cfg(test)]
#[allow(clippy::unwrap_used)] mod markdown_line_map_tests {
use super::extract_markdown_headers_with_spans;
use crate::{Language, Parser};
fn parse_and_extract_line_map(source: &str) -> Vec<usize> {
let mut parser = Parser::new(Language::Markdown).unwrap();
let tree = parser.parse(source).unwrap();
let (_text, _spans, line_map) =
extract_markdown_headers_with_spans(source, &tree, 1, 6).unwrap();
line_map
}
#[test]
fn test_markdown_line_map_non_contiguous_headers() {
let source = "# Title\n\nSome text.\n\n## Section\n\nMore text.\n\n### Sub\n";
let line_map = parse_and_extract_line_map(source);
assert_eq!(
line_map.len(),
3,
"expected 3 output lines (one per header), got {:?}",
line_map
);
let mut sorted = line_map.clone();
sorted.sort_unstable();
assert_eq!(
sorted,
vec![1, 5, 9],
"line map must contain source lines {{1, 5, 9}}, got {:?}",
line_map
);
assert!(
line_map.contains(&9),
"### Sub is at source line 9 — must appear in line map, got {:?}",
line_map
);
assert!(
line_map.contains(&5),
"## Section is at source line 5 — must appear in line map, got {:?}",
line_map
);
assert!(
line_map.contains(&1),
"# Title is at source line 1 — must appear in line map, got {:?}",
line_map
);
}
#[test]
fn test_markdown_line_map_consecutive_headers() {
let source = "# H1\n## H2\n### H3\n";
let line_map = parse_and_extract_line_map(source);
assert_eq!(line_map.len(), 3, "expected 3 headers, got {:?}", line_map);
assert_eq!(
line_map[0], 3,
"### H3 is at source line 3 and is collected first (DFS stack reversal), \
got line_map[0] = {}. A sequential-index implementation would give 1 here.",
line_map[0]
);
assert_eq!(
line_map[1], 2,
"## H2 is at source line 2, got {}",
line_map[1]
);
assert_eq!(
line_map[2], 1,
"# H1 is at source line 1, got {}",
line_map[2]
);
}
#[test]
fn test_markdown_line_map_single_deep_header() {
let source = "Intro line.\n\nAnother para.\n\n# Deep Header\n\nTrailing text.\n";
let line_map = parse_and_extract_line_map(source);
assert_eq!(line_map.len(), 1, "expected 1 header, got {:?}", line_map);
assert_eq!(
line_map[0], 5,
"# Deep Header is on source line 5, got {}. \
A sequential-position implementation would give 1, not 5.",
line_map[0]
);
}
}