use std::io::{self, Write};
use crate::error::PdfError;
#[derive(Clone, Debug)]
pub enum Object {
Null,
Bool(bool),
Integer(i64),
Real(f64),
Name(String),
LiteralString(Vec<u8>),
HexString(Vec<u8>),
Array(Vec<Object>),
Dict(Dict),
Reference(ObjectId),
Stream(Stream),
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub struct ObjectId {
pub number: u32,
pub generation: u16,
}
impl ObjectId {
pub const fn new(number: u32) -> Self {
Self {
number,
generation: 0,
}
}
}
#[derive(Clone, Debug, Default)]
pub struct Dict {
entries: Vec<(String, Object)>,
}
impl Dict {
pub fn new() -> Self {
Self::default()
}
pub fn set(&mut self, key: &str, value: Object) -> &mut Self {
if let Some(slot) = self.entries.iter_mut().find(|(k, _)| k == key) {
slot.1 = value;
} else {
self.entries.push((key.to_owned(), value));
}
self
}
pub fn with(mut self, key: &str, value: Object) -> Self {
self.set(key, value);
self
}
pub fn entries(&self) -> &[(String, Object)] {
&self.entries
}
pub fn is_empty(&self) -> bool {
self.entries.is_empty()
}
}
#[derive(Clone, Debug)]
pub struct Stream {
pub dict: Dict,
pub data: Vec<u8>,
}
impl Stream {
pub fn new(dict: Dict, data: Vec<u8>) -> Self {
Self { dict, data }
}
}
#[derive(Clone, Debug)]
pub struct IndirectObject {
pub id: ObjectId,
pub object: Object,
}
#[derive(Default)]
pub struct Document {
objects: Vec<IndirectObject>,
next_id: u32,
pub root: Option<ObjectId>,
pub info: Option<ObjectId>,
}
impl Document {
pub fn new() -> Self {
Self {
objects: Vec::new(),
next_id: 1,
root: None,
info: None,
}
}
pub fn allocate_id(&mut self) -> ObjectId {
let id = ObjectId::new(self.next_id);
self.next_id += 1;
id
}
pub fn add_object(&mut self, id: ObjectId, object: Object) {
self.objects.push(IndirectObject { id, object });
}
pub fn add(&mut self, object: Object) -> ObjectId {
let id = self.allocate_id();
self.add_object(id, object);
id
}
pub fn object_count(&self) -> usize {
self.objects.len()
}
pub fn write_to(&self, out: &mut Vec<u8>) -> Result<(), PdfError> {
let root = self
.root
.ok_or_else(|| PdfError::other("Document::write_to: missing /Root reference"))?;
out.extend_from_slice(b"%PDF-1.4\n");
out.extend_from_slice(b"%\xE2\xE3\xCF\xD3\n");
let mut sorted = self.objects.iter().collect::<Vec<_>>();
sorted.sort_by_key(|o| o.id.number);
let max_id = sorted.last().map(|o| o.id.number as usize).unwrap_or(0);
let mut offsets: Vec<u64> = vec![0; max_id + 1];
for ind in &sorted {
let off = out.len() as u64;
offsets[ind.id.number as usize] = off;
write_indirect(out, ind).map_err(PdfError::Io)?;
}
let xref_off = out.len() as u64;
out.extend_from_slice(b"xref\n");
let header_line = format!("0 {}\n", max_id + 1);
out.extend_from_slice(header_line.as_bytes());
out.extend_from_slice(b"0000000000 65535 f \n");
for offset in offsets.iter().skip(1) {
let line = format!("{:010} {:05} n \n", offset, 0);
out.extend_from_slice(line.as_bytes());
}
out.extend_from_slice(b"trailer\n");
let mut trailer_dict = Dict::new()
.with("Size", Object::Integer((max_id + 1) as i64))
.with("Root", Object::Reference(root));
if let Some(info_id) = self.info {
trailer_dict.set("Info", Object::Reference(info_id));
}
let trailer = Object::Dict(trailer_dict);
write_object(out, &trailer).map_err(PdfError::Io)?;
out.extend_from_slice(b"\nstartxref\n");
out.extend_from_slice(format!("{}\n", xref_off).as_bytes());
out.extend_from_slice(b"%%EOF\n");
Ok(())
}
}
fn write_indirect(out: &mut Vec<u8>, ind: &IndirectObject) -> io::Result<()> {
let header = format!("{} {} obj\n", ind.id.number, ind.id.generation);
out.write_all(header.as_bytes())?;
write_object(out, &ind.object)?;
out.write_all(b"\nendobj\n")?;
Ok(())
}
fn write_object(out: &mut Vec<u8>, obj: &Object) -> io::Result<()> {
match obj {
Object::Null => out.write_all(b"null"),
Object::Bool(b) => out.write_all(if *b { b"true" } else { b"false" }),
Object::Integer(n) => out.write_all(format!("{}", n).as_bytes()),
Object::Real(f) => out.write_all(format_real(*f).as_bytes()),
Object::Name(s) => {
out.write_all(b"/")?;
for &b in s.as_bytes() {
let needs_escape = matches!(
b,
0x00..=0x20 | 0x23 | 0x25 | 0x28 | 0x29 | 0x2F | 0x3C | 0x3E | 0x5B | 0x5D
| 0x7B | 0x7D | 0x7F..=0xFF
);
if needs_escape {
out.write_all(format!("#{:02X}", b).as_bytes())?;
} else {
out.write_all(&[b])?;
}
}
Ok(())
}
Object::LiteralString(bytes) => {
out.write_all(b"(")?;
for &b in bytes {
match b {
b'\\' => out.write_all(br"\\")?,
b'(' => out.write_all(br"\(")?,
b')' => out.write_all(br"\)")?,
b'\n' => out.write_all(br"\n")?,
b'\r' => out.write_all(br"\r")?,
b'\t' => out.write_all(br"\t")?,
_ => out.write_all(&[b])?,
}
}
out.write_all(b")")
}
Object::HexString(bytes) => {
out.write_all(b"<")?;
for b in bytes {
out.write_all(format!("{:02X}", b).as_bytes())?;
}
out.write_all(b">")
}
Object::Array(items) => {
out.write_all(b"[")?;
for (i, it) in items.iter().enumerate() {
if i > 0 {
out.write_all(b" ")?;
}
write_object(out, it)?;
}
out.write_all(b"]")
}
Object::Dict(d) => write_dict(out, d),
Object::Reference(id) => {
out.write_all(format!("{} {} R", id.number, id.generation).as_bytes())
}
Object::Stream(s) => {
let mut d = s.dict.clone();
d.set("Length", Object::Integer(s.data.len() as i64));
write_dict(out, &d)?;
out.write_all(b"\nstream\n")?;
out.write_all(&s.data)?;
out.write_all(b"\nendstream")
}
}
}
fn write_dict(out: &mut Vec<u8>, d: &Dict) -> io::Result<()> {
out.write_all(b"<<")?;
for (k, v) in &d.entries {
out.write_all(b" /")?;
out.write_all(k.as_bytes())?;
out.write_all(b" ")?;
write_object(out, v)?;
}
out.write_all(b" >>")
}
fn format_real(f: f64) -> String {
if !f.is_finite() {
return "0".to_string();
}
if f.fract() == 0.0 && f.abs() < 1e16 {
return format!("{}", f as i64);
}
let s = format!("{:.6}", f);
let trimmed = s.trim_end_matches('0').trim_end_matches('.');
if trimmed.is_empty() || trimmed == "-" {
"0".to_string()
} else {
trimmed.to_string()
}
}
#[cfg(test)]
mod tests {
use super::*;
fn write_one(obj: &Object) -> Vec<u8> {
let mut buf = Vec::new();
write_object(&mut buf, obj).unwrap();
buf
}
#[test]
fn primitives_serialize() {
assert_eq!(write_one(&Object::Null), b"null");
assert_eq!(write_one(&Object::Bool(true)), b"true");
assert_eq!(write_one(&Object::Bool(false)), b"false");
assert_eq!(write_one(&Object::Integer(42)), b"42");
assert_eq!(write_one(&Object::Integer(-7)), b"-7");
}
#[test]
fn real_numbers_have_no_trailing_zeros() {
assert_eq!(write_one(&Object::Real(0.0)), b"0");
assert_eq!(write_one(&Object::Real(1.0)), b"1");
assert_eq!(write_one(&Object::Real(0.5)), b"0.5");
assert_eq!(write_one(&Object::Real(-1.25)), b"-1.25");
assert_eq!(write_one(&Object::Real(2.345678987654)), b"2.345679");
}
#[test]
fn names_are_slash_prefixed() {
assert_eq!(write_one(&Object::Name("Pages".into())), b"/Pages");
let escaped = write_one(&Object::Name("a b".into()));
assert_eq!(escaped, b"/a#20b");
}
#[test]
fn arrays_have_space_separated_items() {
let a = Object::Array(vec![
Object::Integer(1),
Object::Integer(2),
Object::Real(0.5),
]);
assert_eq!(write_one(&a), b"[1 2 0.5]");
}
#[test]
fn dicts_iterate_in_insertion_order() {
let d = Dict::new()
.with("Type", Object::Name("Pages".into()))
.with("Count", Object::Integer(1));
assert_eq!(write_one(&Object::Dict(d)), b"<< /Type /Pages /Count 1 >>");
}
#[test]
fn streams_serialize_with_length() {
let body = b"hello".to_vec();
let s = Stream::new(Dict::new(), body);
let bytes = write_one(&Object::Stream(s));
let needle = b"/Length 5";
assert!(
bytes.windows(needle.len()).any(|w| w == needle),
"expected /Length 5 in {:?}",
String::from_utf8_lossy(&bytes)
);
assert!(bytes.windows(7).any(|w| w == b"stream\n"));
assert!(bytes.windows(9).any(|w| w == b"endstream"));
}
#[test]
fn document_writes_full_pdf_envelope() {
let mut doc = Document::new();
let pages_id = doc.allocate_id();
let catalog = Object::Dict(
Dict::new()
.with("Type", Object::Name("Catalog".into()))
.with("Pages", Object::Reference(pages_id)),
);
let catalog_id = doc.add(catalog);
doc.add_object(
pages_id,
Object::Dict(
Dict::new()
.with("Type", Object::Name("Pages".into()))
.with("Count", Object::Integer(0))
.with("Kids", Object::Array(Vec::new())),
),
);
doc.root = Some(catalog_id);
let mut bytes = Vec::new();
doc.write_to(&mut bytes).unwrap();
assert!(bytes.starts_with(b"%PDF-1.4\n"));
assert!(bytes.ends_with(b"%%EOF\n"));
assert!(bytes.windows(5).any(|w| w == b"xref\n"));
assert!(bytes.windows(8).any(|w| w == b"trailer\n"));
assert!(bytes.windows(10).any(|w| w == b"startxref\n"));
}
}