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, ObjectSource, 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![ObjectSource::Inline(plan.data)],
288        program: Program::new(vec![Op::PackSegments {
289            data_object: 0,
290            items: plan.items,
291        }]),
292        observation_index: None,
293        seek_directory: false,
294        source_sha256: sha256(input),
295        source_len: input.len() as u64,
296    };
297
298    let candidate = Candidate {
299        kind: CandidateKind::PdfLayout,
300        descriptor,
301    };
302
303    // Verify byte-exactness through the normative decoder before returning. An
304    // inexact program must never be emitted.
305    let (encoded, _) = candidate.descriptor.serialize()?;
306    let parsed = match Descriptor::parse(&encoded, limits) {
307        Ok(p) => p,
308        Err(_) => return Ok(None),
309    };
310    let out = match crate::materialize::materialize(&parsed, limits) {
311        Ok(o) => o,
312        Err(_) => return Ok(None),
313    };
314    if out != input {
315        return Ok(None);
316    }
317
318    Ok(Some(candidate))
319}
320
321/// Propose a layout + rANS candidate, or `None`.
322///
323/// Builds the same [`LayoutPlan`] as [`propose_pdf_layout`], then entropy-codes
324/// its parts into two rANS channels: channel `0` carries the plan's literal data
325/// object, and channel `1` carries `encode_items` of the plan's item table. A
326/// single [`Op::PackedChannels`] reconstructs the source from both, so the whole
327/// plan — data *and* item table — pays entropy-coding cost instead of being
328/// stored as literal bytes. Each channel uses its own order-0 byte model
329/// normalized from its own byte histogram at `scale_bits` 12.
330///
331/// Exactly as with the literal layout lane, an end-to-end serialize / parse /
332/// materialize / byte-compare check gates the return: an inexact program yields
333/// `Ok(None)` rather than an inexact candidate.
334#[cfg(feature = "rans")]
335pub fn propose_pdf_layout_rans(input: &[u8], limits: Limits) -> Result<Option<Candidate>> {
336    let plan = match build_layout_plan(input, limits)? {
337        Some(p) => p,
338        None => return Ok(None),
339    };
340
341    let plan_bytes = crate::dra::op::encode_items(&plan.items)?;
342
343    // Channel 0: the literal data object, coded against its own byte histogram.
344    let mut data_counts = [0u64; ALPHABET];
345    for &b in &plan.data {
346        data_counts[b as usize] += 1;
347    }
348    let data_model = EntropyModel::from_counts(&data_counts, 12)?;
349    let data_capsule = encode_channel(&data_model, &plan.data)?;
350
351    // Channel 1: the serialized item table, coded against its own histogram.
352    let mut plan_counts = [0u64; ALPHABET];
353    for &b in &plan_bytes {
354        plan_counts[b as usize] += 1;
355    }
356    let plan_model = EntropyModel::from_counts(&plan_counts, 12)?;
357    let plan_capsule = encode_channel(&plan_model, &plan_bytes)?;
358
359    let data_channel = EntropyChannelDescriptor {
360        coder: CODER_ORDER0_BYTE_RANS,
361        coder_version: CODER_VERSION_1,
362        scale_bits: data_model.scale_bits,
363        lane_count: 1,
364        model_id: 0,
365        symbol_count: data_capsule.symbol_count,
366        decoded_length: data_capsule.decoded_length,
367        initial_state: data_capsule.initial_state,
368        payload: data_capsule.payload,
369    };
370    let plan_channel = EntropyChannelDescriptor {
371        coder: CODER_ORDER0_BYTE_RANS,
372        coder_version: CODER_VERSION_1,
373        scale_bits: plan_model.scale_bits,
374        lane_count: 1,
375        model_id: 1,
376        symbol_count: plan_capsule.symbol_count,
377        decoded_length: plan_capsule.decoded_length,
378        initial_state: plan_capsule.initial_state,
379        payload: plan_capsule.payload,
380    };
381
382    let format_basis = format!(
383        "pdf-layout-rans;objects={};xref_predicted={};xref_literal={}",
384        marked_object_count(&plan.items),
385        plan.xref_predicted,
386        plan.xref_literal
387    );
388
389    let descriptor = Descriptor {
390        universe: UNIVERSE.to_string(),
391        source_format: SOURCE_FORMAT_PDF,
392        format_basis,
393        models: vec![data_model, plan_model],
394        channels: vec![data_channel, plan_channel],
395        objects: vec![],
396        program: Program::new(vec![Op::PackedChannels {
397            data_channel: 0,
398            plan_channel: 1,
399            declared_output_len: input.len() as u64,
400        }]),
401        observation_index: None,
402        seek_directory: false,
403        source_sha256: sha256(input),
404        source_len: input.len() as u64,
405    };
406
407    let candidate = Candidate {
408        kind: CandidateKind::PdfLayoutRans,
409        descriptor,
410    };
411
412    // Verify byte-exactness through the normative decoder before returning. An
413    // inexact program must never be emitted.
414    let (encoded, _) = candidate.descriptor.serialize()?;
415    let parsed = match Descriptor::parse(&encoded, limits) {
416        Ok(p) => p,
417        Err(_) => return Ok(None),
418    };
419    let out = match crate::materialize::materialize(&parsed, limits) {
420        Ok(o) => o,
421        Err(_) => return Ok(None),
422    };
423    if out != input {
424        return Ok(None);
425    }
426
427    Ok(Some(candidate))
428}
429
430/// Number of indirect objects whose introducer was marked in the item table.
431fn marked_object_count(items: &[PackItem]) -> usize {
432    items
433        .iter()
434        .filter(|item| matches!(item, PackItem::Mark { slot } if *slot != XREF_SLOT))
435        .count()
436}
437
438/// Number of emitted offsets in the item table.
439fn emit_count(items: &[PackItem]) -> usize {
440    items
441        .iter()
442        .filter(|item| matches!(item, PackItem::Emit { .. }))
443        .count()
444}
445
446/// Append `bytes` to the packed data object, recording one [`PackItem::Literal`]
447/// when non-empty, and advance the simulated output position. Returns `false`
448/// (so the caller declines) when a single run would not fit a `u32` length.
449fn push_pack_literal(
450    data: &mut Vec<u8>,
451    items: &mut Vec<PackItem>,
452    bytes: &[u8],
453    pos: &mut u64,
454) -> bool {
455    if !push_literal(items, data, bytes) {
456        return false;
457    }
458    *pos += bytes.len() as u64;
459    true
460}
461
462/// Append `bytes` to the packed data object, coalescing them into the immediately
463/// preceding [`PackItem::Literal`] when one is present so that consecutive literal
464/// runs collapse into the fewest possible items. The merge never crosses a
465/// [`PackItem::Mark`] or [`PackItem::Emit`], empty pushes are ignored, and the
466/// combined length must remain expressible as a `u32`. Returns `false` (so the
467/// caller declines) when no safe item shape exists.
468fn push_literal(items: &mut Vec<PackItem>, data: &mut Vec<u8>, bytes: &[u8]) -> bool {
469    if bytes.is_empty() {
470        return true;
471    }
472    let Ok(len) = u32::try_from(bytes.len()) else {
473        return false;
474    };
475    if let Some(PackItem::Literal { len: prev }) = items.last_mut() {
476        let Some(total) = prev.checked_add(len) else {
477            return false;
478        };
479        *prev = total;
480    } else {
481        items.push(PackItem::Literal { len });
482    }
483    data.extend_from_slice(bytes);
484    true
485}
486
487/// Index of the first object whose introducer starts at `start`.
488fn object_index_at(objects: &[PdfObjectSpan], start: u64) -> Option<usize> {
489    objects.iter().position(|o| o.start == start)
490}
491
492/// The slot marking the target object `number` at position `value`, if any
493/// earlier [`Op::MarkOffset`] recorded exactly that position.
494fn predicted_slot(
495    objects: &[PdfObjectSpan],
496    slot_value: &[u64; 256],
497    slot_marked: &[bool; 256],
498    number: u64,
499    value: u64,
500) -> Option<u8> {
501    objects.iter().enumerate().find_map(|(i, o)| {
502        let slot = i as u8;
503        (o.number == number && slot_marked[slot as usize] && slot_value[slot as usize] == value)
504            .then_some(slot)
505    })
506}
507
508/// Predict a whole `startxref` span: the trailing run of decimal digits is the
509/// value. Returns `(prefix_len, width)` when the value equals the marked `xref`
510/// position and its width is in `1..=20`.
511fn predict_startxref(
512    bytes: &[u8],
513    slot_value: &[u64; 256],
514    slot_marked: &[bool; 256],
515) -> Option<(usize, u8)> {
516    if !slot_marked[XREF_SLOT as usize] {
517        return None;
518    }
519    let mut i = bytes.len();
520    while i > 0 && bytes[i - 1].is_ascii_digit() {
521        i -= 1;
522    }
523    let width = bytes.len() - i;
524    if width == 0 || width > 20 {
525        return None;
526    }
527    let value = parse_digits(&bytes[i..])?;
528    if value != slot_value[XREF_SLOT as usize] {
529        return None;
530    }
531    Some((i, width as u8))
532}
533
534/// One ordered slice of a classic `xref` section: either literal bytes or a
535/// 20-byte entry whose offset field may be regenerated.
536enum XrefPiece {
537    /// Verbatim bytes `[start, start + len)` of the section.
538    Literal { start: usize, len: usize },
539    /// A 20-byte entry `[start, start + 20)`.
540    Entry {
541        /// Offset of the entry within the section.
542        start: usize,
543        /// Target object number (`subsection_start + i`).
544        number: u64,
545        /// Parsed value of the 10-digit offset field, if it is all digits.
546        offset: Option<u64>,
547        /// Whether the status byte is `n` (in use).
548        in_use: bool,
549    },
550}
551
552/// Parse a classic cross-reference table from a section's bytes, returning an
553/// ordered tiling of the section. Returns `None` (so the caller emits the whole
554/// section literally) when the bytes do not match the classic grammar.
555///
556/// Grammar accepted: `xref` EOL, then one or more `<start> <count>` EOL headers
557/// each followed by exactly `count` 20-byte entries of the shape
558/// `10-digit-offset SP 5-digit-generation SP status 2-byte-EOL`. The two EOL
559/// bytes may be `CR LF`, `LF CR`, `SP LF`, or `SP CR`.
560fn parse_classic_xref(bytes: &[u8]) -> Option<Vec<XrefPiece>> {
561    if !bytes.starts_with(b"xref") {
562        return None;
563    }
564    let mut pieces = Vec::new();
565    let mut pos = 4usize;
566
567    // EOL after the `xref` keyword.
568    let eol = eol_len(&bytes[pos..])?;
569    pieces.push(XrefPiece::Literal {
570        start: 0,
571        len: pos + eol,
572    });
573    pos += eol;
574
575    let mut any = false;
576    while pos < bytes.len() {
577        // Subsection header: `<start> <count>` EOL.
578        let header_start = pos;
579        let (start, after_start) = parse_uint_at(bytes, pos)?;
580        pos = after_start;
581        let spaces_start = pos;
582        while pos < bytes.len() && bytes[pos] == b' ' {
583            pos += 1;
584        }
585        if pos == spaces_start {
586            return None;
587        }
588        let (count, after_count) = parse_uint_at(bytes, pos)?;
589        pos = after_count;
590        let eol = eol_len(&bytes[pos..])?;
591        let header_end = pos + eol;
592        pieces.push(XrefPiece::Literal {
593            start: header_start,
594            len: header_end - header_start,
595        });
596        pos = header_end;
597
598        for i in 0..count {
599            let end = pos.checked_add(20)?;
600            if end > bytes.len() {
601                return None;
602            }
603            let entry = &bytes[pos..end];
604            if !is_entry_shape(entry) {
605                return None;
606            }
607            let number = start.checked_add(i)?;
608            let in_use = entry[17] == b'n';
609            let offset = parse_digits(&entry[0..10]);
610            pieces.push(XrefPiece::Entry {
611                start: pos,
612                number,
613                offset,
614                in_use,
615            });
616            pos = end;
617        }
618        any = true;
619    }
620
621    if !any || pos != bytes.len() {
622        return None;
623    }
624    Some(pieces)
625}
626
627/// Whether a 20-byte window matches the classic cross-reference entry shape.
628fn is_entry_shape(entry: &[u8]) -> bool {
629    if entry.len() != 20 {
630        return false;
631    }
632    if entry[10] != b' ' || entry[16] != b' ' {
633        return false;
634    }
635    if !entry[11..16].iter().all(u8::is_ascii_digit) {
636        return false;
637    }
638    if entry[17] != b'n' && entry[17] != b'f' {
639        return false;
640    }
641    matches!(
642        (entry[18], entry[19]),
643        (b'\r', b'\n') | (b'\n', b'\r') | (b' ', b'\n') | (b' ', b'\r')
644    )
645}
646
647/// Length of an end-of-line marker at the start of `bytes`, if any.
648fn eol_len(bytes: &[u8]) -> Option<usize> {
649    match bytes {
650        [b'\r', b'\n', ..] => Some(2),
651        [b'\n', ..] | [b'\r', ..] => Some(1),
652        _ => None,
653    }
654}
655
656/// Parse a non-negative decimal integer at `at`, returning `(value, next)`.
657fn parse_uint_at(bytes: &[u8], at: usize) -> Option<(u64, usize)> {
658    let mut i = at;
659    let mut value: u64 = 0;
660    while i < bytes.len() && bytes[i].is_ascii_digit() {
661        value = value
662            .checked_mul(10)?
663            .checked_add(u64::from(bytes[i] - b'0'))?;
664        i += 1;
665    }
666    if i == at {
667        return None;
668    }
669    Some((value, i))
670}
671
672/// Parse an all-digit slice as a non-negative decimal integer.
673fn parse_digits(digits: &[u8]) -> Option<u64> {
674    if digits.is_empty() {
675        return None;
676    }
677    let mut value: u64 = 0;
678    for &b in digits {
679        if !b.is_ascii_digit() {
680            return None;
681        }
682        value = value.checked_mul(10)?.checked_add(u64::from(b - b'0'))?;
683    }
684    Some(value)
685}
686
687#[cfg(test)]
688mod tests {
689    use super::*;
690    use crate::adapter::pdf::samples::{is_negative_control, sample_pdfs};
691    use crate::container::Descriptor;
692
693    fn sample(name: &str) -> Vec<u8> {
694        sample_pdfs()
695            .into_iter()
696            .find(|(n, _)| *n == name)
697            .unwrap_or_else(|| panic!("sample {name} missing"))
698            .1
699    }
700
701    fn basis_field(basis: &str, key: &str) -> Option<u64> {
702        basis.split(';').find_map(|part| {
703            let (k, v) = part.split_once('=')?;
704            (k == key).then(|| v.parse().ok()).flatten()
705        })
706    }
707
708    /// The item table of the single `PackSegments` op this candidate builds.
709    fn pack_items(cand: &Candidate) -> &[PackItem] {
710        cand.descriptor
711            .program
712            .ops
713            .iter()
714            .find_map(|op| match op {
715                Op::PackSegments { items, .. } => Some(items.as_slice()),
716                _ => None,
717            })
718            .expect("layout program is one PackSegments op")
719    }
720
721    /// Number of regenerated offsets in the packed item table.
722    fn pack_emit_count(cand: &Candidate) -> usize {
723        pack_items(cand)
724            .iter()
725            .filter(|item| matches!(item, PackItem::Emit { .. }))
726            .count()
727    }
728
729    fn assert_materializes_exactly(name: &str, bytes: &[u8]) {
730        let cand = propose_pdf_layout(bytes, Limits::DEFAULT)
731            .unwrap()
732            .unwrap_or_else(|| panic!("{name} must propose a layout candidate"));
733        assert_eq!(cand.kind, CandidateKind::PdfLayout);
734        assert_eq!(cand.descriptor.source_format, SOURCE_FORMAT_PDF);
735        assert_eq!(cand.descriptor.source_len, bytes.len() as u64);
736        // Exactly one object: the packed data object holding every literal byte.
737        assert_eq!(
738            cand.descriptor.objects.len(),
739            1,
740            "{name} layout carries one packed data object"
741        );
742        assert!(matches!(
743            cand.descriptor.program.ops.as_slice(),
744            [Op::PackSegments { .. }]
745        ));
746        assert!(cand.descriptor.models.is_empty());
747        assert!(cand.descriptor.channels.is_empty());
748
749        let (encoded, _) = cand.descriptor.serialize().unwrap();
750        let parsed = Descriptor::parse(&encoded, Limits::DEFAULT).unwrap();
751        let out = crate::materialize::materialize(&parsed, Limits::DEFAULT).unwrap();
752        assert_eq!(out, bytes, "{name} layout candidate materializes exactly");
753        assert_eq!(sha256(&out), sha256(bytes), "{name} layout sha");
754    }
755
756    #[test]
757    fn layout_is_exact_on_corpus() {
758        let mut accepted = 0usize;
759        for (name, bytes) in sample_pdfs() {
760            if is_negative_control(name) {
761                assert!(
762                    propose_pdf_layout(&bytes, Limits::DEFAULT)
763                        .unwrap()
764                        .is_none(),
765                    "{name} is not a classic-xref PDF and must decline"
766                );
767                continue;
768            }
769            if propose_pdf_layout(&bytes, Limits::DEFAULT)
770                .unwrap()
771                .is_some()
772            {
773                assert_materializes_exactly(name, &bytes);
774                accepted += 1;
775            }
776        }
777        assert!(
778            accepted >= 3,
779            "expected several accepted classic-xref samples, found {accepted}"
780        );
781    }
782
783    #[test]
784    fn layout_v2_exact() {
785        for name in ["classic.pdf", "many.pdf", "bigtext.pdf"] {
786            assert_materializes_exactly(name, &sample(name));
787        }
788    }
789
790    #[test]
791    fn layout_v2_predicts_many_entries() {
792        let bytes = sample("many.pdf");
793        let cand = propose_pdf_layout(&bytes, Limits::DEFAULT)
794            .unwrap()
795            .unwrap();
796        let emits = pack_emit_count(&cand);
797        assert!(
798            emits >= 100,
799            "many.pdf must predict at least 100 xref offsets, got {emits}"
800        );
801        // Every Emit is either a predicted xref entry or the predicted startxref.
802        let predicted = basis_field(&cand.descriptor.format_basis, "xref_predicted")
803            .expect("format basis must report xref_predicted");
804        let startxref = basis_field(&cand.descriptor.format_basis, "startxref_predicted")
805            .expect("format basis must report startxref_predicted");
806        assert_eq!(predicted + startxref, emits as u64);
807        assert!(predicted >= 100, "many.pdf xref predictions: {predicted}");
808
809        // The classic sample still predicts, and the pack framing stays compact.
810        let classic = propose_pdf_layout(&sample("classic.pdf"), Limits::DEFAULT)
811            .unwrap()
812            .unwrap();
813        assert!(pack_emit_count(&classic) > 0);
814    }
815
816    #[test]
817    fn layout_falls_back_on_bad_offset() {
818        // Object 1's real offset is `off1`, but the xref entry records a wrong
819        // value. The entry must stay literal. `startxref` is still correct, so it
820        // is the only predicted offset in the whole program.
821        let mut b: Vec<u8> = Vec::new();
822        b.extend_from_slice(b"%PDF-1.4\n");
823        let off1 = b.len() as u64;
824        b.extend_from_slice(b"1 0 obj\n<< /Type /Catalog >>\nendobj\n");
825        let xref = b.len() as u64;
826        let wrong = off1 + 3;
827        b.extend_from_slice(
828            format!("xref\n0 2\n0000000000 65535 f \n{wrong:010} 00000 n \n").as_bytes(),
829        );
830        b.extend_from_slice(
831            format!("trailer\n<< /Size 2 /Root 1 0 R >>\nstartxref\n{xref}\n%%EOF\n").as_bytes(),
832        );
833
834        let cand = propose_pdf_layout(&b, Limits::DEFAULT).unwrap().unwrap();
835        assert_eq!(
836            basis_field(&cand.descriptor.format_basis, "xref_predicted"),
837            Some(0),
838            "a mismatched offset must never be predicted"
839        );
840
841        // The only regenerated offset is the correctly predicted `startxref`.
842        let emits = pack_emit_count(&cand);
843        assert_eq!(
844            emits, 1,
845            "only startxref is predicted; bad entry is literal"
846        );
847
848        let (encoded, _) = cand.descriptor.serialize().unwrap();
849        let parsed = Descriptor::parse(&encoded, Limits::DEFAULT).unwrap();
850        let out = crate::materialize::materialize(&parsed, Limits::DEFAULT).unwrap();
851        assert_eq!(out, b, "fallback must still be byte-exact");
852    }
853
854    #[test]
855    fn layout_v2_declines() {
856        // A cross-reference-stream PDF has no classic table to regenerate.
857        assert!(
858            propose_pdf_layout(&sample("xrefstream.pdf"), Limits::DEFAULT)
859                .unwrap()
860                .is_none(),
861            "an xref-stream PDF must decline the layout candidate"
862        );
863        // More objects than the 255 markable slots cannot be expressed exactly
864        // under the slot bound, so the candidate must decline rather than guess.
865        assert!(
866            propose_pdf_layout(&classic_with_objects(256), Limits::DEFAULT)
867                .unwrap()
868                .is_none(),
869            "256 objects exceeds the 255 markable slots"
870        );
871        // 255 objects still fit: indices 0..=254, slot 255 reserved for xref.
872        assert!(
873            propose_pdf_layout(&classic_with_objects(255), Limits::DEFAULT)
874                .unwrap()
875                .is_some(),
876            "255 objects must still be markable"
877        );
878    }
879
880    /// Build a classic-xref PDF with `n` indirect objects and correct offsets.
881    fn classic_with_objects(n: usize) -> Vec<u8> {
882        let mut b: Vec<u8> = Vec::new();
883        b.extend_from_slice(b"%PDF-1.4\n");
884        let mut offsets = Vec::with_capacity(n);
885        for number in 1..=n {
886            offsets.push(b.len() as u64);
887            b.extend_from_slice(format!("{number} 0 obj\n<< >>\nendobj\n").as_bytes());
888        }
889        let xref = b.len() as u64;
890        b.extend_from_slice(format!("xref\n0 {}\n", n + 1).as_bytes());
891        b.extend_from_slice(b"0000000000 65535 f \n");
892        for &off in &offsets {
893            b.extend_from_slice(format!("{off:010} 00000 n \n").as_bytes());
894        }
895        b.extend_from_slice(
896            format!(
897                "trailer\n<< /Size {} /Root 1 0 R >>\nstartxref\n{xref}\n%%EOF\n",
898                n + 1
899            )
900            .as_bytes(),
901        );
902        b
903    }
904
905    #[test]
906    fn layout_v2_deterministic() {
907        for name in ["classic.pdf", "many.pdf", "bigtext.pdf"] {
908            let bytes = sample(name);
909            let a = propose_pdf_layout(&bytes, Limits::DEFAULT)
910                .unwrap()
911                .unwrap()
912                .descriptor
913                .serialize()
914                .unwrap()
915                .0;
916            let b = propose_pdf_layout(&bytes, Limits::DEFAULT)
917                .unwrap()
918                .unwrap()
919                .descriptor
920                .serialize()
921                .unwrap()
922                .0;
923            assert_eq!(a, b, "{name} layout bytes must be deterministic");
924        }
925    }
926
927    /// Force the layout+rANS candidate for `bytes` and assert the full exact
928    /// triple end-to-end through the normative decoder.
929    #[cfg(feature = "rans")]
930    fn assert_rans_materializes_exactly(name: &str, bytes: &[u8]) {
931        let cand = propose_pdf_layout_rans(bytes, Limits::DEFAULT)
932            .unwrap()
933            .unwrap_or_else(|| panic!("{name} must propose a layout-rANS candidate"));
934        assert_eq!(cand.kind, CandidateKind::PdfLayoutRans);
935        assert_eq!(cand.descriptor.source_format, SOURCE_FORMAT_PDF);
936        assert_eq!(cand.descriptor.source_len, bytes.len() as u64);
937        // No literal objects: the data and the plan both travel in channels.
938        assert!(cand.descriptor.objects.is_empty());
939        assert_eq!(cand.descriptor.models.len(), 2);
940        assert_eq!(cand.descriptor.channels.len(), 2);
941        assert_eq!(cand.descriptor.channels[0].model_id, 0);
942        assert_eq!(cand.descriptor.channels[1].model_id, 1);
943        assert!(matches!(
944            cand.descriptor.program.ops.as_slice(),
945            [Op::PackedChannels {
946                data_channel: 0,
947                plan_channel: 1,
948                ..
949            }]
950        ));
951
952        let (encoded, _) = cand.descriptor.serialize().unwrap();
953        let parsed = Descriptor::parse(&encoded, Limits::DEFAULT).unwrap();
954        let out = crate::materialize::materialize(&parsed, Limits::DEFAULT).unwrap();
955        assert_eq!(out, bytes, "{name} layout-rANS materializes exactly");
956        assert_eq!(sha256(&out), sha256(bytes), "{name} layout-rANS sha");
957
958        // The forced lane must survive the court's own decode-before-commit.
959        let (forced, report) =
960            crate::encode::encode_with(bytes, Limits::DEFAULT, Some(CandidateKind::PdfLayoutRans))
961                .unwrap();
962        assert_eq!(report.kind, CandidateKind::PdfLayoutRans);
963        let (forced_out, _) =
964            crate::materialize::decode_to_bytes(&forced, Limits::DEFAULT).unwrap();
965        assert_eq!(forced_out, bytes, "{name} forced layout-rANS bytes");
966        assert_eq!(
967            sha256(&forced_out),
968            sha256(bytes),
969            "{name} forced layout-rANS sha"
970        );
971    }
972
973    #[cfg(feature = "rans")]
974    #[test]
975    fn layout_rans_exact() {
976        for name in ["classic.pdf", "bigtext.pdf", "many.pdf"] {
977            assert_rans_materializes_exactly(name, &sample(name));
978        }
979    }
980
981    #[cfg(feature = "rans")]
982    #[test]
983    fn layout_rans_deterministic() {
984        for name in ["classic.pdf", "bigtext.pdf", "many.pdf"] {
985            let bytes = sample(name);
986            let a = propose_pdf_layout_rans(&bytes, Limits::DEFAULT)
987                .unwrap()
988                .unwrap()
989                .descriptor
990                .serialize()
991                .unwrap()
992                .0;
993            let b = propose_pdf_layout_rans(&bytes, Limits::DEFAULT)
994                .unwrap()
995                .unwrap()
996                .descriptor
997                .serialize()
998                .unwrap()
999                .0;
1000            assert_eq!(a, b, "{name} layout-rANS bytes must be deterministic");
1001        }
1002    }
1003
1004    #[cfg(feature = "rans")]
1005    #[test]
1006    fn layout_rans_declines_on_non_pdf() {
1007        // Non-PDF controls and cross-reference-stream PDFs: nothing to plan, so
1008        // the candidate must decline rather than store a non-plan.
1009        for name in ["notpdf.bin", "malformed.pdf", "xrefstream.pdf"] {
1010            assert!(
1011                propose_pdf_layout_rans(&sample(name), Limits::DEFAULT)
1012                    .unwrap()
1013                    .is_none(),
1014                "{name} must decline the layout-rANS candidate"
1015            );
1016        }
1017    }
1018
1019    #[test]
1020    fn layout_measurements_report() {
1021        for (name, bytes) in sample_pdfs() {
1022            let Some(cand) = propose_pdf_layout(&bytes, Limits::DEFAULT).unwrap() else {
1023                eprintln!("layout[{name}]: declined");
1024                continue;
1025            };
1026            let (layout_bytes, _) = cand.descriptor.serialize().unwrap();
1027            let emits = pack_emit_count(&cand);
1028            eprintln!(
1029                "layout[{name}] source={} items={} emits={emits} layout={}",
1030                bytes.len(),
1031                pack_items(&cand).len(),
1032                layout_bytes.len(),
1033            );
1034        }
1035
1036        for name in ["classic.pdf", "bigtext.pdf", "many.pdf"] {
1037            let bytes = sample(name);
1038            let (layout_bytes, _) =
1039                crate::encode::encode_with(&bytes, Limits::DEFAULT, Some(CandidateKind::PdfLayout))
1040                    .unwrap();
1041            let (raw_bytes, _) =
1042                crate::encode::encode_with(&bytes, Limits::DEFAULT, Some(CandidateKind::Raw))
1043                    .unwrap();
1044            #[cfg(feature = "rans")]
1045            let (byte_rans_bytes, _) =
1046                crate::encode::encode_with(&bytes, Limits::DEFAULT, Some(CandidateKind::ByteRans))
1047                    .unwrap();
1048            #[cfg(not(feature = "rans"))]
1049            let byte_rans_bytes: Vec<u8> = Vec::new();
1050            #[cfg(feature = "rans")]
1051            let (layout_rans_bytes, _) = crate::encode::encode_with(
1052                &bytes,
1053                Limits::DEFAULT,
1054                Some(CandidateKind::PdfLayoutRans),
1055            )
1056            .unwrap();
1057            #[cfg(not(feature = "rans"))]
1058            let layout_rans_bytes: Vec<u8> = Vec::new();
1059            let (_, auto) = crate::encode::encode(&bytes, Limits::DEFAULT).unwrap();
1060            eprintln!(
1061                "sizes[{name}] source={} raw={} byte_rans={} layout={} layout_rans={} auto={}({})",
1062                bytes.len(),
1063                raw_bytes.len(),
1064                byte_rans_bytes.len(),
1065                layout_bytes.len(),
1066                layout_rans_bytes.len(),
1067                auto.kind.name(),
1068                auto.encoded_len,
1069            );
1070
1071            #[cfg(feature = "rans")]
1072            {
1073                let plan = build_layout_plan(&bytes, Limits::DEFAULT).unwrap().unwrap();
1074                let plan_bytes_len = crate::dra::op::encode_items(&plan.items).unwrap().len();
1075                let cand = propose_pdf_layout_rans(&bytes, Limits::DEFAULT)
1076                    .unwrap()
1077                    .unwrap();
1078                let model_bytes: usize = cand
1079                    .descriptor
1080                    .models
1081                    .iter()
1082                    .map(|m| m.encode().unwrap().len())
1083                    .sum();
1084                let payload_bytes: usize = cand
1085                    .descriptor
1086                    .channels
1087                    .iter()
1088                    .map(|c| c.payload.len())
1089                    .sum();
1090                eprintln!(
1091                    "rans[{name}] data={} plan_bytes={} models={} payloads={} total={}",
1092                    plan.data.len(),
1093                    plan_bytes_len,
1094                    model_bytes,
1095                    payload_bytes,
1096                    layout_rans_bytes.len(),
1097                );
1098            }
1099        }
1100    }
1101}