fuel_block_committer_encoding/bundle/
encoder.rs

1use std::{io::Write, str::FromStr};
2
3use anyhow::{bail, Context};
4use flate2::{write::GzEncoder, Compression};
5
6use super::Bundle;
7
8#[derive(Clone, Debug)]
9pub struct Encoder {
10    compression_level: Option<Compression>,
11}
12
13impl Default for Encoder {
14    fn default() -> Self {
15        Self::new(CompressionLevel::default())
16    }
17}
18
19impl Encoder {
20    pub fn new(level: CompressionLevel) -> Self {
21        let level = match level {
22            CompressionLevel::Disabled => None,
23            CompressionLevel::Min => Some(0),
24            CompressionLevel::Level1 => Some(1),
25            CompressionLevel::Level2 => Some(2),
26            CompressionLevel::Level3 => Some(3),
27            CompressionLevel::Level4 => Some(4),
28            CompressionLevel::Level5 => Some(5),
29            CompressionLevel::Level6 => Some(6),
30            CompressionLevel::Level7 => Some(7),
31            CompressionLevel::Level8 => Some(8),
32            CompressionLevel::Level9 => Some(9),
33            CompressionLevel::Max => Some(10),
34        };
35        Self {
36            compression_level: level.map(Compression::new),
37        }
38    }
39
40    pub fn encode(&self, bundle: Bundle) -> anyhow::Result<Vec<u8>> {
41        const VERSION_SIZE: usize = std::mem::size_of::<u16>();
42
43        let Bundle::V1(v1) = bundle;
44
45        let postcard_encoded =
46            postcard::to_allocvec(&v1).context("could not postcard encode BundleV1 contents")?;
47
48        let blocks_encoded = self
49            .compress(&postcard_encoded)
50            .context("could not compress postcard encoded BundleV1")?;
51        let mut bundle_data = vec![
52            0u8;
53            blocks_encoded
54                .len()
55                .checked_add(VERSION_SIZE)
56                .expect("never to encode more than usize::MAX")
57        ];
58
59        let version = 1u16;
60        bundle_data[..VERSION_SIZE].copy_from_slice(&version.to_be_bytes());
61
62        bundle_data[VERSION_SIZE..].copy_from_slice(&blocks_encoded);
63
64        Ok(bundle_data)
65    }
66
67    fn compress(&self, data: &[u8]) -> anyhow::Result<Vec<u8>> {
68        let Some(compression) = self.compression_level else {
69            return Ok(data.to_vec());
70        };
71
72        let mut encoder = GzEncoder::new(Vec::new(), compression);
73        encoder.write_all(data)?;
74
75        Ok(encoder.finish()?)
76    }
77}
78
79#[derive(Default, Debug, Clone, Copy)]
80#[allow(dead_code)]
81pub enum CompressionLevel {
82    Disabled,
83    Min,
84    Level1,
85    Level2,
86    Level3,
87    Level4,
88    Level5,
89    #[default]
90    Level6,
91    Level7,
92    Level8,
93    Level9,
94    Max,
95}
96
97impl<'a> serde::Deserialize<'a> for CompressionLevel {
98    fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
99    where
100        D: serde::Deserializer<'a>,
101    {
102        let as_string = String::deserialize(deserializer)?;
103
104        Self::from_str(&as_string)
105            .map_err(|e| serde::de::Error::custom(format!("Invalid compression level: {e}")))
106    }
107}
108
109impl FromStr for CompressionLevel {
110    type Err = anyhow::Error;
111
112    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
113        match s.to_lowercase().as_str() {
114            "disabled" => Ok(Self::Disabled),
115            "min" => Ok(Self::Min),
116            "level1" => Ok(Self::Level1),
117            "level2" => Ok(Self::Level2),
118            "level3" => Ok(Self::Level3),
119            "level4" => Ok(Self::Level4),
120            "level5" => Ok(Self::Level5),
121            "level6" => Ok(Self::Level6),
122            "level7" => Ok(Self::Level7),
123            "level8" => Ok(Self::Level8),
124            "level9" => Ok(Self::Level9),
125            "max" => Ok(Self::Max),
126            _ => bail!("Invalid compression level: {s}"),
127        }
128    }
129}
130
131#[allow(dead_code)]
132impl CompressionLevel {
133    pub fn levels() -> Vec<Self> {
134        vec![
135            Self::Disabled,
136            Self::Min,
137            Self::Level1,
138            Self::Level2,
139            Self::Level3,
140            Self::Level4,
141            Self::Level5,
142            Self::Level6,
143            Self::Level7,
144            Self::Level8,
145            Self::Level9,
146            Self::Max,
147        ]
148    }
149}
150
151#[cfg(test)]
152mod tests {
153    use crate::bundle::{CompressionLevel, Encoder};
154
155    #[test]
156    fn can_disable_compression() {
157        // given
158        let compressor = Encoder::new(CompressionLevel::Disabled);
159        let data = vec![1, 2, 3];
160
161        // when
162        let compressed = compressor.compress(&data).unwrap();
163
164        // then
165        assert_eq!(data, compressed);
166    }
167
168    #[test]
169    fn all_compression_levels_work() {
170        let data = vec![1, 2, 3];
171        for level in CompressionLevel::levels() {
172            let compressor = Encoder::new(level);
173            compressor.compress(&data).unwrap();
174        }
175    }
176}