rustolio_utils/bytes/
compression.rs1use std::io;
12use std::io::prelude::*;
13
14use flate2::{
15 write::{GzDecoder, GzEncoder},
16 Compression,
17};
18
19use crate::bytes::Bytes;
20
21const LEVEL: u32 = 6;
22
23pub struct Encoder {
24 encoder: GzEncoder<Vec<u8>>,
25 data_size: usize,
26}
27
28impl Default for Encoder {
29 fn default() -> Self {
30 Self::new()
31 }
32}
33
34impl Encoder {
35 pub fn new() -> Self {
36 Self {
37 encoder: GzEncoder::new(Vec::new(), Compression::new(LEVEL)),
38 data_size: 0,
39 }
40 }
41
42 pub fn with_capacity(capacity: usize) -> Self {
43 Self {
44 encoder: GzEncoder::new(Vec::with_capacity(capacity), Compression::new(6)),
45 data_size: 0,
46 }
47 }
48
49 pub fn write(mut self, bytes: impl AsRef<[u8]>) -> io::Result<Self> {
50 let bytes = bytes.as_ref();
51 self.data_size += bytes.len();
52 self.encoder.write_all(bytes)?;
53 Ok(self)
54 }
55
56 pub fn encode(self) -> io::Result<Compressed> {
57 self.encoder
58 .finish()
59 .map(|v| Compressed::new(Bytes::from_owner(v), self.data_size))
60 }
61
62 pub fn encode_once(bytes: impl AsRef<[u8]>) -> io::Result<Compressed> {
63 let bytes = bytes.as_ref();
64 Self::with_capacity(bytes.len()).write(bytes)?.encode()
65 }
66}
67
68#[derive(Debug, Clone, PartialEq, Eq)]
69pub struct Compressed {
70 uncompressed_size: usize,
71 raw: Bytes,
72}
73
74impl Compressed {
75 pub fn new(raw: Bytes, uncompressed_size: usize) -> Self {
76 Self {
77 raw,
78 uncompressed_size,
79 }
80 }
81
82 pub fn decode(self) -> io::Result<Bytes> {
83 let buffer = Vec::with_capacity(self.uncompressed_size);
84 let mut decoder = GzDecoder::new(buffer);
85 decoder.write_all(&self.raw)?;
86 decoder.finish().map(Bytes::from_owner)
87 }
88
89 pub fn raw(&self) -> &[u8] {
90 &self.raw
91 }
92
93 pub fn uncompressed_size(&self) -> usize {
94 self.uncompressed_size
95 }
96}
97
98#[cfg(test)]
99mod tests {
100 use super::*;
101
102 #[test]
103 fn test_compression() {
104 let msg = &[0; 100000];
105
106 let compressed = Encoder::encode_once(msg).unwrap();
107 let decompressed = compressed.decode().unwrap();
108
109 assert_eq!(*decompressed, *msg);
110 }
111}