1pub mod extd;
10pub mod head;
11pub mod json_chunks;
12pub mod layr;
13pub mod lhas;
14pub mod ltbl;
15pub mod preview;
16pub mod voxl;
17pub mod zdic;
18
19pub use extd::Extension;
20pub use head::Head;
21pub use lhas::LayerHashes;
22pub use ltbl::{LayerEntry, LayerTable};
23pub use preview::{Preview, PreviewRole};
24pub use zdic::ZstdDictionary;
25
26use crate::check::Check;
27use crate::error::{Error, Result};
28use zstd::bulk::{Compressor, Decompressor};
29
30pub fn compress(payload: &[u8], level: i32, dictionary: Option<&[u8]>) -> Result<Vec<u8>> {
32 let mut compressor = match dictionary {
33 Some(dictionary) => Compressor::with_dictionary(level, dictionary),
34 None => Compressor::new(level),
35 }
36 .map_err(|e| {
37 Error::new(
38 Check::LayrFrameDecompressedSize,
39 format!("could not start the zstd compressor: {e}"),
40 )
41 })?;
42 compressor.compress(payload).map_err(|e| {
43 Error::new(
44 Check::LayrFrameDecompressedSize,
45 format!("zstd compression failed: {e}"),
46 )
47 })
48}
49
50pub fn decompress(
57 frame: &[u8],
58 uncompressed_size: u64,
59 dictionary: Option<&[u8]>,
60) -> Result<Vec<u8>> {
61 let bound = (frame.len() as u64).saturating_mul(MAX_ZSTD_RATIO as u64);
62 decompress_within(frame, uncompressed_size, dictionary, bound)
63}
64
65pub fn decompress_within(
67 frame: &[u8],
68 uncompressed_size: u64,
69 dictionary: Option<&[u8]>,
70 bound: u64,
71) -> Result<Vec<u8>> {
72 check_allocation(uncompressed_size, bound, frame.len())?;
75 let capacity = usize::try_from(uncompressed_size).map_err(|_| {
76 Error::new(
77 Check::LayrFrameDecompressedSize,
78 format!("declared uncompressed size {uncompressed_size} is not addressable"),
79 )
80 })?;
81 let mut decompressor = match dictionary {
82 Some(dictionary) => Decompressor::with_dictionary(dictionary),
83 None => Decompressor::new(),
84 }
85 .map_err(|e| {
86 Error::new(
87 Check::LayrFrameDecompressedSize,
88 format!("could not start the zstd decompressor: {e}"),
89 )
90 })?;
91 let decoded = decompressor.decompress(frame, capacity).map_err(|e| {
92 Error::new(
93 Check::LayrFrameDecompressedSize,
94 format!("zstd decompression failed: {e}"),
95 )
96 })?;
97 if decoded.len() as u64 != uncompressed_size {
98 return Err(Error::new(
99 Check::LayrFrameDecompressedSize,
100 format!(
101 "the frame expands to {} bytes, not the declared {uncompressed_size}",
102 decoded.len()
103 ),
104 ));
105 }
106 Ok(decoded)
107}
108
109pub fn frame_dict_id(frame: &[u8]) -> Result<u32> {
111 Ok(zstd::zstd_safe::get_dict_id_from_frame(frame)
112 .map(|id| id.get())
113 .unwrap_or(0))
114}
115
116pub fn frame_content_size(frame: &[u8]) -> Result<u64> {
122 match zstd::zstd_safe::get_frame_content_size(frame) {
123 Ok(Some(size)) => Ok(size),
124 Ok(None) => Err(Error::new(
125 Check::LayrContentSizePresent,
126 "the frame carries no content size, so its output cannot be sized",
127 )),
128 Err(e) => Err(Error::new(
129 Check::LayrContentSizePresent,
130 format!("the frame's content size is unreadable: {e}"),
131 )),
132 }
133}
134
135pub fn decompress_frame(frame: &[u8], dictionary: Option<&[u8]>, bound: u64) -> Result<Vec<u8>> {
145 let size = frame_content_size(frame)?;
146 decompress_within(frame, size, dictionary, bound)
147}
148
149pub fn dictionary_id(dictionary: &[u8]) -> u32 {
151 zstd::zstd_safe::get_dict_id_from_dict(dictionary)
152 .map(|id| id.get())
153 .unwrap_or(0)
154}
155
156pub fn train_dictionary(samples: &[&[u8]], max_size: usize) -> Result<Vec<u8>> {
158 zstd::dict::from_samples(samples, max_size).map_err(|e| {
159 Error::new(
160 Check::ZdicDictSize,
161 format!("zstd dictionary training failed: {e}"),
162 )
163 })
164}
165
166pub fn allocation_bound(layers: u64, total_pixels: u32) -> u64 {
172 layers.saturating_mul(u64::from(total_pixels).saturating_mul(5).saturating_add(64))
173}
174
175pub fn check_allocation(uncompressed_size: u64, bound: u64, payload_len: usize) -> Result<()> {
177 if uncompressed_size > bound {
178 return Err(Error::new(
179 Check::FrameAllocationBound,
180 format!(
181 "declared uncompressed size {uncompressed_size} exceeds the {bound} bytes the \
182 structure allows for {payload_len} stored bytes"
183 ),
184 ));
185 }
186 Ok(())
187}
188
189pub const MAX_ZSTD_RATIO: usize = 4096;
192
193#[cfg(test)]
194mod tests {
195 use super::*;
196 use std::io::Write;
197
198 #[test]
199 fn zstd_round_trips_with_and_without_a_dictionary() {
200 let payload: Vec<u8> = (0..4096u32).map(|i| (i % 17) as u8).collect();
201
202 let framed = compress(&payload, 3, None).unwrap();
203 assert_eq!(frame_dict_id(&framed).unwrap(), 0);
204 assert_eq!(
205 decompress(&framed, payload.len() as u64, None).unwrap(),
206 payload
207 );
208
209 let samples: Vec<Vec<u8>> = (0..256u32)
211 .map(|s| (0..256u32).map(|i| ((i + s) % 23) as u8).collect())
212 .collect();
213 let refs: Vec<&[u8]> = samples.iter().map(|s| s.as_slice()).collect();
214 let dictionary = train_dictionary(&refs, 512).unwrap();
215 let dict_id = dictionary_id(&dictionary);
216 assert_ne!(dict_id, 0, "a trained dictionary reports its ID");
217
218 let framed = compress(&payload, 3, Some(&dictionary)).unwrap();
219 assert_eq!(frame_dict_id(&framed).unwrap(), dict_id);
220 assert_eq!(
221 decompress(&framed, payload.len() as u64, Some(&dictionary)).unwrap(),
222 payload
223 );
224
225 let err = decompress(&framed, payload.len() as u64 + 1, Some(&dictionary)).unwrap_err();
227 assert_eq!(err.check(), Check::LayrFrameDecompressedSize);
228 }
229
230 #[test]
231 fn a_writer_frame_declares_the_size_it_expands_to() {
232 let payload: Vec<u8> = (0..4096u32).map(|i| (i % 17) as u8).collect();
233 let framed = compress(&payload, 3, None).unwrap();
234 assert_eq!(frame_content_size(&framed).unwrap(), payload.len() as u64);
235 let bound = (framed.len() as u64).saturating_mul(MAX_ZSTD_RATIO as u64);
238 assert_eq!(decompress_frame(&framed, None, bound).unwrap(), payload);
239
240 let streamed = {
245 let mut encoder = zstd::stream::write::Encoder::new(Vec::new(), 3).unwrap();
246 encoder.write_all(&payload).unwrap();
247 encoder.finish().unwrap()
248 };
249 assert_eq!(
250 frame_content_size(&streamed).unwrap_err().check(),
251 Check::LayrContentSizePresent
252 );
253 }
254
255 #[test]
256 fn check_allocation_bounds_the_declared_size() {
257 let bound = (MAX_ZSTD_RATIO - 1) as u64;
258 assert!(check_allocation(bound, bound, 1).is_ok());
259 assert_eq!(
260 check_allocation(bound + 1, bound, 1).unwrap_err().check(),
261 Check::FrameAllocationBound
262 );
263 assert!(allocation_bound(2, 1000) > 0);
264 }
265
266 #[test]
271 fn a_layr_frame_is_bounded_by_its_slices_not_by_a_ratio() {
272 let payload = vec![0u8; 1 << 20];
273 let frame = compress(&payload, 3, None).unwrap();
274 let ratio = payload.len() as f64 / frame.len() as f64;
275 assert!(
276 ratio > MAX_ZSTD_RATIO as f64,
277 "this test needs a frame past the ratio guard, got {ratio:.0}:1"
278 );
279
280 assert_eq!(
282 decompress_frame(&frame, None, (frame.len() as u64) * MAX_ZSTD_RATIO as u64)
283 .unwrap_err()
284 .check(),
285 Check::FrameAllocationBound
286 );
287
288 let total_pixels = 1 << 20;
291 let bound = allocation_bound(1, total_pixels);
292 assert!(bound >= payload.len() as u64);
293 assert_eq!(decompress_frame(&frame, None, bound).unwrap(), payload);
294 }
295}