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 total_window = dict.content().len().saturating_add(input.len());
227 let params = strategy::level_params_for_size(self.level, total_window)
228 .expect("level validated at construction");
229 self.workspace.prev_ll = dict
230 .ll_table()
231 .map(|(dt, al)| block_encoder::FseEncodeTable::from_decode_table(dt, al, 35));
232 self.workspace.prev_of = dict
233 .of_table()
234 .map(|(dt, al)| block_encoder::FseEncodeTable::from_decode_table(dt, al, 31));
235 self.workspace.prev_ml = dict
236 .ml_table()
237 .map(|(dt, al)| block_encoder::FseEncodeTable::from_decode_table(dt, al, 52));
238 self.workspace.prev_huffman = dict
239 .huf_table()
240 .and_then(|(dt, tl)| HuffmanEncodeTable::from_decode_table(dt, tl));
241 compress_core(
242 input,
243 params,
244 Some(dict.id()),
245 dict.content(),
246 *dict.rep_offsets(),
247 &mut self.hash_table,
248 &mut self.hash_long,
249 &mut self.dict_hash,
250 &mut self.sequences,
251 &mut self.output,
252 &mut self.workspace,
253 &mut self.combined,
254 )?;
255 Ok(self.take_or_borrow_output())
256 }
257
258 fn compress_with_prepared(&mut self, input: &[u8]) -> Result<Cow<'_, [u8]>, CompressError> {
259 let prep = self.prepared.as_ref().unwrap();
260 let total_window = prep.prefix_len + input.len();
261 let mut params = strategy::level_params_for_size(self.level, total_window)
262 .expect("level validated at construction");
263 strategy::apply_raw_literals_size_override(&mut params, input.len());
264
265 let use_attached = !input.is_empty()
266 && input.len() <= ATTACH_THRESHOLD
267 && params.strategy == Strategy::Fast;
268
269 let dict_id = prep.dict_id;
270 let prefix_len = prep.prefix_len;
271 let dict_hash_log = prep.hash_log;
272
273 if !use_attached {
274 let snapshot_matches = match params.strategy {
275 Strategy::Fast => (1usize << params.hash_log) == prep.hash_snapshot.len(),
276 Strategy::DFast => {
277 (1usize << params.chain_log) == prep.hash_snapshot.len()
278 && (1usize << params.hash_log) == prep.hash_long_snapshot.len()
279 }
280 };
281 if !snapshot_matches {
282 return self.compress_with_dict_fallback(input, dict_id, prefix_len);
283 }
284 }
285
286 {
287 let prep = self.prepared.as_mut().unwrap();
288 if !use_attached {
289 self.hash_table.copy_from_slice(&prep.hash_snapshot);
290 if !prep.hash_long_snapshot.is_empty() {
291 self.hash_long.copy_from_slice(&prep.hash_long_snapshot);
292 }
293 }
294 prep.combined.truncate(prep.prefix_len);
295 prep.combined.extend_from_slice(input);
296 }
297
298 let prep = self.prepared.as_ref().unwrap();
299
300 if use_attached {
301 self.workspace.prev_huffman = if params.force_raw_literals {
302 None
303 } else {
304 prep.huf_table.clone()
305 };
306 } else if let Some(ref huf) = prep.huf_table {
307 self.workspace.prev_huffman = Some(huf.clone());
308 } else {
309 self.workspace.prev_huffman = None;
310 }
311 self.workspace.prev_ll = prep.ll_table.clone();
312 self.workspace.prev_of = prep.of_table.clone();
313 self.workspace.prev_ml = prep.ml_table.clone();
314
315 self.output.clear();
316 self.output.reserve(input.len() + 32);
317 write_frame_header(&mut self.output, input.len(), Some(dict_id));
318
319 if input.is_empty() {
320 block_encoder::encode_raw_block(&[], true, &mut self.output);
321 } else if use_attached {
322 let input_hash_log = if input.len() >= 2 {
323 let src_log = 32 - ((input.len() as u32) - 1).leading_zeros();
324 params.hash_log.min(src_log).max(strategy::HASH_LOG_MIN)
325 } else {
326 strategy::HASH_LOG_MIN
327 };
328 let input_hash_size = 1usize << input_hash_log;
329 self.small_hash.resize(input_hash_size, 0);
330 self.small_hash.fill(0);
331
332 let mut rep_offsets = prep.rep_offsets;
333
334 fast::compress_fast_attached(
335 &prep.combined,
336 prefix_len,
337 prefix_len + input.len(),
338 ¶ms,
339 &prep.rep_offsets,
340 &prep.hash_snapshot,
341 dict_hash_log,
342 &mut self.small_hash,
343 input_hash_log,
344 &mut self.sequences,
345 );
346
347 if params.force_raw_literals {
348 block_encoder::encode_compressed_block_raw(
349 input,
350 &self.sequences,
351 &mut rep_offsets,
352 true,
353 &mut self.output,
354 &mut self.workspace,
355 );
356 } else {
357 block_encoder::encode_compressed_block(
358 input,
359 &self.sequences,
360 &mut rep_offsets,
361 true,
362 &mut self.output,
363 &mut self.workspace,
364 strategy::use_custom_sequence_tables(¶ms, input.len()),
365 );
366 }
367 } else {
368 let combined = &prep.combined;
369 let mut rep_offsets = prep.rep_offsets;
370
371 if input.len() <= MAX_BLOCK_SIZE {
372 match params.strategy {
373 Strategy::Fast => {
374 fast::compress_fast_block(
375 combined,
376 prefix_len,
377 prefix_len + input.len(),
378 ¶ms,
379 &rep_offsets,
380 &mut self.hash_table,
381 &mut self.sequences,
382 );
383 }
384 Strategy::DFast => {
385 dfast::compress_dfast_block(
386 combined,
387 prefix_len,
388 prefix_len + input.len(),
389 ¶ms,
390 &rep_offsets,
391 &mut self.hash_table,
392 &mut self.hash_long,
393 &mut self.sequences,
394 );
395 }
396 }
397 if params.force_raw_literals {
398 block_encoder::encode_compressed_block_raw(
399 input,
400 &self.sequences,
401 &mut rep_offsets,
402 true,
403 &mut self.output,
404 &mut self.workspace,
405 );
406 } else {
407 block_encoder::encode_compressed_block(
408 input,
409 &self.sequences,
410 &mut rep_offsets,
411 true,
412 &mut self.output,
413 &mut self.workspace,
414 strategy::use_custom_sequence_tables(¶ms, input.len()),
415 );
416 }
417 } else {
418 let mut offset = 0;
419 while offset < input.len() {
420 let chunk_size = (input.len() - offset).min(MAX_BLOCK_SIZE);
421 let is_last = offset + chunk_size >= input.len();
422 match params.strategy {
423 Strategy::Fast => {
424 fast::compress_fast_block(
425 combined,
426 prefix_len + offset,
427 prefix_len + offset + chunk_size,
428 ¶ms,
429 &rep_offsets,
430 &mut self.hash_table,
431 &mut self.sequences,
432 );
433 }
434 Strategy::DFast => {
435 dfast::compress_dfast_block(
436 combined,
437 prefix_len + offset,
438 prefix_len + offset + chunk_size,
439 ¶ms,
440 &rep_offsets,
441 &mut self.hash_table,
442 &mut self.hash_long,
443 &mut self.sequences,
444 );
445 }
446 }
447 if params.force_raw_literals {
448 block_encoder::encode_compressed_block_raw(
449 &input[offset..offset + chunk_size],
450 &self.sequences,
451 &mut rep_offsets,
452 is_last,
453 &mut self.output,
454 &mut self.workspace,
455 );
456 } else {
457 block_encoder::encode_compressed_block(
458 &input[offset..offset + chunk_size],
459 &self.sequences,
460 &mut rep_offsets,
461 is_last,
462 &mut self.output,
463 &mut self.workspace,
464 strategy::use_custom_sequence_tables(¶ms, input.len()),
465 );
466 }
467 offset += chunk_size;
468 }
469 }
470 }
471
472 let hash = xxh64(input, 0);
473 let checksum = (hash & 0xFFFF_FFFF) as u32;
474 self.output.extend_from_slice(&checksum.to_le_bytes());
475
476 Ok(self.take_or_borrow_output())
477 }
478
479 fn compress_with_dict_fallback(
480 &mut self,
481 input: &[u8],
482 dict_id: u32,
483 prefix_len: usize,
484 ) -> Result<Cow<'_, [u8]>, CompressError> {
485 let prep = self.prepared.as_ref().unwrap();
486 let rep_offsets = prep.rep_offsets;
487 let prefix = &prep.combined[..prefix_len];
488
489 let total_window = prefix_len.saturating_add(input.len());
490 let params = strategy::level_params_for_size(self.level, total_window)
491 .expect("level validated at construction");
492 self.workspace.prev_huffman = prep.huf_table.clone();
493 self.workspace.prev_ll = prep.ll_table.clone();
494 self.workspace.prev_of = prep.of_table.clone();
495 self.workspace.prev_ml = prep.ml_table.clone();
496 compress_core(
497 input,
498 params,
499 Some(dict_id),
500 prefix,
501 rep_offsets,
502 &mut self.hash_table,
503 &mut self.hash_long,
504 &mut self.dict_hash,
505 &mut self.sequences,
506 &mut self.output,
507 &mut self.workspace,
508 &mut self.combined,
509 )?;
510 Ok(self.take_or_borrow_output())
511 }
512
513 fn take_or_borrow_output(&mut self) -> Cow<'_, [u8]> {
514 if self.output.len() >= zrip_core::LARGE_OUTPUT_THRESHOLD {
515 Cow::Owned(core::mem::take(&mut self.output))
516 } else {
517 Cow::Borrowed(&self.output)
518 }
519 }
520}
521
522#[allow(clippy::too_many_arguments, clippy::unnecessary_wraps)]
523fn compress_core(
524 input: &[u8],
525 params: LevelParams,
526 dict_id: Option<u32>,
527 prefix: &[u8],
528 init_rep_offsets: [u32; 3],
529 hash_table: &mut Vec<u32>,
530 hash_long: &mut Vec<u32>,
531 dict_hash: &mut Vec<u32>,
532 sequences: &mut Vec<Sequence>,
533 output: &mut Vec<u8>,
534 workspace: &mut BlockEncodeWorkspace,
535 combined: &mut Vec<u8>,
536) -> Result<(), CompressError> {
537 let mut params = params;
538 strategy::apply_raw_literals_size_override(&mut params, input.len());
539
540 let hash_size = match params.strategy {
541 Strategy::Fast => 1usize << params.hash_log,
542 Strategy::DFast => 1usize << params.chain_log,
543 };
544 let long_size = 1usize << params.hash_log;
545
546 if prefix.is_empty() {
547 workspace.prev_huffman = None;
548 workspace.prev_ll = None;
549 workspace.prev_of = None;
550 workspace.prev_ml = None;
551 }
552
553 output.clear();
554 output.reserve(input.len() + 32);
555 write_frame_header(output, input.len(), dict_id);
556
557 if input.is_empty() {
558 block_encoder::encode_raw_block(&[], true, output);
559 } else {
560 let has_prefix = !prefix.is_empty();
561 let mut rep_offsets = init_rep_offsets;
562 let mut offset = 0;
563
564 if hash_table.len() != hash_size {
565 hash_table.resize(hash_size, 0);
566 }
567
568 match params.strategy {
569 Strategy::Fast => {
570 if has_prefix && input.len() <= MAX_BLOCK_SIZE {
571 if dict_hash.len() != hash_size {
572 dict_hash.resize(hash_size, 0);
573 }
574 fast::compress_fast_with_prefix_reuse(
575 input,
576 ¶ms,
577 &rep_offsets,
578 prefix,
579 dict_hash,
580 hash_table,
581 sequences,
582 combined,
583 );
584 if params.force_raw_literals {
585 block_encoder::encode_compressed_block_raw(
586 input,
587 sequences,
588 &mut rep_offsets,
589 true,
590 output,
591 workspace,
592 );
593 } else {
594 block_encoder::encode_compressed_block(
595 input,
596 sequences,
597 &mut rep_offsets,
598 true,
599 output,
600 workspace,
601 strategy::use_custom_sequence_tables(¶ms, input.len()),
602 );
603 }
604 } else if has_prefix {
605 combined.clear();
606 combined.reserve(prefix.len() + input.len());
607 combined.extend_from_slice(prefix);
608 combined.extend_from_slice(input);
609 let plen = prefix.len();
610 fast::prefill_hash_table(combined, plen, params.hash_log, hash_table);
611
612 while offset < input.len() {
613 let chunk_size = (input.len() - offset).min(MAX_BLOCK_SIZE);
614 let is_last = offset + chunk_size >= input.len();
615 fast::compress_fast_block(
616 combined,
617 plen + offset,
618 plen + offset + chunk_size,
619 ¶ms,
620 &rep_offsets,
621 hash_table,
622 sequences,
623 );
624 if params.force_raw_literals {
625 block_encoder::encode_compressed_block_raw(
626 &input[offset..offset + chunk_size],
627 sequences,
628 &mut rep_offsets,
629 is_last,
630 output,
631 workspace,
632 );
633 } else {
634 block_encoder::encode_compressed_block(
635 &input[offset..offset + chunk_size],
636 sequences,
637 &mut rep_offsets,
638 is_last,
639 output,
640 workspace,
641 strategy::use_custom_sequence_tables(¶ms, input.len()),
642 );
643 }
644 offset += chunk_size;
645 }
646 } else {
647 hash_table.fill(0);
648 while offset < input.len() {
649 let chunk_size = (input.len() - offset).min(MAX_BLOCK_SIZE);
650 let block_end = offset + chunk_size;
651 let is_last = block_end >= input.len();
652
653 if block_looks_incompressible(&input[offset..block_end]) {
654 block_encoder::encode_raw_block(
655 &input[offset..block_end],
656 is_last,
657 output,
658 );
659 } else {
660 fast::compress_fast_block(
661 input,
662 offset,
663 block_end,
664 ¶ms,
665 &rep_offsets,
666 hash_table,
667 sequences,
668 );
669 if params.force_raw_literals {
670 block_encoder::encode_compressed_block_raw(
671 &input[offset..block_end],
672 sequences,
673 &mut rep_offsets,
674 is_last,
675 output,
676 workspace,
677 );
678 } else {
679 block_encoder::encode_compressed_block(
680 &input[offset..block_end],
681 sequences,
682 &mut rep_offsets,
683 is_last,
684 output,
685 workspace,
686 strategy::use_custom_sequence_tables(¶ms, input.len()),
687 );
688 }
689 }
690 offset = block_end;
691 }
692 }
693 }
694 Strategy::DFast => {
695 if hash_long.len() != long_size {
696 hash_long.resize(long_size, 0);
697 }
698 if has_prefix && input.len() <= MAX_BLOCK_SIZE {
699 dfast::compress_dfast_with_prefix_reuse(
700 input,
701 ¶ms,
702 &rep_offsets,
703 prefix,
704 hash_table,
705 hash_long,
706 sequences,
707 combined,
708 );
709 block_encoder::encode_compressed_block(
710 input,
711 sequences,
712 &mut rep_offsets,
713 true,
714 output,
715 workspace,
716 strategy::use_custom_sequence_tables(¶ms, input.len()),
717 );
718 } else if has_prefix {
719 combined.clear();
720 combined.reserve(prefix.len() + input.len());
721 combined.extend_from_slice(prefix);
722 combined.extend_from_slice(input);
723 let plen = prefix.len();
724 dfast::prefill_hash_tables(
725 combined,
726 plen,
727 params.hash_log,
728 params.chain_log,
729 params.min_match,
730 hash_table,
731 hash_long,
732 );
733
734 while offset < input.len() {
735 let chunk_size = (input.len() - offset).min(MAX_BLOCK_SIZE);
736 let is_last = offset + chunk_size >= input.len();
737 dfast::compress_dfast_block(
738 combined,
739 plen + offset,
740 plen + offset + chunk_size,
741 ¶ms,
742 &rep_offsets,
743 hash_table,
744 hash_long,
745 sequences,
746 );
747 block_encoder::encode_compressed_block(
748 &input[offset..offset + chunk_size],
749 sequences,
750 &mut rep_offsets,
751 is_last,
752 output,
753 workspace,
754 strategy::use_custom_sequence_tables(¶ms, input.len()),
755 );
756 offset += chunk_size;
757 }
758 } else {
759 hash_table.fill(0);
760 hash_long.fill(0);
761 while offset < input.len() {
762 let chunk_size = (input.len() - offset).min(MAX_BLOCK_SIZE);
763 let block_end = offset + chunk_size;
764 let is_last = block_end >= input.len();
765
766 if block_looks_incompressible(&input[offset..block_end]) {
767 block_encoder::encode_raw_block(
768 &input[offset..block_end],
769 is_last,
770 output,
771 );
772 } else {
773 dfast::compress_dfast_block(
774 input,
775 offset,
776 block_end,
777 ¶ms,
778 &rep_offsets,
779 hash_table,
780 hash_long,
781 sequences,
782 );
783 block_encoder::encode_compressed_block(
784 &input[offset..block_end],
785 sequences,
786 &mut rep_offsets,
787 is_last,
788 output,
789 workspace,
790 strategy::use_custom_sequence_tables(¶ms, input.len()),
791 );
792 }
793 offset = block_end;
794 }
795 }
796 }
797 }
798 }
799
800 let hash = xxh64(input, 0);
801 let checksum = (hash & 0xFFFF_FFFF) as u32;
802 output.extend_from_slice(&checksum.to_le_bytes());
803
804 Ok(())
805}