Skip to main content

zrip_encode/
lib.rs

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, &params)
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, &params)
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, &params)
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, &params, &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, &params, &rep_offsets, prefix)
334                }
335                Strategy::DFast => {
336                    dfast::compress_dfast_with_prefix(input, &params, &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(&params, 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                            &params,
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(&params, 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                            &params,
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(&params, 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, &params, &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}