parity_db/
compress.rs

1// Copyright 2021-2022 Parity Technologies (UK) Ltd.
2// This file is dual-licensed as Apache-2.0 or MIT.
3
4//! Compression utility and types.
5
6use std::str::FromStr;
7
8use crate::error::Result;
9
10/// Different compression type
11/// allowed and their u8 representation.
12#[derive(Clone, Copy, Debug, Eq, PartialEq)]
13#[repr(u8)]
14pub enum CompressionType {
15	NoCompression = 0,
16	Lz4 = 1,
17	Snappy = 2,
18}
19
20/// Compression implementation.
21#[derive(Debug)]
22pub struct Compress {
23	inner: Compressor,
24	pub threshold: u32,
25}
26
27impl Compress {
28	pub fn new(kind: CompressionType, threshold: u32) -> Self {
29		Compress { inner: kind.into(), threshold }
30	}
31}
32
33pub const NO_COMPRESSION: Compress =
34	Compress { inner: Compressor::NoCompression(NoCompression), threshold: u32::MAX };
35
36#[derive(Debug)]
37enum Compressor {
38	NoCompression(NoCompression),
39	Lz4(lz4::Lz4),
40	Snappy(snappy::Snappy),
41}
42
43impl From<u8> for CompressionType {
44	fn from(comp_type: u8) -> Self {
45		match comp_type {
46			a if a == CompressionType::NoCompression as u8 => CompressionType::NoCompression,
47			a if a == CompressionType::Lz4 as u8 => CompressionType::Lz4,
48			a if a == CompressionType::Snappy as u8 => CompressionType::Snappy,
49			_ => panic!("Unknown compression."),
50		}
51	}
52}
53
54impl From<CompressionType> for Compressor {
55	fn from(comp_type: CompressionType) -> Self {
56		match comp_type {
57			CompressionType::NoCompression => Compressor::NoCompression(NoCompression),
58			CompressionType::Lz4 => Compressor::Lz4(lz4::Lz4::new()),
59			CompressionType::Snappy => Compressor::Snappy(snappy::Snappy::new()),
60			#[allow(unreachable_patterns)]
61			_ => unimplemented!("Missing compression implementation."),
62		}
63	}
64}
65
66impl From<&Compress> for CompressionType {
67	fn from(compression: &Compress) -> Self {
68		match compression.inner {
69			Compressor::NoCompression(_) => CompressionType::NoCompression,
70			Compressor::Lz4(_) => CompressionType::Lz4,
71			Compressor::Snappy(_) => CompressionType::Snappy,
72			#[allow(unreachable_patterns)]
73			_ => unimplemented!("Missing compression implementation."),
74		}
75	}
76}
77
78impl FromStr for CompressionType {
79	type Err = crate::Error;
80
81	fn from_str(s: &str) -> Result<Self> {
82		match s.to_lowercase().as_str() {
83			"none" => Ok(CompressionType::NoCompression),
84			"lz4" => Ok(CompressionType::Lz4),
85			"snappy" => Ok(CompressionType::Snappy),
86			_ => Err(crate::Error::Compression),
87		}
88	}
89}
90
91impl Compress {
92	pub fn compress(&self, buf: &[u8]) -> Vec<u8> {
93		match &self.inner {
94			Compressor::NoCompression(inner) => inner.compress(buf),
95			Compressor::Lz4(inner) => inner.compress(buf),
96			Compressor::Snappy(inner) => inner.compress(buf),
97			#[allow(unreachable_patterns)]
98			_ => unimplemented!("Missing compression implementation."),
99		}
100	}
101
102	pub fn decompress(&self, buf: &[u8]) -> Result<Vec<u8>> {
103		Ok(match &self.inner {
104			Compressor::NoCompression(inner) => inner.decompress(buf)?,
105			Compressor::Lz4(inner) => inner.decompress(buf)?,
106			Compressor::Snappy(inner) => inner.decompress(buf)?,
107			#[allow(unreachable_patterns)]
108			_ => unimplemented!("Missing compression implementation."),
109		})
110	}
111}
112
113#[derive(Debug)]
114struct NoCompression;
115
116impl NoCompression {
117	fn compress(&self, buf: &[u8]) -> Vec<u8> {
118		buf.to_vec()
119	}
120
121	fn decompress(&self, buf: &[u8]) -> Result<Vec<u8>> {
122		Ok(buf.to_vec())
123	}
124}
125
126mod lz4 {
127	use crate::error::{Error, Result};
128
129	#[derive(Debug)]
130	pub(super) struct Lz4;
131
132	impl Lz4 {
133		pub(super) fn new() -> Self {
134			Lz4
135		}
136
137		pub(super) fn compress(&self, buf: &[u8]) -> Vec<u8> {
138			lz4::block::compress(buf, Some(lz4::block::CompressionMode::DEFAULT), true).unwrap()
139		}
140
141		pub(super) fn decompress(&self, buf: &[u8]) -> Result<Vec<u8>> {
142			lz4::block::decompress(buf, None).map_err(|_| Error::Compression)
143		}
144	}
145}
146
147mod snappy {
148	use crate::error::{Error, Result};
149	use std::io::{Read, Write};
150
151	#[derive(Debug)]
152	pub(super) struct Snappy;
153
154	impl Snappy {
155		pub(super) fn new() -> Self {
156			Snappy
157		}
158
159		pub(super) fn compress(&self, value: &[u8]) -> Vec<u8> {
160			let mut buf = Vec::with_capacity(value.len() << 3);
161			{
162				let mut encoder = snap::write::FrameEncoder::new(&mut buf);
163				encoder.write_all(value).expect("Expect in memory write to succeed.");
164			}
165			buf
166		}
167
168		pub(super) fn decompress(&self, value: &[u8]) -> Result<Vec<u8>> {
169			let mut buf = Vec::with_capacity(value.len());
170			let mut decoder = snap::read::FrameDecoder::new(value);
171			decoder.read_to_end(&mut buf).map_err(|_| Error::Compression)?;
172			Ok(buf)
173		}
174	}
175}
176
177#[cfg(test)]
178mod tests {
179	use super::*;
180
181	#[test]
182	fn test_compression_interfaces() {
183		let original = vec![42; 100];
184		let types =
185			vec![CompressionType::NoCompression, CompressionType::Snappy, CompressionType::Lz4];
186
187		for compression_type in types {
188			let compress = Compress::new(compression_type, 0);
189			let v = compress.compress(&original[..]);
190			assert!(v.len() <= 100);
191			let round_tripped = compress.decompress(&v[..]).unwrap();
192			assert_eq!(original, round_tripped);
193		}
194	}
195
196	#[test]
197	fn test_compression_from_str() {
198		let correct_cases = [
199			("lz4", CompressionType::Lz4),
200			("snappy", CompressionType::Snappy),
201			("none", CompressionType::NoCompression),
202			("SNAPPy", CompressionType::Snappy),
203		];
204
205		for (input, expected_compression) in correct_cases {
206			assert_eq!(expected_compression, CompressionType::from_str(input).unwrap());
207		}
208
209		let invalid_cases = ["lz5", "", "cthulhu"];
210
211		for input in invalid_cases {
212			assert!(CompressionType::from_str(input).is_err());
213		}
214	}
215}