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