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