mod citations;
mod figures;
mod headers;
mod headings;
mod lists;
mod paragraphs;
mod tables;
use std::cmp::Ordering;
use kopitiam_pdf::{Page, TextSpan};
use crate::{Block, Document, Heading, Metadata, Paragraph};
pub(crate) use headers::strip_marginalia;
pub(crate) use figures::collapse_figure_regions;
struct Line {
text: String,
y: f32,
font_size: f32,
cells: Vec<Cell>,
}
struct Cell {
text: String,
x: f32,
x_end: f32,
}
const SAME_LINE_Y_TOLERANCE_RATIO: f32 = 0.4;
const WORD_GAP_RATIO: f32 = 0.15;
const COLUMN_GAP_RATIO: f32 = 2.5;
const STRADDLE_LINE_MAX_FRACTION: f32 = 0.15;
const FULL_WIDTH_CELL_MIN_RATIO: f32 = 0.66;
pub fn reconstruct(pages: &[Page]) -> Document {
let pages = strip_marginalia(pages);
let pages = collapse_figure_regions(&pages);
let pages = pages.as_slice();
let body_font_size = estimate_body_font_size(pages);
let mut citations = Vec::new();
let mut pages_blocks: Vec<Vec<Block>> = Vec::with_capacity(pages.len());
for page in pages {
let mut page_blocks = Vec::new();
for column_spans in split_columns(page) {
let lines = group_lines(&column_spans);
for block in build_blocks(&lines, body_font_size) {
if let Block::Paragraph(paragraph) = &block {
citations.extend(citations::detect(¶graph.text));
}
page_blocks.push(block);
}
}
pages_blocks.push(page_blocks);
}
let (blocks, block_pages) = merge_page_breaks(pages_blocks);
Document {
title: infer_title(&blocks),
metadata: Metadata {
source_pages: pages.len(),
},
blocks,
block_pages,
citations,
}
}
pub fn reconstruct_preordered(pages: &[Page]) -> Document {
let pages = strip_marginalia(pages);
let pages = collapse_figure_regions(&pages);
let pages = pages.as_slice();
let body_font_size = estimate_body_font_size(pages);
let mut citations = Vec::new();
let mut pages_blocks: Vec<Vec<Block>> = Vec::with_capacity(pages.len());
for page in pages {
let lines: Vec<Line> = page.spans.iter().map(|span| build_line(&[span])).collect();
let mut page_blocks = Vec::new();
for block in build_blocks(&lines, body_font_size) {
if let Block::Paragraph(paragraph) = &block {
citations.extend(citations::detect(¶graph.text));
}
page_blocks.push(block);
}
pages_blocks.push(page_blocks);
}
let (blocks, block_pages) = merge_page_breaks(pages_blocks);
Document {
title: infer_title(&blocks),
metadata: Metadata {
source_pages: pages.len(),
},
blocks,
block_pages,
citations,
}
}
fn merge_page_breaks(pages_blocks: Vec<Vec<Block>>) -> (Vec<Block>, Vec<usize>) {
let mut blocks: Vec<Block> = Vec::new();
let mut block_pages: Vec<usize> = Vec::new();
for (page_index, page_blocks) in pages_blocks.into_iter().enumerate() {
if page_blocks.is_empty() {
continue;
}
let page = page_index + 1;
let mut page_blocks = page_blocks.into_iter();
let leading = page_blocks.next();
let merged_text = match (blocks.last(), &leading) {
(Some(Block::Paragraph(trailing)), Some(Block::Paragraph(leading_paragraph))) => {
paragraphs::merge_across_page_break(&trailing.text, &leading_paragraph.text)
}
_ => None,
};
match merged_text {
Some(text) => {
*blocks
.last_mut()
.expect("merged_text is only Some when blocks.last() matched") =
Block::Paragraph(Paragraph { text });
}
None => {
if let Some(leading_block) = leading {
blocks.push(leading_block);
block_pages.push(page);
}
}
}
for block in page_blocks {
blocks.push(block);
block_pages.push(page);
}
}
debug_assert_eq!(
blocks.len(),
block_pages.len(),
"block_pages must stay parallel to blocks, or every citation this document produces is wrong"
);
(blocks, block_pages)
}
fn build_blocks(lines: &[Line], body_font_size: f32) -> Vec<Block> {
let mut blocks = Vec::new();
let mut i = 0;
while i < lines.len() {
if let Some((table, consumed)) = tables::try_table(&lines[i..]) {
blocks.push(Block::Table(table));
i += consumed;
continue;
}
if let Some(figure) = figures::try_figure(&lines[i]) {
blocks.push(Block::Figure(figure));
i += 1;
continue;
}
if let Some(level) = headings::heading_level(&lines[i], body_font_size) {
blocks.push(Block::Heading(Heading {
level,
text: lines[i].text.trim().to_string(),
}));
i += 1;
continue;
}
if let Some((list, consumed)) = lists::try_list(&lines[i..]) {
blocks.push(Block::List(list));
i += consumed;
continue;
}
let (paragraph, consumed) = paragraphs::consume_paragraph(&lines[i..]);
blocks.push(Block::Paragraph(paragraph));
i += consumed;
}
blocks
}
fn infer_title(blocks: &[Block]) -> Option<String> {
blocks.iter().find_map(|block| match block {
Block::Heading(Heading { level: 1, text }) => Some(text.clone()),
_ => None,
})
}
fn estimate_body_font_size(pages: &[Page]) -> f32 {
use std::collections::BTreeMap;
let mut counts: BTreeMap<u32, usize> = BTreeMap::new();
for page in pages {
for span in &page.spans {
let bucket = (span.font_size * 2.0).round() as u32;
*counts.entry(bucket).or_default() += 1;
}
}
counts
.into_iter()
.min_by_key(|&(bucket, count)| (std::cmp::Reverse(count), bucket))
.map(|(bucket, _)| bucket as f32 / 2.0)
.unwrap_or(12.0)
}
fn split_columns(page: &Page) -> Vec<Vec<TextSpan>> {
if page.spans.is_empty() {
return vec![Vec::new()];
}
let midpoint = page.width / 2.0;
let full_lines = group_lines(&page.spans);
let straddling = full_lines
.iter()
.filter(|line| is_full_width_line(line, page.width, midpoint))
.count();
let straddle_fraction = straddling as f32 / full_lines.len().max(1) as f32;
if straddle_fraction > STRADDLE_LINE_MAX_FRACTION {
return vec![page.spans.clone()];
}
let mut left = Vec::new();
let mut right = Vec::new();
for span in &page.spans {
let center = span.x + span.width / 2.0;
if center < midpoint {
left.push(span.clone());
} else {
right.push(span.clone());
}
}
if left.is_empty() || right.is_empty() {
return vec![page.spans.clone()];
}
split_two_column_page_into_bands(page, midpoint)
}
fn is_full_width_line(line: &Line, page_width: f32, midpoint: f32) -> bool {
line.cells.iter().any(|cell| {
(cell.x < midpoint && cell.x_end > midpoint)
|| (cell.x_end - cell.x) > page_width * FULL_WIDTH_CELL_MIN_RATIO
})
}
fn split_two_column_page_into_bands(page: &Page, midpoint: f32) -> Vec<Vec<TextSpan>> {
let mut result = Vec::new();
let mut band_left: Vec<TextSpan> = Vec::new();
let mut band_right: Vec<TextSpan> = Vec::new();
let mut band_full_width: Vec<TextSpan> = Vec::new();
for mut group in group_spans_by_baseline(&page.spans) {
group.sort_by(|a, b| a.x.partial_cmp(&b.x).unwrap_or(Ordering::Equal));
let refs: Vec<&TextSpan> = group.iter().collect();
let line = build_line(&refs);
if is_full_width_line(&line, page.width, midpoint) {
if !band_left.is_empty() {
result.push(std::mem::take(&mut band_left));
}
if !band_right.is_empty() {
result.push(std::mem::take(&mut band_right));
}
band_full_width.extend(group);
} else {
if !band_full_width.is_empty() {
result.push(std::mem::take(&mut band_full_width));
}
for span in group {
let center = span.x + span.width / 2.0;
if center < midpoint {
band_left.push(span);
} else {
band_right.push(span);
}
}
}
}
if !band_full_width.is_empty() {
result.push(band_full_width);
}
if !band_left.is_empty() {
result.push(band_left);
}
if !band_right.is_empty() {
result.push(band_right);
}
result
}
fn group_lines(spans: &[TextSpan]) -> Vec<Line> {
group_spans_by_baseline(spans)
.into_iter()
.map(|mut group| {
group.sort_by(|a, b| a.x.partial_cmp(&b.x).unwrap_or(Ordering::Equal));
let refs: Vec<&TextSpan> = group.iter().collect();
build_line(&refs)
})
.collect()
}
fn group_spans_by_baseline(spans: &[TextSpan]) -> Vec<Vec<TextSpan>> {
let mut ordered: Vec<&TextSpan> = spans.iter().collect();
ordered.sort_by(|a, b| b.y.partial_cmp(&a.y).unwrap_or(Ordering::Equal));
let mut groups: Vec<Vec<TextSpan>> = Vec::new();
for span in ordered {
let joins_last = groups.last().is_some_and(|group: &Vec<TextSpan>| {
let anchor = &group[0];
let tolerance = anchor.font_size.max(span.font_size) * SAME_LINE_Y_TOLERANCE_RATIO;
(anchor.y - span.y).abs() <= tolerance
});
if joins_last {
groups.last_mut().unwrap().push(span.clone());
} else {
groups.push(vec![span.clone()]);
}
}
groups
}
fn build_line(spans: &[&TextSpan]) -> Line {
let mut cells: Vec<Cell> = Vec::new();
let mut text = String::new();
let mut prev_end: Option<f32> = None;
for span in spans {
let gap = prev_end.map(|end| span.x - end);
let is_word_gap = gap.is_some_and(|gap| gap > span.font_size * WORD_GAP_RATIO);
let starts_new_cell = match gap {
Some(gap) => gap > span.font_size * COLUMN_GAP_RATIO,
None => true,
};
if starts_new_cell {
cells.push(Cell {
text: span.text.clone(),
x: span.x,
x_end: span.x + span.width,
});
} else if let Some(cell) = cells.last_mut() {
if is_word_gap {
cell.text.push(' ');
}
cell.text.push_str(&span.text);
cell.x_end = span.x + span.width;
}
if is_word_gap {
text.push(' ');
}
text.push_str(&span.text);
prev_end = Some(span.x + span.width);
}
let y = spans.first().map(|s| s.y).unwrap_or(0.0);
let font_size = spans.iter().map(|s| s.font_size).fold(0.0_f32, f32::max);
Line {
text,
y,
font_size,
cells,
}
}
#[cfg(test)]
mod tests {
use super::*;
fn span(text: &str, x: f32, y: f32, width: f32, font_size: f32) -> TextSpan {
TextSpan {
text: text.to_string(),
x,
y,
width,
height: font_size,
font_size,
font_name: None,
..TextSpan::default()
}
}
#[test]
fn build_line_merges_contiguous_glyph_runs_without_a_space() {
let boo = span("Boo", 0.0, 0.0, 18.0, 10.0);
let k = span("k", 18.2, 0.0, 6.0, 10.0);
let line = build_line(&[&boo, &k]);
assert_eq!(line.text, "Book");
}
#[test]
fn build_line_keeps_a_space_for_a_real_word_gap() {
let book = span("Book", 0.0, 0.0, 24.0, 10.0);
let reviews = span("Reviews", 27.0, 0.0, 40.0, 10.0);
let line = build_line(&[&book, &reviews]);
assert_eq!(line.text, "Book Reviews");
}
#[test]
fn build_line_splits_a_wide_gap_into_a_new_cell() {
let metric = span("Metric", 0.0, 0.0, 30.0, 10.0);
let value = span("Value", 60.0, 0.0, 20.0, 10.0);
let line = build_line(&[&metric, &value]);
assert_eq!(line.cells.len(), 2);
}
#[test]
fn single_column_page_is_not_split() {
let page = Page {
number: 1,
width: 600.0,
height: 800.0,
spans: vec![
span(
"Line one spans the whole page width here",
50.0,
700.0,
500.0,
10.0,
),
span(
"Line two also spans the whole page width",
50.0,
686.0,
500.0,
10.0,
),
],
};
assert_eq!(split_columns(&page).len(), 1);
}
#[test]
fn two_column_page_is_split() {
let page = Page {
number: 1,
width: 600.0,
height: 800.0,
spans: vec![
span("Left column text here", 40.0, 700.0, 220.0, 10.0),
span("Left column text here", 40.0, 686.0, 220.0, 10.0),
span("Right column text here", 340.0, 700.0, 220.0, 10.0),
span("Right column text here", 340.0, 686.0, 220.0, 10.0),
],
};
assert_eq!(split_columns(&page).len(), 2);
}
fn two_column_page_with_full_width_interruption() -> Page {
let mut spans = Vec::new();
for (i, y) in [760.0, 748.0, 736.0].into_iter().enumerate() {
spans.push(span(&format!("Top left {i}"), 40.0, y, 220.0, 10.0));
spans.push(span(&format!("Top right {i}"), 340.0, y, 220.0, 10.0));
}
spans.push(span(
"Full width heading spanning both columns",
40.0,
700.0,
520.0,
10.0,
));
for (i, y) in [660.0, 648.0, 636.0].into_iter().enumerate() {
spans.push(span(&format!("Bottom left {i}"), 40.0, y, 220.0, 10.0));
spans.push(span(&format!("Bottom right {i}"), 340.0, y, 220.0, 10.0));
}
Page {
number: 1,
width: 600.0,
height: 800.0,
spans,
}
}
#[test]
fn full_width_element_splits_a_two_column_page_into_bands() {
let page = two_column_page_with_full_width_interruption();
let columns = split_columns(&page);
assert_eq!(columns.len(), 5);
assert!(columns[0].iter().all(|s| s.text.starts_with("Top left")));
assert!(columns[1].iter().all(|s| s.text.starts_with("Top right")));
assert_eq!(columns[2].len(), 1);
assert_eq!(columns[2][0].text, "Full width heading spanning both columns");
assert!(columns[3].iter().all(|s| s.text.starts_with("Bottom left")));
assert!(columns[4].iter().all(|s| s.text.starts_with("Bottom right")));
}
#[test]
fn plain_two_column_page_is_unaffected_by_band_splitting() {
let page = Page {
number: 1,
width: 600.0,
height: 800.0,
spans: vec![
span("Left column text here", 40.0, 700.0, 220.0, 10.0),
span("Left column text here", 40.0, 686.0, 220.0, 10.0),
span("Right column text here", 340.0, 700.0, 220.0, 10.0),
span("Right column text here", 340.0, 686.0, 220.0, 10.0),
],
};
let columns = split_columns(&page);
assert_eq!(columns.len(), 2);
assert!(columns[0].iter().all(|s| s.text == "Left column text here"));
assert!(columns[1].iter().all(|s| s.text == "Right column text here"));
}
#[test]
fn paragraph_split_across_a_page_break_is_merged() {
let page1 = Page {
number: 1,
width: 600.0,
height: 800.0,
spans: vec![span(
"This paragraph is cut off at the bottom of the page and",
50.0,
700.0,
500.0,
10.0,
)],
};
let page2 = Page {
number: 2,
width: 600.0,
height: 800.0,
spans: vec![span(
"continues here after the page break.",
50.0,
700.0,
500.0,
10.0,
)],
};
let document = reconstruct(&[page1, page2]);
assert_eq!(document.blocks.len(), 1);
match &document.blocks[0] {
Block::Paragraph(paragraph) => assert_eq!(
paragraph.text,
"This paragraph is cut off at the bottom of the page and continues here after the page break."
),
other => panic!("expected a merged Paragraph block, got {other:?}"),
}
assert_eq!(document.page_of(0), Some(1));
}
#[test]
fn every_block_records_the_page_it_starts_on() {
let page1 = Page {
number: 1,
width: 600.0,
height: 800.0,
spans: vec![span("Sentence one finishes here.", 50.0, 700.0, 500.0, 10.0)],
};
let page2 = Page {
number: 2,
width: 600.0,
height: 800.0,
spans: vec![span("Sentence two begins the second page.", 50.0, 700.0, 500.0, 10.0)],
};
let document = reconstruct(&[page1, page2]);
assert_eq!(document.blocks.len(), document.block_pages.len());
assert_eq!(document.page_of(0), Some(1));
assert_eq!(document.page_of(1), Some(2));
assert_eq!(document.page_of(99), None);
let paired: Vec<Option<usize>> = document.blocks_with_pages().map(|(_, page)| page).collect();
assert_eq!(paired, vec![Some(1), Some(2)]);
}
#[test]
fn body_font_size_is_deterministic_when_two_sizes_tie() {
let tied = || Page {
number: 1,
width: 600.0,
height: 800.0,
spans: vec![
span("Alpha at ten point", 50.0, 700.0, 400.0, 10.0),
span("Bravo at fourteen", 50.0, 660.0, 400.0, 14.0),
],
};
let first = estimate_body_font_size(&[tied()]);
for _ in 0..64 {
assert_eq!(estimate_body_font_size(&[tied()]), first, "font-size estimate is not deterministic");
}
assert_eq!(first, 10.0);
}
#[test]
fn a_document_with_no_page_information_reports_none_rather_than_guessing() {
let document = Document {
blocks: vec![Block::Paragraph(Paragraph { text: "orphan".to_string() })],
..Document::default()
};
assert_eq!(document.page_of(0), None);
assert_eq!(document.blocks_with_pages().next().unwrap().1, None);
}
#[test]
fn preordered_trusts_span_order_and_does_not_reinterleave_columns() {
let page = Page {
number: 1,
width: 600.0,
height: 800.0,
spans: vec![
span("Left column sentence one.", 40.0, 700.0, 220.0, 10.0),
span("Right column sentence two.", 340.0, 700.0, 220.0, 10.0),
],
};
let document = reconstruct_preordered(&[page]);
let texts: Vec<&str> = document
.blocks
.iter()
.filter_map(|b| match b {
Block::Paragraph(p) => Some(p.text.as_str()),
_ => None,
})
.collect();
assert_eq!(
texts,
vec!["Left column sentence one. Right column sentence two."],
"preordered reconstruction must preserve the given reading order"
);
}
#[test]
fn preordered_still_detects_headings_and_merges_page_breaks() {
let page1 = Page {
number: 1,
width: 600.0,
height: 800.0,
spans: vec![
span("Big Heading", 50.0, 750.0, 200.0, 20.0),
span("A paragraph that runs off the bottom of the first page and", 50.0, 700.0, 500.0, 10.0),
],
};
let page2 = Page {
number: 2,
width: 600.0,
height: 800.0,
spans: vec![span("continues onto the second page.", 50.0, 750.0, 400.0, 10.0)],
};
let document = reconstruct_preordered(&[page1, page2]);
assert!(
matches!(&document.blocks[0], Block::Heading(h) if h.text == "Big Heading"),
"large-font line must still be a heading, got {:?}",
document.blocks[0]
);
let paragraphs: Vec<&str> = document
.blocks
.iter()
.filter_map(|b| match b {
Block::Paragraph(p) => Some(p.text.as_str()),
_ => None,
})
.collect();
assert_eq!(
paragraphs,
vec!["A paragraph that runs off the bottom of the first page and continues onto the second page."]
);
}
fn cell(text: &str, x: f32) -> Cell {
Cell {
text: text.to_string(),
x,
x_end: x + 20.0,
}
}
fn line_with_cells(cells: Vec<Cell>) -> Line {
Line {
text: cells
.iter()
.map(|c| c.text.as_str())
.collect::<Vec<_>>()
.join(" "),
y: 0.0,
font_size: 10.0,
cells,
}
}
#[test]
fn ragged_table_row_truncates_the_table_and_the_remainder_survives() {
let lines = vec![
line_with_cells(vec![cell("Metric", 0.0), cell("Value", 60.0)]),
line_with_cells(vec![cell("Commits", 0.0), cell("282", 60.0)]),
line_with_cells(vec![cell("Outside", 0.0), cell("81", 60.0)]),
line_with_cells(vec![cell("Subtotal across both columns", 0.0)]),
line_with_cells(vec![cell("Reviews", 0.0), cell("7", 60.0)]),
];
let blocks = build_blocks(&lines, 10.0);
assert!(
matches!(&blocks[0], Block::Table(t)
if t.headers == vec!["Metric", "Value"]
&& t.rows == vec![vec!["Commits", "282"], vec!["Outside", "81"]]),
"first block must be the uniform table prefix, got {:?}",
blocks[0]
);
let tail: String = blocks[1..]
.iter()
.filter_map(|b| match b {
Block::Paragraph(p) => Some(p.text.clone()),
_ => None,
})
.collect::<Vec<_>>()
.join(" ");
assert!(
tail.contains("Subtotal across both columns"),
"ragged subtotal row must survive as normal text, got {tail:?}"
);
}
#[test]
fn paragraph_ending_a_sentence_does_not_merge_across_a_page_break() {
let page1 = Page {
number: 1,
width: 600.0,
height: 800.0,
spans: vec![span(
"This sentence finishes cleanly on the first page.",
50.0,
700.0,
500.0,
10.0,
)],
};
let page2 = Page {
number: 2,
width: 600.0,
height: 800.0,
spans: vec![span(
"New paragraph starts capitalized on the next page.",
50.0,
700.0,
500.0,
10.0,
)],
};
let document = reconstruct(&[page1, page2]);
assert_eq!(document.blocks.len(), 2);
for block in &document.blocks {
assert!(matches!(block, Block::Paragraph(_)));
}
}
}