imprimatur 0.6.0 → 0.7.0

data/ext/imprimatur_native/core/src/read/pdf/isolation/worker.rs

@@ -1,43 +1,74 @@
-//! The companion worker sets limits before touching the document.
+//! Apply resource limits before parsing any document bytes.
 
-use std::io::Read;
+use std::io::{Read, Write};
 use std::process::ExitCode;
 
-use super::{protocol, system};
-use crate::read::pdf::render_page;
+use super::protocol::{Operation, Request};
+use super::{protocol, stream, system};
+use crate::read::pdf::{self, bounded};
 use crate::{Error, Result};
 
-/// Internal entry point for the bundled worker executable.
+/// Internal entry point for the bundled PDF worker executable.
 #[doc(hidden)]
 #[must_use]
 pub fn worker_main() -> ExitCode {
-    match protocol::write_reply(std::io::stdout().lock(), run()) {
+    let mut out = std::io::stdout().lock();
+    let result = Request::parse().and_then(|request| run(&request, &mut out));
+    match result.or_else(|error| protocol::write_reply(&mut out, Err(error))) {
         Ok(()) => ExitCode::SUCCESS,
         Err(error) => {
-            eprintln!("imprimatur render worker: {error}");
+            eprintln!("imprimatur PDF worker: {error}");
             ExitCode::FAILURE
         }
     }
 }
 
-fn run() -> Result<Vec<u8>> {
-    let request = protocol::Request::parse()?;
-    let output = request
-        .options
-        .max_output_bytes
-        .max(protocol::MAX_ERROR_BYTES)
-        .checked_add(protocol::HEADER.len())
-        .and_then(|value| u64::try_from(value).ok())
-        .ok_or(Error::InvalidPdfOption("output limit is too large"))?;
+fn run(request: &Request, out: &mut impl Write) -> Result<()> {
+    let output = u64::try_from(request.operation.maximum()?)
+        .map_err(|_| Error::InvalidPdfOption("worker output limit is too large"))?;
     system::apply_limits(request.memory, request.cpu_seconds, output)?;
-    let mut bytes = Vec::new();
+    #[cfg(any(target_os = "linux", target_os = "freebsd"))]
+    if !request.restricted {
+        system::reexec(&request.arguments())?;
+    }
     let maximum = u64::try_from(request.limits.max_input)
         .ok()
         .and_then(|value| value.checked_add(1))
         .ok_or(Error::InvalidPdfOption("input limit is too large"))?;
+    let mut bytes = Vec::new();
     std::io::stdin().lock().take(maximum).read_to_end(&mut bytes)?;
     if bytes.len() > request.limits.max_input {
         return Err(Error::LimitExceeded("input size"));
     }
-    render_page(&bytes, request.number, &request.options, &request.limits).map(|image| image.data)
+    match &request.operation {
+        Operation::Info => protocol::write_reply(
+            out,
+            pdf::info(&bytes, &request.limits).and_then(|info| {
+                serde_json::to_vec(&info)
+                    .map_err(|error| Error::malformed("PDF worker", error.to_string()))
+            }),
+        ),
+        Operation::Render(number, options) => protocol::write_reply(
+            out,
+            pdf::render_page(&bytes, *number, options, &request.limits).map(|image| image.data),
+        ),
+        Operation::Pages(options) => {
+            out.write_all(protocol::HEADER)?;
+            let result = pdf::info(&bytes, &request.limits).and_then(|info| {
+                stream::write(out, &stream::Frame::Info(info))?;
+                bounded::visit_pages(&bytes, options, &request.limits, true, |page| {
+                    stream::write(out, &stream::Frame::Page(page))
+                })?;
+                Ok(())
+            });
+            let frame = match result {
+                Ok(()) => stream::Frame::End,
+                Err(error) => {
+                    let (kind, message) = protocol::failure(&error);
+                    stream::Frame::Error { kind, message }
+                }
+            };
+            stream::write(out, &frame)
+        }
+    }
 }

data/ext/imprimatur_native/core/src/read/pdf/isolation.rs

@@ -1,7 +1,8 @@
-//! Render in a separate, resource-limited process. No untrusted PDF parsing
+//! Read and render in a resource-limited process. No untrusted PDF parsing
 //! occurs in the supervisor.
 
 mod protocol;
+mod stream;
 mod system;
 #[cfg(test)]
 mod tests;
@@ -14,11 +15,47 @@ use std::path::Path;
 use std::process::{Child, Command, ExitStatus, Stdio};
 use std::time::{Duration, Instant};
 
-use super::RenderOptions;
+use super::{Info, Pages, ReadOptions, RenderOptions};
 use crate::read::{Budget, Image, Limits};
 use crate::{Error, Result};
 pub use worker::worker_main;
 
+/// Inspects a PDF in the trusted, resource-limited worker.
+///
+/// # Errors
+/// Invalid options, worker errors, or exceeded input, memory or time limits.
+pub fn info_isolated(
+    worker: &Path,
+    data: &[u8],
+    limits: &Limits,
+    process: &ProcessLimits,
+) -> Result<Info> {
+    let output = invoke(worker, data, protocol::Operation::Info, limits, process)?;
+    let bytes = output.complete(protocol::MAX_ERROR_BYTES)?;
+    serde_json::from_slice(&bytes)
+        .map_err(|error| Error::malformed("PDF worker", error.to_string()))
+}
+
+/// Reads PDF pages in the worker, retaining completed records on resource failure.
+///
+/// Page-local stream limits trigger bounded individual recovery.
+/// All attempts share one external wall-clock and process-memory limit.
+///
+/// # Errors
+/// As [`info_isolated`], and cancellation or an invalid worker reply.
+/// A failure before page metadata is available remains a document error.
+pub fn pages_isolated(
+    worker: &Path,
+    data: &[u8],
+    options: &ReadOptions,
+    limits: &Limits,
+    process: &ProcessLimits,
+) -> Result<Pages> {
+    let output =
+        invoke(worker, data, protocol::Operation::Pages(options.clone()), limits, process)?;
+    stream::decode(output, options)
+}
+
 /// Mandatory process limits, in addition to the reader/render limits.
 #[derive(Debug, Clone)]
 pub struct ProcessLimits {
@@ -51,6 +88,44 @@ pub fn render_page_isolated(
     limits: &Limits,
     process: &ProcessLimits,
 ) -> Result<Image> {
+    let output = invoke(
+        worker,
+        data,
+        protocol::Operation::Render(number, options.clone()),
+        limits,
+        process,
+    )?;
+    let data = output.complete(options.max_output_bytes)?;
+    // The trusted worker writes only PNG; reject malformed/empty replies too.
+    if !data.starts_with(b"\x89PNG\r\n\x1a\n") || data.len() < 26 {
+        return Err(Error::malformed("PDF worker", "invalid PNG response"));
+    }
+    let dimension = |at| u32::from_be_bytes(data[at..at + 4].try_into().unwrap_or([0; 4]));
+    let (width, height) = (dimension(16), dimension(20));
+    options.image(data, number, (width, height))
+}
+
+struct Output {
+    data: Vec<u8>,
+    failure: Option<Error>,
+}
+
+impl Output {
+    fn complete(self, maximum: usize) -> Result<Vec<u8>> {
+        if let Some(error) = self.failure {
+            return Err(error);
+        }
+        protocol::decode(self.data, maximum)
+    }
+}
+
+fn invoke(
+    worker: &Path,
+    data: &[u8],
+    operation: protocol::Operation,
+    limits: &Limits,
+    process: &ProcessLimits,
+) -> Result<Output> {
     let budget = Budget::new(limits, data.len())?;
     budget.check()?;
     let deadline = Instant::now()
@@ -59,10 +134,12 @@ pub fn render_page_isolated(
     if !worker.is_absolute() || !worker.is_file() {
         return Err(Error::InvalidPdfOption("an absolute path to the render worker is required"));
     }
-    if process.memory_bytes == 0 || process.wall_time.is_zero() || number == 0 {
-        return Err(Error::InvalidPdfOption("process limits and page number must be positive"));
+    if process.memory_bytes == 0 || process.wall_time.is_zero() {
+        return Err(Error::InvalidPdfOption("process limits must be positive"));
     }
-    let request = protocol::Request::new(number, options, limits, process)?;
+    let request = protocol::Request::new(operation, limits, process)?;
+    let maximum = u64::try_from(request.operation.maximum()?)
+        .map_err(|_| Error::InvalidPdfOption("worker output limit is too large"))?;
     let directory = tempfile::tempdir()?;
     let input_path = directory.path().join("input.pdf");
     File::create(&input_path)?.write_all(data)?;
@@ -77,30 +154,20 @@ pub fn render_page_isolated(
         .process_group(0)
         .spawn()?;
     let mut child = Running { child, finished: false };
-    let status = supervise(&mut child, &budget, deadline, process.memory_bytes)?;
-    if !status.success() {
-        return Err(Error::LimitExceeded("render worker resources"));
-    }
-    let maximum = options
-        .max_output_bytes
-        .max(protocol::MAX_ERROR_BYTES)
-        .checked_add(protocol::HEADER.len())
-        .ok_or(Error::InvalidPdfOption("output limit is too large"))?;
-    let maximum = u64::try_from(maximum).unwrap_or(u64::MAX);
+    let failure = match supervise(&mut child, &budget, deadline, process.memory_bytes) {
+        Ok(status) if status.success() => check_finished(&budget, deadline).err(),
+        Ok(_) => Some(Error::LimitExceeded("PDF worker resources")),
+        Err(Error::Cancelled) => return Err(Error::Cancelled),
+        Err(error) => Some(error),
+    };
+    // Stop and reap an unfinished worker before inspecting its flushed output.
+    drop(child);
     if fs::metadata(&output_path)?.len() > maximum {
-        return Err(Error::LimitExceeded("render worker output"));
+        return Err(Error::LimitExceeded("PDF worker output"));
     }
     let mut output = Vec::new();
     File::open(output_path)?.take(maximum).read_to_end(&mut output)?;
-    check_finished(&budget, deadline)?;
-    let data = protocol::decode(output, options.max_output_bytes)?;
-    // The trusted worker writes only PNG; also reject malformed/empty replies.
-    if !data.starts_with(b"\x89PNG\r\n\x1a\n") || data.len() < 26 {
-        return Err(Error::malformed("render worker", "invalid PNG response"));
-    }
-    let dimension = |at| u32::from_be_bytes(data[at..at + 4].try_into().unwrap_or([0; 4]));
-    let (width, height) = (dimension(16), dimension(20));
-    options.image(data, number, (width, height))
+    Ok(Output { data: output, failure })
 }
 
 fn check(budget: &Budget<'_>, deadline: Instant) -> Result<()> {
@@ -155,12 +222,34 @@ fn supervise(
         check(budget, deadline)?;
         #[cfg(target_os = "macos")]
         if let Err(error) = system::check_memory(child.child.id(), memory) {
-            // It may have exited between try_wait and the memory query.
-            if let Some(status) = child.poll()? {
-                return Ok(status);
+            if matches!(error, Error::Io(_)) {
+                return finish_monitor_error(child, budget, deadline, error);
             }
             return Err(error);
         }
         std::thread::sleep(Duration::from_millis(5));
     }
 }
+
+#[cfg(target_os = "macos")]
+fn finish_monitor_error(
+    child: &mut Running,
+    budget: &Budget<'_>,
+    deadline: Instant,
+    error: Error,
+) -> Result<ExitStatus> {
+    // proc_pidinfo can stop seeing an exiting task before waitpid can reap it.
+    // Never bypass monitoring for a live worker: allow only a bounded exit
+    // grace, still respecting cancellation and the original overall deadline.
+    let grace = Instant::now() + Duration::from_millis(50);
+    loop {
+        if let Some(status) = child.poll()? {
+            return Ok(status);
+        }
+        check(budget, deadline)?;
+        if Instant::now() >= grace {
+            return Err(error);
+        }
+        std::thread::sleep(Duration::from_millis(5));
+    }
+}

data/ext/imprimatur_native/core/src/read/pdf/layout.rs

@@ -20,6 +20,7 @@ use crate::script::Counts;
 const EDGE: f64 = 0.09;
 /// A gap wider than this share of the font size separates words.
 const WORD_GAP: f64 = 0.18;
+const POSITION_BUCKET: f64 = 3.0;
 
 /// A line of a page, its text a span of the page's buffer.
 #[derive(Debug, Clone)]
@@ -36,6 +37,9 @@ pub struct Line {
     /// In the top or bottom margin of an upright page: a hash of the text,
     /// digits aside, to spot running headers and footers. 0 elsewhere.
     pub margin: u64,
+    /// A hash of the text and its bucketed position and direction, to spot
+    /// anything else that repeats at the same spot on most pages.
+    pub repeat: u64,
 }
 
 /// A page's lines and their text (one buffer for all of them), its height
@@ -82,14 +86,9 @@ pub fn lines(page: &PageText) -> PageLines {
         glyphs.sort_by(|a, b| b.1.total_cmp(&a.1).then(a.0.total_cmp(&b.0)));
         let mut start = 0;
         while start < glyphs.len() {
-            let (_, baseline, first) = glyphs[start];
-            let mut end = start + 1;
-            while end < glyphs.len() {
-                let (_, y, glyph) = glyphs[end];
-                if (baseline - y).abs() > 0.45 * first.size.max(glyph.size) {
-                    break;
-                }
-                end += 1;
+            let mut end = row_end(&glyphs, start);
+            if let Some(merged) = attach_superscript(&glyphs, start, end) {
+                end = merged;
             }
             let row = &mut glyphs[start..end];
             row.sort_by(|a, b| a.0.total_cmp(&b.0));
@@ -108,6 +107,54 @@ pub fn lines(page: &PageText) -> PageLines {
     out
 }
 
+fn row_end(glyphs: &[(f64, f64, &Glyph)], start: usize) -> usize {
+    let (_, baseline, first) = glyphs[start];
+    let mut end = start + 1;
+    while end < glyphs.len() {
+        let (_, y, glyph) = glyphs[end];
+        if (baseline - y).abs() > 0.45 * first.size.max(glyph.size) {
+            break;
+        }
+        end += 1;
+    }
+    end
+}
+
+const SUPERSCRIPT_GLYPHS: usize = 3;
+const SUPERSCRIPT_SIZE_RATIO: f64 = 0.85;
+
+fn attach_superscript(glyphs: &[(f64, f64, &Glyph)], start: usize, end: usize) -> Option<usize> {
+    if end - start > SUPERSCRIPT_GLYPHS || end >= glyphs.len() {
+        return None;
+    }
+    let row = &glyphs[start..end];
+    let (_, baseline, _) = glyphs[start];
+    let (next_x, next_y, next) = glyphs[end];
+    let base_size = next.size;
+    let candidate_size = row.iter().map(|&(_, _, glyph)| glyph.size).fold(0.0, f64::max);
+    let raised = baseline - next_y;
+    if raised <= 0.0 || raised > base_size.max(candidate_size) {
+        return None;
+    }
+    if row.iter().any(|&(_, _, glyph)| glyph.size > SUPERSCRIPT_SIZE_RATIO * base_size) {
+        return None;
+    }
+    let row_left = row.iter().map(|&(x, _, _)| x).fold(f64::MAX, f64::min);
+    let row_right = row.iter().map(|&(x, _, glyph)| x + glyph.width).fold(f64::MIN, f64::max);
+    let next_right = next_x + next.width;
+    let gap = if row_left >= next_right {
+        row_left - next_right
+    } else if next_x >= row_right {
+        next_x - row_right
+    } else {
+        return None;
+    };
+    if gap > WORD_GAP * base_size.max(1.0) {
+        return None;
+    }
+    Some(row_end(glyphs, end))
+}
+
 /// Whether a page reads right to left: more of its letters, those in the
 /// top and bottom margins of upright text aside, are right to left.
 fn right_to_left(page: &PageText) -> bool {
@@ -179,7 +226,9 @@ fn line<'p>(
     let (left, y, _) = row[0];
     let span = span::end(&out.text[..start])..span::end(&out.text);
     let bold = bold * 2 > count;
-    out.lines.push(Line { span, y, left, right, size, bold, turns, margin: 0 });
+    let text = span::slice(&out.text, &span);
+    let repeat = repeat_key(text, left, y, turns);
+    out.lines.push(Line { span, y, left, right, size, bold, turns, margin: 0, repeat });
 }
 
 /// A hash of a line's text with digits as `#` and ASCII letters lowercased
@@ -192,18 +241,39 @@ fn margin_key(text: &str) -> u64 {
     hash.max(1)
 }
 
+#[expect(clippy::cast_possible_truncation, reason = "page coordinates fit in i64")]
+fn repeat_key(text: &str, left: f64, y: f64, turns: u8) -> u64 {
+    let bucket = |value: f64| (value / POSITION_BUCKET).round() as i64;
+    let mut hash = margin_key(text);
+    for part in [u64::from(turns), bucket(left).cast_unsigned(), bucket(y).cast_unsigned()] {
+        hash = (hash ^ part).wrapping_mul(0x0100_0000_01b3);
+    }
+    hash.max(1)
+}
+
 /// The document's pages, from their lines: blocks of the lines that are
-/// not running headers or footers.
-pub fn pages(pages: &[PageLines]) -> Vec<Page> {
-    let boilerplate = running_lines(pages);
-    let kept = |line: &&Line| line.margin == 0 || boilerplate.binary_search(&line.margin).is_err();
-    let body = body_size(pages, &kept);
-    let levels = heading_sizes(pages, &kept, body);
+/// not running headers or footers, and, with `drop_repeated_anywhere`, not
+/// text repeated at the same spot on most pages anywhere on the page.
+pub fn pages(
+    pages: &[PageLines],
+    labels: Option<&[Option<String>]>,
+    drop_repeated_anywhere: bool,
+) -> Vec<Page> {
+    let boilerplate = repeated_keys(pages, |line| line.margin);
+    let margin_kept =
+        |line: &&Line| line.margin == 0 || boilerplate.binary_search(&line.margin).is_err();
+    let repeated =
+        if drop_repeated_anywhere { repeated_keys(pages, |line| line.repeat) } else { Vec::new() };
+    let kept = |line: &&Line| {
+        margin_kept(line) && (repeated.is_empty() || repeated.binary_search(&line.repeat).is_err())
+    };
+    let body = body_size(pages, &margin_kept);
+    let levels = heading_sizes(pages, &margin_kept, body);
     let mut out = Vec::with_capacity(pages.len());
     let mut paragraph: Vec<&Line> = Vec::new();
     // Each page's blocks, collected here, then moved out at their size.
     let mut blocks = Vec::new();
-    for (number, page) in (1..).zip(pages) {
+    for (index, (number, page)) in (1..).zip(pages).enumerate() {
         paragraph.clear();
         for line in page.lines.iter().filter(kept) {
             if let Some(last) = paragraph.last()
@@ -223,11 +293,39 @@ pub fn pages(pages: &[PageLines]) -> Vec<Page> {
             reason = "`blocks` keeps its buffer for the next page"
         )]
         let blocks = blocks.drain(..).collect();
-        out.push(Page { number, blocks, images: page.images });
+        let label = labels.map_or_else(
+            || footer_label(page, &boilerplate),
+            |labels| labels.get(index).cloned().flatten(),
+        );
+        out.push(Page { number, blocks, images: page.images, label });
     }
     out
 }
 
+fn footer_label(page: &PageLines, boilerplate: &[u64]) -> Option<String> {
+    let mut candidates = page.lines.iter().filter(|line: &&Line| {
+        line.margin != 0
+            && line.y < page.height * EDGE
+            && boilerplate.binary_search(&line.margin).is_err()
+    });
+    let line = candidates.next()?;
+    if candidates.next().is_some() {
+        return None;
+    }
+    let text = page.text(line).trim();
+    looks_like_page_number(text).then(|| text.to_owned())
+}
+
+fn looks_like_page_number(text: &str) -> bool {
+    let core = text.trim_matches(|ch: char| !ch.is_alphanumeric());
+    !core.is_empty()
+        && core.len() <= 8
+        && (core.chars().all(|ch| ch.is_ascii_digit())
+            || core.chars().all(|ch| {
+                matches!(ch.to_ascii_lowercase(), 'i' | 'v' | 'x' | 'l' | 'c' | 'd' | 'm')
+            }))
+}
+
 /// Whether `next` carries on the paragraph `last` is in: same direction,
 /// similar size, and no more than a line and a half below it.
 fn continues(last: &Line, next: &Line) -> bool {
@@ -245,41 +343,79 @@ fn continues(last: &Line, next: &Line) -> bool {
 fn block(page: &PageLines, lines: &[&Line], body: f64, levels: &[f64]) -> Placed {
     let first = lines[0];
     let length: usize = lines.iter().map(|line| line.span.len()).sum();
-    let level = if lines.len() <= 3 && length <= 200 {
-        levels.iter().position(|&size| first.size >= size - 0.25)
-    } else {
-        None
+    let short = lines.len() <= 3 && length <= 200;
+    let text = paragraph_text(page, lines);
+    let level =
+        if short { levels.iter().position(|&size| first.size >= size - 0.25) } else { None };
+    let sized = level
+        .filter(|_| first.size >= body * 1.15)
+        .map(|level| u8::try_from(level.min(2) + 1).unwrap_or(3));
+    let numbered = (short && first.bold && first.size >= body - 0.25)
+        .then(|| numbered_heading_level(&text))
+        .flatten();
+    let block = match sized.or(numbered) {
+        Some(level) => Block::Heading { level, text: Text::plain(text) },
+        None => Block::Paragraph { text: Text::plain(text) },
     };
+    Placed { block, bounds: bounds(page, lines) }
+}
+
+fn paragraph_text(page: &PageLines, lines: &[&Line]) -> String {
+    let length: usize = lines.iter().map(|line| line.span.len()).sum();
     let mut text = String::with_capacity(length + lines.len());
-    let block = if let Some(level) = level.filter(|_| first.size >= body * 1.15) {
-        for (index, line) in lines.iter().enumerate() {
-            if index > 0 {
-                text.push(' ');
+    for (index, line) in lines.iter().enumerate() {
+        let line = page.text(line);
+        if index > 0 {
+            let joined = text.ends_with('-') && {
+                let mut before = text.chars().rev().skip(1);
+                before.next().is_some_and(char::is_lowercase)
+                    && line.chars().next().is_some_and(char::is_lowercase)
+            };
+            if joined {
+                text.pop();
+            } else {
+                text.push('\n');
             }
-            text.push_str(page.text(line));
         }
-        let level = u8::try_from(level.min(2) + 1).unwrap_or(3);
-        Block::Heading { level, text: Text::plain(text) }
-    } else {
-        for (index, line) in lines.iter().enumerate() {
-            let line = page.text(line);
-            if index > 0 {
-                let joined = text.ends_with('-') && {
-                    let mut before = text.chars().rev().skip(1);
-                    before.next().is_some_and(char::is_lowercase)
-                        && line.chars().next().is_some_and(char::is_lowercase)
-                };
-                if joined {
-                    text.pop();
-                } else {
-                    text.push('\n');
-                }
-            }
-            text.push_str(line);
+        text.push_str(line);
+    }
+    text
+}
+
+fn numbered_heading_level(text: &str) -> Option<u8> {
+    let (depth, rest) = numbered_prefix(text)?;
+    if depth < 2 || !(rest.is_empty() || rest.starts_with(' ')) {
+        return None;
+    }
+    let title = rest.trim_start();
+    if title.contains("....") || title.starts_with(|ch: char| !ch.is_alphanumeric()) {
+        return None;
+    }
+    Some(u8::try_from(depth).unwrap_or(6).min(6))
+}
+
+fn numbered_prefix(text: &str) -> Option<(usize, &str)> {
+    let bytes = text.as_bytes();
+    let mut groups = 0_usize;
+    let mut pos = 0_usize;
+    loop {
+        let start = pos;
+        while bytes.get(pos).is_some_and(u8::is_ascii_digit) {
+            pos += 1;
         }
-        Block::Paragraph { text: Text::plain(text) }
-    };
-    Placed { block, bounds: bounds(page, lines) }
+        if pos == start {
+            break;
+        }
+        groups += 1;
+        if groups >= 4
+            || bytes.get(pos) != Some(&b'.')
+            || !bytes.get(pos + 1).is_some_and(u8::is_ascii_digit)
+        {
+            break;
+        }
+        pos += 1;
+    }
+    (groups > 0).then(|| (groups, &text[pos..]))
 }
 
 /// Where upright lines lie on the page, from its top-left corner: from the
@@ -296,9 +432,9 @@ fn bounds(page: &PageLines, lines: &[&Line]) -> Option<Bounds> {
     Some(Bounds { left, top, right, bottom })
 }
 
-/// Text that repeats in the margins of many pages (page numbers, running
-/// titles), digits aside: their margin keys, sorted.
-fn running_lines(pages: &[PageLines]) -> Vec<u64> {
+/// The keys `key` returns for lines that repeat on enough pages to be
+/// boilerplate rather than coincidence, sorted.
+fn repeated_keys(pages: &[PageLines], key: impl Fn(&Line) -> u64) -> Vec<u64> {
     if pages.len() < 3 {
         return Vec::new();
     }
@@ -306,7 +442,7 @@ fn running_lines(pages: &[PageLines]) -> Vec<u64> {
     let mut seen: Vec<u64> = Vec::new();
     for page in pages {
         seen.clear();
-        seen.extend(page.lines.iter().map(|line| line.margin).filter(|&key| key != 0));
+        seen.extend(page.lines.iter().map(&key).filter(|&key| key != 0));
         seen.sort_unstable();
         seen.dedup();
         for &key in &seen {
@@ -314,10 +450,10 @@ fn running_lines(pages: &[PageLines]) -> Vec<u64> {
         }
     }
     let threshold = (pages.len() * 2 / 5).max(3);
-    let mut running: Vec<u64> =
+    let mut repeated: Vec<u64> =
         counts.into_iter().filter(|&(_, count)| count >= threshold).map(|(key, _)| key).collect();
-    running.sort_unstable();
-    running
+    repeated.sort_unstable();
+    repeated
 }
 
 /// Every page's lines that `kept` keeps.

data/ext/imprimatur_native/core/src/read/pdf.rs

@@ -57,5 +57,6 @@ pub(super) fn content(data: &[u8], budget: &Budget<'_>) -> Result<Vec<Page>> {
             Err(_) => lines.push(layout::PageLines::default()),
         }
     }
-    Ok(layout::pages(&lines))
+    let labels = document.page_labels(lines.len())?;
+    Ok(layout::pages(&lines, labels.as_deref(), budget.drop_repeated()))
 }

data/ext/imprimatur_native/ffi/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "imprimatur_native"
-version = "0.6.0"
+version = "0.7.0"
 description = "Ruby bindings for the imprimatur crate"
 publish = false
 edition.workspace = true

data/ext/imprimatur_native/ffi/src/lib.rs

@@ -326,8 +326,8 @@ fn init(ruby: &Ruby) -> Result<(), Error> {
     native.define_singleton_method("markdown", function!(read::markdown_native, 4))?;
     native.define_singleton_method("content", function!(read::content_native, 5))?;
     native.define_singleton_method("images", function!(read::images_native, 4))?;
-    native.define_singleton_method("pdf_info", function!(pdf::info, 2))?;
-    native.define_singleton_method("pdf_pages", function!(pdf::pages, 3))?;
+    native.define_singleton_method("pdf_info", function!(pdf::info, 4))?;
+    native.define_singleton_method("pdf_pages", function!(pdf::pages, 5))?;
     native.define_singleton_method("render_pdf_page", function!(pdf::render, 5))?;
     native.define_singleton_method("sheets", function!(sheets::sheets_native, 5))?;
     let zip_writer = native.define_class("ZipWriter", ruby.class_object())?;