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
241 if block_looks_incompressible(&input[offset..block_end]) {
242 block_encoder::encode_raw_block(
243 &input[offset..block_end],
244 is_last,
245 output,
246 )?;
247 } else {
248 #[cfg(feature = "ldm")]
249 let used_ldm = if let Some(ref mut ldm) = ldm_state {
250 let mut empty = Vec::new();
251 ldm.compress_block(
252 input,
253 offset,
254 block_end,
255 params,
256 &rep_offsets,
257 &mut hash_table,
258 &mut empty,
259 &mut sequences,
260 );
261 true
262 } else {
263 false
264 };
265 #[cfg(not(feature = "ldm"))]
266 let used_ldm = false;
267
268 if !used_ldm {
269 fast::compress_fast_block(
270 input,
271 offset,
272 block_end,
273 params,
274 &rep_offsets,
275 &mut hash_table,
276 &mut sequences,
277 );
278 }
279 if params.force_raw_literals {
280 block_encoder::encode_compressed_block_raw(
281 &input[offset..block_end],
282 &sequences,
283 &mut rep_offsets,
284 is_last,
285 output,
286 &mut workspace,
287 )?;
288 } else {
289 block_encoder::encode_compressed_block(
290 &input[offset..block_end],
291 &sequences,
292 &mut rep_offsets,
293 is_last,
294 output,
295 &mut workspace,
296 strategy::use_custom_sequence_tables(params, input.len()),
297 )?;
298 }
299 }
300 offset = block_end;
301 }
302 }
303 Strategy::DFast => {
304 let short_size = 1usize << params.chain_log;
305 let long_size = 1usize << params.hash_log;
306 let mut hash_short = vec![0u32; short_size];
307 let mut hash_long = vec![0u32; long_size];
308 while offset < input.len() {
309 let chunk_size = (input.len() - offset).min(MAX_BLOCK_SIZE);
310 let block_end = offset + chunk_size;
311 let is_last = block_end >= input.len();
312
313 if block_looks_incompressible(&input[offset..block_end]) {
314 block_encoder::encode_raw_block(
315 &input[offset..block_end],
316 is_last,
317 output,
318 )?;
319 } else {
320 #[cfg(feature = "ldm")]
321 let used_ldm = if let Some(ref mut ldm) = ldm_state {
322 ldm.compress_block(
323 input,
324 offset,
325 block_end,
326 params,
327 &rep_offsets,
328 &mut hash_short,
329 &mut hash_long,
330 &mut sequences,
331 );
332 true
333 } else {
334 false
335 };
336 #[cfg(not(feature = "ldm"))]
337 let used_ldm = false;
338
339 if !used_ldm {
340 dfast::compress_dfast_block(
341 input,
342 offset,
343 block_end,
344 params,
345 &rep_offsets,
346 &mut hash_short,
347 &mut hash_long,
348 &mut sequences,
349 );
350 }
351 block_encoder::encode_compressed_block(
352 &input[offset..block_end],
353 &sequences,
354 &mut rep_offsets,
355 is_last,
356 output,
357 &mut workspace,
358 strategy::use_custom_sequence_tables(params, input.len()),
359 )?;
360 }
361 offset = block_end;
362 }
363 }
364 }
365 }
366
367 let hash = xxh64(input, 0);
368 let checksum = (hash & 0xFFFF_FFFF) as u32;
369 output.extend_from_slice(&checksum.to_le_bytes())?;
370 Ok(())
371}
372
373pub fn compress_with_dict(
374 input: &[u8],
375 level: i32,
376 dict: &zrip_core::dict::Dictionary,
377) -> Result<Vec<u8>, CompressError> {
378 let total_window = dict.content().len() + input.len();
379 let mut params = strategy::level_params_for_size(level, total_window)
380 .ok_or(CompressError::InvalidLevel(level))?;
381 strategy::apply_raw_literals_size_override(&mut params, input.len());
382
383 let mut output = Vec::with_capacity(input.len() + 32);
384 write_frame_header(&mut output, input.len(), Some(dict.id()), params.window_log)?;
385
386 if input.is_empty() {
387 block_encoder::encode_raw_block(&[], true, &mut output)?;
388 } else {
389 let prefix = dict.content();
390 let mut rep_offsets = *dict.rep_offsets();
391 let mut workspace = block_encoder::BlockEncodeWorkspace::new();
392
393 workspace.prev_ll = dict
394 .ll_table()
395 .map(|(dt, al)| block_encoder::FseEncodeTable::from_decode_table(dt, al, 35));
396 workspace.prev_of = dict
397 .of_table()
398 .map(|(dt, al)| block_encoder::FseEncodeTable::from_decode_table(dt, al, 31));
399 workspace.prev_ml = dict
400 .ml_table()
401 .map(|(dt, al)| block_encoder::FseEncodeTable::from_decode_table(dt, al, 52));
402 workspace.prev_huffman = dict.huf_table().and_then(|(dt, tl)| {
403 zrip_core::huffman::encode::HuffmanEncodeTable::from_decode_table(dt, tl)
404 });
405
406 if input.len() <= MAX_BLOCK_SIZE {
407 let sequences = match params.strategy {
408 Strategy::Fast => {
409 fast::compress_fast_with_prefix(input, ¶ms, &rep_offsets, prefix)
410 }
411 Strategy::DFast => {
412 dfast::compress_dfast_with_prefix(input, ¶ms, &rep_offsets, prefix)
413 }
414 };
415 if params.force_raw_literals {
416 block_encoder::encode_compressed_block_raw(
417 input,
418 &sequences,
419 &mut rep_offsets,
420 true,
421 &mut output,
422 &mut workspace,
423 )?;
424 } else {
425 block_encoder::encode_compressed_block(
426 input,
427 &sequences,
428 &mut rep_offsets,
429 true,
430 &mut output,
431 &mut workspace,
432 strategy::use_custom_sequence_tables(¶ms, input.len()),
433 )?;
434 }
435 } else {
436 let mut combined = Vec::with_capacity(prefix.len() + input.len());
437 combined.extend_from_slice(prefix);
438 combined.extend_from_slice(input);
439 let plen = prefix.len();
440 let hash_size = 1usize << params.hash_log;
441 let mut sequences = Vec::new();
442
443 match params.strategy {
444 Strategy::Fast => {
445 let mut hash_table = vec![0u32; hash_size];
446 fast::prefill_hash_table(&combined, plen, params.hash_log, &mut hash_table);
447 let mut offset = 0;
448 while offset < input.len() {
449 let chunk_size = (input.len() - offset).min(MAX_BLOCK_SIZE);
450 let is_last = offset + chunk_size >= input.len();
451 fast::compress_fast_block(
452 &combined,
453 plen + offset,
454 plen + offset + chunk_size,
455 ¶ms,
456 &rep_offsets,
457 &mut hash_table,
458 &mut sequences,
459 );
460 if params.force_raw_literals {
461 block_encoder::encode_compressed_block_raw(
462 &input[offset..offset + chunk_size],
463 &sequences,
464 &mut rep_offsets,
465 is_last,
466 &mut output,
467 &mut workspace,
468 )?;
469 } else {
470 block_encoder::encode_compressed_block(
471 &input[offset..offset + chunk_size],
472 &sequences,
473 &mut rep_offsets,
474 is_last,
475 &mut output,
476 &mut workspace,
477 strategy::use_custom_sequence_tables(¶ms, input.len()),
478 )?;
479 }
480 offset += chunk_size;
481 }
482 }
483 Strategy::DFast => {
484 let short_size = 1usize << params.chain_log;
485 let long_size = 1usize << params.hash_log;
486 let mut hash_short = vec![0u32; short_size];
487 let mut hash_long = vec![0u32; long_size];
488 dfast::prefill_hash_tables(
489 &combined,
490 plen,
491 params.hash_log,
492 params.chain_log,
493 params.min_match,
494 &mut hash_short,
495 &mut hash_long,
496 );
497 let mut offset = 0;
498 while offset < input.len() {
499 let chunk_size = (input.len() - offset).min(MAX_BLOCK_SIZE);
500 let is_last = offset + chunk_size >= input.len();
501 dfast::compress_dfast_block(
502 &combined,
503 plen + offset,
504 plen + offset + chunk_size,
505 ¶ms,
506 &rep_offsets,
507 &mut hash_short,
508 &mut hash_long,
509 &mut sequences,
510 );
511 block_encoder::encode_compressed_block(
512 &input[offset..offset + chunk_size],
513 &sequences,
514 &mut rep_offsets,
515 is_last,
516 &mut output,
517 &mut workspace,
518 strategy::use_custom_sequence_tables(¶ms, input.len()),
519 )?;
520 offset += chunk_size;
521 }
522 }
523 }
524 }
525 }
526
527 let hash = xxh64(input, 0);
528 let checksum = (hash & 0xFFFF_FFFF) as u32;
529 output.extend_from_slice(&checksum.to_le_bytes());
530
531 Ok(output)
532}
533
534pub fn compress_into(input: &[u8], output: &mut [u8], level: i32) -> Result<usize, CompressError> {
535 let mut params = strategy::level_params_for_size(level, input.len())
536 .ok_or(CompressError::InvalidLevel(level))?;
537 strategy::apply_raw_literals_size_override(&mut params, input.len());
538 let mut sink = SliceSink::new(output);
539 compress_frame(input, ¶ms, &mut sink)?;
540 Ok(sink.pos())
541}
542
543#[cfg(test)]
544mod tests {
545 use super::*;
546 use zrip_core::frame::header::parse_frame_header;
547
548 #[test]
549 fn clamp_params_normalizes_public_log_values() {
550 let mut params = strategy::level_params(1).unwrap();
551 params.hash_log = 0;
552 params.chain_log = 40;
553 params.window_log = 40;
554
555 clamp_params_to_src_size(&mut params, usize::MAX);
556
557 assert_eq!(params.hash_log, strategy::HASH_LOG_MIN);
558 assert_eq!(params.chain_log, strategy::HASH_LOG_MAX);
559 assert_eq!(params.window_log, strategy::WINDOW_LOG_MAX);
560 }
561
562 #[test]
563 fn options_clamp_window_log_before_ldm_defaults() {
564 let mut params = strategy::level_params(1).unwrap();
565 let opts = strategy::Options::default().window_log(0);
566
567 strategy::apply_options(&mut params, &opts);
568
569 assert_eq!(params.window_log, strategy::WINDOW_LOG_MIN);
570 #[cfg(feature = "ldm")]
571 {
572 let mut params = strategy::level_params(1).unwrap();
573 let opts = strategy::Options::default().window_log(0).ldm(true);
574 strategy::apply_options(&mut params, &opts);
575
576 let ldm = params.ldm_params.unwrap();
577 assert!(ldm.hash_log >= ldm.bucket_size_log);
578 }
579 }
580
581 #[test]
582 fn small_plain_frame_uses_single_segment_header() {
583 let mut output = Vec::new();
584
585 write_frame_header(&mut output, 12, None, 19).unwrap();
586 let header = parse_frame_header(&output).unwrap();
587
588 assert!(header.single_segment);
589 assert_eq!(header.frame_content_size, Some(12));
590 assert_eq!(header.window_size, 12);
591 assert_eq!(header.dict_id, None);
592 assert!(header.content_checksum);
593 assert_eq!(header.header_size, 6);
594 }
595
596 #[test]
597 fn large_plain_frame_uses_bounded_window_descriptor() {
598 let mut output = Vec::new();
599 let content_size = MAX_WINDOW_SIZE as usize + 1;
600
601 write_frame_header(&mut output, content_size, None, 19).unwrap();
602 let header = parse_frame_header(&output).unwrap();
603
604 assert!(!header.single_segment);
605 assert_eq!(header.frame_content_size, Some(content_size as u64));
606 assert_eq!(header.window_size, 1 << 19);
607 assert_eq!(header.dict_id, None);
608 assert!(header.content_checksum);
609 assert_eq!(header.header_size, 10);
610 }
611
612 #[test]
613 fn dict_frame_uses_window_descriptor_even_when_small() {
614 let mut output = Vec::new();
615
616 write_frame_header(&mut output, 12, Some(0x1234), 10).unwrap();
617 let header = parse_frame_header(&output).unwrap();
618
619 assert!(!header.single_segment);
620 assert_eq!(header.frame_content_size, Some(12));
621 assert_eq!(header.window_size, 1 << 10);
622 assert_eq!(header.dict_id, Some(0x1234));
623 assert!(header.content_checksum);
624 assert_eq!(header.header_size, 12);
625 }
626
627 #[test]
628 fn no_fcs_frame_uses_window_descriptor() {
629 let mut output = Vec::new();
630
631 write_frame_header_without_content_size(&mut output, None, 19).unwrap();
632 let header = parse_frame_header(&output).unwrap();
633
634 assert!(!header.single_segment);
635 assert_eq!(header.frame_content_size, None);
636 assert_eq!(header.window_size, 1 << 19);
637 assert_eq!(header.dict_id, None);
638 assert!(header.content_checksum);
639 assert_eq!(header.header_size, 6);
640 }
641}