use crate::SOURCE_FORMAT_PDF;
use crate::container::{Descriptor, UNIVERSE};
use crate::dra::op::PackItem;
use crate::dra::{Op, Program};
use crate::encode::candidates::{Candidate, CandidateKind};
use crate::error::Result;
use crate::integrity::sha256;
use crate::limits::Limits;
use super::physical::{ObjRole, PdfObjectSpan, PhysicalKind, scan};
pub const XREF_SLOT: u8 = u8::MAX;
pub const MAX_MARKED_OBJECTS: usize = XREF_SLOT as usize;
pub fn propose_pdf_layout(input: &[u8], limits: Limits) -> Result<Option<Candidate>> {
let physical = match scan(input, limits) {
Ok(p) => p,
Err(_) => return Ok(None),
};
if !physical
.spans
.iter()
.any(|s| s.kind == PhysicalKind::XrefSection)
{
return Ok(None);
}
if physical
.objects
.iter()
.any(|o| o.role == ObjRole::XRefStream)
{
return Ok(None);
}
if physical.objects.len() > MAX_MARKED_OBJECTS {
return Ok(None);
}
let mut data: Vec<u8> = Vec::new();
let mut items: Vec<PackItem> = Vec::new();
let mut pos: u64 = 0;
let mut slot_value = [0u64; 256];
let mut slot_marked = [false; 256];
let mut xref_predicted: u64 = 0;
let mut xref_literal: u64 = 0;
let mut startxref_predicted: u64 = 0;
for span in &physical.spans {
let start = span.start as usize;
let end = start + span.len as usize;
let bytes = &input[start..end];
if pos != span.start {
return Ok(None);
}
match span.kind {
PhysicalKind::ObjHeader => {
if let Some(idx) = object_index_at(&physical.objects, span.start) {
let slot = idx as u8;
items.push(PackItem::Mark { slot });
slot_value[slot as usize] = pos;
slot_marked[slot as usize] = true;
}
if !push_pack_literal(&mut data, &mut items, bytes, &mut pos) {
return Ok(None);
}
}
PhysicalKind::XrefSection => {
items.push(PackItem::Mark { slot: XREF_SLOT });
slot_value[XREF_SLOT as usize] = pos;
slot_marked[XREF_SLOT as usize] = true;
match parse_classic_xref(bytes) {
Some(pieces) => {
for piece in pieces {
match piece {
XrefPiece::Literal { start, len } => {
if !push_pack_literal(
&mut data,
&mut items,
&bytes[start..start + len],
&mut pos,
) {
return Ok(None);
}
}
XrefPiece::Entry {
start,
number,
offset,
in_use,
} => {
let slot = if in_use {
offset.and_then(|value| {
predicted_slot(
&physical.objects,
&slot_value,
&slot_marked,
number,
value,
)
})
} else {
None
};
match slot {
Some(slot) => {
items.push(PackItem::Emit { slot, width: 10 });
pos += 10;
if !push_pack_literal(
&mut data,
&mut items,
&bytes[start + 10..start + 20],
&mut pos,
) {
return Ok(None);
}
xref_predicted += 1;
}
None => {
if !push_pack_literal(
&mut data,
&mut items,
&bytes[start..start + 20],
&mut pos,
) {
return Ok(None);
}
xref_literal += 1;
}
}
}
}
}
}
None => {
if !push_pack_literal(&mut data, &mut items, bytes, &mut pos) {
return Ok(None);
}
}
}
}
PhysicalKind::StartXref => match predict_startxref(bytes, &slot_value, &slot_marked) {
Some((prefix_len, width)) => {
if !push_pack_literal(&mut data, &mut items, &bytes[..prefix_len], &mut pos) {
return Ok(None);
}
items.push(PackItem::Emit {
slot: XREF_SLOT,
width,
});
pos += width as u64;
startxref_predicted += 1;
}
None => {
if !push_pack_literal(&mut data, &mut items, bytes, &mut pos) {
return Ok(None);
}
}
},
_ => {
if !push_pack_literal(&mut data, &mut items, bytes, &mut pos) {
return Ok(None);
}
}
}
}
let format_basis = format!(
"pdf-layout;objects={};xref_predicted={};xref_literal={};startxref_predicted={}",
physical.objects.len(),
xref_predicted,
xref_literal,
startxref_predicted
);
let descriptor = Descriptor {
universe: UNIVERSE.to_string(),
source_format: SOURCE_FORMAT_PDF,
format_basis,
models: vec![],
channels: vec![],
objects: vec![data],
program: Program::new(vec![Op::PackSegments {
data_object: 0,
items,
}]),
source_sha256: sha256(input),
source_len: input.len() as u64,
};
let candidate = Candidate {
kind: CandidateKind::PdfLayout,
descriptor,
};
let (encoded, _) = candidate.descriptor.serialize()?;
let parsed = match Descriptor::parse(&encoded, limits) {
Ok(p) => p,
Err(_) => return Ok(None),
};
let out = match crate::materialize::materialize(&parsed, limits) {
Ok(o) => o,
Err(_) => return Ok(None),
};
if out != input {
return Ok(None);
}
Ok(Some(candidate))
}
fn push_pack_literal(
data: &mut Vec<u8>,
items: &mut Vec<PackItem>,
bytes: &[u8],
pos: &mut u64,
) -> bool {
if !push_literal(items, data, bytes) {
return false;
}
*pos += bytes.len() as u64;
true
}
fn push_literal(items: &mut Vec<PackItem>, data: &mut Vec<u8>, bytes: &[u8]) -> bool {
if bytes.is_empty() {
return true;
}
let Ok(len) = u32::try_from(bytes.len()) else {
return false;
};
if let Some(PackItem::Literal { len: prev }) = items.last_mut() {
let Some(total) = prev.checked_add(len) else {
return false;
};
*prev = total;
} else {
items.push(PackItem::Literal { len });
}
data.extend_from_slice(bytes);
true
}
fn object_index_at(objects: &[PdfObjectSpan], start: u64) -> Option<usize> {
objects.iter().position(|o| o.start == start)
}
fn predicted_slot(
objects: &[PdfObjectSpan],
slot_value: &[u64; 256],
slot_marked: &[bool; 256],
number: u64,
value: u64,
) -> Option<u8> {
objects.iter().enumerate().find_map(|(i, o)| {
let slot = i as u8;
(o.number == number && slot_marked[slot as usize] && slot_value[slot as usize] == value)
.then_some(slot)
})
}
fn predict_startxref(
bytes: &[u8],
slot_value: &[u64; 256],
slot_marked: &[bool; 256],
) -> Option<(usize, u8)> {
if !slot_marked[XREF_SLOT as usize] {
return None;
}
let mut i = bytes.len();
while i > 0 && bytes[i - 1].is_ascii_digit() {
i -= 1;
}
let width = bytes.len() - i;
if width == 0 || width > 20 {
return None;
}
let value = parse_digits(&bytes[i..])?;
if value != slot_value[XREF_SLOT as usize] {
return None;
}
Some((i, width as u8))
}
enum XrefPiece {
Literal { start: usize, len: usize },
Entry {
start: usize,
number: u64,
offset: Option<u64>,
in_use: bool,
},
}
fn parse_classic_xref(bytes: &[u8]) -> Option<Vec<XrefPiece>> {
if !bytes.starts_with(b"xref") {
return None;
}
let mut pieces = Vec::new();
let mut pos = 4usize;
let eol = eol_len(&bytes[pos..])?;
pieces.push(XrefPiece::Literal {
start: 0,
len: pos + eol,
});
pos += eol;
let mut any = false;
while pos < bytes.len() {
let header_start = pos;
let (start, after_start) = parse_uint_at(bytes, pos)?;
pos = after_start;
let spaces_start = pos;
while pos < bytes.len() && bytes[pos] == b' ' {
pos += 1;
}
if pos == spaces_start {
return None;
}
let (count, after_count) = parse_uint_at(bytes, pos)?;
pos = after_count;
let eol = eol_len(&bytes[pos..])?;
let header_end = pos + eol;
pieces.push(XrefPiece::Literal {
start: header_start,
len: header_end - header_start,
});
pos = header_end;
for i in 0..count {
let end = pos.checked_add(20)?;
if end > bytes.len() {
return None;
}
let entry = &bytes[pos..end];
if !is_entry_shape(entry) {
return None;
}
let number = start.checked_add(i)?;
let in_use = entry[17] == b'n';
let offset = parse_digits(&entry[0..10]);
pieces.push(XrefPiece::Entry {
start: pos,
number,
offset,
in_use,
});
pos = end;
}
any = true;
}
if !any || pos != bytes.len() {
return None;
}
Some(pieces)
}
fn is_entry_shape(entry: &[u8]) -> bool {
if entry.len() != 20 {
return false;
}
if entry[10] != b' ' || entry[16] != b' ' {
return false;
}
if !entry[11..16].iter().all(u8::is_ascii_digit) {
return false;
}
if entry[17] != b'n' && entry[17] != b'f' {
return false;
}
matches!(
(entry[18], entry[19]),
(b'\r', b'\n') | (b'\n', b'\r') | (b' ', b'\n') | (b' ', b'\r')
)
}
fn eol_len(bytes: &[u8]) -> Option<usize> {
match bytes {
[b'\r', b'\n', ..] => Some(2),
[b'\n', ..] | [b'\r', ..] => Some(1),
_ => None,
}
}
fn parse_uint_at(bytes: &[u8], at: usize) -> Option<(u64, usize)> {
let mut i = at;
let mut value: u64 = 0;
while i < bytes.len() && bytes[i].is_ascii_digit() {
value = value
.checked_mul(10)?
.checked_add(u64::from(bytes[i] - b'0'))?;
i += 1;
}
if i == at {
return None;
}
Some((value, i))
}
fn parse_digits(digits: &[u8]) -> Option<u64> {
if digits.is_empty() {
return None;
}
let mut value: u64 = 0;
for &b in digits {
if !b.is_ascii_digit() {
return None;
}
value = value.checked_mul(10)?.checked_add(u64::from(b - b'0'))?;
}
Some(value)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::adapter::pdf::samples::{is_negative_control, sample_pdfs};
use crate::container::Descriptor;
fn sample(name: &str) -> Vec<u8> {
sample_pdfs()
.into_iter()
.find(|(n, _)| *n == name)
.unwrap_or_else(|| panic!("sample {name} missing"))
.1
}
fn basis_field(basis: &str, key: &str) -> Option<u64> {
basis.split(';').find_map(|part| {
let (k, v) = part.split_once('=')?;
(k == key).then(|| v.parse().ok()).flatten()
})
}
fn pack_items(cand: &Candidate) -> &[PackItem] {
cand.descriptor
.program
.ops
.iter()
.find_map(|op| match op {
Op::PackSegments { items, .. } => Some(items.as_slice()),
_ => None,
})
.expect("layout program is one PackSegments op")
}
fn pack_emit_count(cand: &Candidate) -> usize {
pack_items(cand)
.iter()
.filter(|item| matches!(item, PackItem::Emit { .. }))
.count()
}
fn assert_materializes_exactly(name: &str, bytes: &[u8]) {
let cand = propose_pdf_layout(bytes, Limits::DEFAULT)
.unwrap()
.unwrap_or_else(|| panic!("{name} must propose a layout candidate"));
assert_eq!(cand.kind, CandidateKind::PdfLayout);
assert_eq!(cand.descriptor.source_format, SOURCE_FORMAT_PDF);
assert_eq!(cand.descriptor.source_len, bytes.len() as u64);
assert_eq!(
cand.descriptor.objects.len(),
1,
"{name} layout carries one packed data object"
);
assert!(matches!(
cand.descriptor.program.ops.as_slice(),
[Op::PackSegments { .. }]
));
assert!(cand.descriptor.models.is_empty());
assert!(cand.descriptor.channels.is_empty());
let (encoded, _) = cand.descriptor.serialize().unwrap();
let parsed = Descriptor::parse(&encoded, Limits::DEFAULT).unwrap();
let out = crate::materialize::materialize(&parsed, Limits::DEFAULT).unwrap();
assert_eq!(out, bytes, "{name} layout candidate materializes exactly");
assert_eq!(sha256(&out), sha256(bytes), "{name} layout sha");
}
#[test]
fn layout_is_exact_on_corpus() {
let mut accepted = 0usize;
for (name, bytes) in sample_pdfs() {
if is_negative_control(name) {
assert!(
propose_pdf_layout(&bytes, Limits::DEFAULT)
.unwrap()
.is_none(),
"{name} is not a classic-xref PDF and must decline"
);
continue;
}
if propose_pdf_layout(&bytes, Limits::DEFAULT)
.unwrap()
.is_some()
{
assert_materializes_exactly(name, &bytes);
accepted += 1;
}
}
assert!(
accepted >= 3,
"expected several accepted classic-xref samples, found {accepted}"
);
}
#[test]
fn layout_v2_exact() {
for name in ["classic.pdf", "many.pdf", "bigtext.pdf"] {
assert_materializes_exactly(name, &sample(name));
}
}
#[test]
fn layout_v2_predicts_many_entries() {
let bytes = sample("many.pdf");
let cand = propose_pdf_layout(&bytes, Limits::DEFAULT)
.unwrap()
.unwrap();
let emits = pack_emit_count(&cand);
assert!(
emits >= 100,
"many.pdf must predict at least 100 xref offsets, got {emits}"
);
let predicted = basis_field(&cand.descriptor.format_basis, "xref_predicted")
.expect("format basis must report xref_predicted");
let startxref = basis_field(&cand.descriptor.format_basis, "startxref_predicted")
.expect("format basis must report startxref_predicted");
assert_eq!(predicted + startxref, emits as u64);
assert!(predicted >= 100, "many.pdf xref predictions: {predicted}");
let classic = propose_pdf_layout(&sample("classic.pdf"), Limits::DEFAULT)
.unwrap()
.unwrap();
assert!(pack_emit_count(&classic) > 0);
}
#[test]
fn layout_falls_back_on_bad_offset() {
let mut b: Vec<u8> = Vec::new();
b.extend_from_slice(b"%PDF-1.4\n");
let off1 = b.len() as u64;
b.extend_from_slice(b"1 0 obj\n<< /Type /Catalog >>\nendobj\n");
let xref = b.len() as u64;
let wrong = off1 + 3;
b.extend_from_slice(
format!("xref\n0 2\n0000000000 65535 f \n{wrong:010} 00000 n \n").as_bytes(),
);
b.extend_from_slice(
format!("trailer\n<< /Size 2 /Root 1 0 R >>\nstartxref\n{xref}\n%%EOF\n").as_bytes(),
);
let cand = propose_pdf_layout(&b, Limits::DEFAULT).unwrap().unwrap();
assert_eq!(
basis_field(&cand.descriptor.format_basis, "xref_predicted"),
Some(0),
"a mismatched offset must never be predicted"
);
let emits = pack_emit_count(&cand);
assert_eq!(
emits, 1,
"only startxref is predicted; bad entry is literal"
);
let (encoded, _) = cand.descriptor.serialize().unwrap();
let parsed = Descriptor::parse(&encoded, Limits::DEFAULT).unwrap();
let out = crate::materialize::materialize(&parsed, Limits::DEFAULT).unwrap();
assert_eq!(out, b, "fallback must still be byte-exact");
}
#[test]
fn layout_v2_declines() {
assert!(
propose_pdf_layout(&sample("xrefstream.pdf"), Limits::DEFAULT)
.unwrap()
.is_none(),
"an xref-stream PDF must decline the layout candidate"
);
assert!(
propose_pdf_layout(&classic_with_objects(256), Limits::DEFAULT)
.unwrap()
.is_none(),
"256 objects exceeds the 255 markable slots"
);
assert!(
propose_pdf_layout(&classic_with_objects(255), Limits::DEFAULT)
.unwrap()
.is_some(),
"255 objects must still be markable"
);
}
fn classic_with_objects(n: usize) -> Vec<u8> {
let mut b: Vec<u8> = Vec::new();
b.extend_from_slice(b"%PDF-1.4\n");
let mut offsets = Vec::with_capacity(n);
for number in 1..=n {
offsets.push(b.len() as u64);
b.extend_from_slice(format!("{number} 0 obj\n<< >>\nendobj\n").as_bytes());
}
let xref = b.len() as u64;
b.extend_from_slice(format!("xref\n0 {}\n", n + 1).as_bytes());
b.extend_from_slice(b"0000000000 65535 f \n");
for &off in &offsets {
b.extend_from_slice(format!("{off:010} 00000 n \n").as_bytes());
}
b.extend_from_slice(
format!(
"trailer\n<< /Size {} /Root 1 0 R >>\nstartxref\n{xref}\n%%EOF\n",
n + 1
)
.as_bytes(),
);
b
}
#[test]
fn layout_v2_deterministic() {
for name in ["classic.pdf", "many.pdf", "bigtext.pdf"] {
let bytes = sample(name);
let a = propose_pdf_layout(&bytes, Limits::DEFAULT)
.unwrap()
.unwrap()
.descriptor
.serialize()
.unwrap()
.0;
let b = propose_pdf_layout(&bytes, Limits::DEFAULT)
.unwrap()
.unwrap()
.descriptor
.serialize()
.unwrap()
.0;
assert_eq!(a, b, "{name} layout bytes must be deterministic");
}
}
#[test]
fn layout_measurements_report() {
for (name, bytes) in sample_pdfs() {
let Some(cand) = propose_pdf_layout(&bytes, Limits::DEFAULT).unwrap() else {
eprintln!("layout[{name}]: declined");
continue;
};
let (layout_bytes, _) = cand.descriptor.serialize().unwrap();
let emits = pack_emit_count(&cand);
eprintln!(
"layout[{name}] source={} items={} emits={emits} layout={}",
bytes.len(),
pack_items(&cand).len(),
layout_bytes.len(),
);
}
for name in ["classic.pdf", "bigtext.pdf", "many.pdf"] {
let bytes = sample(name);
let (layout_bytes, _) =
crate::encode::encode_with(&bytes, Limits::DEFAULT, Some(CandidateKind::PdfLayout))
.unwrap();
let (raw_bytes, _) =
crate::encode::encode_with(&bytes, Limits::DEFAULT, Some(CandidateKind::Raw))
.unwrap();
#[cfg(feature = "rans")]
let (byte_rans_bytes, _) =
crate::encode::encode_with(&bytes, Limits::DEFAULT, Some(CandidateKind::ByteRans))
.unwrap();
#[cfg(not(feature = "rans"))]
let byte_rans_bytes: Vec<u8> = Vec::new();
let (_, auto) = crate::encode::encode(&bytes, Limits::DEFAULT).unwrap();
eprintln!(
"sizes[{name}] source={} layout_v2={} raw={} byte_rans={} auto={}({})",
bytes.len(),
layout_bytes.len(),
raw_bytes.len(),
byte_rans_bytes.len(),
auto.kind.name(),
auto.encoded_len,
);
}
}
}