use std::io::{Read, Write};
use std::process::{Command, Stdio};
use std::sync::mpsc;
use std::thread;
use std::time::Duration;
use crate::config::CodeLang;
use crate::parser::check_pragma;
use crate::sentence::SentenceSplitter;
pub const FORMATTER_TIMEOUT_SECS: u64 = 30;
pub fn reflow_code_body(
lang: &str,
body: &str,
cfg: &CodeLang,
splitter: &dyn SentenceSplitter,
format_code: bool,
) -> String {
#[cfg(feature = "treesitter")]
let body: &str = &{
let frozen = frozen_lines(body, cfg);
crate::ts_comments::reflow_grammar_comments(lang, body, cfg, splitter, &frozen)
.unwrap_or_else(|| body.to_string())
};
#[cfg(not(feature = "treesitter"))]
let _ = lang;
let after_comment_reflow = reflow_comments(body, cfg, splitter);
if format_code {
if let Some(ref argv) = cfg.formatter {
match run_formatter(&after_comment_reflow, argv) {
Ok(out) => return out,
Err(diag) => {
eprintln!("snapper: {diag}");
return after_comment_reflow;
}
}
}
}
after_comment_reflow
}
fn reflow_comments(body: &str, cfg: &CodeLang, splitter: &dyn SentenceSplitter) -> String {
let lines = crate::parser::iter_lines(body);
let mut out = String::with_capacity(body.len());
let mut i = 0;
let mut pragma_off = false;
while i < lines.len() {
let line = lines[i];
let slice = &body[line.start..line.end];
if let Some(on) = check_pragma_for(line.text, cfg) {
pragma_off = !on;
out.push_str(slice);
i += 1;
continue;
}
if pragma_off {
out.push_str(slice);
i += 1;
continue;
}
if let Some(ref pair) = cfg.block_comment {
let [open, close] = [pair[0].as_str(), pair[1].as_str()];
if !open.is_empty() {
if let Some((indent, after_open)) = split_at_marker(line.text, open) {
let trimmed_after = after_open.trim_start();
if !close.is_empty() {
if let Some(idx) = find_close(trimmed_after, close, cfg) {
let interior = &trimmed_after[..idx];
emit_block_comment(&mut out, indent, open, close, interior, splitter);
i += 1;
continue;
}
}
let mut interior = after_open.to_string();
let mut close_indent: Option<String> = None;
let mut closed_at = None;
for (j, next) in lines.iter().enumerate().skip(i + 1) {
if let Some(idx) = find_close(next.text, close, cfg) {
let pre = &next.text[..idx];
let pre_trim = pre.trim();
if !pre_trim.is_empty() {
if !interior.is_empty() && !interior.ends_with(' ') {
interior.push(' ');
}
interior.push_str(pre_trim);
}
close_indent = Some(
next.text[..next.text.len() - next.text.trim_start().len()]
.to_string(),
);
closed_at = Some(j);
break;
}
let stripped = next.text.trim_start();
let stripped = stripped
.strip_prefix("* ")
.or_else(|| stripped.strip_prefix('*'))
.unwrap_or(stripped);
if !interior.is_empty() && !interior.ends_with(' ') {
interior.push(' ');
}
interior.push_str(stripped.trim());
}
if let Some(j) = closed_at {
let ci = close_indent.unwrap_or_else(|| indent.to_string());
emit_block_comment_multi(
&mut out,
indent,
open,
close,
&ci,
interior.trim(),
splitter,
);
i = j + 1;
continue;
}
for keep in &lines[i..] {
out.push_str(&body[keep.start..keep.end]);
}
break;
}
}
}
if let Some(ref marker) = cfg.line_comment {
if let Some((indent, marker, rest)) = strip_line_comment(line.text, marker) {
let prose = rest.trim();
if prose.is_empty() {
out.push_str(slice);
i += 1;
continue;
}
let sentences = splitter.split(prose);
if sentences.len() <= 1 {
out.push_str(slice);
i += 1;
continue;
}
for (k, s) in sentences.iter().enumerate() {
out.push_str(indent);
out.push_str(marker);
out.push(' ');
out.push_str(s);
if k + 1 < sentences.len() {
out.push('\n');
} else {
out.push_str(
&body[line.terminator_span().start..line.terminator_span().end],
);
}
}
i += 1;
continue;
}
}
if let Some(ref marker) = cfg.line_comment {
if let Some(at) = trailing_comment_at(line.text, marker, cfg) {
if let Some(rewritten) = rewrite_trailing(line.text, at, marker, splitter) {
out.push_str(&rewritten);
out.push_str(&body[line.terminator_span().start..line.terminator_span().end]);
i += 1;
continue;
}
}
}
out.push_str(slice);
i += 1;
}
out
}
#[cfg(feature = "treesitter")]
fn frozen_lines(body: &str, cfg: &CodeLang) -> std::collections::HashSet<usize> {
let mut frozen = std::collections::HashSet::new();
let mut off = false;
for (i, line) in body.lines().enumerate() {
if let Some(on) = check_pragma_for(line, cfg) {
frozen.insert(i);
off = !on;
continue;
}
if off {
frozen.insert(i);
}
}
frozen
}
fn check_pragma_for(line: &str, cfg: &CodeLang) -> Option<bool> {
if let Some(b) = check_pragma(line) {
return Some(b);
}
let trimmed = line.trim();
if let Some(ref marker) = cfg.line_comment {
if let Some(rest) = trimmed.strip_prefix(marker.as_str()) {
let rest = rest.trim();
if rest == "snapper:off" {
return Some(false);
}
if rest == "snapper:on" {
return Some(true);
}
}
}
None
}
fn trailing_comment_at(line: &str, marker: &str, cfg: &CodeLang) -> Option<usize> {
if marker.is_empty() {
return None;
}
let quotes = cfg.quote_chars();
let escape = cfg.escape_char();
let block_open = cfg
.block_comment
.as_ref()
.map(|pair| pair[0].as_str())
.unwrap_or("");
let bytes = line.as_bytes();
let mut in_string: Option<char> = None;
let mut i = 0;
let mut seen_code = false;
while i < bytes.len() {
let rest = &line[i..];
let ch = rest.chars().next()?;
match in_string {
Some(delim) => {
if ch == escape {
i += ch.len_utf8();
if let Some(next) = line[i..].chars().next() {
i += next.len_utf8();
}
continue;
}
if ch == delim {
in_string = None;
}
}
None => {
if quotes.contains(&ch) {
in_string = Some(ch);
seen_code = true;
} else if !block_open.is_empty() && rest.starts_with(block_open) {
return None;
} else if rest.starts_with(marker) {
return if seen_code { Some(i) } else { None };
} else if !ch.is_whitespace() {
seen_code = true;
}
}
}
i += ch.len_utf8();
}
None
}
fn rewrite_trailing(
line: &str,
at: usize,
marker: &str,
splitter: &dyn SentenceSplitter,
) -> Option<String> {
let (code, comment) = line.split_at(at);
let (_, found, rest) = strip_line_comment(comment, marker)?;
let prose = rest.trim();
if prose.is_empty() {
return None;
}
let sentences = splitter.split(prose);
if sentences.len() < 2 {
return None;
}
let pad: String = code
.chars()
.map(|c| if c == '\t' { '\t' } else { ' ' })
.collect();
let mut out = String::with_capacity(line.len() + sentences.len() * (pad.len() + 4));
for (i, sentence) in sentences.iter().enumerate() {
if i == 0 {
out.push_str(code);
} else {
out.push('\n');
out.push_str(&pad);
}
out.push_str(found);
out.push(' ');
out.push_str(sentence);
}
Some(out)
}
fn find_close(line: &str, close: &str, cfg: &CodeLang) -> Option<usize> {
if close.is_empty() {
return None;
}
let quotes = cfg.quote_chars();
if close.chars().all(|c| quotes.contains(&c)) {
return line.find(close);
}
let escape = cfg.escape_char();
let bytes = line.as_bytes();
let mut in_string: Option<char> = None;
let mut i = 0;
while i < bytes.len() {
let rest = &line[i..];
let ch = rest.chars().next()?;
match in_string {
Some(delim) => {
if ch == escape {
i += ch.len_utf8();
if let Some(next) = line[i..].chars().next() {
i += next.len_utf8();
}
continue;
}
if ch == delim {
in_string = None;
}
}
None => {
if rest.starts_with(close) {
return Some(i);
}
if quotes.contains(&ch) {
in_string = Some(ch);
}
}
}
i += ch.len_utf8();
}
None
}
fn split_at_marker<'a>(line: &'a str, marker: &str) -> Option<(&'a str, &'a str)> {
let leading = line.len() - line.trim_start().len();
let (indent, rest) = line.split_at(leading);
rest.strip_prefix(marker).map(|after| (indent, after))
}
fn strip_line_comment<'a>(line: &'a str, marker: &str) -> Option<(&'a str, &'a str, &'a str)> {
let leading = line.len() - line.trim_start().len();
let (indent, rest) = line.split_at(leading);
rest.strip_prefix(marker)?;
let mut end = marker.len();
if let Some(last) = marker.chars().last() {
let bytes = rest.as_bytes();
while end < bytes.len() && bytes[end] == last as u8 {
end += 1;
}
if marker.len() > 1 && end < bytes.len() && bytes[end] == b'!' {
end += 1;
}
}
let (found, after) = rest.split_at(end);
let after = after.strip_prefix(' ').unwrap_or(after);
Some((indent, found, after))
}
fn emit_block_comment(
out: &mut String,
indent: &str,
open: &str,
close: &str,
interior: &str,
splitter: &dyn SentenceSplitter,
) {
out.push_str(indent);
out.push_str(open);
out.push('\n');
let sentences = splitter.split(interior.trim());
for s in &sentences {
out.push_str(indent);
out.push(' ');
out.push_str(s);
out.push('\n');
}
out.push_str(indent);
out.push_str(close);
out.push('\n');
}
fn emit_block_comment_multi(
out: &mut String,
indent: &str,
open: &str,
close: &str,
close_indent: &str,
interior: &str,
splitter: &dyn SentenceSplitter,
) {
out.push_str(indent);
out.push_str(open);
out.push('\n');
let sentences = splitter.split(interior);
for s in &sentences {
out.push_str(indent);
out.push(' ');
out.push_str(s);
out.push('\n');
}
out.push_str(close_indent);
out.push_str(close);
out.push('\n');
}
pub fn run_formatter(body: &str, argv: &[String]) -> Result<String, String> {
if argv.is_empty() {
return Err("formatter argv is empty".to_string());
}
let mut cmd = Command::new(&argv[0]);
cmd.args(&argv[1..])
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::piped());
let mut child = match cmd.spawn() {
Ok(c) => c,
Err(e) => {
if e.kind() == std::io::ErrorKind::NotFound {
return Err(format!("formatter not found: {}", argv[0]));
}
return Err(format!("formatter spawn failed: {}: {e}", argv[0]));
}
};
if let Some(mut stdin) = child.stdin.take() {
let body_owned = body.to_string();
let _ = thread::spawn(move || {
let _ = stdin.write_all(body_owned.as_bytes());
});
}
let (done_tx, done_rx) = mpsc::channel::<()>();
let child_id = child.id();
let watchdog = thread::spawn(move || {
match done_rx.recv_timeout(Duration::from_secs(FORMATTER_TIMEOUT_SECS)) {
Ok(_) | Err(mpsc::RecvTimeoutError::Disconnected) => {
}
Err(mpsc::RecvTimeoutError::Timeout) => {
#[cfg(unix)]
unsafe {
libc_kill(child_id as i32);
}
#[cfg(not(unix))]
{
let _ = std::process::Command::new("taskkill")
.args(["/F", "/PID", &child_id.to_string()])
.output();
}
}
}
});
let output = child.wait_with_output();
let _ = done_tx.send(());
let _ = watchdog.join();
let output = match output {
Ok(o) => o,
Err(e) => return Err(format!("formatter wait failed: {}: {e}", argv[0])),
};
if !output.status.success() {
let stderr = String::from_utf8_lossy(&output.stderr);
return Err(format!(
"formatter {} exited non-zero (status {:?}): {}",
argv[0],
output.status.code(),
stderr.trim()
));
}
String::from_utf8(output.stdout)
.map_err(|e| format!("formatter {} produced non-UTF-8 output: {e}", argv[0]))
}
#[cfg(unix)]
unsafe fn libc_kill(pid: i32) {
unsafe extern "C" {
fn kill(pid: i32, sig: i32) -> i32;
}
const SIGKILL: i32 = 9;
unsafe {
let _ = kill(pid, SIGKILL);
}
}
#[doc(hidden)]
pub fn read_to_string(mut r: impl Read) -> std::io::Result<String> {
let mut s = String::new();
r.read_to_string(&mut s)?;
Ok(s)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::sentence::unicode::UnicodeSentenceSplitter;
fn rust_cfg() -> CodeLang {
CodeLang {
line_comment: Some("//".to_string()),
block_comment: Some(["/*".to_string(), "*/".to_string()]),
..Default::default()
}
}
#[test]
fn line_comment_two_sentences_split() {
let body = "// First sentence. Second sentence.\nfn main() {}\n";
let out = reflow_comments(body, &rust_cfg(), &UnicodeSentenceSplitter::new());
assert_eq!(
out,
"// First sentence.\n// Second sentence.\nfn main() {}\n"
);
}
#[test]
fn indented_comment_preserved() {
let body = " // First. Second.\n fn x() {}\n";
let out = reflow_comments(body, &rust_cfg(), &UnicodeSentenceSplitter::new());
assert_eq!(out, " // First.\n // Second.\n fn x() {}\n");
}
#[test]
fn non_comment_passes_through() {
let body = "fn main() { println!(\"hi\"); }\n";
let out = reflow_comments(body, &rust_cfg(), &UnicodeSentenceSplitter::new());
assert_eq!(out, body);
}
#[test]
fn block_comment_one_liner_splits() {
let body = "/* First. Second. */\n";
let out = reflow_comments(body, &rust_cfg(), &UnicodeSentenceSplitter::new());
assert_eq!(out, "/*\n First.\n Second.\n*/\n");
}
#[test]
fn pragma_freezes_run() {
let body = "// snapper:off\n// Long.\n// Off.\n// snapper:on\n// Reflow this. Now.\n";
let out = reflow_comments(body, &rust_cfg(), &UnicodeSentenceSplitter::new());
let expected = concat!(
"// snapper:off\n",
"// Long.\n",
"// Off.\n",
"// snapper:on\n",
"// Reflow this.\n",
"// Now.\n",
);
assert_eq!(out, expected);
}
}