1#![cfg_attr(not(feature = "std"), no_std)]
2#![cfg_attr(feature = "nightly", feature(optimize_attribute))]
3#![cfg_attr(feature = "paranoid", forbid(unsafe_code))]
4
5#[cfg(feature = "alloc")]
6extern crate alloc;
7
8pub(crate) mod block_encoder;
9#[cfg(feature = "std")]
10pub mod context;
11pub(crate) mod dfast;
12pub(crate) mod fast;
13#[cfg(feature = "ldm")]
14pub(crate) mod ldm;
15pub(crate) mod primitives;
16pub mod strategy;
17#[cfg(feature = "std")]
18pub mod streaming;
19
20#[cfg(feature = "alloc")]
21use alloc::vec;
22#[cfg(feature = "alloc")]
23use alloc::vec::Vec;
24
25use crate::strategy::Strategy;
26use zrip_core::error::CompressError;
27use zrip_core::frame::{MAX_BLOCK_SIZE, ZSTD_MAGIC};
28use zrip_core::xxhash::xxh64;
29
30pub(crate) fn write_frame_header(output: &mut Vec<u8>, content_size: usize, dict_id: Option<u32>) {
31 output.extend_from_slice(&ZSTD_MAGIC.to_le_bytes());
32
33 let fcs_size = if content_size <= 255 {
34 1
35 } else if content_size <= 0xFFFF + 256 {
36 2
37 } else if content_size <= 0xFFFF_FFFF {
38 4
39 } else {
40 8
41 };
42 let fcs_flag: u8 = match fcs_size {
43 1 => 0,
44 2 => 1,
45 4 => 2,
46 _ => 3,
47 };
48
49 let dict_id_flag: u8 = match dict_id {
50 None => 0,
51 Some(id) if id <= 0xFF => 1,
52 Some(id) if id <= 0xFFFF => 2,
53 Some(_) => 3,
54 };
55
56 let descriptor = 0x20 | 0x04 | (fcs_flag << 6) | dict_id_flag;
57 output.push(descriptor);
58
59 match dict_id {
60 Some(id) if id <= 0xFF => output.push(id as u8),
61 Some(id) if id <= 0xFFFF => output.extend_from_slice(&(id as u16).to_le_bytes()),
62 Some(id) => output.extend_from_slice(&id.to_le_bytes()),
63 None => {}
64 }
65
66 match fcs_size {
67 1 => output.push(content_size as u8),
68 2 => {
69 let v = (content_size - 256) as u16;
70 output.extend_from_slice(&v.to_le_bytes());
71 }
72 4 => output.extend_from_slice(&(content_size as u32).to_le_bytes()),
73 _ => output.extend_from_slice(&(content_size as u64).to_le_bytes()),
74 }
75}
76
77pub(crate) fn block_looks_incompressible(data: &[u8]) -> bool {
78 const SAMPLE: usize = 1024;
79 const DISTINCT_THRESHOLD: u32 = 200;
80 const MAX_FREQ_DENOM: u32 = 24;
81 if data.len() < SAMPLE {
82 return false;
83 }
84 let mut counts = [0u16; 256];
85 for &b in &data[..SAMPLE] {
86 counts[b as usize] += 1;
87 }
88 let mut distinct: u32 = 0;
89 let mut max_freq: u16 = 0;
90 for &c in &counts {
91 distinct += (c > 0) as u32;
92 max_freq = max_freq.max(c);
93 }
94 distinct >= DISTINCT_THRESHOLD && (max_freq as u32) <= SAMPLE as u32 / MAX_FREQ_DENOM
95}
96
97pub(crate) fn clamp_params_to_src_size(params: &mut strategy::LevelParams, src_len: usize) {
98 params.hash_log = params
99 .hash_log
100 .clamp(strategy::HASH_LOG_MIN, strategy::HASH_LOG_MAX);
101 params.chain_log = params
102 .chain_log
103 .clamp(strategy::HASH_LOG_MIN, strategy::HASH_LOG_MAX);
104 params.window_log = params
105 .window_log
106 .clamp(strategy::WINDOW_LOG_MIN, strategy::WINDOW_LOG_MAX);
107 if src_len >= 2 {
108 let src_log = 32 - ((src_len as u32) - 1).leading_zeros();
109 params.hash_log = params.hash_log.min(src_log).max(strategy::HASH_LOG_MIN);
110 params.chain_log = params.chain_log.min(src_log).max(strategy::HASH_LOG_MIN);
111 params.window_log = params.window_log.min(src_log);
112 }
113}
114
115pub fn compress_with_params(
116 input: &[u8],
117 params: &strategy::LevelParams,
118) -> Result<Vec<u8>, CompressError> {
119 let mut params = *params;
120 clamp_params_to_src_size(&mut params, input.len());
121 compress_inner(input, ¶ms)
122}
123
124pub fn compress(input: &[u8], level: i32) -> Result<Vec<u8>, CompressError> {
125 let params = strategy::level_params_for_size(level, input.len())
126 .ok_or(CompressError::InvalidLevel(level))?;
127 compress_inner(input, ¶ms)
128}
129
130pub fn compress_opts(
131 input: &[u8],
132 level: i32,
133 opts: &strategy::Options,
134) -> Result<Vec<u8>, CompressError> {
135 let mut params = strategy::level_params_for_size(level, input.len())
136 .ok_or(CompressError::InvalidLevel(level))?;
137 strategy::apply_options(&mut params, opts);
138 clamp_params_to_src_size(&mut params, input.len());
139 compress_inner(input, ¶ms)
140}
141
142#[allow(clippy::unnecessary_wraps)]
143fn compress_inner(input: &[u8], params: &strategy::LevelParams) -> Result<Vec<u8>, CompressError> {
144 let mut params = *params;
145 strategy::apply_raw_literals_size_override(&mut params, input.len());
146 let mut output = Vec::with_capacity(input.len() + 32);
147 compress_frame(input, ¶ms, &mut output);
148 Ok(output)
149}
150
151fn compress_frame(input: &[u8], params: &strategy::LevelParams, output: &mut Vec<u8>) {
152 write_frame_header(output, input.len(), None);
153
154 if input.is_empty() {
155 block_encoder::encode_raw_block(&[], true, output);
156 } else {
157 let mut rep_offsets = [1u32, 4, 8];
158 let mut offset = 0;
159 let mut sequences = Vec::with_capacity(MAX_BLOCK_SIZE / 8);
160 let mut workspace = block_encoder::BlockEncodeWorkspace::new();
161
162 #[cfg(feature = "ldm")]
163 let mut ldm_state = params.ldm_params.as_ref().map(ldm::LdmState::new);
164
165 match params.strategy {
166 Strategy::Fast => {
167 let hash_size = 1usize << params.hash_log;
168 let mut hash_table = vec![0u32; hash_size];
169 while offset < input.len() {
170 let chunk_size = (input.len() - offset).min(MAX_BLOCK_SIZE);
171 let block_end = offset + chunk_size;
172 let is_last = block_end >= input.len();
173
174 if block_looks_incompressible(&input[offset..block_end]) {
175 block_encoder::encode_raw_block(&input[offset..block_end], is_last, output);
176 } else {
177 #[cfg(feature = "ldm")]
178 let used_ldm = if let Some(ref mut ldm) = ldm_state {
179 let mut empty = Vec::new();
180 ldm.compress_block(
181 input,
182 offset,
183 block_end,
184 params,
185 &rep_offsets,
186 &mut hash_table,
187 &mut empty,
188 &mut sequences,
189 );
190 true
191 } else {
192 false
193 };
194 #[cfg(not(feature = "ldm"))]
195 let used_ldm = false;
196
197 if !used_ldm {
198 fast::compress_fast_block(
199 input,
200 offset,
201 block_end,
202 params,
203 &rep_offsets,
204 &mut hash_table,
205 &mut sequences,
206 );
207 }
208 if params.force_raw_literals {
209 block_encoder::encode_compressed_block_raw(
210 &input[offset..block_end],
211 &sequences,
212 &mut rep_offsets,
213 is_last,
214 output,
215 &mut workspace,
216 );
217 } else {
218 block_encoder::encode_compressed_block(
219 &input[offset..block_end],
220 &sequences,
221 &mut rep_offsets,
222 is_last,
223 output,
224 &mut workspace,
225 strategy::use_custom_sequence_tables(params, input.len()),
226 );
227 }
228 }
229 offset = block_end;
230 }
231 }
232 Strategy::DFast => {
233 let short_size = 1usize << params.chain_log;
234 let long_size = 1usize << params.hash_log;
235 let mut hash_short = vec![0u32; short_size];
236 let mut hash_long = vec![0u32; long_size];
237 while offset < input.len() {
238 let chunk_size = (input.len() - offset).min(MAX_BLOCK_SIZE);
239 let block_end = offset + chunk_size;
240 let is_last = block_end >= input.len();
241
242 if block_looks_incompressible(&input[offset..block_end]) {
243 block_encoder::encode_raw_block(&input[offset..block_end], is_last, output);
244 } else {
245 #[cfg(feature = "ldm")]
246 let used_ldm = if let Some(ref mut ldm) = ldm_state {
247 ldm.compress_block(
248 input,
249 offset,
250 block_end,
251 params,
252 &rep_offsets,
253 &mut hash_short,
254 &mut hash_long,
255 &mut sequences,
256 );
257 true
258 } else {
259 false
260 };
261 #[cfg(not(feature = "ldm"))]
262 let used_ldm = false;
263
264 if !used_ldm {
265 dfast::compress_dfast_block(
266 input,
267 offset,
268 block_end,
269 params,
270 &rep_offsets,
271 &mut hash_short,
272 &mut hash_long,
273 &mut sequences,
274 );
275 }
276 block_encoder::encode_compressed_block(
277 &input[offset..block_end],
278 &sequences,
279 &mut rep_offsets,
280 is_last,
281 output,
282 &mut workspace,
283 strategy::use_custom_sequence_tables(params, input.len()),
284 );
285 }
286 offset = block_end;
287 }
288 }
289 }
290 }
291
292 let hash = xxh64(input, 0);
293 let checksum = (hash & 0xFFFF_FFFF) as u32;
294 output.extend_from_slice(&checksum.to_le_bytes());
295}
296
297pub fn compress_with_dict(
298 input: &[u8],
299 level: i32,
300 dict: &zrip_core::dict::Dictionary,
301) -> Result<Vec<u8>, CompressError> {
302 let total_window = dict.content().len() + input.len();
303 let mut params = strategy::level_params_for_size(level, total_window)
304 .ok_or(CompressError::InvalidLevel(level))?;
305 strategy::apply_raw_literals_size_override(&mut params, input.len());
306
307 let mut output = Vec::with_capacity(input.len() + 32);
308 write_frame_header(&mut output, input.len(), Some(dict.id()));
309
310 if input.is_empty() {
311 block_encoder::encode_raw_block(&[], true, &mut output);
312 } else {
313 let prefix = dict.content();
314 let mut rep_offsets = *dict.rep_offsets();
315 let mut workspace = block_encoder::BlockEncodeWorkspace::new();
316
317 workspace.prev_ll = dict
318 .ll_table()
319 .map(|(dt, al)| block_encoder::FseEncodeTable::from_decode_table(dt, al, 35));
320 workspace.prev_of = dict
321 .of_table()
322 .map(|(dt, al)| block_encoder::FseEncodeTable::from_decode_table(dt, al, 31));
323 workspace.prev_ml = dict
324 .ml_table()
325 .map(|(dt, al)| block_encoder::FseEncodeTable::from_decode_table(dt, al, 52));
326 workspace.prev_huffman = dict.huf_table().and_then(|(dt, tl)| {
327 zrip_core::huffman::encode::HuffmanEncodeTable::from_decode_table(dt, tl)
328 });
329
330 if input.len() <= MAX_BLOCK_SIZE {
331 let sequences = match params.strategy {
332 Strategy::Fast => {
333 fast::compress_fast_with_prefix(input, ¶ms, &rep_offsets, prefix)
334 }
335 Strategy::DFast => {
336 dfast::compress_dfast_with_prefix(input, ¶ms, &rep_offsets, prefix)
337 }
338 };
339 if params.force_raw_literals {
340 block_encoder::encode_compressed_block_raw(
341 input,
342 &sequences,
343 &mut rep_offsets,
344 true,
345 &mut output,
346 &mut workspace,
347 );
348 } else {
349 block_encoder::encode_compressed_block(
350 input,
351 &sequences,
352 &mut rep_offsets,
353 true,
354 &mut output,
355 &mut workspace,
356 strategy::use_custom_sequence_tables(¶ms, input.len()),
357 );
358 }
359 } else {
360 let mut combined = Vec::with_capacity(prefix.len() + input.len());
361 combined.extend_from_slice(prefix);
362 combined.extend_from_slice(input);
363 let plen = prefix.len();
364 let hash_size = 1usize << params.hash_log;
365 let mut sequences = Vec::new();
366
367 match params.strategy {
368 Strategy::Fast => {
369 let mut hash_table = vec![0u32; hash_size];
370 fast::prefill_hash_table(&combined, plen, params.hash_log, &mut hash_table);
371 let mut offset = 0;
372 while offset < input.len() {
373 let chunk_size = (input.len() - offset).min(MAX_BLOCK_SIZE);
374 let is_last = offset + chunk_size >= input.len();
375 fast::compress_fast_block(
376 &combined,
377 plen + offset,
378 plen + offset + chunk_size,
379 ¶ms,
380 &rep_offsets,
381 &mut hash_table,
382 &mut sequences,
383 );
384 if params.force_raw_literals {
385 block_encoder::encode_compressed_block_raw(
386 &input[offset..offset + chunk_size],
387 &sequences,
388 &mut rep_offsets,
389 is_last,
390 &mut output,
391 &mut workspace,
392 );
393 } else {
394 block_encoder::encode_compressed_block(
395 &input[offset..offset + chunk_size],
396 &sequences,
397 &mut rep_offsets,
398 is_last,
399 &mut output,
400 &mut workspace,
401 strategy::use_custom_sequence_tables(¶ms, input.len()),
402 );
403 }
404 offset += chunk_size;
405 }
406 }
407 Strategy::DFast => {
408 let short_size = 1usize << params.chain_log;
409 let long_size = 1usize << params.hash_log;
410 let mut hash_short = vec![0u32; short_size];
411 let mut hash_long = vec![0u32; long_size];
412 dfast::prefill_hash_tables(
413 &combined,
414 plen,
415 params.hash_log,
416 params.chain_log,
417 params.min_match,
418 &mut hash_short,
419 &mut hash_long,
420 );
421 let mut offset = 0;
422 while offset < input.len() {
423 let chunk_size = (input.len() - offset).min(MAX_BLOCK_SIZE);
424 let is_last = offset + chunk_size >= input.len();
425 dfast::compress_dfast_block(
426 &combined,
427 plen + offset,
428 plen + offset + chunk_size,
429 ¶ms,
430 &rep_offsets,
431 &mut hash_short,
432 &mut hash_long,
433 &mut sequences,
434 );
435 block_encoder::encode_compressed_block(
436 &input[offset..offset + chunk_size],
437 &sequences,
438 &mut rep_offsets,
439 is_last,
440 &mut output,
441 &mut workspace,
442 strategy::use_custom_sequence_tables(¶ms, input.len()),
443 );
444 offset += chunk_size;
445 }
446 }
447 }
448 }
449 }
450
451 let hash = xxh64(input, 0);
452 let checksum = (hash & 0xFFFF_FFFF) as u32;
453 output.extend_from_slice(&checksum.to_le_bytes());
454
455 Ok(output)
456}
457
458pub fn compress_into(input: &[u8], output: &mut [u8], level: i32) -> Result<usize, CompressError> {
459 let mut params = strategy::level_params_for_size(level, input.len())
460 .ok_or(CompressError::InvalidLevel(level))?;
461 strategy::apply_raw_literals_size_override(&mut params, input.len());
462 let mut buf = Vec::with_capacity(output.len());
463 compress_frame(input, ¶ms, &mut buf);
464 if buf.len() > output.len() {
465 return Err(CompressError::OutputTooSmall);
466 }
467 output[..buf.len()].copy_from_slice(&buf);
468 Ok(buf.len())
469}
470
471#[cfg(test)]
472mod tests {
473 use super::*;
474
475 #[test]
476 fn clamp_params_normalizes_public_log_values() {
477 let mut params = strategy::level_params(1).unwrap();
478 params.hash_log = 0;
479 params.chain_log = 40;
480 params.window_log = 40;
481
482 clamp_params_to_src_size(&mut params, usize::MAX);
483
484 assert_eq!(params.hash_log, strategy::HASH_LOG_MIN);
485 assert_eq!(params.chain_log, strategy::HASH_LOG_MAX);
486 assert_eq!(params.window_log, strategy::WINDOW_LOG_MAX);
487 }
488
489 #[test]
490 fn options_clamp_window_log_before_ldm_defaults() {
491 let mut params = strategy::level_params(1).unwrap();
492 let opts = strategy::Options::default().window_log(0);
493
494 strategy::apply_options(&mut params, &opts);
495
496 assert_eq!(params.window_log, strategy::WINDOW_LOG_MIN);
497 #[cfg(feature = "ldm")]
498 {
499 let mut params = strategy::level_params(1).unwrap();
500 let opts = strategy::Options::default().window_log(0).ldm(true);
501 strategy::apply_options(&mut params, &opts);
502
503 let ldm = params.ldm_params.unwrap();
504 assert!(ldm.hash_log >= ldm.bucket_size_log);
505 }
506 }
507}