use crate::finding::section::Section;
use crate::finding::test_file_item::TestFileItem;
use crate::finding::test_file_parser::TestFileParser;
use crate::source_file::SourceFile;
use std::collections::BTreeMap;
pub struct TestFileRewriter;
impl TestFileRewriter {
pub const ORDER: [Section; 4] = [
Section::Imports,
Section::Constants,
Section::Helpers,
Section::Tests,
];
pub const REGISTRIES: [&'static str; 2] = ["all_tests.rs", "mod.rs"];
pub const TESTS_ROOT: &'static str = "tests/";
pub fn rewrite(file: &SourceFile) -> Option<String> {
if !Self::governs(file) {
return None;
}
let items = TestFileParser::parse(file)?;
if items.is_empty() {
return None;
}
let rewritten = Self::assemble(file, &items);
(rewritten != file.contents()).then_some(rewritten)
}
fn governs(file: &SourceFile) -> bool {
let path = file.relative_path();
path.starts_with(Self::TESTS_ROOT)
&& !path
.rsplit('/')
.next()
.is_some_and(|name| Self::REGISTRIES.contains(&name))
}
fn assemble(file: &SourceFile, items: &[TestFileItem]) -> String {
let grouped = Self::grouped(items);
let chunks: Vec<String> = Self::ORDER
.iter()
.filter_map(|section| {
let members = grouped.get(section)?;
let gap = "\n".repeat(section.blank_lines_between_entries() + 1);
Some(
members
.iter()
.map(|item| Self::block(file, item))
.collect::<Vec<String>>()
.join(&gap),
)
})
.collect();
let mut out = Self::preamble(file, items);
out.push_str(&chunks.join("\n\n"));
out.push('\n');
let tail = Self::tail(file, items);
if !tail.is_empty() {
out.push('\n');
out.push_str(&tail);
}
out
}
fn preamble(file: &SourceFile, items: &[TestFileItem]) -> String {
let first = items.iter().map(|item| item.first_line).min().unwrap_or(1);
let kept: Vec<&String> = file.lines().iter().take(first - 1).collect();
let trimmed: Vec<&&String> = kept
.iter()
.rev()
.skip_while(|line| line.trim().is_empty())
.collect();
if trimmed.is_empty() {
return String::new();
}
let mut preamble: Vec<String> = trimmed.iter().rev().map(|line| (**line).clone()).collect();
preamble.push(String::new());
preamble.join("\n") + "\n"
}
fn tail(file: &SourceFile, items: &[TestFileItem]) -> String {
let last = items.iter().map(|item| item.last_line).max().unwrap_or(0);
let rest: Vec<String> = file
.lines()
.iter()
.skip(last)
.skip_while(|line| line.trim().is_empty())
.cloned()
.collect();
let Some(end) = rest.iter().rposition(|line| !line.trim().is_empty()) else {
return String::new();
};
rest[..=end].join("\n") + "\n"
}
fn block(file: &SourceFile, item: &TestFileItem) -> String {
file.lines()
.iter()
.skip(item.first_line - 1)
.take(item.last_line + 1 - item.first_line)
.cloned()
.collect::<Vec<String>>()
.join("\n")
}
fn grouped(items: &[TestFileItem]) -> BTreeMap<Section, Vec<&TestFileItem>> {
let mut grouped: BTreeMap<Section, Vec<&TestFileItem>> = BTreeMap::new();
for item in items {
grouped.entry(item.section).or_default().push(item);
}
for (section, members) in grouped.iter_mut() {
if *section != Section::Imports {
members.sort_by_key(|item| item.sort_key());
}
}
grouped
}
}