use std::collections::HashMap;
use crate::nlp::{
contains_page_number_token, ends_with_block, ends_with_hyphen, is_page_number_line,
is_terminal_punct, last_line_is_list_item, last_significant_char, looks_like_boundary_header,
looks_like_numbered_running_header, looks_like_running_header, starts_with_block,
};
use crate::prompt::Exclude;
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub enum StitchAction {
None,
Space,
Hyphen,
}
pub struct PostProcessState {
header_counts: HashMap<String, usize>,
footer_counts: HashMap<String, usize>,
}
impl PostProcessState {
pub fn new() -> Self {
Self {
header_counts: HashMap::new(),
footer_counts: HashMap::new(),
}
}
pub fn process(&mut self, text: &str, excludes: &[Exclude]) -> String {
let lines: Vec<String> = text.split('\n').map(|l| l.to_string()).collect();
if lines.is_empty() {
return String::new();
}
let mut keep = vec![true; lines.len()];
let first_non_empty = lines.iter().position(|l| !l.trim().is_empty());
let last_non_empty = lines.iter().rposition(|l| !l.trim().is_empty());
if has_exclude(excludes, Exclude::PageNumbers) {
for i in 0..lines.len() {
if !keep[i] {
continue;
}
let trimmed = lines[i].trim();
if trimmed.is_empty() {
continue;
}
if is_page_number_line(trimmed) {
let prev_blank = i == 0 || lines[i - 1].trim().is_empty();
let next_blank = i + 1 >= lines.len() || lines[i + 1].trim().is_empty();
let is_boundary = Some(i) == first_non_empty || Some(i) == last_non_empty;
if is_boundary || (prev_blank && next_blank) {
keep[i] = false;
}
}
}
}
remove_standalone_running_artifacts(&lines, &mut keep, excludes);
let mut first_kept = None;
let mut last_kept = None;
for i in 0..lines.len() {
if keep[i] && !lines[i].trim().is_empty() {
if first_kept.is_none() {
first_kept = Some(i);
}
last_kept = Some(i);
}
}
let header_candidate = first_kept
.map(|i| lines[i].trim().to_string())
.filter(|s| !s.is_empty());
let footer_candidate = last_kept
.map(|i| lines[i].trim().to_string())
.filter(|s| !s.is_empty());
if has_exclude(excludes, Exclude::Headers)
&& let (Some(i), Some(candidate)) = (first_kept, header_candidate.as_ref())
&& (self.header_counts.get(candidate).copied().unwrap_or(0) >= 1
|| contains_page_number_token(candidate)
|| looks_like_running_header(candidate))
{
keep[i] = false;
}
if has_exclude(excludes, Exclude::Footers)
&& let (Some(i), Some(candidate)) = (last_kept, footer_candidate.as_ref())
&& (self.footer_counts.get(candidate).copied().unwrap_or(0) >= 1
|| contains_page_number_token(candidate))
{
keep[i] = false;
}
if let Some(candidate) = header_candidate {
*self.header_counts.entry(candidate).or_insert(0) += 1;
}
if let Some(candidate) = footer_candidate {
*self.footer_counts.entry(candidate).or_insert(0) += 1;
}
let mut output = String::new();
for (i, line) in lines.iter().enumerate() {
if !keep[i] {
continue;
}
if !output.is_empty() {
output.push('\n');
}
output.push_str(line.trim_end());
}
output
}
}
pub fn stitch_decision(prev: &str, next: &str) -> StitchAction {
let prev_trim = prev.trim_end();
let next_trim = next.trim_start();
if prev_trim.is_empty() || next_trim.is_empty() {
return StitchAction::None;
}
if starts_with_block(next_trim) {
return StitchAction::None;
}
if ends_with_block(prev_trim) {
return StitchAction::None;
}
if ends_with_hyphen(prev_trim) && !last_line_is_list_item(prev_trim) {
return StitchAction::Hyphen;
}
let prev_last = last_significant_char(prev_trim).unwrap_or(' ');
if is_terminal_punct(prev_last) {
return StitchAction::None;
}
let next_first = next_trim.chars().next().unwrap_or(' ');
if next_first.is_lowercase() {
return StitchAction::Space;
}
if matches!(prev_last, ',' | ';' | ':' | '—' | '–') {
return StitchAction::Space;
}
if matches!(next_first, ')' | ']' | '}' | ',' | '.' | ';' | ':') {
return StitchAction::Space;
}
StitchAction::None
}
pub fn merge_stitched(prev: &str, next: &str, action: StitchAction) -> String {
let prev_trim = prev.trim_end();
let next_trim = next.trim_start();
match action {
StitchAction::None => prev_trim.to_string(),
StitchAction::Space => format!("{} {}", prev_trim, next_trim),
StitchAction::Hyphen => {
let trimmed = if let Some(stripped) = prev_trim.strip_suffix('-') {
stripped
} else if let Some(stripped) = prev_trim.strip_suffix('‐') {
stripped
} else if let Some(stripped) = prev_trim.strip_suffix('‑') {
stripped
} else {
prev_trim
};
format!("{}{}", trimmed, next_trim)
}
}
}
fn has_exclude(excludes: &[Exclude], target: Exclude) -> bool {
excludes.contains(&target)
}
fn remove_standalone_running_artifacts(lines: &[String], keep: &mut [bool], excludes: &[Exclude]) {
let remove_headers = has_exclude(excludes, Exclude::Headers);
let remove_page_numbers = has_exclude(excludes, Exclude::PageNumbers);
if !remove_headers && !remove_page_numbers {
return;
}
for i in 0..lines.len() {
if !keep[i] {
continue;
}
let line = lines[i].trim();
if line.is_empty() || starts_with_block(line) {
continue;
}
let prev_blank = i == 0 || lines[i - 1].trim().is_empty();
let next_blank = i + 1 >= lines.len() || lines[i + 1].trim().is_empty();
if !(prev_blank && next_blank) {
continue;
}
if remove_page_numbers && looks_like_numbered_running_header(line) {
keep[i] = false;
continue;
}
if remove_headers && looks_like_running_header(line) && contains_page_number_token(line) {
keep[i] = false;
}
}
}
pub fn strip_boundary_artifacts(prev: &str, next: &str, _excludes: &[Exclude]) -> String {
let prev_trim = prev.trim_end();
let Some(prev_last) = last_significant_char(prev_trim) else {
return next.to_string();
};
if is_terminal_punct(prev_last) {
return next.to_string();
}
let mut lines: Vec<String> = next.split('\n').map(|line| line.to_string()).collect();
let Some(first_non_empty) = lines.iter().position(|line| !line.trim().is_empty()) else {
return String::new();
};
let candidate = lines[first_non_empty].trim();
if !looks_like_boundary_header(candidate) {
return next.to_string();
}
let Some(next_non_empty) = lines
.iter()
.skip(first_non_empty + 1)
.position(|line| !line.trim().is_empty())
.map(|idx| idx + first_non_empty + 1)
else {
return next.to_string();
};
let next_first = lines[next_non_empty]
.trim_start()
.chars()
.next()
.unwrap_or(' ');
if !next_first.is_lowercase() {
return next.to_string();
}
lines.remove(first_non_empty);
while lines.first().is_some_and(|line| line.trim().is_empty()) {
lines.remove(0);
}
lines.join("\n")
}
#[cfg(test)]
mod tests {
use super::{
PostProcessState, StitchAction, merge_stitched, stitch_decision, strip_boundary_artifacts,
};
use crate::prompt::Exclude;
#[test]
fn removes_standalone_roman_page_numbers() {
let mut state = PostProcessState::new();
let input = "Paragraph one.\n\nvii\n\nParagraph two.";
let output = state.process(input, &[Exclude::PageNumbers]);
assert!(!output.contains("\nvii\n"));
assert!(output.contains("Paragraph one."));
assert!(output.contains("Paragraph two."));
}
#[test]
fn removes_numbered_running_header_lines() {
let mut state = PostProcessState::new();
let input = "Paragraph one.\n\nFREDERICK THE GREAT AND NAPOLEON 9\n\nParagraph two.";
let output = state.process(input, &[Exclude::Headers, Exclude::PageNumbers]);
assert!(!output.contains("FREDERICK THE GREAT AND NAPOLEON 9"));
}
#[test]
fn strips_boundary_header_inside_sentence_flow() {
let prev = "find the measure";
let next = "Introduction\n\nof their dwelling. If man is dwelling, then...";
let output = strip_boundary_artifacts(prev, next, &[]);
assert!(!output.contains("Introduction"));
assert!(output.starts_with("of their dwelling."));
}
#[test]
fn keeps_boundary_heading_after_sentence_end() {
let prev = "This section ends cleanly.";
let next = "Introduction\n\nThis chapter opens with context.";
let output = strip_boundary_artifacts(prev, next, &[]);
assert_eq!(output, next);
}
#[test]
fn stitch_space_when_next_starts_lowercase() {
let decision = stitch_decision("find the measure", "of their dwelling");
assert_eq!(decision, StitchAction::Space);
let merged = merge_stitched("find the measure", "of their dwelling", decision);
assert_eq!(merged, "find the measure of their dwelling");
}
#[test]
fn stitch_none_when_prev_empty() {
assert_eq!(stitch_decision("", "next text"), StitchAction::None);
assert_eq!(stitch_decision(" ", "next text"), StitchAction::None);
}
#[test]
fn stitch_none_when_next_empty() {
assert_eq!(stitch_decision("prev text", ""), StitchAction::None);
assert_eq!(stitch_decision("prev text", " "), StitchAction::None);
}
#[test]
fn stitch_none_when_next_starts_with_block() {
assert_eq!(
stitch_decision("some text", "# Heading"),
StitchAction::None
);
assert_eq!(stitch_decision("some text", "```code"), StitchAction::None);
assert_eq!(stitch_decision("some text", "> quote"), StitchAction::None);
assert_eq!(
stitch_decision("some text", "- list item"),
StitchAction::None
);
assert_eq!(
stitch_decision("some text", "1. numbered"),
StitchAction::None
);
}
#[test]
fn stitch_none_when_prev_ends_with_block() {
assert_eq!(
stitch_decision("text</pre>", "next text"),
StitchAction::None
);
assert_eq!(
stitch_decision("text</table>", "next text"),
StitchAction::None
);
assert_eq!(stitch_decision("---", "next text"), StitchAction::None);
}
#[test]
fn stitch_hyphen_ascii() {
let decision = stitch_decision("compre-", "hending");
assert_eq!(decision, StitchAction::Hyphen);
let merged = merge_stitched("compre-", "hending", decision);
assert_eq!(merged, "comprehending");
}
#[test]
fn stitch_hyphen_unicode() {
let decision = stitch_decision("compre\u{2010}", "hending");
assert_eq!(decision, StitchAction::Hyphen);
let merged = merge_stitched("compre\u{2010}", "hending", decision);
assert_eq!(merged, "comprehending");
let decision = stitch_decision("compre\u{2011}", "hending");
assert_eq!(decision, StitchAction::Hyphen);
let merged = merge_stitched("compre\u{2011}", "hending", decision);
assert_eq!(merged, "comprehending");
}
#[test]
fn stitch_no_hyphen_for_list_items() {
assert_eq!(stitch_decision("- item-", "next text"), StitchAction::Space);
assert_eq!(stitch_decision("- item-", "Next text"), StitchAction::None);
}
#[test]
fn stitch_none_on_terminal_punct() {
assert_eq!(
stitch_decision("sentence ends.", "Next sentence"),
StitchAction::None
);
assert_eq!(stitch_decision("question?", "Answer"), StitchAction::None);
assert_eq!(stitch_decision("exclaim!", "More text"), StitchAction::None);
assert_eq!(
stitch_decision("trailing\u{2026}", "More text"),
StitchAction::None
);
}
#[test]
fn stitch_none_terminal_punct_through_quotes() {
assert_eq!(stitch_decision("end.\"", "Next line"), StitchAction::None);
assert_eq!(stitch_decision("end.)", "Next line"), StitchAction::None);
}
#[test]
fn stitch_space_on_continuation_punct() {
assert_eq!(stitch_decision("clause,", "More text"), StitchAction::Space);
assert_eq!(stitch_decision("clause;", "More text"), StitchAction::Space);
assert_eq!(stitch_decision("clause:", "More text"), StitchAction::Space);
assert_eq!(
stitch_decision("clause\u{2014}", "More text"),
StitchAction::Space
); assert_eq!(
stitch_decision("clause\u{2013}", "More text"),
StitchAction::Space
); }
#[test]
fn stitch_space_on_closing_delimiters_next() {
assert_eq!(stitch_decision("some text", ") rest"), StitchAction::Space);
assert_eq!(stitch_decision("some text", "] rest"), StitchAction::Space);
assert_eq!(stitch_decision("some text", "} rest"), StitchAction::Space);
}
#[test]
fn stitch_none_both_capitalized_no_continuation() {
assert_eq!(
stitch_decision("End of paragraph", "Start of next"),
StitchAction::None
);
}
#[test]
fn merge_none_returns_prev_only() {
let result = merge_stitched("first page", "second page", StitchAction::None);
assert_eq!(result, "first page");
}
#[test]
fn merge_space_trims() {
let result = merge_stitched("first ", " second", StitchAction::Space);
assert_eq!(result, "first second");
}
#[test]
fn process_strips_repeated_header() {
let mut state = PostProcessState::new();
let excludes = [Exclude::Headers];
let page1 = "JOURNAL OF SCIENCE\n\nContent of page one.";
let out1 = state.process(page1, &excludes);
assert!(!out1.contains("JOURNAL OF SCIENCE"));
let page2 = "JOURNAL OF SCIENCE\n\nContent of page two.";
let out2 = state.process(page2, &excludes);
assert!(!out2.contains("JOURNAL OF SCIENCE"));
assert!(out2.contains("Content of page two."));
}
#[test]
fn process_strips_repeated_footer() {
let mut state = PostProcessState::new();
let excludes = [Exclude::Footers];
let page1 = "Content one.\n\nCopyright 2024 Press";
let _out1 = state.process(page1, &excludes);
let page2 = "Content two.\n\nCopyright 2024 Press";
let out2 = state.process(page2, &excludes);
assert!(!out2.contains("Copyright 2024 Press"));
assert!(out2.contains("Content two."));
}
#[test]
fn process_keeps_unique_headers() {
let mut state = PostProcessState::new();
let excludes = [Exclude::Headers];
let page1 = "Chapter One\n\nContent.";
let out1 = state.process(page1, &excludes);
assert!(out1.contains("Chapter One"));
let page2 = "Chapter Two\n\nMore content.";
let out2 = state.process(page2, &excludes);
assert!(out2.contains("Chapter Two")); }
#[test]
fn process_page_number_at_top() {
let mut state = PostProcessState::new();
let input = "42\n\nContent here.";
let output = state.process(input, &[Exclude::PageNumbers]);
assert!(!output.contains("42"));
assert!(output.contains("Content here."));
}
#[test]
fn process_page_number_at_bottom() {
let mut state = PostProcessState::new();
let input = "Content here.\n\n7";
let output = state.process(input, &[Exclude::PageNumbers]);
assert!(!output.contains("\n7"));
assert!(output.contains("Content here."));
}
#[test]
fn process_page_number_inline_not_removed() {
let mut state = PostProcessState::new();
let input = "Line one\n42\nLine three";
let output = state.process(input, &[Exclude::PageNumbers]);
assert!(output.contains("42"));
}
#[test]
fn process_empty_input() {
let mut state = PostProcessState::new();
let output = state.process("", &[Exclude::Headers, Exclude::PageNumbers]);
assert!(output.is_empty());
}
#[test]
fn process_no_excludes_keeps_everything() {
let mut state = PostProcessState::new();
let input = "42\n\nContent\n\nJOURNAL\n\n7";
let output = state.process(input, &[]);
assert!(output.contains("42"));
assert!(output.contains("JOURNAL"));
assert!(output.contains("7"));
}
#[test]
fn strip_boundary_keeps_header_when_followed_by_uppercase() {
let prev = "find the measure";
let next = "Introduction\n\nThe next chapter begins.";
let output = strip_boundary_artifacts(prev, next, &[]);
assert_eq!(output, next);
}
#[test]
fn strip_boundary_keeps_non_title_case() {
let prev = "find the measure";
let next = "not a header\n\nof their dwelling.";
let output = strip_boundary_artifacts(prev, next, &[]);
assert_eq!(output, next);
}
#[test]
fn strip_boundary_all_empty_next() {
let prev = "find the measure";
let next = "\n\n\n";
let output = strip_boundary_artifacts(prev, next, &[]);
assert!(output.is_empty());
}
#[test]
fn strip_boundary_prev_is_only_quotes() {
let prev = "\"'";
let next = "Introduction\n\nof their dwelling.";
let output = strip_boundary_artifacts(prev, next, &[]);
assert_eq!(output, next);
}
#[test]
fn strip_boundary_candidate_with_digits_kept() {
let prev = "find the measure";
let next = "Chapter 3\n\nof their dwelling.";
let output = strip_boundary_artifacts(prev, next, &[]);
assert_eq!(output, next); }
#[test]
fn strip_boundary_header_only_no_content_after() {
let prev = "find the measure";
let next = "Introduction";
let output = strip_boundary_artifacts(prev, next, &[]);
assert_eq!(output, next);
}
}