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