Skip to main content

vole_document/adapter/pdf/
layout.rs

1//! PDF layout candidate: regenerate classic cross-reference entry offsets and
2//! the `startxref` value from positions marked during materialization.
3//!
4//! This is the first Phase-5 candidate that replaces literal structural bytes
5//! with *procedurally determined* ones. The mechanism is deliberately narrow and
6//! conservative:
7//!
8//! * It applies only to files that already carry a classic cross-reference
9//!   section ([`PhysicalKind::XrefSection`]) and contain no
10//!   [`ObjRole::XRefStream`] object. Anything else declines (`Ok(None)`).
11//! * Every literal byte is accumulated into a single data object and the whole
12//!   reconstruction is one compact [`Op::PackSegments`] item table, so the
13//!   per-segment framing is paid once rather than once per span or xref entry.
14//! * It records each indirect object's introducer offset with a
15//!   [`PackItem::Mark`] (slot = the object's index in [`PdfPhysical::objects`])
16//!   and, for every `n`-status xref entry whose 10-digit offset field equals the
17//!   marked position of its target object, emits that field with a
18//!   [`PackItem::Emit`] rather than storing the digits literally.
19//! * The most recent `xref` section start is marked in the reserved slot
20//!   [`XREF_SLOT`] (`255`); each `startxref` value is regenerated from it only
21//!   when the emitted position equals the source value.
22//! * Whenever a precondition fails — a non-standard offset field, a mismatched
23//!   position, a malformed table, too many objects — the site (or the whole
24//!   section) falls back to a literal [`PackItem::Literal`]. Prediction never
25//!   invents bytes: a fallback is always byte-exact, and a prediction is only
26//!   emitted when it reproduces the source digits exactly.
27//!
28//! After building the program the candidate is verified end-to-end (serialize,
29//! parse, materialize, byte-compare) before it is returned; if that round trip
30//! is not exact, the proposal declines rather than emitting an inexact
31//! candidate.
32
33use crate::SOURCE_FORMAT_PDF;
34use crate::container::{Descriptor, UNIVERSE};
35use crate::dra::op::PackItem;
36use crate::dra::{Op, Program};
37use crate::encode::candidates::{Candidate, CandidateKind};
38use crate::error::Result;
39use crate::integrity::sha256;
40use crate::limits::Limits;
41
42#[cfg(feature = "rans")]
43use crate::entropy::{
44    ALPHABET, CODER_ORDER0_BYTE_RANS, CODER_VERSION_1, EntropyChannelDescriptor, EntropyModel,
45    encode_channel,
46};
47
48use super::physical::{ObjRole, PdfObjectSpan, PhysicalKind, scan};
49
50/// Reserved slot index for the most recent classic `xref` section start.
51pub const XREF_SLOT: u8 = u8::MAX;
52
53/// Largest number of indirect objects that can be marked (indices `0..=254`),
54/// leaving slot `255` free for [`XREF_SLOT`].
55pub const MAX_MARKED_OBJECTS: usize = XREF_SLOT as usize;
56
57/// The structural layout plan for a classic-cross-reference PDF: the ordered
58/// packed item table, the single literal data object it consumes, and the
59/// prediction counters used to describe the plan.
60pub struct LayoutPlan {
61    /// Ordered reconstruction items (literal runs, position marks, emitted
62    /// offsets), as consumed by [`Op::PackSegments`] and [`Op::PackedChannels`].
63    pub items: Vec<PackItem>,
64    /// Every literal byte, in item order. Must be consumed exactly by `items`.
65    pub data: Vec<u8>,
66    /// Number of xref entry offsets regenerated from a marked position.
67    pub xref_predicted: usize,
68    /// Number of xref entry offsets stored literally (precondition failed).
69    pub xref_literal: usize,
70    /// Whether at least one `startxref` value was regenerated.
71    pub startxref_predicted: bool,
72}
73
74/// Build the structural layout plan for `input`, or `None` when the input is not
75/// a classic-cross-reference PDF this mechanism can express exactly.
76///
77/// Declines (`Ok(None)`) whenever a precondition fails — no classic `xref`
78/// section, a cross-reference stream present, too many objects, a non-contiguous
79/// span cover, or a literal run that cannot fit a `u32`. See the module
80/// documentation for the algorithm. The caller is responsible for the final
81/// byte-exactness check through the normative decoder.
82pub fn build_layout_plan(input: &[u8], limits: Limits) -> Result<Option<LayoutPlan>> {
83    let physical = match scan(input, limits) {
84        Ok(p) => p,
85        Err(_) => return Ok(None),
86    };
87
88    // Precondition: a classic cross-reference section must exist, and no
89    // cross-reference stream may be present.
90    if !physical
91        .spans
92        .iter()
93        .any(|s| s.kind == PhysicalKind::XrefSection)
94    {
95        return Ok(None);
96    }
97    if physical
98        .objects
99        .iter()
100        .any(|o| o.role == ObjRole::XRefStream)
101    {
102        return Ok(None);
103    }
104    if physical.objects.len() > MAX_MARKED_OBJECTS {
105        return Ok(None);
106    }
107
108    // The data object carries every literal byte; the compact item table
109    // interleaves literal runs, position marks, and regenerated offsets so the
110    // per-segment framing is paid once rather than once per span/entry.
111    let mut data: Vec<u8> = Vec::new();
112    let mut items: Vec<PackItem> = Vec::new();
113    // Simulated output position. The physical cover is contiguous, so this
114    // tracks the source offset of the next byte exactly: literals add their
115    // length, emits add their width, and marks add nothing.
116    let mut pos: u64 = 0;
117    let mut slot_value = [0u64; 256];
118    let mut slot_marked = [false; 256];
119    let mut xref_predicted: usize = 0;
120    let mut xref_literal: usize = 0;
121    let mut startxref_predicted = false;
122
123    for span in &physical.spans {
124        let start = span.start as usize;
125        let end = start + span.len as usize;
126        let bytes = &input[start..end];
127
128        if pos != span.start {
129            // A non-contiguous simulation would break the offset contract; bail.
130            return Ok(None);
131        }
132
133        match span.kind {
134            PhysicalKind::ObjHeader => {
135                if let Some(idx) = object_index_at(&physical.objects, span.start) {
136                    let slot = idx as u8;
137                    items.push(PackItem::Mark { slot });
138                    slot_value[slot as usize] = pos;
139                    slot_marked[slot as usize] = true;
140                }
141                if !push_pack_literal(&mut data, &mut items, bytes, &mut pos) {
142                    return Ok(None);
143                }
144            }
145            PhysicalKind::XrefSection => {
146                // Mark this section's start in the reserved slot, then emit.
147                items.push(PackItem::Mark { slot: XREF_SLOT });
148                slot_value[XREF_SLOT as usize] = pos;
149                slot_marked[XREF_SLOT as usize] = true;
150
151                match parse_classic_xref(bytes) {
152                    Some(pieces) => {
153                        for piece in pieces {
154                            match piece {
155                                XrefPiece::Literal { start, len } => {
156                                    if !push_pack_literal(
157                                        &mut data,
158                                        &mut items,
159                                        &bytes[start..start + len],
160                                        &mut pos,
161                                    ) {
162                                        return Ok(None);
163                                    }
164                                }
165                                XrefPiece::Entry {
166                                    start,
167                                    number,
168                                    offset,
169                                    in_use,
170                                } => {
171                                    let slot = if in_use {
172                                        offset.and_then(|value| {
173                                            predicted_slot(
174                                                &physical.objects,
175                                                &slot_value,
176                                                &slot_marked,
177                                                number,
178                                                value,
179                                            )
180                                        })
181                                    } else {
182                                        None
183                                    };
184                                    match slot {
185                                        Some(slot) => {
186                                            items.push(PackItem::Emit { slot, width: 10 });
187                                            pos += 10;
188                                            if !push_pack_literal(
189                                                &mut data,
190                                                &mut items,
191                                                &bytes[start + 10..start + 20],
192                                                &mut pos,
193                                            ) {
194                                                return Ok(None);
195                                            }
196                                            xref_predicted += 1;
197                                        }
198                                        None => {
199                                            if !push_pack_literal(
200                                                &mut data,
201                                                &mut items,
202                                                &bytes[start..start + 20],
203                                                &mut pos,
204                                            ) {
205                                                return Ok(None);
206                                            }
207                                            xref_literal += 1;
208                                        }
209                                    }
210                                }
211                            }
212                        }
213                    }
214                    None => {
215                        // Not a classic table we understand: literal whole section.
216                        if !push_pack_literal(&mut data, &mut items, bytes, &mut pos) {
217                            return Ok(None);
218                        }
219                    }
220                }
221            }
222            PhysicalKind::StartXref => match predict_startxref(bytes, &slot_value, &slot_marked) {
223                Some((prefix_len, width)) => {
224                    if !push_pack_literal(&mut data, &mut items, &bytes[..prefix_len], &mut pos) {
225                        return Ok(None);
226                    }
227                    items.push(PackItem::Emit {
228                        slot: XREF_SLOT,
229                        width,
230                    });
231                    pos += width as u64;
232                    startxref_predicted = true;
233                }
234                None => {
235                    if !push_pack_literal(&mut data, &mut items, bytes, &mut pos) {
236                        return Ok(None);
237                    }
238                }
239            },
240            _ => {
241                if !push_pack_literal(&mut data, &mut items, bytes, &mut pos) {
242                    return Ok(None);
243                }
244            }
245        }
246    }
247
248    Ok(Some(LayoutPlan {
249        items,
250        data,
251        xref_predicted,
252        xref_literal,
253        startxref_predicted,
254    }))
255}
256
257/// Propose a layout candidate that regenerates xref offsets / startxref, or
258/// `None`.
259///
260/// Wraps [`build_layout_plan`] into a single [`Op::PackSegments`] program over
261/// one literal data object. Declines (`Ok(None)`) whenever the plan cannot be
262/// built or the assembled program does not materialize byte-for-byte. See the
263/// module documentation for the algorithm.
264pub fn propose_pdf_layout(input: &[u8], limits: Limits) -> Result<Option<Candidate>> {
265    let plan = match build_layout_plan(input, limits)? {
266        Some(p) => p,
267        None => return Ok(None),
268    };
269
270    // `startxref` is predicted once per `StartXref` span; every `Emit` is either
271    // a predicted xref entry or a predicted startxref, so the count is exact.
272    let startxref_predicted = emit_count(&plan.items).saturating_sub(plan.xref_predicted);
273    let format_basis = format!(
274        "pdf-layout;objects={};xref_predicted={};xref_literal={};startxref_predicted={}",
275        marked_object_count(&plan.items),
276        plan.xref_predicted,
277        plan.xref_literal,
278        startxref_predicted
279    );
280
281    let descriptor = Descriptor {
282        universe: UNIVERSE.to_string(),
283        source_format: SOURCE_FORMAT_PDF,
284        format_basis,
285        models: vec![],
286        channels: vec![],
287        objects: vec![plan.data],
288        program: Program::new(vec![Op::PackSegments {
289            data_object: 0,
290            items: plan.items,
291        }]),
292        source_sha256: sha256(input),
293        source_len: input.len() as u64,
294    };
295
296    let candidate = Candidate {
297        kind: CandidateKind::PdfLayout,
298        descriptor,
299    };
300
301    // Verify byte-exactness through the normative decoder before returning. An
302    // inexact program must never be emitted.
303    let (encoded, _) = candidate.descriptor.serialize()?;
304    let parsed = match Descriptor::parse(&encoded, limits) {
305        Ok(p) => p,
306        Err(_) => return Ok(None),
307    };
308    let out = match crate::materialize::materialize(&parsed, limits) {
309        Ok(o) => o,
310        Err(_) => return Ok(None),
311    };
312    if out != input {
313        return Ok(None);
314    }
315
316    Ok(Some(candidate))
317}
318
319/// Propose a layout + rANS candidate, or `None`.
320///
321/// Builds the same [`LayoutPlan`] as [`propose_pdf_layout`], then entropy-codes
322/// its parts into two rANS channels: channel `0` carries the plan's literal data
323/// object, and channel `1` carries `encode_items` of the plan's item table. A
324/// single [`Op::PackedChannels`] reconstructs the source from both, so the whole
325/// plan — data *and* item table — pays entropy-coding cost instead of being
326/// stored as literal bytes. Each channel uses its own order-0 byte model
327/// normalized from its own byte histogram at `scale_bits` 12.
328///
329/// Exactly as with the literal layout lane, an end-to-end serialize / parse /
330/// materialize / byte-compare check gates the return: an inexact program yields
331/// `Ok(None)` rather than an inexact candidate.
332#[cfg(feature = "rans")]
333pub fn propose_pdf_layout_rans(input: &[u8], limits: Limits) -> Result<Option<Candidate>> {
334    let plan = match build_layout_plan(input, limits)? {
335        Some(p) => p,
336        None => return Ok(None),
337    };
338
339    let plan_bytes = crate::dra::op::encode_items(&plan.items)?;
340
341    // Channel 0: the literal data object, coded against its own byte histogram.
342    let mut data_counts = [0u64; ALPHABET];
343    for &b in &plan.data {
344        data_counts[b as usize] += 1;
345    }
346    let data_model = EntropyModel::from_counts(&data_counts, 12)?;
347    let data_capsule = encode_channel(&data_model, &plan.data)?;
348
349    // Channel 1: the serialized item table, coded against its own histogram.
350    let mut plan_counts = [0u64; ALPHABET];
351    for &b in &plan_bytes {
352        plan_counts[b as usize] += 1;
353    }
354    let plan_model = EntropyModel::from_counts(&plan_counts, 12)?;
355    let plan_capsule = encode_channel(&plan_model, &plan_bytes)?;
356
357    let data_channel = EntropyChannelDescriptor {
358        coder: CODER_ORDER0_BYTE_RANS,
359        coder_version: CODER_VERSION_1,
360        scale_bits: data_model.scale_bits,
361        lane_count: 1,
362        model_id: 0,
363        symbol_count: data_capsule.symbol_count,
364        decoded_length: data_capsule.decoded_length,
365        initial_state: data_capsule.initial_state,
366        payload: data_capsule.payload,
367    };
368    let plan_channel = EntropyChannelDescriptor {
369        coder: CODER_ORDER0_BYTE_RANS,
370        coder_version: CODER_VERSION_1,
371        scale_bits: plan_model.scale_bits,
372        lane_count: 1,
373        model_id: 1,
374        symbol_count: plan_capsule.symbol_count,
375        decoded_length: plan_capsule.decoded_length,
376        initial_state: plan_capsule.initial_state,
377        payload: plan_capsule.payload,
378    };
379
380    let format_basis = format!(
381        "pdf-layout-rans;objects={};xref_predicted={};xref_literal={}",
382        marked_object_count(&plan.items),
383        plan.xref_predicted,
384        plan.xref_literal
385    );
386
387    let descriptor = Descriptor {
388        universe: UNIVERSE.to_string(),
389        source_format: SOURCE_FORMAT_PDF,
390        format_basis,
391        models: vec![data_model, plan_model],
392        channels: vec![data_channel, plan_channel],
393        objects: vec![],
394        program: Program::new(vec![Op::PackedChannels {
395            data_channel: 0,
396            plan_channel: 1,
397            declared_output_len: input.len() as u64,
398        }]),
399        source_sha256: sha256(input),
400        source_len: input.len() as u64,
401    };
402
403    let candidate = Candidate {
404        kind: CandidateKind::PdfLayoutRans,
405        descriptor,
406    };
407
408    // Verify byte-exactness through the normative decoder before returning. An
409    // inexact program must never be emitted.
410    let (encoded, _) = candidate.descriptor.serialize()?;
411    let parsed = match Descriptor::parse(&encoded, limits) {
412        Ok(p) => p,
413        Err(_) => return Ok(None),
414    };
415    let out = match crate::materialize::materialize(&parsed, limits) {
416        Ok(o) => o,
417        Err(_) => return Ok(None),
418    };
419    if out != input {
420        return Ok(None);
421    }
422
423    Ok(Some(candidate))
424}
425
426/// Number of indirect objects whose introducer was marked in the item table.
427fn marked_object_count(items: &[PackItem]) -> usize {
428    items
429        .iter()
430        .filter(|item| matches!(item, PackItem::Mark { slot } if *slot != XREF_SLOT))
431        .count()
432}
433
434/// Number of emitted offsets in the item table.
435fn emit_count(items: &[PackItem]) -> usize {
436    items
437        .iter()
438        .filter(|item| matches!(item, PackItem::Emit { .. }))
439        .count()
440}
441
442/// Append `bytes` to the packed data object, recording one [`PackItem::Literal`]
443/// when non-empty, and advance the simulated output position. Returns `false`
444/// (so the caller declines) when a single run would not fit a `u32` length.
445fn push_pack_literal(
446    data: &mut Vec<u8>,
447    items: &mut Vec<PackItem>,
448    bytes: &[u8],
449    pos: &mut u64,
450) -> bool {
451    if !push_literal(items, data, bytes) {
452        return false;
453    }
454    *pos += bytes.len() as u64;
455    true
456}
457
458/// Append `bytes` to the packed data object, coalescing them into the immediately
459/// preceding [`PackItem::Literal`] when one is present so that consecutive literal
460/// runs collapse into the fewest possible items. The merge never crosses a
461/// [`PackItem::Mark`] or [`PackItem::Emit`], empty pushes are ignored, and the
462/// combined length must remain expressible as a `u32`. Returns `false` (so the
463/// caller declines) when no safe item shape exists.
464fn push_literal(items: &mut Vec<PackItem>, data: &mut Vec<u8>, bytes: &[u8]) -> bool {
465    if bytes.is_empty() {
466        return true;
467    }
468    let Ok(len) = u32::try_from(bytes.len()) else {
469        return false;
470    };
471    if let Some(PackItem::Literal { len: prev }) = items.last_mut() {
472        let Some(total) = prev.checked_add(len) else {
473            return false;
474        };
475        *prev = total;
476    } else {
477        items.push(PackItem::Literal { len });
478    }
479    data.extend_from_slice(bytes);
480    true
481}
482
483/// Index of the first object whose introducer starts at `start`.
484fn object_index_at(objects: &[PdfObjectSpan], start: u64) -> Option<usize> {
485    objects.iter().position(|o| o.start == start)
486}
487
488/// The slot marking the target object `number` at position `value`, if any
489/// earlier [`Op::MarkOffset`] recorded exactly that position.
490fn predicted_slot(
491    objects: &[PdfObjectSpan],
492    slot_value: &[u64; 256],
493    slot_marked: &[bool; 256],
494    number: u64,
495    value: u64,
496) -> Option<u8> {
497    objects.iter().enumerate().find_map(|(i, o)| {
498        let slot = i as u8;
499        (o.number == number && slot_marked[slot as usize] && slot_value[slot as usize] == value)
500            .then_some(slot)
501    })
502}
503
504/// Predict a whole `startxref` span: the trailing run of decimal digits is the
505/// value. Returns `(prefix_len, width)` when the value equals the marked `xref`
506/// position and its width is in `1..=20`.
507fn predict_startxref(
508    bytes: &[u8],
509    slot_value: &[u64; 256],
510    slot_marked: &[bool; 256],
511) -> Option<(usize, u8)> {
512    if !slot_marked[XREF_SLOT as usize] {
513        return None;
514    }
515    let mut i = bytes.len();
516    while i > 0 && bytes[i - 1].is_ascii_digit() {
517        i -= 1;
518    }
519    let width = bytes.len() - i;
520    if width == 0 || width > 20 {
521        return None;
522    }
523    let value = parse_digits(&bytes[i..])?;
524    if value != slot_value[XREF_SLOT as usize] {
525        return None;
526    }
527    Some((i, width as u8))
528}
529
530/// One ordered slice of a classic `xref` section: either literal bytes or a
531/// 20-byte entry whose offset field may be regenerated.
532enum XrefPiece {
533    /// Verbatim bytes `[start, start + len)` of the section.
534    Literal { start: usize, len: usize },
535    /// A 20-byte entry `[start, start + 20)`.
536    Entry {
537        /// Offset of the entry within the section.
538        start: usize,
539        /// Target object number (`subsection_start + i`).
540        number: u64,
541        /// Parsed value of the 10-digit offset field, if it is all digits.
542        offset: Option<u64>,
543        /// Whether the status byte is `n` (in use).
544        in_use: bool,
545    },
546}
547
548/// Parse a classic cross-reference table from a section's bytes, returning an
549/// ordered tiling of the section. Returns `None` (so the caller emits the whole
550/// section literally) when the bytes do not match the classic grammar.
551///
552/// Grammar accepted: `xref` EOL, then one or more `<start> <count>` EOL headers
553/// each followed by exactly `count` 20-byte entries of the shape
554/// `10-digit-offset SP 5-digit-generation SP status 2-byte-EOL`. The two EOL
555/// bytes may be `CR LF`, `LF CR`, `SP LF`, or `SP CR`.
556fn parse_classic_xref(bytes: &[u8]) -> Option<Vec<XrefPiece>> {
557    if !bytes.starts_with(b"xref") {
558        return None;
559    }
560    let mut pieces = Vec::new();
561    let mut pos = 4usize;
562
563    // EOL after the `xref` keyword.
564    let eol = eol_len(&bytes[pos..])?;
565    pieces.push(XrefPiece::Literal {
566        start: 0,
567        len: pos + eol,
568    });
569    pos += eol;
570
571    let mut any = false;
572    while pos < bytes.len() {
573        // Subsection header: `<start> <count>` EOL.
574        let header_start = pos;
575        let (start, after_start) = parse_uint_at(bytes, pos)?;
576        pos = after_start;
577        let spaces_start = pos;
578        while pos < bytes.len() && bytes[pos] == b' ' {
579            pos += 1;
580        }
581        if pos == spaces_start {
582            return None;
583        }
584        let (count, after_count) = parse_uint_at(bytes, pos)?;
585        pos = after_count;
586        let eol = eol_len(&bytes[pos..])?;
587        let header_end = pos + eol;
588        pieces.push(XrefPiece::Literal {
589            start: header_start,
590            len: header_end - header_start,
591        });
592        pos = header_end;
593
594        for i in 0..count {
595            let end = pos.checked_add(20)?;
596            if end > bytes.len() {
597                return None;
598            }
599            let entry = &bytes[pos..end];
600            if !is_entry_shape(entry) {
601                return None;
602            }
603            let number = start.checked_add(i)?;
604            let in_use = entry[17] == b'n';
605            let offset = parse_digits(&entry[0..10]);
606            pieces.push(XrefPiece::Entry {
607                start: pos,
608                number,
609                offset,
610                in_use,
611            });
612            pos = end;
613        }
614        any = true;
615    }
616
617    if !any || pos != bytes.len() {
618        return None;
619    }
620    Some(pieces)
621}
622
623/// Whether a 20-byte window matches the classic cross-reference entry shape.
624fn is_entry_shape(entry: &[u8]) -> bool {
625    if entry.len() != 20 {
626        return false;
627    }
628    if entry[10] != b' ' || entry[16] != b' ' {
629        return false;
630    }
631    if !entry[11..16].iter().all(u8::is_ascii_digit) {
632        return false;
633    }
634    if entry[17] != b'n' && entry[17] != b'f' {
635        return false;
636    }
637    matches!(
638        (entry[18], entry[19]),
639        (b'\r', b'\n') | (b'\n', b'\r') | (b' ', b'\n') | (b' ', b'\r')
640    )
641}
642
643/// Length of an end-of-line marker at the start of `bytes`, if any.
644fn eol_len(bytes: &[u8]) -> Option<usize> {
645    match bytes {
646        [b'\r', b'\n', ..] => Some(2),
647        [b'\n', ..] | [b'\r', ..] => Some(1),
648        _ => None,
649    }
650}
651
652/// Parse a non-negative decimal integer at `at`, returning `(value, next)`.
653fn parse_uint_at(bytes: &[u8], at: usize) -> Option<(u64, usize)> {
654    let mut i = at;
655    let mut value: u64 = 0;
656    while i < bytes.len() && bytes[i].is_ascii_digit() {
657        value = value
658            .checked_mul(10)?
659            .checked_add(u64::from(bytes[i] - b'0'))?;
660        i += 1;
661    }
662    if i == at {
663        return None;
664    }
665    Some((value, i))
666}
667
668/// Parse an all-digit slice as a non-negative decimal integer.
669fn parse_digits(digits: &[u8]) -> Option<u64> {
670    if digits.is_empty() {
671        return None;
672    }
673    let mut value: u64 = 0;
674    for &b in digits {
675        if !b.is_ascii_digit() {
676            return None;
677        }
678        value = value.checked_mul(10)?.checked_add(u64::from(b - b'0'))?;
679    }
680    Some(value)
681}
682
683#[cfg(test)]
684mod tests {
685    use super::*;
686    use crate::adapter::pdf::samples::{is_negative_control, sample_pdfs};
687    use crate::container::Descriptor;
688
689    fn sample(name: &str) -> Vec<u8> {
690        sample_pdfs()
691            .into_iter()
692            .find(|(n, _)| *n == name)
693            .unwrap_or_else(|| panic!("sample {name} missing"))
694            .1
695    }
696
697    fn basis_field(basis: &str, key: &str) -> Option<u64> {
698        basis.split(';').find_map(|part| {
699            let (k, v) = part.split_once('=')?;
700            (k == key).then(|| v.parse().ok()).flatten()
701        })
702    }
703
704    /// The item table of the single `PackSegments` op this candidate builds.
705    fn pack_items(cand: &Candidate) -> &[PackItem] {
706        cand.descriptor
707            .program
708            .ops
709            .iter()
710            .find_map(|op| match op {
711                Op::PackSegments { items, .. } => Some(items.as_slice()),
712                _ => None,
713            })
714            .expect("layout program is one PackSegments op")
715    }
716
717    /// Number of regenerated offsets in the packed item table.
718    fn pack_emit_count(cand: &Candidate) -> usize {
719        pack_items(cand)
720            .iter()
721            .filter(|item| matches!(item, PackItem::Emit { .. }))
722            .count()
723    }
724
725    fn assert_materializes_exactly(name: &str, bytes: &[u8]) {
726        let cand = propose_pdf_layout(bytes, Limits::DEFAULT)
727            .unwrap()
728            .unwrap_or_else(|| panic!("{name} must propose a layout candidate"));
729        assert_eq!(cand.kind, CandidateKind::PdfLayout);
730        assert_eq!(cand.descriptor.source_format, SOURCE_FORMAT_PDF);
731        assert_eq!(cand.descriptor.source_len, bytes.len() as u64);
732        // Exactly one object: the packed data object holding every literal byte.
733        assert_eq!(
734            cand.descriptor.objects.len(),
735            1,
736            "{name} layout carries one packed data object"
737        );
738        assert!(matches!(
739            cand.descriptor.program.ops.as_slice(),
740            [Op::PackSegments { .. }]
741        ));
742        assert!(cand.descriptor.models.is_empty());
743        assert!(cand.descriptor.channels.is_empty());
744
745        let (encoded, _) = cand.descriptor.serialize().unwrap();
746        let parsed = Descriptor::parse(&encoded, Limits::DEFAULT).unwrap();
747        let out = crate::materialize::materialize(&parsed, Limits::DEFAULT).unwrap();
748        assert_eq!(out, bytes, "{name} layout candidate materializes exactly");
749        assert_eq!(sha256(&out), sha256(bytes), "{name} layout sha");
750    }
751
752    #[test]
753    fn layout_is_exact_on_corpus() {
754        let mut accepted = 0usize;
755        for (name, bytes) in sample_pdfs() {
756            if is_negative_control(name) {
757                assert!(
758                    propose_pdf_layout(&bytes, Limits::DEFAULT)
759                        .unwrap()
760                        .is_none(),
761                    "{name} is not a classic-xref PDF and must decline"
762                );
763                continue;
764            }
765            if propose_pdf_layout(&bytes, Limits::DEFAULT)
766                .unwrap()
767                .is_some()
768            {
769                assert_materializes_exactly(name, &bytes);
770                accepted += 1;
771            }
772        }
773        assert!(
774            accepted >= 3,
775            "expected several accepted classic-xref samples, found {accepted}"
776        );
777    }
778
779    #[test]
780    fn layout_v2_exact() {
781        for name in ["classic.pdf", "many.pdf", "bigtext.pdf"] {
782            assert_materializes_exactly(name, &sample(name));
783        }
784    }
785
786    #[test]
787    fn layout_v2_predicts_many_entries() {
788        let bytes = sample("many.pdf");
789        let cand = propose_pdf_layout(&bytes, Limits::DEFAULT)
790            .unwrap()
791            .unwrap();
792        let emits = pack_emit_count(&cand);
793        assert!(
794            emits >= 100,
795            "many.pdf must predict at least 100 xref offsets, got {emits}"
796        );
797        // Every Emit is either a predicted xref entry or the predicted startxref.
798        let predicted = basis_field(&cand.descriptor.format_basis, "xref_predicted")
799            .expect("format basis must report xref_predicted");
800        let startxref = basis_field(&cand.descriptor.format_basis, "startxref_predicted")
801            .expect("format basis must report startxref_predicted");
802        assert_eq!(predicted + startxref, emits as u64);
803        assert!(predicted >= 100, "many.pdf xref predictions: {predicted}");
804
805        // The classic sample still predicts, and the pack framing stays compact.
806        let classic = propose_pdf_layout(&sample("classic.pdf"), Limits::DEFAULT)
807            .unwrap()
808            .unwrap();
809        assert!(pack_emit_count(&classic) > 0);
810    }
811
812    #[test]
813    fn layout_falls_back_on_bad_offset() {
814        // Object 1's real offset is `off1`, but the xref entry records a wrong
815        // value. The entry must stay literal. `startxref` is still correct, so it
816        // is the only predicted offset in the whole program.
817        let mut b: Vec<u8> = Vec::new();
818        b.extend_from_slice(b"%PDF-1.4\n");
819        let off1 = b.len() as u64;
820        b.extend_from_slice(b"1 0 obj\n<< /Type /Catalog >>\nendobj\n");
821        let xref = b.len() as u64;
822        let wrong = off1 + 3;
823        b.extend_from_slice(
824            format!("xref\n0 2\n0000000000 65535 f \n{wrong:010} 00000 n \n").as_bytes(),
825        );
826        b.extend_from_slice(
827            format!("trailer\n<< /Size 2 /Root 1 0 R >>\nstartxref\n{xref}\n%%EOF\n").as_bytes(),
828        );
829
830        let cand = propose_pdf_layout(&b, Limits::DEFAULT).unwrap().unwrap();
831        assert_eq!(
832            basis_field(&cand.descriptor.format_basis, "xref_predicted"),
833            Some(0),
834            "a mismatched offset must never be predicted"
835        );
836
837        // The only regenerated offset is the correctly predicted `startxref`.
838        let emits = pack_emit_count(&cand);
839        assert_eq!(
840            emits, 1,
841            "only startxref is predicted; bad entry is literal"
842        );
843
844        let (encoded, _) = cand.descriptor.serialize().unwrap();
845        let parsed = Descriptor::parse(&encoded, Limits::DEFAULT).unwrap();
846        let out = crate::materialize::materialize(&parsed, Limits::DEFAULT).unwrap();
847        assert_eq!(out, b, "fallback must still be byte-exact");
848    }
849
850    #[test]
851    fn layout_v2_declines() {
852        // A cross-reference-stream PDF has no classic table to regenerate.
853        assert!(
854            propose_pdf_layout(&sample("xrefstream.pdf"), Limits::DEFAULT)
855                .unwrap()
856                .is_none(),
857            "an xref-stream PDF must decline the layout candidate"
858        );
859        // More objects than the 255 markable slots cannot be expressed exactly
860        // under the slot bound, so the candidate must decline rather than guess.
861        assert!(
862            propose_pdf_layout(&classic_with_objects(256), Limits::DEFAULT)
863                .unwrap()
864                .is_none(),
865            "256 objects exceeds the 255 markable slots"
866        );
867        // 255 objects still fit: indices 0..=254, slot 255 reserved for xref.
868        assert!(
869            propose_pdf_layout(&classic_with_objects(255), Limits::DEFAULT)
870                .unwrap()
871                .is_some(),
872            "255 objects must still be markable"
873        );
874    }
875
876    /// Build a classic-xref PDF with `n` indirect objects and correct offsets.
877    fn classic_with_objects(n: usize) -> Vec<u8> {
878        let mut b: Vec<u8> = Vec::new();
879        b.extend_from_slice(b"%PDF-1.4\n");
880        let mut offsets = Vec::with_capacity(n);
881        for number in 1..=n {
882            offsets.push(b.len() as u64);
883            b.extend_from_slice(format!("{number} 0 obj\n<< >>\nendobj\n").as_bytes());
884        }
885        let xref = b.len() as u64;
886        b.extend_from_slice(format!("xref\n0 {}\n", n + 1).as_bytes());
887        b.extend_from_slice(b"0000000000 65535 f \n");
888        for &off in &offsets {
889            b.extend_from_slice(format!("{off:010} 00000 n \n").as_bytes());
890        }
891        b.extend_from_slice(
892            format!(
893                "trailer\n<< /Size {} /Root 1 0 R >>\nstartxref\n{xref}\n%%EOF\n",
894                n + 1
895            )
896            .as_bytes(),
897        );
898        b
899    }
900
901    #[test]
902    fn layout_v2_deterministic() {
903        for name in ["classic.pdf", "many.pdf", "bigtext.pdf"] {
904            let bytes = sample(name);
905            let a = propose_pdf_layout(&bytes, Limits::DEFAULT)
906                .unwrap()
907                .unwrap()
908                .descriptor
909                .serialize()
910                .unwrap()
911                .0;
912            let b = propose_pdf_layout(&bytes, Limits::DEFAULT)
913                .unwrap()
914                .unwrap()
915                .descriptor
916                .serialize()
917                .unwrap()
918                .0;
919            assert_eq!(a, b, "{name} layout bytes must be deterministic");
920        }
921    }
922
923    /// Force the layout+rANS candidate for `bytes` and assert the full exact
924    /// triple end-to-end through the normative decoder.
925    #[cfg(feature = "rans")]
926    fn assert_rans_materializes_exactly(name: &str, bytes: &[u8]) {
927        let cand = propose_pdf_layout_rans(bytes, Limits::DEFAULT)
928            .unwrap()
929            .unwrap_or_else(|| panic!("{name} must propose a layout-rANS candidate"));
930        assert_eq!(cand.kind, CandidateKind::PdfLayoutRans);
931        assert_eq!(cand.descriptor.source_format, SOURCE_FORMAT_PDF);
932        assert_eq!(cand.descriptor.source_len, bytes.len() as u64);
933        // No literal objects: the data and the plan both travel in channels.
934        assert!(cand.descriptor.objects.is_empty());
935        assert_eq!(cand.descriptor.models.len(), 2);
936        assert_eq!(cand.descriptor.channels.len(), 2);
937        assert_eq!(cand.descriptor.channels[0].model_id, 0);
938        assert_eq!(cand.descriptor.channels[1].model_id, 1);
939        assert!(matches!(
940            cand.descriptor.program.ops.as_slice(),
941            [Op::PackedChannels {
942                data_channel: 0,
943                plan_channel: 1,
944                ..
945            }]
946        ));
947
948        let (encoded, _) = cand.descriptor.serialize().unwrap();
949        let parsed = Descriptor::parse(&encoded, Limits::DEFAULT).unwrap();
950        let out = crate::materialize::materialize(&parsed, Limits::DEFAULT).unwrap();
951        assert_eq!(out, bytes, "{name} layout-rANS materializes exactly");
952        assert_eq!(sha256(&out), sha256(bytes), "{name} layout-rANS sha");
953
954        // The forced lane must survive the court's own decode-before-commit.
955        let (forced, report) =
956            crate::encode::encode_with(bytes, Limits::DEFAULT, Some(CandidateKind::PdfLayoutRans))
957                .unwrap();
958        assert_eq!(report.kind, CandidateKind::PdfLayoutRans);
959        let (forced_out, _) =
960            crate::materialize::decode_to_bytes(&forced, Limits::DEFAULT).unwrap();
961        assert_eq!(forced_out, bytes, "{name} forced layout-rANS bytes");
962        assert_eq!(
963            sha256(&forced_out),
964            sha256(bytes),
965            "{name} forced layout-rANS sha"
966        );
967    }
968
969    #[cfg(feature = "rans")]
970    #[test]
971    fn layout_rans_exact() {
972        for name in ["classic.pdf", "bigtext.pdf", "many.pdf"] {
973            assert_rans_materializes_exactly(name, &sample(name));
974        }
975    }
976
977    #[cfg(feature = "rans")]
978    #[test]
979    fn layout_rans_deterministic() {
980        for name in ["classic.pdf", "bigtext.pdf", "many.pdf"] {
981            let bytes = sample(name);
982            let a = propose_pdf_layout_rans(&bytes, Limits::DEFAULT)
983                .unwrap()
984                .unwrap()
985                .descriptor
986                .serialize()
987                .unwrap()
988                .0;
989            let b = propose_pdf_layout_rans(&bytes, Limits::DEFAULT)
990                .unwrap()
991                .unwrap()
992                .descriptor
993                .serialize()
994                .unwrap()
995                .0;
996            assert_eq!(a, b, "{name} layout-rANS bytes must be deterministic");
997        }
998    }
999
1000    #[cfg(feature = "rans")]
1001    #[test]
1002    fn layout_rans_declines_on_non_pdf() {
1003        // Non-PDF controls and cross-reference-stream PDFs: nothing to plan, so
1004        // the candidate must decline rather than store a non-plan.
1005        for name in ["notpdf.bin", "malformed.pdf", "xrefstream.pdf"] {
1006            assert!(
1007                propose_pdf_layout_rans(&sample(name), Limits::DEFAULT)
1008                    .unwrap()
1009                    .is_none(),
1010                "{name} must decline the layout-rANS candidate"
1011            );
1012        }
1013    }
1014
1015    #[test]
1016    fn layout_measurements_report() {
1017        for (name, bytes) in sample_pdfs() {
1018            let Some(cand) = propose_pdf_layout(&bytes, Limits::DEFAULT).unwrap() else {
1019                eprintln!("layout[{name}]: declined");
1020                continue;
1021            };
1022            let (layout_bytes, _) = cand.descriptor.serialize().unwrap();
1023            let emits = pack_emit_count(&cand);
1024            eprintln!(
1025                "layout[{name}] source={} items={} emits={emits} layout={}",
1026                bytes.len(),
1027                pack_items(&cand).len(),
1028                layout_bytes.len(),
1029            );
1030        }
1031
1032        for name in ["classic.pdf", "bigtext.pdf", "many.pdf"] {
1033            let bytes = sample(name);
1034            let (layout_bytes, _) =
1035                crate::encode::encode_with(&bytes, Limits::DEFAULT, Some(CandidateKind::PdfLayout))
1036                    .unwrap();
1037            let (raw_bytes, _) =
1038                crate::encode::encode_with(&bytes, Limits::DEFAULT, Some(CandidateKind::Raw))
1039                    .unwrap();
1040            #[cfg(feature = "rans")]
1041            let (byte_rans_bytes, _) =
1042                crate::encode::encode_with(&bytes, Limits::DEFAULT, Some(CandidateKind::ByteRans))
1043                    .unwrap();
1044            #[cfg(not(feature = "rans"))]
1045            let byte_rans_bytes: Vec<u8> = Vec::new();
1046            #[cfg(feature = "rans")]
1047            let (layout_rans_bytes, _) = crate::encode::encode_with(
1048                &bytes,
1049                Limits::DEFAULT,
1050                Some(CandidateKind::PdfLayoutRans),
1051            )
1052            .unwrap();
1053            #[cfg(not(feature = "rans"))]
1054            let layout_rans_bytes: Vec<u8> = Vec::new();
1055            let (_, auto) = crate::encode::encode(&bytes, Limits::DEFAULT).unwrap();
1056            eprintln!(
1057                "sizes[{name}] source={} raw={} byte_rans={} layout={} layout_rans={} auto={}({})",
1058                bytes.len(),
1059                raw_bytes.len(),
1060                byte_rans_bytes.len(),
1061                layout_bytes.len(),
1062                layout_rans_bytes.len(),
1063                auto.kind.name(),
1064                auto.encoded_len,
1065            );
1066
1067            #[cfg(feature = "rans")]
1068            {
1069                let plan = build_layout_plan(&bytes, Limits::DEFAULT).unwrap().unwrap();
1070                let plan_bytes_len = crate::dra::op::encode_items(&plan.items).unwrap().len();
1071                let cand = propose_pdf_layout_rans(&bytes, Limits::DEFAULT)
1072                    .unwrap()
1073                    .unwrap();
1074                let model_bytes: usize = cand
1075                    .descriptor
1076                    .models
1077                    .iter()
1078                    .map(|m| m.encode().unwrap().len())
1079                    .sum();
1080                let payload_bytes: usize = cand
1081                    .descriptor
1082                    .channels
1083                    .iter()
1084                    .map(|c| c.payload.len())
1085                    .sum();
1086                eprintln!(
1087                    "rans[{name}] data={} plan_bytes={} models={} payloads={} total={}",
1088                    plan.data.len(),
1089                    plan_bytes_len,
1090                    model_bytes,
1091                    payload_bytes,
1092                    layout_rans_bytes.len(),
1093                );
1094            }
1095        }
1096    }
1097}