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