paperforge-pdf 0.1.0

use crate::error::PdfResult;

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum StreamFilter {
    ASCIIHex,
    ASCII85,
    Flate,
    LZW,
    RunLength,
    CCITTFax,
    JBIG2,
    DCT,
    JPX,
}

pub struct StreamDecoder;

impl StreamDecoder {
    pub fn new() -> Self {
        Self
    }

    pub fn decode(&self, data: &[u8], filter: StreamFilter) -> PdfResult<Vec<u8>> {
        self.decode_with_limit(data, filter, usize::MAX)
    }

    /// Decode with a hard cap on the produced bytes (guards decompression bombs).
    pub fn decode_with_limit(
        &self,
        data: &[u8],
        filter: StreamFilter,
        limit: usize,
    ) -> PdfResult<Vec<u8>> {
        match filter {
            StreamFilter::Flate => self.decode_flate(data, limit),
            _ => Err(crate::error::PdfError::NotImplemented(format!(
                "filter {:?} not yet implemented",
                filter
            ))),
        }
    }

    fn decode_flate(&self, data: &[u8], limit: usize) -> PdfResult<Vec<u8>> {
        use std::io::Read;
        let mut decoder = flate2::read::ZlibDecoder::new(data);
        let mut result = Vec::new();
        decoder
            .by_ref()
            .take(limit.saturating_add(1) as u64)
            .read_to_end(&mut result)?;
        Ok(result)
    }
}

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

/// Flate (zlib) stream encoder.
///
/// The encoder keeps the underlying [`flate2::Compress`] state across calls and
/// resets it per stream, so a single instance can compress many small streams
/// (one per content stream, for instance) without paying the zlib state
/// initialisation cost every time. This is roughly 4-6x faster than creating a
/// fresh encoder per stream on typical tiny PDF content streams.
pub struct StreamEncoder {
    /// `(zlib level, compressor)`; the compressor is recreated only when the
    /// requested level changes between calls.
    compressor: Option<(u32, flate2::Compress)>,
}

impl StreamEncoder {
    pub fn new() -> Self {
        Self { compressor: None }
    }

    pub fn encode(&mut self, data: &[u8], filter: StreamFilter) -> PdfResult<Vec<u8>> {
        self.encode_with_compression(data, filter, 6)
    }

    /// Like [`StreamEncoder::encode`] but with an explicit zlib level (0-9)
    /// for the flate filter; other filters ignore the level.
    pub fn encode_with_compression(
        &mut self,
        data: &[u8],
        filter: StreamFilter,
        level: u32,
    ) -> PdfResult<Vec<u8>> {
        let level = level.clamp(0, 9);
        match filter {
            StreamFilter::Flate => self.encode_flate(data, level),
            _ => Err(crate::error::PdfError::NotImplemented(format!(
                "filter {:?} not yet implemented",
                filter
            ))),
        }
    }

    fn encode_flate(&mut self, data: &[u8], level: u32) -> PdfResult<Vec<u8>> {
        if self
            .compressor
            .as_ref()
            .map(|(l, _)| *l != level)
            .unwrap_or(true)
        {
            self.compressor = Some((
                level,
                flate2::Compress::new(flate2::Compression::new(level), true),
            ));
        }
        let (_, compressor) = self.compressor.as_mut().expect("initialized above");
        compressor.reset();

        // Pre-size the output; PDF content streams are small, so a single
        // allocation usually suffices. `compress_vec` writes into the Vec's
        // spare capacity (appending, never overwriting), so on a full buffer we
        // simply reserve more space and call again until the stream finishes.
        let mut out = Vec::with_capacity(data.len().saturating_add(data.len() / 2).max(128));
        loop {
            let before = compressor.total_out();
            match compressor.compress_vec(data, &mut out, flate2::FlushCompress::Finish) {
                Ok(flate2::Status::StreamEnd) => break,
                Ok(_) | Err(_) => {
                    // compress_vec only uses spare capacity; grow it and
                    // continue. If the stream makes no progress at all, bail
                    // out instead of looping forever.
                    if compressor.total_out() == before {
                        return Err(crate::error::PdfError::NotImplemented(
                            "flate compression made no progress".to_string(),
                        ));
                    }
                    out.reserve(data.len().max(64));
                }
            }
        }
        Ok(out)
    }
}

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