notescraft/crates/engine/tests/common/pdf.rs

//! Recovers the text a produced PDF draws: the writer stores text as glyph id
//! runs inside flate-compressed page streams, so a test that wants to assert
//! on rendered frontmatter properties must inflate the streams and decode the
//! runs through the fonts' `ToUnicode` maps — the same round-trip the PDF
//! crate's own integration tests rely on.

use flate2::read::ZlibDecoder;
use std::collections::BTreeMap;
use std::io::Read as _;

/// The text every page of `pdf` draws, joined in page order.
pub fn page_text(pdf: &[u8]) -> String {
    let objects = objects(pdf);
    let mut out = String::new();
    for number in 1..objects.len() {
        let body = body_of(&objects, number);
        if !body.contains("/Type /Page ") {
            continue;
        }
        let names = font_names(&body);
        let maps = font_maps(&objects, &names);
        let Some(stream) = reference(&body, "/Contents").and_then(|target| inflate(raw_of(&objects, target))) else {
            continue;
        };
        out.push_str(&drawn(&stream, &maps));
        out.push('\n');
    }
    out
}

/// Walks one content stream: `/F<n> size Tf` selects the font, `<hex> Tj`
/// draws the run through that font's glyph map.
fn drawn(stream: &str, maps: &BTreeMap<String, BTreeMap<u16, String>>) -> String {
    let mut out = String::new();
    let mut font = String::from("/F1");
    let mut hex: Option<&str> = None;
    for token in stream.split_whitespace() {
        if let Some(name) = token.strip_prefix('/') {
            if name.starts_with('F') && name.get(1..).is_some_and(|rest| !rest.is_empty() && rest.chars().all(|ch| ch.is_ascii_digit())) {
                font = format!("/{name}");
            }
            continue;
        }
        match token {
            "Tj" => {
                if let Some(map) = maps.get(&font) {
                    for word in hex.unwrap_or_default().as_bytes().chunks(4) {
                        let Some(word) = std::str::from_utf8(word).ok().and_then(|text| u16::from_str_radix(text, 16).ok()) else { continue };
                        if let Some(text) = map.get(&word) {
                            out.push_str(text);
                        }
                    }
                    out.push(' ');
                }
            }
            _ => {
                hex = token.strip_prefix('<').map(|rest| rest.trim_end_matches('>'));
            }
        }
    }
    out.trim().to_owned()
}

/// The `ToUnicode` glyph id to text maps of one page's font resources, keyed
/// by the `/F<n>` name the content stream uses.
fn font_maps(objects: &[Vec<u8>], names: &BTreeMap<usize, String>) -> BTreeMap<String, BTreeMap<u16, String>> {
    let mut out = BTreeMap::new();
    for (number, name) in names {
        let object = body_of(objects, *number);
        if !object.contains("/Subtype /Type0") {
            continue;
        }
        let Some(target) = reference(&object, "/ToUnicode") else { continue };
        let Some(text) = inflate(raw_of(objects, target)) else { continue };
        out.insert(name.clone(), parse_cmap(&text));
    }
    out
}

/// The whole `/Font << /F1 5 0 R .. >>` mapping of one page: object number
/// to resource name.
fn font_names(page: &str) -> BTreeMap<usize, String> {
    let mut out = BTreeMap::new();
    let Some(dict) = page.split("/Font <<").nth(1).and_then(|rest| rest.split(">>").next()) else { return out };
    for entry in dict.split('/').skip(1) {
        let mut parts = entry.split_whitespace();
        let Some(name) = parts.next() else { continue };
        let Some(target) = parts.next().and_then(|value| value.parse::<usize>().ok()) else { continue };
        out.insert(target, format!("/{name}"));
    }
    out
}

/// Reads a `bfchar`-style `ToUnicode` `CMap` into a glyph id to text map.
fn parse_cmap(cmap: &str) -> BTreeMap<u16, String> {
    let mut out = BTreeMap::new();
    let mut rest = cmap;
    while let Some(found) = rest.find('<') {
        let Some((gid, after)) = hex_word(rest.get(found + 1..).unwrap_or_default()) else { break };
        let Some((text, after)) = hex_word(after.trim_start()) else { break };
        if let Ok(gid) = u16::from_str_radix(gid, 16) {
            out.insert(gid, decode_utf16(text));
        }
        rest = after;
    }
    out
}

/// One `<hex>` value and the text after it.
fn hex_word(text: &str) -> Option<(&str, &str)> {
    let text = text.strip_prefix('<').unwrap_or(text).trim_start();
    let end = text.find('>')?;
    Some((text.get(..end)?, text.get(end + 1..)?))
}

/// UTF-16BE hex (with or without a BOM) as text.
fn decode_utf16(hex: &str) -> String {
    let hex = hex.strip_prefix("FEFF").unwrap_or(hex);
    let units: Vec<u16> = hex.as_bytes().chunks(4).filter_map(|chunk| u16::from_str_radix(std::str::from_utf8(chunk).unwrap_or_default(), 16).ok()).collect();
    String::from_utf16_lossy(&units)
}

/// The inflated text of a stream object body.
fn inflate(object: &[u8]) -> Option<String> {
    let data = stream_data(object)?;
    let mut out = Vec::new();
    ZlibDecoder::new(data).read_to_end(&mut out).ok()?;
    String::from_utf8(out).ok()
}

/// The compressed bytes between `stream\n` and the last `\nendstream`.
fn stream_data(object: &[u8]) -> Option<&[u8]> {
    let start = find(object, b"stream\n", 0)? + b"stream\n".len();
    let end = object.windows(10).rposition(|probe| probe == b"\nendstream")?;
    object.get(start..end)
}

/// The file offsets of every object, read from the xref table rather than by
/// scanning for `N 0 obj` (binary font payloads hold those bytes by chance).
fn offsets(pdf: &[u8]) -> Vec<usize> {
    let Some(start) = find(pdf, b"startxref", 0) else { return Vec::new() };
    let rest = pdf.get(start + b"startxref\n".len()..).unwrap_or_default();
    let line_end = rest.iter().position(|byte| byte == &b'\n').unwrap_or(rest.len());
    let Ok(table) = std::str::from_utf8(rest.get(..line_end).unwrap_or_default()).unwrap_or_default().trim().parse::<usize>() else {
        return Vec::new();
    };
    let head = pdf.get(table + 5..).unwrap_or_default();
    let head_end = head.iter().position(|byte| byte == &b'\n').unwrap_or(head.len());
    let count = std::str::from_utf8(head.get(..head_end).unwrap_or_default())
        .unwrap_or_default()
        .split_whitespace()
        .nth(1)
        .and_then(|value| value.parse::<usize>().ok())
        .unwrap_or(0);
    let entries = table + 5 + head_end + 1 + 20;
    (1..count)
        .filter_map(|number| {
            let cell = pdf.get(entries + (number - 1) * 20..entries + (number - 1) * 20 + 10)?;
            std::str::from_utf8(cell).unwrap_or_default().trim().parse::<usize>().ok()
        })
        .collect()
}

/// Every object body of the document as raw bytes, indexed by number (the
/// table is 1-based, entry 0 is the free entry).
fn objects(pdf: &[u8]) -> Vec<Vec<u8>> {
    let mut out = vec![Vec::new()];
    for offset in offsets(pdf) {
        let body = pdf.get(offset..).unwrap_or_default();
        let head = find(body, b" obj", 0).map_or(body.len(), |at| at + 4);
        let dict_end = find(body, b"endobj", 0);
        let stream_at = find(body, b"stream\n", 0).filter(|at| dict_end.is_none_or(|end| at < &end));
        // A stream object ends after `endstream`; a binary payload can hold
        // the bytes of `endobj` by chance.
        let end = match stream_at {
            Some(at) => find(body, b"\nendstream", at).map_or(body.len(), |stop| stop + 10),
            None => dict_end.map_or(body.len(), |at| at + 6),
        };
        out.push(body.get(head..end).unwrap_or_default().to_vec());
    }
    out
}

/// First position of `needle` in `hay` at or after `from`.
fn find(hay: &[u8], needle: &[u8], from: usize) -> Option<usize> {
    hay.get(from..)?.windows(needle.len()).position(|window| window == needle).map(|at| at + from)
}

/// The body of one object as text; entry `n` holds object `n 0 obj`.
fn body_of(objects: &[Vec<u8>], number: usize) -> String {
    String::from_utf8_lossy(raw_of(objects, number)).to_string()
}

/// The raw bytes of one object body.
fn raw_of(objects: &[Vec<u8>], number: usize) -> &[u8] {
    objects.get(number).map(Vec::as_slice).unwrap_or_default()
}

/// The object number a `/Key N 0 R` entry points at.
fn reference(text: &str, key: &str) -> Option<usize> {
    let tail = text.split(key).nth(1)?.trim_start();
    let digits: String = tail.chars().take_while(char::is_ascii_digit).collect();
    digits.parse().ok()
}