1#[cfg(feature = "alloc")]
2use alloc::borrow::Cow;
3#[cfg(feature = "alloc")]
4use alloc::vec;
5#[cfg(feature = "alloc")]
6use alloc::vec::Vec;
7
8use crate::block_encoder::{self, BlockEncodeWorkspace};
9use crate::strategy::{self, LevelParams, Strategy};
10use crate::{block_looks_incompressible, dfast, fast, write_frame_header};
11use zrip_core::Sequence;
12use zrip_core::dict::Dictionary;
13use zrip_core::error::CompressError;
14use zrip_core::frame::MAX_BLOCK_SIZE;
15use zrip_core::huffman::encode::HuffmanEncodeTable;
16use zrip_core::xxhash::xxh64;
17
18const ATTACH_THRESHOLD: usize = 16384;
24
25pub(crate) struct PreparedDict {
26 combined: Vec<u8>,
27 hash_snapshot: Vec<u32>,
28 hash_long_snapshot: Vec<u32>,
29 hash_log: u32,
30 prefix_len: usize,
31 rep_offsets: [u32; 3],
32 dict_id: u32,
33 huf_table: Option<HuffmanEncodeTable>,
34 ll_table: Option<block_encoder::FseEncodeTable>,
35 of_table: Option<block_encoder::FseEncodeTable>,
36 ml_table: Option<block_encoder::FseEncodeTable>,
37}
38
39impl PreparedDict {
40 pub fn new(dict: &Dictionary, params: &LevelParams) -> Self {
41 let prefix = dict.content();
42 let prefix_len = prefix.len();
43
44 let mut combined = Vec::with_capacity(prefix_len + MAX_BLOCK_SIZE);
45 combined.extend_from_slice(prefix);
46
47 let (hash_snapshot, hash_long_snapshot) = match params.strategy {
48 Strategy::Fast => {
49 let hash_size = 1usize << params.hash_log;
50 let mut hash_table = vec![0u32; hash_size];
51 fast::prefill_hash_table(&combined, prefix_len, params.hash_log, &mut hash_table);
52 (hash_table, Vec::new())
53 }
54 Strategy::DFast => {
55 let short_size = 1usize << params.chain_log;
56 let long_size = 1usize << params.hash_log;
57 let mut hash_short = vec![0u32; short_size];
58 let mut hash_long = vec![0u32; long_size];
59 dfast::prefill_hash_tables(
60 &combined,
61 prefix_len,
62 params.hash_log,
63 params.chain_log,
64 params.min_match,
65 &mut hash_short,
66 &mut hash_long,
67 );
68 (hash_short, hash_long)
69 }
70 };
71
72 let huf_table = dict
73 .huf_table()
74 .and_then(|(dt, tl)| HuffmanEncodeTable::from_decode_table(dt, tl));
75
76 let ll_table = dict
77 .ll_table()
78 .map(|(dt, al)| block_encoder::FseEncodeTable::from_decode_table(dt, al, 35));
79 let of_table = dict
80 .of_table()
81 .map(|(dt, al)| block_encoder::FseEncodeTable::from_decode_table(dt, al, 31));
82 let ml_table = dict
83 .ml_table()
84 .map(|(dt, al)| block_encoder::FseEncodeTable::from_decode_table(dt, al, 52));
85
86 Self {
87 combined,
88 hash_snapshot,
89 hash_long_snapshot,
90 hash_log: params.hash_log,
91 prefix_len,
92 rep_offsets: *dict.rep_offsets(),
93 dict_id: dict.id(),
94 huf_table,
95 ll_table,
96 of_table,
97 ml_table,
98 }
99 }
100}
101
102pub struct CompressContext {
116 level: i32,
117 prepared: Option<PreparedDict>,
118 hash_table: Vec<u32>,
119 hash_long: Vec<u32>,
120 dict_hash: Vec<u32>,
121 small_hash: Vec<u32>,
122 sequences: Vec<Sequence>,
123 output: Vec<u8>,
124 workspace: BlockEncodeWorkspace,
125 combined: Vec<u8>,
126}
127
128impl CompressContext {
129 pub fn new(level: i32) -> Result<Self, CompressError> {
131 let params = strategy::level_params(level).ok_or(CompressError::InvalidLevel(level))?;
132 let max_log = strategy::max_hash_log(level).expect("level validated above");
133 let alloc_size = 1usize << max_log;
134 let (hash_table, hash_long) = match params.strategy {
135 Strategy::Fast => (vec![0u32; alloc_size], Vec::new()),
136 Strategy::DFast => (vec![0u32; alloc_size], vec![0u32; alloc_size]),
137 };
138 Ok(Self {
139 level,
140 prepared: None,
141 hash_table,
142 hash_long,
143 dict_hash: Vec::new(),
144 small_hash: Vec::new(),
145 sequences: Vec::new(),
146 output: Vec::new(),
147 workspace: BlockEncodeWorkspace::new(),
148 combined: Vec::new(),
149 })
150 }
151
152 pub fn with_dict(level: i32, dict: Dictionary) -> Result<Self, CompressError> {
162 Self::with_dict_for_size(level, dict, usize::MAX)
163 }
164
165 pub fn with_dict_for_size(
172 level: i32,
173 dict: Dictionary,
174 expected_size: usize,
175 ) -> Result<Self, CompressError> {
176 let total_window = dict.content().len().saturating_add(expected_size);
177 let params = strategy::level_params_for_size(level, total_window)
178 .ok_or(CompressError::InvalidLevel(level))?;
179 let prepared = PreparedDict::new(&dict, ¶ms);
180 let hash_table = vec![0u32; prepared.hash_snapshot.len()];
181 let hash_long = vec![0u32; prepared.hash_long_snapshot.len()];
182 Ok(Self {
183 level,
184 prepared: Some(prepared),
185 hash_table,
186 hash_long,
187 dict_hash: Vec::new(),
188 small_hash: Vec::new(),
189 sequences: Vec::new(),
190 output: Vec::new(),
191 workspace: BlockEncodeWorkspace::new(),
192 combined: Vec::new(),
193 })
194 }
195
196 pub fn compress(&mut self, input: &[u8]) -> Result<Cow<'_, [u8]>, CompressError> {
198 if self.prepared.is_some() {
199 return self.compress_with_prepared(input);
200 }
201 let params = strategy::level_params_for_size(self.level, input.len())
202 .expect("level validated at construction");
203 compress_core(
204 input,
205 params,
206 None,
207 &[],
208 [1u32, 4, 8],
209 &mut self.hash_table,
210 &mut self.hash_long,
211 &mut self.dict_hash,
212 &mut self.sequences,
213 &mut self.output,
214 &mut self.workspace,
215 &mut self.combined,
216 )?;
217 Ok(self.take_or_borrow_output())
218 }
219
220 pub fn compress_with_dict(
222 &mut self,
223 input: &[u8],
224 dict: &Dictionary,
225 ) -> Result<Cow<'_, [u8]>, CompressError> {
226 let params = strategy::level_params_for_size(self.level, input.len())
227 .expect("level validated at construction");
228 compress_core(
229 input,
230 params,
231 Some(dict.id()),
232 dict.content(),
233 *dict.rep_offsets(),
234 &mut self.hash_table,
235 &mut self.hash_long,
236 &mut self.dict_hash,
237 &mut self.sequences,
238 &mut self.output,
239 &mut self.workspace,
240 &mut self.combined,
241 )?;
242 Ok(self.take_or_borrow_output())
243 }
244
245 fn compress_with_prepared(&mut self, input: &[u8]) -> Result<Cow<'_, [u8]>, CompressError> {
246 let prep = self.prepared.as_ref().unwrap();
247 let total_window = prep.prefix_len + input.len();
248 let mut params = strategy::level_params_for_size(self.level, total_window)
249 .expect("level validated at construction");
250 strategy::apply_raw_literals_size_override(&mut params, input.len());
251
252 let use_attached = !input.is_empty()
253 && input.len() <= ATTACH_THRESHOLD
254 && params.strategy == Strategy::Fast;
255
256 let dict_id = prep.dict_id;
257 let prefix_len = prep.prefix_len;
258 let dict_hash_log = prep.hash_log;
259
260 if !use_attached {
261 let snapshot_matches = match params.strategy {
262 Strategy::Fast => (1usize << params.hash_log) == prep.hash_snapshot.len(),
263 Strategy::DFast => {
264 (1usize << params.chain_log) == prep.hash_snapshot.len()
265 && (1usize << params.hash_log) == prep.hash_long_snapshot.len()
266 }
267 };
268 if !snapshot_matches {
269 return self.compress_with_dict_fallback(input, dict_id, prefix_len);
270 }
271 }
272
273 {
274 let prep = self.prepared.as_mut().unwrap();
275 if !use_attached {
276 self.hash_table.copy_from_slice(&prep.hash_snapshot);
277 if !prep.hash_long_snapshot.is_empty() {
278 self.hash_long.copy_from_slice(&prep.hash_long_snapshot);
279 }
280 }
281 prep.combined.truncate(prep.prefix_len);
282 prep.combined.extend_from_slice(input);
283 }
284
285 let prep = self.prepared.as_ref().unwrap();
286
287 if use_attached {
288 self.workspace.prev_huffman = if params.force_raw_literals {
289 None
290 } else {
291 prep.huf_table.clone()
292 };
293 } else if let Some(ref huf) = prep.huf_table {
294 self.workspace.prev_huffman = Some(huf.clone());
295 } else {
296 self.workspace.prev_huffman = None;
297 }
298 self.workspace.prev_ll = prep.ll_table.clone();
299 self.workspace.prev_of = prep.of_table.clone();
300 self.workspace.prev_ml = prep.ml_table.clone();
301
302 self.output.clear();
303 self.output.reserve(input.len() + 32);
304 write_frame_header(&mut self.output, input.len(), Some(dict_id));
305
306 if input.is_empty() {
307 block_encoder::encode_raw_block(&[], true, &mut self.output);
308 } else if use_attached {
309 let input_hash_log = if input.len() >= 2 {
310 let src_log = 32 - ((input.len() as u32) - 1).leading_zeros();
311 params.hash_log.min(src_log).max(strategy::HASH_LOG_MIN)
312 } else {
313 strategy::HASH_LOG_MIN
314 };
315 let input_hash_size = 1usize << input_hash_log;
316 self.small_hash.resize(input_hash_size, 0);
317 self.small_hash.fill(0);
318
319 let mut rep_offsets = prep.rep_offsets;
320
321 fast::compress_fast_attached(
322 &prep.combined,
323 prefix_len,
324 prefix_len + input.len(),
325 ¶ms,
326 &prep.rep_offsets,
327 &prep.hash_snapshot,
328 dict_hash_log,
329 &mut self.small_hash,
330 input_hash_log,
331 &mut self.sequences,
332 );
333
334 if params.force_raw_literals {
335 block_encoder::encode_compressed_block_raw(
336 input,
337 &self.sequences,
338 &mut rep_offsets,
339 true,
340 &mut self.output,
341 &mut self.workspace,
342 );
343 } else {
344 block_encoder::encode_compressed_block(
345 input,
346 &self.sequences,
347 &mut rep_offsets,
348 true,
349 &mut self.output,
350 &mut self.workspace,
351 strategy::use_custom_sequence_tables(¶ms, input.len()),
352 );
353 }
354 } else {
355 let combined = &prep.combined;
356 let mut rep_offsets = prep.rep_offsets;
357
358 if input.len() <= MAX_BLOCK_SIZE {
359 match params.strategy {
360 Strategy::Fast => {
361 fast::compress_fast_block(
362 combined,
363 prefix_len,
364 prefix_len + input.len(),
365 ¶ms,
366 &rep_offsets,
367 &mut self.hash_table,
368 &mut self.sequences,
369 );
370 }
371 Strategy::DFast => {
372 dfast::compress_dfast_block(
373 combined,
374 prefix_len,
375 prefix_len + input.len(),
376 ¶ms,
377 &rep_offsets,
378 &mut self.hash_table,
379 &mut self.hash_long,
380 &mut self.sequences,
381 );
382 }
383 }
384 if params.force_raw_literals {
385 block_encoder::encode_compressed_block_raw(
386 input,
387 &self.sequences,
388 &mut rep_offsets,
389 true,
390 &mut self.output,
391 &mut self.workspace,
392 );
393 } else {
394 block_encoder::encode_compressed_block(
395 input,
396 &self.sequences,
397 &mut rep_offsets,
398 true,
399 &mut self.output,
400 &mut self.workspace,
401 strategy::use_custom_sequence_tables(¶ms, input.len()),
402 );
403 }
404 } else {
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 match params.strategy {
410 Strategy::Fast => {
411 fast::compress_fast_block(
412 combined,
413 prefix_len + offset,
414 prefix_len + offset + chunk_size,
415 ¶ms,
416 &rep_offsets,
417 &mut self.hash_table,
418 &mut self.sequences,
419 );
420 }
421 Strategy::DFast => {
422 dfast::compress_dfast_block(
423 combined,
424 prefix_len + offset,
425 prefix_len + offset + chunk_size,
426 ¶ms,
427 &rep_offsets,
428 &mut self.hash_table,
429 &mut self.hash_long,
430 &mut self.sequences,
431 );
432 }
433 }
434 if params.force_raw_literals {
435 block_encoder::encode_compressed_block_raw(
436 &input[offset..offset + chunk_size],
437 &self.sequences,
438 &mut rep_offsets,
439 is_last,
440 &mut self.output,
441 &mut self.workspace,
442 );
443 } else {
444 block_encoder::encode_compressed_block(
445 &input[offset..offset + chunk_size],
446 &self.sequences,
447 &mut rep_offsets,
448 is_last,
449 &mut self.output,
450 &mut self.workspace,
451 strategy::use_custom_sequence_tables(¶ms, input.len()),
452 );
453 }
454 offset += chunk_size;
455 }
456 }
457 }
458
459 let hash = xxh64(input, 0);
460 let checksum = (hash & 0xFFFF_FFFF) as u32;
461 self.output.extend_from_slice(&checksum.to_le_bytes());
462
463 Ok(self.take_or_borrow_output())
464 }
465
466 fn compress_with_dict_fallback(
467 &mut self,
468 input: &[u8],
469 dict_id: u32,
470 prefix_len: usize,
471 ) -> Result<Cow<'_, [u8]>, CompressError> {
472 let prep = self.prepared.as_ref().unwrap();
473 let rep_offsets = prep.rep_offsets;
474 let prefix = &prep.combined[..prefix_len];
475
476 let params = strategy::level_params_for_size(self.level, input.len())
477 .expect("level validated at construction");
478 compress_core(
479 input,
480 params,
481 Some(dict_id),
482 prefix,
483 rep_offsets,
484 &mut self.hash_table,
485 &mut self.hash_long,
486 &mut self.dict_hash,
487 &mut self.sequences,
488 &mut self.output,
489 &mut self.workspace,
490 &mut self.combined,
491 )?;
492 Ok(self.take_or_borrow_output())
493 }
494
495 fn take_or_borrow_output(&mut self) -> Cow<'_, [u8]> {
496 if self.output.len() >= zrip_core::LARGE_OUTPUT_THRESHOLD {
497 Cow::Owned(core::mem::take(&mut self.output))
498 } else {
499 Cow::Borrowed(&self.output)
500 }
501 }
502}
503
504#[allow(clippy::too_many_arguments, clippy::unnecessary_wraps)]
505fn compress_core(
506 input: &[u8],
507 params: LevelParams,
508 dict_id: Option<u32>,
509 prefix: &[u8],
510 init_rep_offsets: [u32; 3],
511 hash_table: &mut Vec<u32>,
512 hash_long: &mut Vec<u32>,
513 dict_hash: &mut Vec<u32>,
514 sequences: &mut Vec<Sequence>,
515 output: &mut Vec<u8>,
516 workspace: &mut BlockEncodeWorkspace,
517 combined: &mut Vec<u8>,
518) -> Result<(), CompressError> {
519 let mut params = params;
520 strategy::apply_raw_literals_size_override(&mut params, input.len());
521
522 let hash_size = match params.strategy {
523 Strategy::Fast => 1usize << params.hash_log,
524 Strategy::DFast => 1usize << params.chain_log,
525 };
526 let long_size = 1usize << params.hash_log;
527
528 workspace.prev_huffman = None;
529
530 output.clear();
531 output.reserve(input.len() + 32);
532 write_frame_header(output, input.len(), dict_id);
533
534 if input.is_empty() {
535 block_encoder::encode_raw_block(&[], true, output);
536 } else {
537 let has_prefix = !prefix.is_empty();
538 let mut rep_offsets = init_rep_offsets;
539 let mut offset = 0;
540
541 if hash_table.len() != hash_size {
542 hash_table.resize(hash_size, 0);
543 }
544
545 match params.strategy {
546 Strategy::Fast => {
547 if has_prefix && input.len() <= MAX_BLOCK_SIZE {
548 if dict_hash.len() != hash_size {
549 dict_hash.resize(hash_size, 0);
550 }
551 fast::compress_fast_with_prefix_reuse(
552 input,
553 ¶ms,
554 &rep_offsets,
555 prefix,
556 dict_hash,
557 hash_table,
558 sequences,
559 combined,
560 );
561 if params.force_raw_literals {
562 block_encoder::encode_compressed_block_raw(
563 input,
564 sequences,
565 &mut rep_offsets,
566 true,
567 output,
568 workspace,
569 );
570 } else {
571 block_encoder::encode_compressed_block(
572 input,
573 sequences,
574 &mut rep_offsets,
575 true,
576 output,
577 workspace,
578 strategy::use_custom_sequence_tables(¶ms, input.len()),
579 );
580 }
581 } else if has_prefix {
582 combined.clear();
583 combined.reserve(prefix.len() + input.len());
584 combined.extend_from_slice(prefix);
585 combined.extend_from_slice(input);
586 let plen = prefix.len();
587 fast::prefill_hash_table(combined, plen, params.hash_log, hash_table);
588
589 while offset < input.len() {
590 let chunk_size = (input.len() - offset).min(MAX_BLOCK_SIZE);
591 let is_last = offset + chunk_size >= input.len();
592 fast::compress_fast_block(
593 combined,
594 plen + offset,
595 plen + offset + chunk_size,
596 ¶ms,
597 &rep_offsets,
598 hash_table,
599 sequences,
600 );
601 if params.force_raw_literals {
602 block_encoder::encode_compressed_block_raw(
603 &input[offset..offset + chunk_size],
604 sequences,
605 &mut rep_offsets,
606 is_last,
607 output,
608 workspace,
609 );
610 } else {
611 block_encoder::encode_compressed_block(
612 &input[offset..offset + chunk_size],
613 sequences,
614 &mut rep_offsets,
615 is_last,
616 output,
617 workspace,
618 strategy::use_custom_sequence_tables(¶ms, input.len()),
619 );
620 }
621 offset += chunk_size;
622 }
623 } else {
624 hash_table.fill(0);
625 while offset < input.len() {
626 let chunk_size = (input.len() - offset).min(MAX_BLOCK_SIZE);
627 let block_end = offset + chunk_size;
628 let is_last = block_end >= input.len();
629
630 if block_looks_incompressible(&input[offset..block_end]) {
631 block_encoder::encode_raw_block(
632 &input[offset..block_end],
633 is_last,
634 output,
635 );
636 } else {
637 fast::compress_fast_block(
638 input,
639 offset,
640 block_end,
641 ¶ms,
642 &rep_offsets,
643 hash_table,
644 sequences,
645 );
646 if params.force_raw_literals {
647 block_encoder::encode_compressed_block_raw(
648 &input[offset..block_end],
649 sequences,
650 &mut rep_offsets,
651 is_last,
652 output,
653 workspace,
654 );
655 } else {
656 block_encoder::encode_compressed_block(
657 &input[offset..block_end],
658 sequences,
659 &mut rep_offsets,
660 is_last,
661 output,
662 workspace,
663 strategy::use_custom_sequence_tables(¶ms, input.len()),
664 );
665 }
666 }
667 offset = block_end;
668 }
669 }
670 }
671 Strategy::DFast => {
672 if hash_long.len() != long_size {
673 hash_long.resize(long_size, 0);
674 }
675 if has_prefix && input.len() <= MAX_BLOCK_SIZE {
676 dfast::compress_dfast_with_prefix_reuse(
677 input,
678 ¶ms,
679 &rep_offsets,
680 prefix,
681 hash_table,
682 hash_long,
683 sequences,
684 combined,
685 );
686 block_encoder::encode_compressed_block(
687 input,
688 sequences,
689 &mut rep_offsets,
690 true,
691 output,
692 workspace,
693 strategy::use_custom_sequence_tables(¶ms, input.len()),
694 );
695 } else if has_prefix {
696 combined.clear();
697 combined.reserve(prefix.len() + input.len());
698 combined.extend_from_slice(prefix);
699 combined.extend_from_slice(input);
700 let plen = prefix.len();
701 dfast::prefill_hash_tables(
702 combined,
703 plen,
704 params.hash_log,
705 params.chain_log,
706 params.min_match,
707 hash_table,
708 hash_long,
709 );
710
711 while offset < input.len() {
712 let chunk_size = (input.len() - offset).min(MAX_BLOCK_SIZE);
713 let is_last = offset + chunk_size >= input.len();
714 dfast::compress_dfast_block(
715 combined,
716 plen + offset,
717 plen + offset + chunk_size,
718 ¶ms,
719 &rep_offsets,
720 hash_table,
721 hash_long,
722 sequences,
723 );
724 block_encoder::encode_compressed_block(
725 &input[offset..offset + chunk_size],
726 sequences,
727 &mut rep_offsets,
728 is_last,
729 output,
730 workspace,
731 strategy::use_custom_sequence_tables(¶ms, input.len()),
732 );
733 offset += chunk_size;
734 }
735 } else {
736 hash_table.fill(0);
737 hash_long.fill(0);
738 while offset < input.len() {
739 let chunk_size = (input.len() - offset).min(MAX_BLOCK_SIZE);
740 let block_end = offset + chunk_size;
741 let is_last = block_end >= input.len();
742
743 if block_looks_incompressible(&input[offset..block_end]) {
744 block_encoder::encode_raw_block(
745 &input[offset..block_end],
746 is_last,
747 output,
748 );
749 } else {
750 dfast::compress_dfast_block(
751 input,
752 offset,
753 block_end,
754 ¶ms,
755 &rep_offsets,
756 hash_table,
757 hash_long,
758 sequences,
759 );
760 block_encoder::encode_compressed_block(
761 &input[offset..block_end],
762 sequences,
763 &mut rep_offsets,
764 is_last,
765 output,
766 workspace,
767 strategy::use_custom_sequence_tables(¶ms, input.len()),
768 );
769 }
770 offset = block_end;
771 }
772 }
773 }
774 }
775 }
776
777 let hash = xxh64(input, 0);
778 let checksum = (hash & 0xFFFF_FFFF) as u32;
779 output.extend_from_slice(&checksum.to_le_bytes());
780
781 Ok(())
782}