Skip to main content

justpdf_core/writer/
encode.rs

1use flate2::Compression;
2use flate2::write::ZlibEncoder;
3use std::io::Write;
4
5use crate::error::Result;
6use crate::object::{PdfDict, PdfObject};
7
8/// Compress data using zlib (FlateDecode) at default compression level.
9pub fn encode_flate(data: &[u8]) -> Result<Vec<u8>> {
10    let mut encoder = ZlibEncoder::new(Vec::new(), Compression::default());
11    encoder.write_all(data)?;
12    let compressed = encoder.finish()?;
13    Ok(compressed)
14}
15
16/// Compress data using zlib (FlateDecode) at maximum compression level.
17pub fn encode_flate_best(data: &[u8]) -> Result<Vec<u8>> {
18    let mut encoder = ZlibEncoder::new(Vec::new(), Compression::best());
19    encoder.write_all(data)?;
20    let compressed = encoder.finish()?;
21    Ok(compressed)
22}
23
24/// Create a stream dictionary and encoded data.
25///
26/// If `compress` is true, the data is compressed with FlateDecode and the
27/// `/Filter` entry is set. The returned dict always contains `/Length`.
28///
29/// Returns `(dict, encoded_data)`.
30pub fn make_stream(data: &[u8], compress: bool) -> (PdfDict, Vec<u8>) {
31    if compress {
32        let compressed = encode_flate(data).unwrap_or_else(|_| data.to_vec());
33        let mut dict = PdfDict::new();
34        dict.insert(
35            b"Length".to_vec(),
36            PdfObject::Integer(compressed.len() as i64),
37        );
38        dict.insert(
39            b"Filter".to_vec(),
40            PdfObject::Name(b"FlateDecode".to_vec()),
41        );
42        (dict, compressed)
43    } else {
44        let mut dict = PdfDict::new();
45        dict.insert(
46            b"Length".to_vec(),
47            PdfObject::Integer(data.len() as i64),
48        );
49        (dict, data.to_vec())
50    }
51}
52
53#[cfg(test)]
54mod tests {
55    use super::*;
56    use flate2::read::ZlibDecoder;
57    use std::io::Read;
58
59    #[test]
60    fn test_flate_roundtrip() {
61        let original = b"Hello, World! This is a test of PDF stream encoding.";
62        let compressed = encode_flate(original).unwrap();
63        assert_ne!(compressed, original);
64
65        let mut decoder = ZlibDecoder::new(&compressed[..]);
66        let mut decompressed = Vec::new();
67        decoder.read_to_end(&mut decompressed).unwrap();
68        assert_eq!(decompressed, original);
69    }
70
71    #[test]
72    fn test_flate_empty() {
73        let compressed = encode_flate(b"").unwrap();
74        let mut decoder = ZlibDecoder::new(&compressed[..]);
75        let mut decompressed = Vec::new();
76        decoder.read_to_end(&mut decompressed).unwrap();
77        assert_eq!(decompressed, b"");
78    }
79
80    #[test]
81    fn test_make_stream_uncompressed() {
82        let data = b"raw content";
83        let (dict, out_data) = make_stream(data, false);
84
85        assert_eq!(out_data, data);
86        assert_eq!(dict.get_i64(b"Length"), Some(data.len() as i64));
87        assert!(dict.get(b"Filter").is_none());
88    }
89
90    #[test]
91    fn test_make_stream_compressed() {
92        let data = b"some content to compress";
93        let (dict, out_data) = make_stream(data, true);
94
95        assert_eq!(
96            dict.get_i64(b"Length"),
97            Some(out_data.len() as i64)
98        );
99        assert_eq!(dict.get_name(b"Filter"), Some(b"FlateDecode".as_slice()));
100
101        // Verify decompression gives back original
102        let mut decoder = ZlibDecoder::new(&out_data[..]);
103        let mut decompressed = Vec::new();
104        decoder.read_to_end(&mut decompressed).unwrap();
105        assert_eq!(decompressed, data);
106    }
107}