use tower_lsp::lsp_types::{FoldingRange, FoldingRangeKind};
#[must_use]
pub fn folding_ranges(text: &str) -> Vec<FoldingRange> {
let lines: Vec<&str> = text.lines().collect();
if lines.is_empty() {
return Vec::new();
}
let mut ranges = Vec::new();
collect_indentation_folds(&lines, &mut ranges);
collect_document_section_folds(&lines, &mut ranges);
collect_comment_block_folds(&lines, &mut ranges);
ranges
}
struct OpenRegion {
start_line: usize,
indent: usize,
}
fn collect_indentation_folds(lines: &[&str], ranges: &mut Vec<FoldingRange>) {
let mut stack: Vec<OpenRegion> = Vec::new();
for (i, line) in lines.iter().enumerate() {
let trimmed = line.trim();
if trimmed.is_empty() || trimmed == "---" || trimmed == "..." {
continue;
}
if trimmed.starts_with('#') {
let indent = line.len() - line.trim_start().len();
close_regions_at_or_above(&mut stack, indent, i, lines, ranges);
continue;
}
let indent = line.len() - line.trim_start().len();
close_regions_at_or_above(&mut stack, indent, i, lines, ranges);
if starts_fold_region(trimmed) {
stack.push(OpenRegion {
start_line: i,
indent,
});
}
}
let total = lines.len();
stack.into_iter().rev().for_each(|region| {
let end = find_last_content_line(lines, region.start_line, total);
if end > region.start_line {
push_fold(ranges, region.start_line, end, None);
}
});
}
fn close_regions_at_or_above(
stack: &mut Vec<OpenRegion>,
indent: usize,
current_line: usize,
lines: &[&str],
ranges: &mut Vec<FoldingRange>,
) {
while let Some(top) = stack.last() {
if top.indent >= indent {
let Some(region) = stack.pop() else { break };
let end = find_last_content_line(lines, region.start_line, current_line);
if end > region.start_line {
push_fold(ranges, region.start_line, end, None);
}
} else {
break;
}
}
}
fn starts_fold_region(trimmed: &str) -> bool {
if trimmed.ends_with(':') {
return true;
}
if let Some(colon_pos) = find_mapping_colon(trimmed) {
let after_colon = trimmed[colon_pos + 1..].trim();
if after_colon.is_empty() {
return true;
}
if is_block_scalar_indicator(after_colon) {
return true;
}
}
false
}
fn is_block_scalar_indicator(value: &str) -> bool {
let mut chars = value.chars();
let Some(first) = chars.next() else {
return false;
};
if first != '|' && first != '>' {
return false;
}
for ch in chars {
match ch {
'+' | '-' | '0'..='9' => {}
' ' | '\t' | '#' => return true, _ => return false, }
}
true
}
fn find_mapping_colon(line: &str) -> Option<usize> {
let mut in_single_quote = false;
let mut in_double_quote = false;
for (i, ch) in line.char_indices() {
match ch {
'\'' if !in_double_quote => in_single_quote = !in_single_quote,
'"' if !in_single_quote => in_double_quote = !in_double_quote,
':' if !in_single_quote && !in_double_quote => {
let rest = &line[i + 1..];
if rest.is_empty() || rest.starts_with(' ') || rest.starts_with('\t') {
return Some(i);
}
}
_ => {}
}
}
None
}
fn find_last_content_line(lines: &[&str], start: usize, before: usize) -> usize {
((start + 1)..before)
.rev()
.find(|&i| {
lines
.get(i)
.is_some_and(|l| !l.trim().is_empty() && l.trim() != "---" && l.trim() != "...")
})
.unwrap_or(start)
}
fn collect_document_section_folds(lines: &[&str], ranges: &mut Vec<FoldingRange>) {
let separator_positions: Vec<usize> = lines
.iter()
.enumerate()
.filter(|(_, l)| l.trim() == "---")
.map(|(i, _)| i)
.collect();
if separator_positions.is_empty() {
return;
}
if let Some(&first_sep) = separator_positions.first()
&& first_sep > 0
{
let end = find_last_content_line_in_range(lines, 0, first_sep);
if let Some(end) = end
&& end > 0
{
push_fold(ranges, 0, end, Some(FoldingRangeKind::Region));
}
}
for window in separator_positions.windows(2) {
if let [start_sep, before] = window {
let start = start_sep + 1;
if start < *before {
let end = find_last_content_line_in_range(lines, start, *before);
if let Some(end) = end
&& end > start
{
push_fold(ranges, start, end, Some(FoldingRangeKind::Region));
}
}
}
}
let Some(&last_sep) = separator_positions.last() else {
return;
};
let start = last_sep + 1;
if start < lines.len() {
let end = find_last_content_line_in_range(lines, start, lines.len());
if let Some(end) = end
&& end > start
{
push_fold(ranges, start, end, Some(FoldingRangeKind::Region));
}
}
}
fn find_last_content_line_in_range(lines: &[&str], from: usize, before: usize) -> Option<usize> {
(from..before).rev().find(|&i| {
lines
.get(i)
.is_some_and(|l| !l.trim().is_empty() && l.trim() != "---" && l.trim() != "...")
})
}
fn collect_comment_block_folds(lines: &[&str], ranges: &mut Vec<FoldingRange>) {
let mut comment_start: Option<usize> = None;
for (i, line) in lines.iter().enumerate() {
let trimmed = line.trim();
if trimmed.starts_with('#') {
if comment_start.is_none() {
comment_start = Some(i);
}
} else {
if let Some(start) = comment_start
&& i - 1 > start
{
push_fold(ranges, start, i - 1, Some(FoldingRangeKind::Comment));
}
comment_start = None;
}
}
if let Some(start) = comment_start {
let end = lines.len() - 1;
if end > start {
push_fold(ranges, start, end, Some(FoldingRangeKind::Comment));
}
}
}
fn push_fold(
ranges: &mut Vec<FoldingRange>,
start: usize,
end: usize,
kind: Option<FoldingRangeKind>,
) {
#[expect(
clippy::cast_possible_truncation,
reason = "LSP line/col are u32; always fits"
)]
ranges.push(FoldingRange {
start_line: start as u32,
start_character: None,
end_line: end as u32,
end_character: None,
kind,
collapsed_text: None,
});
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use super::*;
use tower_lsp::lsp_types::FoldingRangeKind;
fn ranges_as_tuples(ranges: &[FoldingRange]) -> Vec<(u32, u32)> {
ranges.iter().map(|r| (r.start_line, r.end_line)).collect()
}
#[rstest]
#[case::mapping_with_nested_content("server:\n host: localhost\n port: 8080\n", (0, 2))]
#[case::sequence("items:\n - one\n - two\n - three\n", (0, 3))]
#[case::sequence_of_mappings("users:\n - name: Alice\n age: 30\n - name: Bob\n age: 25\n", (0, 4))]
#[case::literal_block_scalar("description: |\n This is a\n multi-line\n description\n", (0, 3))]
#[case::folded_block_scalar("summary: >\n This is a\n folded\n paragraph\n", (0, 3))]
#[case::comment_within_server_fold("server:\n # This is a comment\n host: localhost\n port: 8080\n", (0, 3))]
#[case::blank_lines_within_fold("server:\n host: localhost\n\n port: 8080\n", (0, 3))]
#[case::mixed_content_types("config:\n name: app\n ports:\n - 80\n - 443\n description: |\n A multi-line\n description\n", (0, 7))]
#[case::block_scalar_with_indentation_indicator("text: |2\n indented content\n more content\n", (0, 2))]
fn folding_ranges_contains_range(#[case] text: &str, #[case] expected: (u32, u32)) {
let result = folding_ranges(text);
let tuples = ranges_as_tuples(&result);
assert!(
tuples.contains(&expected),
"expected tuple {expected:?} in folding ranges, got: {tuples:?}"
);
}
#[rstest]
#[case::empty_document("")]
#[case::single_line_document("key: value")]
#[case::mapping_with_inline_value_only("name: Alice\nage: 30\n")]
#[case::gt_or_pipe_in_value_not_block_scalar("condition: a > b\nresult: true\n")]
fn folding_ranges_returns_empty(#[case] text: &str) {
let result = folding_ranges(text);
assert!(result.is_empty(), "expected empty result, got: {result:?}");
}
#[test]
fn should_fold_document_sections() {
let text = "key1: val1\nkey2: val2\n---\nkey3: val3\nkey4: val4\n";
let result = folding_ranges(text);
assert!(
result.len() >= 2,
"should have at least 2 folding ranges for 2 document sections, got: {}",
result.len()
);
}
#[test]
fn should_fold_document_sections_with_nested_content() {
let text = "doc1:\n key: val\n---\ndoc2:\n key: val\n";
let result = folding_ranges(text);
let tuples = ranges_as_tuples(&result);
assert!(
tuples.contains(&(0, 1)),
"should fold doc1 mapping (lines 0-1), got: {tuples:?}"
);
assert!(
tuples.contains(&(3, 4)),
"should fold doc2 mapping (lines 3-4), got: {tuples:?}"
);
}
#[test]
fn should_fold_consecutive_comment_block() {
let text = "# Header comment\n# continues here\n# and here\nkey: value\n";
let result = folding_ranges(text);
let comment_folds: Vec<&FoldingRange> = result
.iter()
.filter(|r| r.kind == Some(FoldingRangeKind::Comment))
.collect();
let region_folds: Vec<&FoldingRange> = result
.iter()
.filter(|r| r.kind == Some(FoldingRangeKind::Region) || r.kind.is_none())
.collect();
for fold in ®ion_folds {
assert!(
fold.start_line > 2,
"comment block should not produce a Region fold, got region fold starting at line {}",
fold.start_line
);
}
if !comment_folds.is_empty() {
let tuples: Vec<(u32, u32)> = comment_folds
.iter()
.map(|r| (r.start_line, r.end_line))
.collect();
assert!(
tuples.contains(&(0, 2)),
"comment fold should span lines 0-2, got: {tuples:?}"
);
}
}
#[test]
fn should_return_empty_for_comment_only_document() {
let text = "# just a comment\n";
let result = folding_ranges(text);
for fold in &result {
assert!(
fold.kind == Some(FoldingRangeKind::Comment) || fold.kind.is_none(),
"comment-only document should not produce Region folds"
);
}
}
#[test]
fn should_fold_three_document_sections() {
let text = "a: 1\nb: 2\n---\nc: 3\nd: 4\n---\ne: 5\nf: 6\n";
let result = folding_ranges(text);
let region_folds: Vec<_> = result
.iter()
.filter(|r| r.kind == Some(FoldingRangeKind::Region))
.collect();
assert!(
region_folds.len() >= 3,
"three document sections should produce at least 3 region folds, got: {region_folds:?}"
);
}
#[test]
fn should_fold_last_section_after_final_separator() {
let text = "---\nkey1: val1\nkey2: val2\n";
let result = folding_ranges(text);
assert!(
result
.iter()
.any(|r| r.kind == Some(FoldingRangeKind::Region)),
"content after separator should produce a region fold, got: {result:?}"
);
}
#[test]
fn should_fold_comment_block_at_end_of_file() {
let text = "key: value\n# comment line 1\n# comment line 2\n# comment line 3\n";
let result = folding_ranges(text);
let comment_folds: Vec<_> = result
.iter()
.filter(|r| r.kind == Some(FoldingRangeKind::Comment))
.collect();
assert!(
!comment_folds.is_empty(),
"comment block at end of file should produce a Comment fold, got: {result:?}"
);
let tuples: Vec<(u32, u32)> = comment_folds
.iter()
.map(|r| (r.start_line, r.end_line))
.collect();
assert!(
tuples.contains(&(1, 3)),
"comment fold should span lines 1-3, got: {tuples:?}"
);
}
#[test]
fn should_fold_block_scalar_with_chomping_indicator() {
let text =
"a: |-\n content line 1\n content line 2\nb: >+\n folded line 1\n folded line 2\n";
let result = folding_ranges(text);
let tuples = ranges_as_tuples(&result);
assert!(
tuples.contains(&(0, 2)),
"should fold '|-' block scalar (lines 0-2), got: {tuples:?}"
);
assert!(
tuples.contains(&(3, 5)),
"should fold '>+' block scalar (lines 3-5), got: {tuples:?}"
);
}
#[test]
fn should_not_fold_section_consisting_only_of_blank_lines() {
let text = "a: 1\n---\n\n\n---\nb: 2\n";
let result = folding_ranges(text);
let region_folds: Vec<_> = result
.iter()
.filter(|r| r.kind == Some(FoldingRangeKind::Region))
.collect();
for fold in ®ion_folds {
assert!(
fold.start_line != 2 && fold.start_line != 3,
"blank-only section should not produce a region fold, got: {fold:?}"
);
}
}
#[test]
fn should_not_fold_single_comment_line() {
let text = "# only one comment\nkey: value\n";
let result = folding_ranges(text);
let comment_folds: Vec<_> = result
.iter()
.filter(|r| r.kind == Some(FoldingRangeKind::Comment))
.collect();
assert!(
comment_folds.is_empty(),
"single comment line should not fold, got: {comment_folds:?}"
);
}
}