use super::Language;
use super::symbols::{extract_symbols, find_symbol, full_symbol_span};
#[derive(Debug)]
pub struct WrapResult {
pub content: String,
}
pub fn wrap_code(
source: &str,
symbols_arg: Option<&[String]>,
lines_arg: Option<&str>,
wrapper: &str,
preamble: Option<&str>,
lang: Language,
) -> anyhow::Result<WrapResult> {
let mode_count = symbols_arg.is_some() as u8 + lines_arg.is_some() as u8;
if mode_count != 1 {
anyhow::bail!("exactly one of 'symbols' or 'lines' must be specified");
}
let eol = crate::write::detect_eol(source);
let source_lines: Vec<&str> = source.lines().collect();
let (start_idx, end_idx) = if let Some(symbol_names) = symbols_arg {
find_symbol_range(source, symbol_names, lang)?
} else {
parse_line_range_to_indices(
lines_arg.expect("mode_count==1 guarantees Some"),
source_lines.len(),
)?
};
let target_lines = &source_lines[start_idx..end_idx];
let base_indent = if start_idx < source_lines.len() {
let line = source_lines[start_idx];
let trimmed = line.trim_start();
&line[..line.len() - trimmed.len()]
} else {
""
};
let mut wrapped = String::new();
let wrapper_opening = format_wrapper_opening(wrapper);
wrapped.push_str(base_indent);
wrapped.push_str(&wrapper_opening);
wrapped.push_str(eol);
if let Some(pre) = preamble {
for pline in pre.lines() {
if pline.trim().is_empty() {
wrapped.push_str(eol);
} else {
wrapped.push_str(base_indent);
wrapped.push_str(" ");
wrapped.push_str(pline.trim());
wrapped.push_str(eol);
}
}
wrapped.push_str(eol);
}
for line in target_lines {
if line.trim().is_empty() {
wrapped.push_str(eol);
} else {
wrapped.push_str(" ");
wrapped.push_str(line);
wrapped.push_str(eol);
}
}
wrapped.push_str(base_indent);
wrapped.push('}');
wrapped.push_str(eol);
let mut result = String::new();
for line in &source_lines[..start_idx] {
result.push_str(line);
result.push_str(eol);
}
result.push_str(&wrapped);
for line in &source_lines[end_idx..] {
result.push_str(line);
result.push_str(eol);
}
if !source.ends_with('\n') && result.ends_with('\n') {
result.truncate(result.len() - eol.len());
}
Ok(WrapResult { content: result })
}
fn format_wrapper_opening(wrapper: &str) -> String {
let trimmed = wrapper.trim();
if trimmed.ends_with('{') {
trimmed.to_string()
} else {
format!("{trimmed} {{")
}
}
fn find_symbol_range(
source: &str,
symbol_names: &[String],
lang: Language,
) -> anyhow::Result<(usize, usize)> {
if symbol_names.is_empty() {
anyhow::bail!("symbols list must not be empty");
}
let symbols = extract_symbols(source, lang);
let mut min_start = usize::MAX;
let mut max_end = 0usize;
for name in symbol_names {
let sym = find_symbol(&symbols, name)
.ok_or_else(|| anyhow::anyhow!("symbol '{}' not found", name))?;
let (full_start, full_end) = full_symbol_span(source, sym, lang);
let start = full_start.saturating_sub(1);
let end = full_end;
if start < min_start {
min_start = start;
}
if end > max_end {
max_end = end;
}
}
Ok((min_start, max_end))
}
fn parse_line_range_to_indices(spec: &str, total_lines: usize) -> anyhow::Result<(usize, usize)> {
let parts: Vec<&str> = if spec.contains(':') {
spec.splitn(2, ':').collect()
} else if spec.contains('-') && !spec.starts_with('-') {
spec.splitn(2, '-').collect()
} else {
anyhow::bail!("invalid line range format: '{spec}' (expected START:END)");
};
if parts.len() != 2 {
anyhow::bail!("invalid line range: '{spec}'");
}
let start: usize = parts[0]
.parse()
.map_err(|_| anyhow::anyhow!("invalid start line: {}", parts[0]))?;
if start == 0 {
anyhow::bail!("line numbers are 1-based, got 0");
}
let start_idx = start - 1;
if parts[1].is_empty() {
return Ok((start_idx, total_lines));
}
let end: usize = parts[1]
.parse()
.map_err(|_| anyhow::anyhow!("invalid end line: {}", parts[1]))?;
if end < start {
anyhow::bail!("end line {end} is before start line {start}");
}
let end_idx = end.min(total_lines);
Ok((start_idx, end_idx))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn wrap_symbols_in_module() {
let source = "fn helper_a() {}\n\nfn helper_b() {}\n";
let result = wrap_code(
source,
Some(&["helper_a".into(), "helper_b".into()]),
None,
"pub(crate) mod helpers",
None,
Language::Rust,
)
.unwrap();
assert!(result.content.contains("pub(crate) mod helpers {"));
assert!(result.content.contains(" fn helper_a()"));
assert!(result.content.contains(" fn helper_b()"));
assert!(
result.content.ends_with("}\n") || result.content.ends_with('}'),
"module wrapper should close with }}: {}",
result.content
);
}
#[test]
fn wrap_with_preamble() {
let source = "fn test_a() {}\n\nfn test_b() {}\n";
let result = wrap_code(
source,
Some(&["test_a".into(), "test_b".into()]),
None,
"mod tests",
Some("use super::*;"),
Language::Rust,
)
.unwrap();
assert!(result.content.contains("mod tests {"));
assert!(result.content.contains(" use super::*;"));
let use_pos = result.content.find("use super::*").unwrap();
let fn_pos = result.content.find("fn test_a").unwrap();
assert!(use_pos < fn_pos);
}
#[test]
fn wrap_lines_range() {
let source = "line1\nline2\nline3\nline4\nline5\n";
let result = wrap_code(
source,
None,
Some("2:4"),
"#[cfg(test)]",
None,
Language::Rust,
)
.unwrap();
assert!(result.content.contains("#[cfg(test)] {"));
assert!(result.content.contains(" line2"));
assert!(result.content.contains(" line3"));
assert!(result.content.contains(" line4"));
assert!(result.content.starts_with("line1\n"));
assert!(result.content.contains("}\nline5"));
}
#[test]
fn wrap_preserves_indentation() {
let source = " fn inner_a() {}\n fn inner_b() {}\n";
let result = wrap_code(
source,
Some(&["inner_a".into(), "inner_b".into()]),
None,
"mod nested",
None,
Language::Rust,
)
.unwrap();
assert!(result.content.contains(" mod nested {"));
assert!(result.content.contains(" fn inner_a()"));
}
#[test]
fn wrap_symbol_not_found() {
let source = "fn foo() {}\n";
let result = wrap_code(
source,
Some(&["nonexistent".into()]),
None,
"mod wrapper",
None,
Language::Rust,
);
assert!(result.is_err());
assert!(result.unwrap_err().to_string().contains("not found"));
}
#[test]
fn wrap_requires_exactly_one_mode() {
let source = "fn foo() {}\n";
let result = wrap_code(source, None, None, "mod x", None, Language::Rust);
assert!(result.is_err());
let result = wrap_code(
source,
Some(&["foo".into()]),
Some("1:1"),
"mod x",
None,
Language::Rust,
);
assert!(result.is_err());
}
#[test]
fn wrap_includes_doc_comments_and_attributes() {
let source = "/// Helper docs\n#[inline]\nfn helper() {\n 42\n}\n";
let result = wrap_code(
source,
Some(&["helper".into()]),
None,
"mod utils",
None,
Language::Rust,
)
.unwrap();
let wrapper_pos = result.content.find("mod utils {").unwrap();
let doc_pos = result.content.find("/// Helper docs").unwrap();
let attr_pos = result.content.find("#[inline]").unwrap();
assert!(
doc_pos > wrapper_pos,
"doc comment should be inside wrapper: {}",
result.content
);
assert!(
attr_pos > wrapper_pos,
"attribute should be inside wrapper: {}",
result.content
);
}
#[test]
fn wrap_preserves_crlf_line_endings() {
let source = "fn helper_a() {}\r\n\r\nfn helper_b() {}\r\n";
let result = wrap_code(
source,
Some(&["helper_a".into(), "helper_b".into()]),
None,
"mod helpers",
None,
Language::Rust,
)
.unwrap();
let bytes = result.content.as_bytes();
for (i, &b) in bytes.iter().enumerate() {
if b == b'\n' {
assert!(
i > 0 && bytes[i - 1] == b'\r',
"bare LF at byte {i} in: {:?}",
result.content
);
}
}
}
#[test]
fn wrap_cfg_attribute() {
let source = "fn feature_a() {\n // experimental\n}\n";
let result = wrap_code(
source,
Some(&["feature_a".into()]),
None,
"#[cfg(feature = \"experimental\")]",
None,
Language::Rust,
)
.unwrap();
assert!(
result
.content
.contains("#[cfg(feature = \"experimental\")] {")
);
}
}