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