paperforge_pdf/

object.rs

use std::fmt;

use rustc_hash::FxHashMap;

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct ObjectId {
    pub number: u32,
    pub generation: u16,
}

impl ObjectId {
    pub fn new(number: u32, generation: u16) -> Self {
        Self { number, generation }
    }
}

impl fmt::Display for ObjectId {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{} {} R", self.number, self.generation)
    }
}

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct PdfName(pub String);

impl PdfName {
    pub fn new(s: &str) -> Self {
        Self(s.to_string())
    }

    pub fn as_str(&self) -> &str {
        &self.0
    }
}

impl fmt::Display for PdfName {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "/{}", self.0)
    }
}

#[derive(Debug, Clone, PartialEq)]
pub struct PdfString(pub Vec<u8>);

impl PdfString {
    pub fn from_literal(s: &str) -> Self {
        Self(s.as_bytes().to_vec())
    }

    pub fn from_bytes(bytes: &[u8]) -> Self {
        Self(bytes.to_vec())
    }

    pub fn as_bytes(&self) -> &[u8] {
        &self.0
    }
}

/// Escapes raw bytes as the body of a PDF literal string (no surrounding parens),
/// appending the escaped bytes to `out`. The output is always pure ASCII.
///
/// Parens, backslashes and EOL / control characters must be escaped, otherwise a
/// round-trip through the parser would change the byte sequence (e.g. an unescaped
/// `)` would close the string early). Bytes outside the printable ASCII range are
/// emitted as 3-digit octal escapes so the output stays byte-exact.
pub fn escape_literal_string_into(out: &mut Vec<u8>, bytes: &[u8]) {
    for &b in bytes {
        match b {
            b'(' => out.extend_from_slice(b"\\("),
            b')' => out.extend_from_slice(b"\\)"),
            b'\\' => out.extend_from_slice(b"\\\\"),
            b'\n' => out.extend_from_slice(b"\\n"),
            b'\r' => out.extend_from_slice(b"\\r"),
            b'\t' => out.extend_from_slice(b"\\t"),
            8 => out.extend_from_slice(b"\\b"),
            12 => out.extend_from_slice(b"\\f"),
            0x20..=0x7e => out.push(b),
            b => {
                // 3-digit octal: \000 .. \377
                out.push(b'\\');
                out.push(b'0' + (b >> 6));
                out.push(b'0' + ((b >> 3) & 7));
                out.push(b'0' + (b & 7));
            }
        }
    }
}

/// Escapes raw bytes as the body of a PDF literal string (no surrounding parens).
/// See [`escape_literal_string_into`] for the escaping rules.
pub fn escape_literal_string(bytes: &[u8]) -> String {
    let mut out = Vec::with_capacity(bytes.len() + 8);
    escape_literal_string_into(&mut out, bytes);
    // The escaped output is always pure ASCII, so this conversion cannot fail.
    String::from_utf8(out).expect("escaped literal string is ASCII")
}

impl fmt::Display for PdfString {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "({})", escape_literal_string(&self.0))
    }
}

#[derive(Debug, Clone, PartialEq)]
pub struct PdfArray(pub Vec<PdfObject>);

impl PdfArray {
    pub fn new() -> Self {
        Self(Vec::new())
    }

    pub fn push(&mut self, obj: PdfObject) {
        self.0.push(obj);
    }

    pub fn get(&self, index: usize) -> Option<&PdfObject> {
        self.0.get(index)
    }

    pub fn len(&self) -> usize {
        self.0.len()
    }

    pub fn is_empty(&self) -> bool {
        self.0.is_empty()
    }
}

impl Default for PdfArray {
    fn default() -> Self {
        Self::new()
    }
}

/// A PDF dictionary backed by an [`FxHashMap`]. PDF names are short strings, so
/// the faster non-cryptographic hasher costs little in practice while speeding
/// up both parsing and serialization. The hasher is not collision-resistant,
/// but dictionaries parsed from untrusted input are bounded by the parser's
/// `max_dict_entries` limit, which caps the worst-case collision cost.
#[derive(Debug, Clone, PartialEq)]
pub struct PdfDictionary {
    entries: FxHashMap<PdfName, PdfObject>,
}

impl PdfDictionary {
    pub fn new() -> Self {
        Self {
            entries: FxHashMap::default(),
        }
    }

    pub fn insert(&mut self, key: &str, value: PdfObject) {
        self.entries.insert(PdfName::new(key), value);
    }

    pub fn get(&self, key: &str) -> Option<&PdfObject> {
        self.entries.get(&PdfName::new(key))
    }

    pub fn get_name(&self, key: &str) -> Option<&PdfName> {
        match self.get(key) {
            Some(PdfObject::Name(n)) => Some(n),
            _ => None,
        }
    }

    pub fn get_integer(&self, key: &str) -> Option<i64> {
        match self.get(key) {
            Some(PdfObject::Integer(i)) => Some(*i),
            _ => None,
        }
    }

    pub fn get_array(&self, key: &str) -> Option<&PdfArray> {
        match self.get(key) {
            Some(PdfObject::Array(a)) => Some(a),
            _ => None,
        }
    }

    pub fn get_dict(&self, key: &str) -> Option<&PdfDictionary> {
        match self.get(key) {
            Some(PdfObject::Dictionary(d)) => Some(d),
            _ => None,
        }
    }

    /// Returns the raw bytes of a string entry, if the value is a `PdfString`.
    pub fn get_string_bytes(&self, key: &str) -> Option<&[u8]> {
        self.get(key)
            .and_then(|o| o.as_string())
            .map(PdfString::as_bytes)
    }

    pub fn len(&self) -> usize {
        self.entries.len()
    }

    pub fn is_empty(&self) -> bool {
        self.entries.is_empty()
    }

    pub fn iter(&self) -> impl Iterator<Item = (&PdfName, &PdfObject)> {
        self.entries.iter()
    }
}

impl Default for PdfDictionary {
    fn default() -> Self {
        Self::new()
    }
}

#[derive(Debug, Clone, PartialEq)]
pub struct PdfStream {
    pub dictionary: PdfDictionary,
    pub data: Vec<u8>,
}

impl PdfStream {
    pub fn new(data: Vec<u8>) -> Self {
        Self {
            dictionary: PdfDictionary::new(),
            data,
        }
    }

    pub fn with_dict(dictionary: PdfDictionary, data: Vec<u8>) -> Self {
        Self { dictionary, data }
    }

    pub fn length(&self) -> usize {
        self.data.len()
    }
}

#[derive(Debug, Clone, PartialEq)]
pub enum PdfObject {
    Null,
    Boolean(bool),
    Integer(i64),
    Real(f64),
    Name(PdfName),
    String(PdfString),
    Array(PdfArray),
    Dictionary(PdfDictionary),
    Stream(PdfStream),
    Reference(ObjectId),
}

impl PdfObject {
    pub fn is_null(&self) -> bool {
        matches!(self, Self::Null)
    }

    pub fn as_bool(&self) -> Option<bool> {
        match self {
            Self::Boolean(b) => Some(*b),
            _ => None,
        }
    }

    pub fn as_integer(&self) -> Option<i64> {
        match self {
            Self::Integer(i) => Some(*i),
            _ => None,
        }
    }

    pub fn as_real(&self) -> Option<f64> {
        match self {
            Self::Real(r) => Some(*r),
            Self::Integer(i) => Some(*i as f64),
            _ => None,
        }
    }

    pub fn as_name(&self) -> Option<&PdfName> {
        match self {
            Self::Name(n) => Some(n),
            _ => None,
        }
    }

    pub fn as_string(&self) -> Option<&PdfString> {
        match self {
            Self::String(s) => Some(s),
            _ => None,
        }
    }

    pub fn as_array(&self) -> Option<&PdfArray> {
        match self {
            Self::Array(a) => Some(a),
            _ => None,
        }
    }

    pub fn as_dict(&self) -> Option<&PdfDictionary> {
        match self {
            Self::Dictionary(d) => Some(d),
            _ => None,
        }
    }

    pub fn as_stream(&self) -> Option<&PdfStream> {
        match self {
            Self::Stream(s) => Some(s),
            _ => None,
        }
    }

    pub fn as_reference(&self) -> Option<ObjectId> {
        match self {
            Self::Reference(id) => Some(*id),
            _ => None,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn escape_literal_handles_all_special_bytes() {
        let input: &[u8] = b"a(b)c\\d\ne\rf\t";
        let escaped = escape_literal_string(input);
        assert_eq!(escaped, "a\\(b\\)c\\\\d\\ne\\rf\\t");
    }

    #[test]
    fn escape_literal_uses_octal_for_control_bytes() {
        assert_eq!(
            escape_literal_string(&[0x01, 0x07, 0x0b]),
            "\\001\\007\\013"
        );
    }

    #[test]
    fn escape_literal_octal_escapes_non_ascii() {
        // Non-ASCII bytes are emitted as 3-digit octal escapes so the output is
        // pure ASCII and round-trips byte-exactly through any PDF reader.
        assert_eq!(escape_literal_string(&[0x80]), "\\200");
        assert_eq!(escape_literal_string(&[0xc3, 0xa9]), "\\303\\251");
        assert_eq!(escape_literal_string(&[0xff]), "\\377");
    }

    #[test]
    fn string_display_escapes_parens_and_backslash() {
        let s = PdfString::from_literal("(a) \\ b");
        assert_eq!(s.to_string(), "(\\(a\\) \\\\ b)");
    }

    #[test]
    fn string_display_escapes_line_breaks() {
        let s = PdfString::from_bytes(b"line1\nline2\r\nline3");
        assert_eq!(s.to_string(), "(line1\\nline2\\r\\nline3)");
    }
}