use super::Language;
use super::symbols::{
SymbolKind, extract_symbol_text, extract_symbols, find_symbol, full_symbol_span,
};
#[derive(Debug)]
pub struct ExtractResult {
pub source_content: String,
pub target_content: String,
pub extracted_lines: usize,
}
pub fn extract_to_file(
source: &str,
symbol: &str,
replacement: Option<&str>,
unwrap: bool,
prepend: Option<&str>,
lang: Language,
) -> anyhow::Result<ExtractResult> {
let eol = crate::write::detect_eol(source);
let symbols = extract_symbols(source, lang);
let sym = find_symbol(&symbols, symbol).ok_or_else(|| {
anyhow::Error::new(crate::exit::NoMatchError {
msg: format!("symbol '{symbol}' not found"),
})
})?;
let (full_start, full_end) = full_symbol_span(source, sym, lang);
let lines: Vec<&str> = source.lines().collect();
let start_0 = full_start.saturating_sub(1);
let end_0 = full_end.min(lines.len());
let target_body = if unwrap && sym.kind == SymbolKind::Module {
unwrap_module_body(&lines, sym.start_line.saturating_sub(1), end_0, eol)
} else {
let text = extract_symbol_text(source, sym, lang);
text.to_string()
};
let extracted_lines = target_body.lines().count();
let mut target_content = String::new();
if let Some(pre) = prepend {
target_content.push_str(pre);
if !pre.ends_with('\n') {
target_content.push_str(eol);
}
target_content.push_str(eol);
}
target_content.push_str(&target_body);
if !target_content.ends_with('\n') {
target_content.push_str(eol);
}
let mut source_lines: Vec<String> = lines.iter().map(|l| l.to_string()).collect();
let mut remove_end = end_0;
if remove_end < source_lines.len() && source_lines[remove_end].trim().is_empty() {
remove_end += 1;
}
source_lines.drain(start_0..remove_end.min(source_lines.len()));
if let Some(repl) = replacement {
let repl_lines: Vec<String> = repl.lines().map(String::from).collect();
for (i, line) in repl_lines.iter().enumerate() {
source_lines.insert(start_0 + i, line.clone());
}
}
let mut source_content = source_lines.join(eol);
if source.ends_with('\n') && !source_content.ends_with('\n') {
source_content.push_str(eol);
}
Ok(ExtractResult {
source_content,
target_content,
extracted_lines,
})
}
fn unwrap_module_body(lines: &[&str], sym_start_0: usize, sym_end_0: usize, eol: &str) -> String {
let mut body_start = sym_start_0;
let mut brace_line_tail: Option<&str> = None;
for (i, line) in lines.iter().enumerate().take(sym_end_0).skip(sym_start_0) {
let trimmed = line.trim();
if trimmed.contains('{') {
if let Some(open) = line.find('{') {
let after = line[open + 1..].trim();
if !after.is_empty() && !after.starts_with('}') {
brace_line_tail = Some(line[open + 1..].trim_end());
}
}
body_start = i + 1;
break;
}
}
let body_end = if sym_end_0 > 0 {
sym_end_0 - 1
} else {
sym_end_0
};
if body_start >= body_end {
if body_start > 0 {
let brace_line = lines[body_start - 1];
if let Some(open) = brace_line.find('{') {
let after_open = &brace_line[open + 1..];
if let Some(close) = after_open.rfind('}') {
let inner = after_open[..close].trim();
if !inner.is_empty() {
return inner.to_string();
}
}
}
}
return String::new();
}
let body_lines = &lines[body_start..body_end];
let min_indent = body_lines
.iter()
.filter(|l| !l.trim().is_empty())
.map(|l| l.len() - l.trim_start().len())
.min()
.unwrap_or(0);
let mut parts: Vec<String> = Vec::new();
if let Some(tail) = brace_line_tail {
let trimmed = tail.trim();
parts.push(trimmed.to_string());
}
parts.extend(body_lines.iter().map(|line| {
if line.trim().is_empty() {
String::new()
} else if line.len() >= min_indent {
line[min_indent..].to_string()
} else {
line.to_string()
}
}));
parts.join(eol)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn extract_basic_function() {
let source = "fn foo() {\n 42\n}\n\nfn bar() {}\n";
let result = extract_to_file(source, "foo", None, false, None, Language::Rust).unwrap();
assert!(!result.source_content.contains("fn foo"));
assert!(result.source_content.contains("fn bar"));
assert!(result.target_content.contains("fn foo()"));
assert!(
result.target_content.contains(" 42"),
"target should contain the function body: {}",
result.target_content
);
}
#[test]
fn extract_with_replacement() {
let source = "mod tests {\n fn test_a() {}\n}\n\nfn other() {}\n";
let replacement = "#[path = \"tests.rs\"]\nmod tests;";
let result = extract_to_file(
source,
"tests",
Some(replacement),
true,
None,
Language::Rust,
)
.unwrap();
assert!(result.source_content.contains("#[path = \"tests.rs\"]"));
assert!(result.source_content.contains("mod tests;"));
assert!(!result.source_content.contains("fn test_a"));
}
#[test]
fn extract_unwrap_module() {
let source = "mod tests {\n fn test_a() {\n assert!(true);\n }\n}\n";
let result = extract_to_file(source, "tests", None, true, None, Language::Rust).unwrap();
assert!(result.target_content.contains("fn test_a()"));
assert!(result.target_content.contains(" assert!(true);"));
assert!(!result.target_content.contains("mod tests"));
}
#[test]
fn extract_with_prepend() {
let source = "fn foo() {}\n\nfn bar() {}\n";
let result = extract_to_file(
source,
"foo",
None,
false,
Some("use super::*;"),
Language::Rust,
)
.unwrap();
assert!(result.target_content.starts_with("use super::*;"));
assert!(result.target_content.contains("fn foo()"));
}
#[test]
fn extract_no_unwrap() {
let source = "mod tests {\n fn test_a() {}\n}\n";
let result = extract_to_file(source, "tests", None, false, None, Language::Rust).unwrap();
assert!(result.target_content.contains("mod tests {"));
}
#[test]
fn extract_symbol_not_found() {
let source = "fn foo() {}\n";
let result = extract_to_file(source, "nonexistent", None, false, None, Language::Rust);
assert!(result.is_err(), "expected error, got Ok: {result:?}");
}
#[test]
fn extract_preserves_attributes() {
let source = "/// Doc comment.\n#[cfg(test)]\nmod tests {\n fn t() {}\n}\n";
let result = extract_to_file(source, "tests", None, false, None, Language::Rust).unwrap();
assert!(result.target_content.contains("/// Doc comment."));
assert!(result.target_content.contains("#[cfg(test)]"));
}
#[test]
fn unwrap_module_with_trailing_comment_on_brace_line() {
let source =
"mod tests { // test module\n fn test_a() {\n assert!(true);\n }\n}\n";
let result = extract_to_file(source, "tests", None, true, None, Language::Rust).unwrap();
assert!(
!result.target_content.contains("mod tests"),
"mod declaration should not appear in unwrapped body: {}",
result.target_content
);
assert!(result.target_content.contains("fn test_a()"));
}
#[test]
fn extract_preserves_crlf_line_endings() {
let source = "fn foo() {\r\n 42\r\n}\r\n\r\nfn bar() {}\r\n";
let result = extract_to_file(source, "foo", None, false, None, Language::Rust).unwrap();
let bytes = result.source_content.as_bytes();
for (i, &b) in bytes.iter().enumerate() {
if b == b'\n' {
assert!(
i > 0 && bytes[i - 1] == b'\r',
"bare LF in source at byte {i}: {:?}",
result.source_content
);
}
}
let target_bytes = result.target_content.as_bytes();
for (i, &b) in target_bytes.iter().enumerate() {
if b == b'\n' {
assert!(
i > 0 && target_bytes[i - 1] == b'\r',
"bare LF in target at byte {i}: {:?}",
result.target_content
);
}
}
}
#[test]
fn unwrap_single_line_module() {
let source = "mod foo { fn bar() {} }\n";
let result = extract_to_file(source, "foo", None, true, None, Language::Rust).unwrap();
assert!(
result.target_content.contains("fn bar()"),
"single-line module body should be extracted, got: {:?}",
result.target_content
);
}
#[test]
fn unwrap_module_preserves_brace_line_content() {
let source = "mod foo { use bar::*;\n fn baz() {}\n}\n";
let result = extract_to_file(source, "foo", None, true, None, Language::Rust).unwrap();
assert!(
result.target_content.contains("use bar::*;"),
"content on the brace line should be preserved, got: {:?}",
result.target_content
);
assert!(
result.target_content.contains("fn baz()"),
"body content should also be present, got: {:?}",
result.target_content
);
}
}