kopitiam_document/reconstruction/

mod.rs

mod citations;
mod figures;
mod headings;
mod lists;
mod paragraphs;
mod tables;

use std::cmp::Ordering;

use kopitiam_pdf::{Page, TextSpan};

use crate::{Block, Document, Heading, Metadata};

/// One visual line of text on a page: spans grouped by shared baseline and
/// sorted left to right.
struct Line {
    text: String,
    y: f32,
    font_size: f32,
    /// Sub-runs of this line separated by a gap wide enough to suggest a
    /// column boundary; used by table detection.
    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;

/// Turn a page's raw text spans into the semantic `Document` AST: split each
/// page into reading-order columns, group spans into lines, then classify
/// each line (or run of lines) as a heading, list, table, figure caption, or
/// paragraph.
pub fn reconstruct(pages: &[Page]) -> Document {
    let body_font_size = estimate_body_font_size(pages);
    let mut blocks = Vec::new();
    let mut citations = Vec::new();

    for page in pages {
        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(&paragraph.text));
                }
                blocks.push(block);
            }
        }
    }

    Document {
        title: infer_title(&blocks),
        metadata: Metadata {
            source_pages: pages.len(),
        },
        blocks,
        citations,
    }
}

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,
    })
}

/// The most common font size across the document, used as the "body text"
/// baseline that heading detection compares against.
fn estimate_body_font_size(pages: &[Page]) -> f32 {
    use std::collections::HashMap;

    let mut counts: HashMap<u32, usize> = HashMap::new();
    for page in pages {
        for span in &page.spans {
            // Bucket to the nearest half-point to absorb float noise.
            let bucket = (span.font_size * 2.0).round() as u32;
            *counts.entry(bucket).or_default() += 1;
        }
    }

    counts
        .into_iter()
        .max_by_key(|&(_, count)| count)
        .map(|(bucket, _)| bucket as f32 / 2.0)
        .unwrap_or(12.0)
}

/// Splits a page's spans into left-to-right reading-order column groups.
///
/// A single-column page with normal margins routinely has lines whose text
/// crosses the page's geometric midpoint (most paragraph lines are wider
/// than half the page) -- so "spans exist on both sides of the midpoint" is
/// true for nearly every document and cannot be the two-column test.
///
/// Real two-column layouts instead have a genuine empty gutter at the
/// midpoint. But naively grouping all of a page's spans into y-based lines
/// first (as `group_lines` does) can still merge left- and right-column text
/// that happens to share a baseline (common: columns are typeset on a shared
/// line grid) into one "line" whose overall bounding box crosses the
/// midpoint -- even though neither column's text actually does. So the test
/// is per *cell*, not per line's overall extent: a cell is one uninterrupted
/// glyph run (see `build_line`'s gap detection), so a cell crossing the
/// midpoint means continuous prose was actually typeset across it, whereas
/// two same-baseline column fragments merged by `group_lines` show up as two
/// separate cells that individually stay on one side.
///
/// Does not handle a full-width element (e.g. a spanning figure or table)
/// that interrupts a two-column region partway down the page -- see
/// kopitiam-zay.
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| {
            line.cells
                .iter()
                .any(|cell| cell.x < midpoint && cell.x_end > 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() {
        vec![page.spans.clone()]
    } else {
        vec![left, right]
    }
}

fn group_lines(spans: &[TextSpan]) -> Vec<Line> {
    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);
        } else {
            groups.push(vec![span]);
        }
    }

    groups
        .into_iter()
        .map(|mut group| {
            group.sort_by(|a, b| a.x.partial_cmp(&b.x).unwrap_or(Ordering::Equal));
            build_line(&group)
        })
        .collect()
}

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);

        // A real inter-word space is a much smaller gap than a column/cell
        // boundary. Below `WORD_GAP_RATIO` the spans are contiguous glyphs
        // (e.g. an OCR text layer that split one word into several spans)
        // and must be concatenated with no space, or every such split would
        // otherwise render as a broken word ("Boo k" instead of "Book").
        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,
        }
    }

    #[test]
    fn build_line_merges_contiguous_glyph_runs_without_a_space() {
        // "Boo" then "k" with almost no gap simulates an OCR text layer that
        // split one word into two spans; it must read back as "Book".
        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() {
        // Every line's text spans the full page width, crossing the
        // midpoint -- this must not be mistaken for a two-column layout.
        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);
    }
}