use std::collections::HashSet;
use tree_sitter::Parser;
use crate::config::CodeLang;
use crate::sentence::SentenceSplitter;
fn language_for(lang: &str) -> Option<tree_sitter::Language> {
let key = lang.trim().to_ascii_lowercase();
let language = match key.as_str() {
"rust" | "rs" => tree_sitter_rust::LANGUAGE.into(),
"python" | "py" => tree_sitter_python::LANGUAGE.into(),
"javascript" | "js" | "jsx" => tree_sitter_javascript::LANGUAGE.into(),
"c" => tree_sitter_c::LANGUAGE.into(),
"cpp" | "c++" | "cxx" | "cc" => tree_sitter_cpp::LANGUAGE.into(),
"go" | "golang" => tree_sitter_go::LANGUAGE.into(),
"bash" | "sh" | "shell" | "zsh" => tree_sitter_bash::LANGUAGE.into(),
"html" => tree_sitter_html::LANGUAGE.into(),
_ => return None,
};
Some(language)
}
pub(crate) fn reflow_grammar_comments(
lang: &str,
body: &str,
cfg: &CodeLang,
splitter: &dyn SentenceSplitter,
frozen: &HashSet<usize>,
) -> Option<String> {
let language = language_for(lang)?;
let mut parser = Parser::new();
parser.set_language(&language).ok()?;
let tree = parser.parse(body, None)?;
let block_open = cfg
.block_comment
.as_ref()
.map(|pair| pair[0].clone())
.unwrap_or_default();
let mut spans: Vec<(usize, usize)> = Vec::new();
let mut cursor = tree.walk();
let mut stack = vec![tree.root_node()];
while let Some(node) = stack.pop() {
if node.kind().contains("comment") {
let text = &body[node.byte_range()];
let is_block = node.kind().contains("block")
|| (!block_open.is_empty() && text.starts_with(&block_open));
if !is_block && !frozen.contains(&node.start_position().row) {
spans.push((node.start_byte(), node.end_byte()));
}
continue;
}
for child in node.children(&mut cursor) {
stack.push(child);
}
}
spans.sort_unstable();
let mut out = body.to_string();
for (start, end) in spans.into_iter().rev() {
let text = &body[start..end];
if text.contains('\n') {
continue;
}
let Some(replacement) = rewrite_one(body, start, text, splitter) else {
continue;
};
out.replace_range(start..end, &replacement);
}
Some(out)
}
fn rewrite_one(
body: &str,
start: usize,
text: &str,
splitter: &dyn SentenceSplitter,
) -> Option<String> {
let marker_len = text
.find(|c: char| c.is_whitespace() || c.is_alphanumeric())
.unwrap_or(text.len());
if marker_len == 0 {
return None;
}
let marker = &text[..marker_len];
let prose = text[marker_len..].trim();
if prose.is_empty() {
return None;
}
let sentences = splitter.split(prose);
if sentences.len() < 2 {
return None;
}
let line_start = body[..start].rfind('\n').map_or(0, |i| i + 1);
let pad: String = body[line_start..start]
.chars()
.map(|c| if c == '\t' { '\t' } else { ' ' })
.collect();
let mut replacement = String::with_capacity(text.len() + sentences.len() * (pad.len() + 4));
for (i, sentence) in sentences.iter().enumerate() {
if i > 0 {
replacement.push('\n');
replacement.push_str(&pad);
}
replacement.push_str(marker);
replacement.push(' ');
replacement.push_str(sentence);
}
Some(replacement)
}