Skip to main content

rustolio_utils/bytes/
compression.rs

1//
2// SPDX-License-Identifier: MPL-2.0
3//
4// Copyright (c) 2026 Tobias Binnewies. All rights reserved.
5//
6// This Source Code Form is subject to the terms of the Mozilla Public
7// License, v. 2.0. If a copy of the MPL was not distributed with this
8// file, You can obtain one at http://mozilla.org/MPL/2.0/.
9//
10
11use 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}