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