notescraft/crates/pdf/tests/support/mod.rs

//! Shared helpers of the integration tests: the fonts, a structural PDF
//! checker and a content-stream text extractor.
//!
//! Each test binary uses a subset of this module, so unused helpers here are
//! expected rather than dead code.
#![allow(dead_code)]
//!
//! The extractor is deliberately independent of the writer: it reads the
//! xref table, walks every object, inflates the `ToUnicode` `CMap`s of the
//! fonts and rebuilds the text of each page from the `Tj` operands, so the
//! round-trip tests prove the document is really searchable text.

use std::collections::BTreeMap;
use std::io::Read as _;
use std::sync::Arc;

use notescraft_pdf::{ExportOptions, FontSet};

/// The exporter's fonts, read from the repository assets.
pub fn fonts() -> FontSet {
    FontSet {
        regular: Arc::from(&include_bytes!("../../../../assets/fonts/Inter-Regular.ttf")[..]),
        bold: Arc::from(&include_bytes!("../../../../assets/fonts/Inter-SemiBold.ttf")[..]),
        mono: Arc::from(&include_bytes!("../../../../assets/fonts/JetBrainsMono-Regular.ttf")[..]),
        italic: None,
    }
}

/// Export options with a fixed creation date, for deterministic bytes.
pub fn options() -> ExportOptions {
    let mut opts = ExportOptions::new(fonts());
    opts.created = Some(String::from("D:20260101000000"));
    opts
}

/// A structural check of a produced document: the xref offsets point at
/// their objects, every stream length matches its data and `startxref` is
/// real. Panics with a message when the file is malformed.
pub fn structural_check(pdf: &[u8]) {
    assert!(pdf.starts_with(b"%PDF-1.7\n"), "header");
    assert!(pdf.ends_with(b"%%EOF\n"), "trailer");
    let start = find(pdf, b"startxref", 0).unwrap_or_else(|| panic!("startxref marker"));
    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 offset: usize = std::str::from_utf8(rest.get(..line_end).unwrap_or_default())
        .unwrap_or_default()
        .trim()
        .parse()
        .unwrap_or_else(|_| panic!("startxref offset is a number"));
    assert_eq!(pdf.get(offset..offset + 4), Some(&b"xref"[..]), "startxref points at the table");
    let head = pdf.get(offset + 5..).unwrap_or_default();
    let head_end = head.iter().position(|byte| *byte == b'\n').unwrap_or(head.len());
    let count_line = std::str::from_utf8(head.get(..head_end).unwrap_or_default()).unwrap_or_default();
    let count: usize = count_line.split_whitespace().nth(1).and_then(|value| value.parse().ok()).unwrap_or(0);
    assert!(count > 1, "the xref declares its objects: {count_line:?}");
    let entries = offset + 5 + head_end + 1 + 20;
    for number in 1..count {
        let cell = pdf.get(entries + (number - 1) * 20..entries + (number - 1) * 20 + 10).unwrap_or_default();
        let text = std::str::from_utf8(cell).unwrap_or_default();
        let offset: usize = text.trim().parse().unwrap_or_else(|_| panic!("xref entry {number} holds an offset"));
        let expected = format!("{number} 0 obj");
        assert_eq!(pdf.get(offset..offset + expected.len()), Some(expected.as_bytes()), "object {number} sits at the offset its xref entry claims");
    }
    for (dict_end, data_start) in streams(pdf) {
        let declared = declared_length(pdf, data_start).unwrap_or(0);
        assert!(declared > 0, "every stream declares a positive /Length");
        let end = data_start + declared;
        assert_eq!(pdf.get(end..end + b"\nendstream".len()), Some(&b"\nendstream"[..]), "the stream of {declared} bytes ends at endstream");
        assert!(dict_end >= end, "endstream follows the stream data");
    }
}

/// The `/Length` declared by the dictionary right before `data_start`.
fn declared_length(pdf: &[u8], data_start: usize) -> Option<usize> {
    let window = pdf.get(..data_start)?;
    let at = window.windows(8).rposition(|probe| probe == b"/Length ")?;
    let tail = window.get(at + 8..)?.get(..12)?;
    let digits: String = tail.iter().map(|byte| char::from(*byte)).take_while(char::is_ascii_digit).collect();
    digits.parse().ok()
}

/// The `(endstream offset, data start)` of every stream object.
fn streams(pdf: &[u8]) -> Vec<(usize, usize)> {
    let mut out = Vec::new();
    let mut at = 0usize;
    while let Some(found) = find(pdf, b"stream\n", at) {
        let data_start = found + b"stream\n".len();
        let Some(end) = find(pdf, b"endstream", data_start) else { break };
        out.push((end, data_start));
        at = end + b"endstream".len();
    }
    out
}

/// First position of `needle` in `hay` at or after `from`.
fn find(hay: &[u8], needle: &[u8], from: usize) -> Option<usize> {
    let first = *needle.first()?;
    if from > hay.len() {
        return None;
    }
    hay.get(from..)?
        .iter()
        .enumerate()
        .find(|(at, byte)| **byte == first && hay.get(from + at..from + at + needle.len()) == Some(needle))
        .map(|(at, _)| from + at)
}

/// One page of extracted text.
pub struct PageText {
    /// The page number, 1-based.
    pub number: usize,
    /// The text in drawing order, runs separated by a space.
    pub text: String,
    /// The link URIs of the page, in order.
    pub links: Vec<String>,
}

/// Rebuilds the text of every page of a produced document.
pub fn pages(pdf: &[u8]) -> Vec<PageText> {
    let objects = objects(pdf);
    let shared = objects.clone();
    let cmaps = cmaps(&objects);
    let pages: Vec<usize> = (0..objects.len()).filter(|number| body(&objects, *number).contains("/Type /Page ")).collect();
    pages
        .iter()
        .enumerate()
        .map(|(index, number)| {
            let body = body(&objects, *number);
            let content = reference(&body, "/Contents");
            let stream = content.and_then(|target| inflate(raw(&objects, target)));
            let text = stream.as_deref().map(|ops| extract(ops, &cmaps)).unwrap_or_default();
            PageText { number: index + 1, text, links: uris(&shared, &body) }
        })
        .collect()
}

/// The whole text of a document, page by page.
pub fn text(pdf: &[u8]) -> String {
    pages(pdf).iter().map(|page| page.text.clone()).collect::<Vec<_>>().join("\n")
}

/// The inflated content stream of every page.
pub fn contents(pdf: &[u8]) -> Vec<String> {
    let objects = objects(pdf);
    let pages: Vec<usize> = (0..objects.len()).filter(|number| body(&objects, *number).contains("/Type /Page ")).collect();
    pages.iter().filter_map(|number| reference(&body(&objects, *number), "/Contents")).filter_map(|target| inflate(raw(&objects, target))).collect()
}

/// The URIs the annotations of one page point at.
fn uris(objects: &[Vec<u8>], page: &str) -> Vec<String> {
    let Some(array) = page.split("/Annots [").nth(1).and_then(|rest| rest.split(']').next()) else {
        return Vec::new();
    };
    array
        .split_whitespace()
        .filter_map(|token| token.parse::<usize>().ok())
        .filter_map(|number| {
            let uri = body(objects, number);
            let start = uri.find("/URI (")? + 6;
            let tail = uri.get(start..)?;
            let end = tail.find(')')?;
            tail.get(..end).map(|text| text.replace("\\(", "(").replace("\\)", ")"))
        })
        .collect()
}

/// Every object body of the document as raw bytes, indexed by object number.
///
/// The objects are located through the xref table rather than by scanning
/// for `N 0 obj`: compressed font programs are binary and can contain that
/// byte sequence by chance.
fn objects(pdf: &[u8]) -> Vec<Vec<u8>> {
    let mut out: Vec<Vec<u8>> = vec![Vec::new()];
    for offset in offsets(pdf) {
        let Some(body) = pdf.get(offset..) else { continue };
        let head = find(body, b" obj", 0).unwrap_or(body.len()) + 4;
        // A stream object ends at the first `endstream` after its data; a
        // binary payload can contain the bytes of `endobj` by chance.
        let dict_end = find(body, b"endobj", 0);
        let stream_at = find(body, b"stream\n", 0).filter(|at| dict_end.is_none_or(|dict_end| *at < dict_end));
        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),
        };
        let number = out.len();
        if out.len() < number {
            out.resize(number, Vec::new());
        }
        out.push(body.get(head..end).unwrap_or_default().to_vec());
    }
    out
}

/// The file offsets of every object, read from the xref table.
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: usize = std::str::from_utf8(head.get(..head_end).unwrap_or_default())
        .unwrap_or_default()
        .split_whitespace()
        .nth(1)
        .and_then(|value| value.parse().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()
}

/// The body of one object as text, or an empty string. The table is padded
/// so that entry `n` holds object `n 0 obj`.
fn body(objects: &[Vec<u8>], number: usize) -> String {
    String::from_utf8_lossy(raw(objects, number)).to_string()
}

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

/// The object number a `/Name 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()
}

/// The `ToUnicode` glyph id to text maps of every font resource.
fn cmaps(objects: &[Vec<u8>]) -> BTreeMap<String, BTreeMap<u16, String>> {
    let names = resource_names(objects);
    let mut out = BTreeMap::new();
    for number in 1..objects.len() {
        let object = body(objects, number);
        if !object.contains("/Subtype /Type0") {
            continue;
        }
        let Some(target) = reference(&object, "/ToUnicode") else { continue };
        let Some(text) = inflate(raw(objects, target)) else { continue };
        out.insert(names.get(&number).cloned().unwrap_or_default(), parse_cmap(&text));
    }
    out
}

/// The `/F<n>` name each font object is registered under in the page
/// resource dictionaries.
fn resource_names(objects: &[Vec<u8>]) -> BTreeMap<usize, String> {
    let mut out = BTreeMap::new();
    for number in 1..objects.len() {
        let object = body(objects, number);
        if !object.contains("/Type /Page ") {
            continue;
        }
        let Some(dict) = object.split("/Font <<").nth(1).and_then(|rest| rest.split(">>").next()) else {
            continue;
        };
        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 bytes of a stream object body.
fn inflate(object: &[u8]) -> Option<String> {
    let start = find(object, b"stream\n", 0)? + b"stream\n".len();
    let end = object.windows(10).rposition(|probe| probe == b"\nendstream")?;
    let data = object.get(start..end)?;
    let mut out = Vec::new();
    flate2::read::ZlibDecoder::new(data).read_to_end(&mut out).ok()?;
    String::from_utf8(out).ok()
}

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