1use std::cmp::Reverse;
21use std::collections::BinaryHeap;
22use std::fmt;
23
24pub const DEFAULT_MAX_OUTPUT: usize = 256 * 1024 * 1024;
28
29#[derive(Debug, Clone, Copy, PartialEq, Eq)]
32pub enum InflateError {
33 Truncated,
35 InvalidBlockType,
37 InvalidCode,
40 InvalidCodeLengths,
43 DistanceTooFar,
45 LengthMismatch,
48 OutputTooLarge,
50 InvalidHeader,
53 ChecksumMismatch,
55 TrailingData,
57}
58
59impl fmt::Display for InflateError {
60 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
61 f.write_str(match self {
62 Self::Truncated => "compressed stream ended before its final block",
63 Self::InvalidBlockType => "compressed block uses the reserved block type",
64 Self::InvalidCode => "compressed block contains an invalid Huffman code",
65 Self::InvalidCodeLengths => "compressed block header has malformed code lengths",
66 Self::DistanceTooFar => "compressed block refers to data before the start of the output",
67 Self::LengthMismatch => "compressed stream's length fields disagree with its contents",
68 Self::OutputTooLarge => "decompressed output exceeds the permitted size",
69 Self::InvalidHeader => "compressed stream has an invalid or unsupported header",
70 Self::ChecksumMismatch => "compressed stream's checksum does not match its contents",
71 Self::TrailingData => "unexpected data after the end of the compressed stream",
72 })
73 }
74}
75
76impl std::error::Error for InflateError {}
77
78type Result<T> = std::result::Result<T, InflateError>;
79
80const MAX_BITS: usize = 15;
85const MAX_CODE_LENGTH_BITS: usize = 7;
87
88const LITLEN_SYMBOLS: usize = 286;
91const END_OF_BLOCK: u16 = 256;
92const DIST_SYMBOLS: usize = 30;
93const CODE_LENGTH_SYMBOLS: usize = 19;
94
95const MIN_MATCH: usize = 3;
96const MAX_MATCH: usize = 258;
97const WINDOW_SIZE: usize = 32 * 1024;
98
99const LENGTH_BASE: [u16; 29] = [
102 3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 15, 17, 19, 23, 27, 31, 35, 43, 51, 59, 67, 83, 99, 115, 131, 163, 195,
103 227, 258,
104];
105const LENGTH_EXTRA: [u8; 29] =
106 [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 4, 5, 5, 5, 5, 0];
107
108const DIST_BASE: [u16; 30] = [
110 1, 2, 3, 4, 5, 7, 9, 13, 17, 25, 33, 49, 65, 97, 129, 193, 257, 385, 513, 769, 1025, 1537, 2049, 3073,
111 4097, 6145, 8193, 12289, 16385, 24577,
112];
113const DIST_EXTRA: [u8; 30] =
114 [0, 0, 0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7, 8, 8, 9, 9, 10, 10, 11, 11, 12, 12, 13, 13];
115
116const CODE_LENGTH_ORDER: [usize; 19] = [16, 17, 18, 0, 8, 7, 9, 6, 10, 5, 11, 4, 12, 3, 13, 2, 14, 1, 15];
120
121const LENGTH_CODE: [u8; MAX_MATCH - MIN_MATCH + 1] = build_length_codes();
126
127const fn build_length_codes() -> [u8; MAX_MATCH - MIN_MATCH + 1] {
128 let mut table = [0u8; MAX_MATCH - MIN_MATCH + 1];
129 let mut code = 0;
130 while code < LENGTH_BASE.len() {
131 let base = LENGTH_BASE[code] as usize;
132 let span = 1usize << LENGTH_EXTRA[code];
133 let mut length = base;
134 while length < base + span && length <= MAX_MATCH {
135 table[length - MIN_MATCH] = code as u8;
136 length += 1;
137 }
138 code += 1;
139 }
140 table
141}
142
143fn length_code(length: usize) -> usize {
144 usize::from(LENGTH_CODE[length - MIN_MATCH])
145}
146
147fn dist_code(dist: usize) -> usize {
148 DIST_BASE.partition_point(|&base| usize::from(base) <= dist) - 1
150}
151
152fn fixed_litlen_lengths() -> [u8; 288] {
155 let mut lengths = [8u8; 288];
156 lengths[144..256].fill(9);
157 lengths[256..280].fill(7);
158 lengths
159}
160
161struct BitReader<'a> {
172 data: &'a [u8],
173 pos: usize,
174 bits: u64,
175 count: u32,
176}
177
178impl<'a> BitReader<'a> {
179 fn new(data: &'a [u8]) -> Self {
180 Self { data, pos: 0, bits: 0, count: 0 }
181 }
182
183 fn read(&mut self, n: u32) -> Result<u32> {
184 while self.count < n {
185 let byte = *self.data.get(self.pos).ok_or(InflateError::Truncated)?;
186 self.bits |= u64::from(byte) << self.count;
187 self.count += 8;
188 self.pos += 1;
189 }
190 let value = (self.bits & ((1u64 << n) - 1)) as u32;
191 self.bits >>= n;
192 self.count -= n;
193 Ok(value)
194 }
195
196 fn read_bit(&mut self) -> Result<u32> {
197 self.read(1)
198 }
199
200 fn align_to_byte(&mut self) {
204 self.bits = 0;
205 self.count = 0;
206 }
207
208 fn take_bytes(&mut self, n: usize) -> Result<&'a [u8]> {
211 debug_assert_eq!(self.count, 0, "take_bytes needs a byte-aligned reader");
212 let end = self.pos.checked_add(n).ok_or(InflateError::Truncated)?;
213 let bytes = self.data.get(self.pos..end).ok_or(InflateError::Truncated)?;
214 self.pos = end;
215 Ok(bytes)
216 }
217
218 fn position(&self) -> usize {
221 self.pos
222 }
223}
224
225struct Decoder {
235 count: [u16; MAX_BITS + 1],
236 symbol: Vec<u16>,
237}
238
239impl Decoder {
240 fn new(lengths: &[u8]) -> Result<Self> {
248 let mut count = [0u16; MAX_BITS + 1];
249 for &len in lengths {
250 count[usize::from(len)] += 1;
251 }
252 let used = lengths.len() - usize::from(count[0]);
253 let mut left: i32 = 1;
255 for &c in &count[1..] {
256 left = (left << 1) - i32::from(c);
257 if left < 0 {
258 return Err(InflateError::InvalidCode);
259 }
260 }
261 if left > 0 && used > 1 {
262 return Err(InflateError::InvalidCode);
263 }
264 let mut offsets = [0u16; MAX_BITS + 2];
266 for len in 1..=MAX_BITS {
267 offsets[len + 1] = offsets[len] + count[len];
268 }
269 let mut symbol = vec![0u16; used];
270 for (sym, &len) in lengths.iter().enumerate() {
271 if len != 0 {
272 let slot = &mut offsets[usize::from(len)];
273 symbol[usize::from(*slot)] = sym as u16;
274 *slot += 1;
275 }
276 }
277 Ok(Self { count, symbol })
278 }
279
280 fn decode(&self, reader: &mut BitReader<'_>) -> Result<u16> {
281 let mut code: i32 = 0;
285 let mut first: i32 = 0;
286 let mut index: i32 = 0;
287 for len in 1..=MAX_BITS {
288 code |= reader.read_bit()? as i32;
289 let count = i32::from(self.count[len]);
290 if code - first < count {
291 return Ok(self.symbol[(index + code - first) as usize]);
292 }
293 index += count;
294 first = (first + count) << 1;
295 code <<= 1;
296 }
297 Err(InflateError::InvalidCode)
298 }
299}
300
301pub fn decompress(input: &[u8]) -> Result<Vec<u8>> {
308 decompress_with_limit(input, DEFAULT_MAX_OUTPUT)
309}
310
311pub fn decompress_with_limit(input: &[u8], max_out: usize) -> Result<Vec<u8>> {
316 let (output, consumed) = inflate(input, max_out)?;
317 if consumed != input.len() {
318 return Err(InflateError::TrailingData);
319 }
320 Ok(output)
321}
322
323pub(super) fn inflate(input: &[u8], max_out: usize) -> Result<(Vec<u8>, usize)> {
327 let mut reader = BitReader::new(input);
328 let mut output = Vec::new();
329 loop {
330 let is_final = reader.read_bit()? == 1;
331 match reader.read(2)? {
332 0b00 => inflate_stored(&mut reader, &mut output, max_out)?,
333 0b01 => {
334 let litlen = Decoder::new(&fixed_litlen_lengths())?;
337 let dist = Decoder::new(&[5u8; 32])?;
338 inflate_codes(&mut reader, &mut output, max_out, &litlen, &dist)?;
339 }
340 0b10 => {
341 let (litlen, dist) = read_dynamic_codes(&mut reader)?;
342 inflate_codes(&mut reader, &mut output, max_out, &litlen, &dist)?;
343 }
344 _ => return Err(InflateError::InvalidBlockType),
345 }
346 if is_final {
347 break;
348 }
349 }
350 reader.align_to_byte();
351 Ok((output, reader.position()))
352}
353
354fn inflate_stored(reader: &mut BitReader<'_>, output: &mut Vec<u8>, max_out: usize) -> Result<()> {
357 reader.align_to_byte();
358 let len = reader.read(16)? as usize;
359 let nlen = reader.read(16)? as usize;
360 if len != !nlen & 0xFFFF {
361 return Err(InflateError::LengthMismatch);
362 }
363 let bytes = reader.take_bytes(len)?;
364 if output.len() + len > max_out {
365 return Err(InflateError::OutputTooLarge);
366 }
367 output.extend_from_slice(bytes);
368 Ok(())
369}
370
371fn read_dynamic_codes(reader: &mut BitReader<'_>) -> Result<(Decoder, Decoder)> {
374 let hlit = reader.read(5)? as usize + 257;
375 let hdist = reader.read(5)? as usize + 1;
376 let hclen = reader.read(4)? as usize + 4;
377 if hlit > LITLEN_SYMBOLS || hdist > DIST_SYMBOLS {
378 return Err(InflateError::InvalidCode);
379 }
380
381 let mut code_lengths = [0u8; CODE_LENGTH_SYMBOLS];
382 for &symbol in &CODE_LENGTH_ORDER[..hclen] {
383 code_lengths[symbol] = reader.read(3)? as u8;
384 }
385 let code_length_decoder = Decoder::new(&code_lengths)?;
386
387 let mut lengths = vec![0u8; hlit + hdist];
390 let mut index = 0;
391 while index < lengths.len() {
392 let symbol = code_length_decoder.decode(reader)?;
393 let (value, repeat) = match symbol {
394 0..=15 => (symbol as u8, 1),
395 16 => {
396 if index == 0 {
397 return Err(InflateError::InvalidCodeLengths);
398 }
399 (lengths[index - 1], 3 + reader.read(2)? as usize)
400 }
401 17 => (0, 3 + reader.read(3)? as usize),
402 _ => (0, 11 + reader.read(7)? as usize),
403 };
404 if index + repeat > lengths.len() {
405 return Err(InflateError::InvalidCodeLengths);
406 }
407 lengths[index..index + repeat].fill(value);
408 index += repeat;
409 }
410
411 if lengths[usize::from(END_OF_BLOCK)] == 0 {
413 return Err(InflateError::InvalidCode);
414 }
415 let litlen = Decoder::new(&lengths[..hlit])?;
416 let dist = Decoder::new(&lengths[hlit..])?;
417 Ok((litlen, dist))
418}
419
420fn inflate_codes(
423 reader: &mut BitReader<'_>,
424 output: &mut Vec<u8>,
425 max_out: usize,
426 litlen: &Decoder,
427 dist: &Decoder,
428) -> Result<()> {
429 loop {
430 let symbol = litlen.decode(reader)?;
431 if symbol < END_OF_BLOCK {
432 if output.len() >= max_out {
433 return Err(InflateError::OutputTooLarge);
434 }
435 output.push(symbol as u8);
436 continue;
437 }
438 if symbol == END_OF_BLOCK {
439 return Ok(());
440 }
441 let code = usize::from(symbol - 257);
442 if code >= LENGTH_BASE.len() {
443 return Err(InflateError::InvalidCode);
444 }
445 let length = usize::from(LENGTH_BASE[code]) + reader.read(u32::from(LENGTH_EXTRA[code]))? as usize;
446
447 let code = usize::from(dist.decode(reader)?);
448 if code >= DIST_BASE.len() {
449 return Err(InflateError::InvalidCode);
450 }
451 let distance = usize::from(DIST_BASE[code]) + reader.read(u32::from(DIST_EXTRA[code]))? as usize;
452 if distance > output.len() {
453 return Err(InflateError::DistanceTooFar);
454 }
455 if output.len() + length > max_out {
456 return Err(InflateError::OutputTooLarge);
457 }
458 let start = output.len() - distance;
461 for i in 0..length {
462 output.push(output[start + i]);
463 }
464 }
465}
466
467struct BitWriter {
471 out: Vec<u8>,
472 bits: u64,
473 count: u32,
474}
475
476impl BitWriter {
477 fn new() -> Self {
478 Self { out: Vec::new(), bits: 0, count: 0 }
479 }
480
481 fn write(&mut self, value: u32, n: u32) {
485 self.bits |= u64::from(value) << self.count;
486 self.count += n;
487 while self.count >= 8 {
488 self.out.push(self.bits as u8);
489 self.bits >>= 8;
490 self.count -= 8;
491 }
492 }
493
494 fn write_code(&mut self, code: &Code) {
498 self.write(u32::from(code.bits), u32::from(code.len));
499 }
500
501 fn align_to_byte(&mut self) {
502 if self.count > 0 {
503 self.out.push(self.bits as u8);
504 self.bits = 0;
505 self.count = 0;
506 }
507 }
508
509 fn write_bytes(&mut self, bytes: &[u8]) {
510 debug_assert_eq!(self.count, 0, "write_bytes needs a byte-aligned writer");
511 self.out.extend_from_slice(bytes);
512 }
513
514 fn finish(mut self) -> Vec<u8> {
515 self.align_to_byte();
516 self.out
517 }
518}
519
520#[derive(Clone, Copy, Default)]
523struct Code {
524 bits: u16,
525 len: u8,
526}
527
528fn assign_codes(lengths: &[u8]) -> Vec<Code> {
533 let mut count = [0u16; MAX_BITS + 1];
534 for &len in lengths {
535 count[usize::from(len)] += 1;
536 }
537 count[0] = 0;
538 let mut next_code = [0u16; MAX_BITS + 1];
539 let mut code = 0u16;
540 for len in 1..=MAX_BITS {
541 code = (code + count[len - 1]) << 1;
542 next_code[len] = code;
543 }
544 lengths
545 .iter()
546 .map(|&len| {
547 if len == 0 {
548 return Code::default();
549 }
550 let code = next_code[usize::from(len)];
551 next_code[usize::from(len)] += 1;
552 Code { bits: code.reverse_bits() >> (16 - len), len }
553 })
554 .collect()
555}
556
557fn build_lengths(freqs: &[u32], max_bits: u8) -> Vec<u8> {
568 let mut lengths = vec![0u8; freqs.len()];
569 let used: Vec<usize> = (0..freqs.len()).filter(|&i| freqs[i] > 0).collect();
570 match used.len() {
571 0 => return lengths,
572 1 => {
575 lengths[used[0]] = 1;
576 return lengths;
577 }
578 _ => {}
579 }
580
581 let mut weights: Vec<u64> = freqs.iter().map(|&f| u64::from(f)).collect();
582 loop {
583 let mut parent = vec![usize::MAX; freqs.len() * 2];
586 let mut heap: BinaryHeap<Reverse<(u64, usize)>> =
587 used.iter().map(|&i| Reverse((weights[i], i))).collect();
588 let mut next = freqs.len();
589 while heap.len() > 1 {
590 let Reverse((w1, a)) = heap.pop().expect("heap has at least two entries");
591 let Reverse((w2, b)) = heap.pop().expect("heap has at least two entries");
592 parent[a] = next;
593 parent[b] = next;
594 heap.push(Reverse((w1 + w2, next)));
595 next += 1;
596 }
597 let mut too_deep = false;
598 for &sym in &used {
599 let mut depth = 0u8;
600 let mut node = sym;
601 while parent[node] != usize::MAX {
602 node = parent[node];
603 depth += 1;
604 }
605 lengths[sym] = depth;
606 too_deep |= depth > max_bits;
607 }
608 if !too_deep {
609 return lengths;
610 }
611 for w in &mut weights {
612 if *w > 0 {
613 *w = w.div_ceil(2);
614 }
615 }
616 }
617}
618
619struct MatchFinder {
631 head: Vec<u32>,
632 prev: Vec<u32>,
633 next_insert: usize,
636}
637
638const HASH_BITS: u32 = 15;
639const HASH_SIZE: usize = 1 << HASH_BITS;
640const MAX_CHAIN: usize = 128;
641const LAZY_MATCH_LIMIT: usize = 32;
645const NO_POSITION: u32 = u32::MAX;
646
647impl MatchFinder {
648 fn new() -> Self {
649 Self { head: vec![NO_POSITION; HASH_SIZE], prev: vec![NO_POSITION; WINDOW_SIZE], next_insert: 0 }
650 }
651
652 fn hash(input: &[u8], pos: usize) -> usize {
654 let key =
655 (u32::from(input[pos]) << 16) | (u32::from(input[pos + 1]) << 8) | u32::from(input[pos + 2]);
656 (key.wrapping_mul(0x9E37_79B1) >> (32 - HASH_BITS)) as usize
657 }
658
659 fn insert_through(&mut self, input: &[u8], end: usize) {
661 while self.next_insert < end {
662 let pos = self.next_insert;
663 if pos + MIN_MATCH <= input.len() {
664 let hash = Self::hash(input, pos);
665 self.prev[pos & (WINDOW_SIZE - 1)] = self.head[hash];
666 self.head[hash] = pos as u32;
667 }
668 self.next_insert += 1;
669 }
670 }
671
672 fn longest_match(&self, input: &[u8], pos: usize) -> (usize, usize) {
676 if pos + MIN_MATCH > input.len() {
677 return (0, 0);
678 }
679 let max_len = MAX_MATCH.min(input.len() - pos);
680 let mut best_len = MIN_MATCH - 1;
681 let mut best_dist = 0;
682 let mut candidate = self.head[Self::hash(input, pos)];
683 let mut remaining = MAX_CHAIN;
684 while candidate != NO_POSITION && remaining > 0 {
685 let start = candidate as usize;
686 let dist = pos - start;
687 if dist > WINDOW_SIZE {
688 break;
689 }
690 if input[start + best_len] == input[pos + best_len] {
693 let len = (0..max_len).take_while(|&i| input[start + i] == input[pos + i]).count();
694 if len > best_len {
695 best_len = len;
696 best_dist = dist;
697 if len == max_len {
698 break;
699 }
700 }
701 }
702 candidate = self.prev[start & (WINDOW_SIZE - 1)];
703 remaining -= 1;
704 }
705 if best_len >= MIN_MATCH { (best_len, best_dist) } else { (0, 0) }
706 }
707}
708
709#[derive(Debug, Clone, Copy, PartialEq, Eq)]
711enum Token {
712 Literal(u8),
713 Match { len: u16, dist: u16 },
714}
715
716const BLOCK_TOKENS: usize = 16 * 1024;
721
722pub fn compress(input: &[u8]) -> Vec<u8> {
726 let mut writer = BitWriter::new();
727 let mut finder = MatchFinder::new();
728 let mut tokens = Vec::with_capacity(BLOCK_TOKENS.min(input.len() + 1));
729 let mut block_start = 0;
730 let mut pos = 0;
731 let mut current = finder.longest_match(input, 0);
732
733 while pos < input.len() {
734 let (len, dist) = current;
735 if (MIN_MATCH..LAZY_MATCH_LIMIT).contains(&len) && pos + 1 < input.len() {
739 finder.insert_through(input, pos + 1);
740 let next = finder.longest_match(input, pos + 1);
741 if next.0 > len {
742 tokens.push(Token::Literal(input[pos]));
743 pos += 1;
744 current = next;
745 continue;
746 }
747 }
748 if len >= MIN_MATCH {
749 tokens.push(Token::Match { len: len as u16, dist: dist as u16 });
750 pos += len;
751 } else {
752 tokens.push(Token::Literal(input[pos]));
753 pos += 1;
754 }
755 finder.insert_through(input, pos);
756 if pos < input.len() {
757 current = finder.longest_match(input, pos);
758 }
759 if tokens.len() >= BLOCK_TOKENS && pos < input.len() {
760 write_block(&mut writer, &input[block_start..pos], &tokens, false);
761 tokens.clear();
762 block_start = pos;
763 }
764 }
765 write_block(&mut writer, &input[block_start..], &tokens, true);
766 writer.finish()
767}
768
769#[cfg(test)]
772fn tokenize(input: &[u8]) -> Vec<Token> {
773 let mut finder = MatchFinder::new();
774 let mut tokens = Vec::new();
775 let mut pos = 0;
776 while pos < input.len() {
777 let (len, dist) = finder.longest_match(input, pos);
778 if len >= MIN_MATCH {
779 tokens.push(Token::Match { len: len as u16, dist: dist as u16 });
780 pos += len;
781 } else {
782 tokens.push(Token::Literal(input[pos]));
783 pos += 1;
784 }
785 finder.insert_through(input, pos);
786 }
787 tokens
788}
789
790struct DynamicHeader {
792 litlen_lengths: Vec<u8>,
793 dist_lengths: Vec<u8>,
794 hlit: usize,
795 hdist: usize,
796 hclen: usize,
797 code_length_lengths: Vec<u8>,
798 sequence: Vec<(u8, u8, u8)>,
800}
801
802fn write_block(writer: &mut BitWriter, raw: &[u8], tokens: &[Token], is_final: bool) {
804 let mut litlen_freq = [0u32; LITLEN_SYMBOLS];
805 let mut dist_freq = [0u32; DIST_SYMBOLS];
806 let mut extra_bits = 0usize;
807 for token in tokens {
808 match *token {
809 Token::Literal(byte) => litlen_freq[usize::from(byte)] += 1,
810 Token::Match { len, dist } => {
811 let lc = length_code(usize::from(len));
812 let dc = dist_code(usize::from(dist));
813 litlen_freq[257 + lc] += 1;
814 dist_freq[dc] += 1;
815 extra_bits += usize::from(LENGTH_EXTRA[lc]) + usize::from(DIST_EXTRA[dc]);
816 }
817 }
818 }
819 litlen_freq[usize::from(END_OF_BLOCK)] += 1;
820
821 let stored_bits = 3 + 7 + raw.len().div_ceil(u16::MAX as usize).max(1) * 32 + raw.len() * 8;
825
826 let fixed_lengths = fixed_litlen_lengths();
827 let fixed_bits = 3
828 + extra_bits
829 + litlen_freq
830 .iter()
831 .zip(fixed_lengths.iter())
832 .map(|(&f, &l)| f as usize * usize::from(l))
833 .sum::<usize>()
834 + dist_freq.iter().map(|&f| f as usize * 5).sum::<usize>();
835
836 let dynamic = build_dynamic_header(&litlen_freq, &dist_freq);
837 let dynamic_bits = 3
838 + 14
839 + dynamic.hclen * 3
840 + dynamic
841 .sequence
842 .iter()
843 .map(|&(sym, _, extra)| {
844 usize::from(dynamic.code_length_lengths[usize::from(sym)]) + usize::from(extra)
845 })
846 .sum::<usize>()
847 + extra_bits
848 + litlen_freq
849 .iter()
850 .zip(dynamic.litlen_lengths.iter())
851 .map(|(&f, &l)| f as usize * usize::from(l))
852 .sum::<usize>()
853 + dist_freq
854 .iter()
855 .zip(dynamic.dist_lengths.iter())
856 .map(|(&f, &l)| f as usize * usize::from(l))
857 .sum::<usize>();
858
859 if stored_bits <= fixed_bits && stored_bits <= dynamic_bits {
860 write_stored(writer, raw, is_final);
861 } else if fixed_bits <= dynamic_bits {
862 writer.write(u32::from(is_final), 1);
863 writer.write(0b01, 2);
864 let litlen = assign_codes(&fixed_lengths);
865 let dist = assign_codes(&[5u8; 32]);
866 write_tokens(writer, tokens, &litlen, &dist);
867 } else {
868 writer.write(u32::from(is_final), 1);
869 writer.write(0b10, 2);
870 write_dynamic_header(writer, &dynamic);
871 let litlen = assign_codes(&dynamic.litlen_lengths);
872 let dist = assign_codes(&dynamic.dist_lengths);
873 write_tokens(writer, tokens, &litlen, &dist);
874 }
875}
876
877fn write_stored(writer: &mut BitWriter, raw: &[u8], is_final: bool) {
880 let mut chunks = raw.chunks(u16::MAX as usize).peekable();
881 if chunks.peek().is_none() {
882 write_stored_chunk(writer, &[], is_final);
883 return;
884 }
885 while let Some(chunk) = chunks.next() {
886 write_stored_chunk(writer, chunk, is_final && chunks.peek().is_none());
887 }
888}
889
890fn write_stored_chunk(writer: &mut BitWriter, chunk: &[u8], is_final: bool) {
891 writer.write(u32::from(is_final), 1);
892 writer.write(0b00, 2);
893 writer.align_to_byte();
894 writer.write(chunk.len() as u32, 16);
895 writer.write(!(chunk.len() as u32) & 0xFFFF, 16);
896 writer.write_bytes(chunk);
897}
898
899fn write_tokens(writer: &mut BitWriter, tokens: &[Token], litlen: &[Code], dist: &[Code]) {
900 for token in tokens {
901 match *token {
902 Token::Literal(byte) => writer.write_code(&litlen[usize::from(byte)]),
903 Token::Match { len, dist: distance } => {
904 let len = usize::from(len);
905 let distance = usize::from(distance);
906 let lc = length_code(len);
907 writer.write_code(&litlen[257 + lc]);
908 writer.write((len - usize::from(LENGTH_BASE[lc])) as u32, u32::from(LENGTH_EXTRA[lc]));
909 let dc = dist_code(distance);
910 writer.write_code(&dist[dc]);
911 writer.write((distance - usize::from(DIST_BASE[dc])) as u32, u32::from(DIST_EXTRA[dc]));
912 }
913 }
914 }
915 writer.write_code(&litlen[usize::from(END_OF_BLOCK)]);
916}
917
918fn build_dynamic_header(litlen_freq: &[u32], dist_freq: &[u32]) -> DynamicHeader {
921 let litlen_lengths = build_lengths(litlen_freq, MAX_BITS as u8);
922 let dist_lengths = build_lengths(dist_freq, MAX_BITS as u8);
923
924 let hlit = litlen_lengths.iter().rposition(|&l| l != 0).map_or(257, |i| i + 1).max(257);
929 let hdist = dist_lengths.iter().rposition(|&l| l != 0).map_or(1, |i| i + 1);
930
931 let all: Vec<u8> = litlen_lengths[..hlit].iter().chain(&dist_lengths[..hdist]).copied().collect();
935 let mut sequence = Vec::new();
936 let mut i = 0;
937 while i < all.len() {
938 let value = all[i];
939 let mut run = all[i..].iter().take_while(|&&l| l == value).count();
940 i += run;
941 if value == 0 {
942 while run >= 11 {
943 let n = run.min(138);
944 sequence.push((18, (n - 11) as u8, 7));
945 run -= n;
946 }
947 if run >= 3 {
948 sequence.push((17, (run - 3) as u8, 3));
949 run = 0;
950 }
951 sequence.extend(std::iter::repeat_n((0, 0, 0), run));
952 } else {
953 sequence.push((value, 0, 0));
954 run -= 1;
955 while run >= 3 {
956 let n = run.min(6);
957 sequence.push((16, (n - 3) as u8, 2));
958 run -= n;
959 }
960 sequence.extend(std::iter::repeat_n((value, 0, 0), run));
961 }
962 }
963
964 let mut code_length_freq = [0u32; CODE_LENGTH_SYMBOLS];
965 for &(sym, _, _) in &sequence {
966 code_length_freq[usize::from(sym)] += 1;
967 }
968 let code_length_lengths = build_lengths(&code_length_freq, MAX_CODE_LENGTH_BITS as u8);
969 let hclen =
970 CODE_LENGTH_ORDER.iter().rposition(|&sym| code_length_lengths[sym] != 0).map_or(4, |i| i + 1).max(4);
971
972 DynamicHeader { litlen_lengths, dist_lengths, hlit, hdist, hclen, code_length_lengths, sequence }
973}
974
975fn write_dynamic_header(writer: &mut BitWriter, header: &DynamicHeader) {
976 writer.write((header.hlit - 257) as u32, 5);
977 writer.write((header.hdist - 1) as u32, 5);
978 writer.write((header.hclen - 4) as u32, 4);
979 for &sym in &CODE_LENGTH_ORDER[..header.hclen] {
980 writer.write(u32::from(header.code_length_lengths[sym]), 3);
981 }
982 let codes = assign_codes(&header.code_length_lengths);
983 for &(sym, extra_value, extra_bits) in &header.sequence {
984 writer.write_code(&codes[usize::from(sym)]);
985 writer.write(u32::from(extra_value), u32::from(extra_bits));
986 }
987}
988
989#[cfg(test)]
990mod tests {
991 use super::*;
992
993 fn noise(len: usize, mut seed: u32) -> Vec<u8> {
996 (0..len)
997 .map(|_| {
998 seed ^= seed << 13;
999 seed ^= seed >> 17;
1000 seed ^= seed << 5;
1001 (seed >> 24) as u8
1002 })
1003 .collect()
1004 }
1005
1006 fn fox_text() -> Vec<u8> {
1009 (0..40)
1010 .map(|i| format!("line {i}: the quick brown fox jumps over the lazy dog {}\n", i * i))
1011 .collect::<String>()
1012 .into_bytes()
1013 }
1014
1015 fn block_type(stream: &[u8]) -> u32 {
1016 (u32::from(stream[0]) >> 1) & 0b11
1017 }
1018
1019 fn round_trip(input: &[u8]) -> Vec<u8> {
1020 let compressed = compress(input);
1021 let output = decompress(&compressed).expect("our own output must inflate");
1022 assert_eq!(output, input);
1023 compressed
1024 }
1025
1026 #[test]
1027 fn round_trips_empty_input() {
1028 let compressed = round_trip(b"");
1029 assert_eq!(compressed, [0x03, 0x00]);
1031 }
1032
1033 #[test]
1034 fn round_trips_one_byte() {
1035 round_trip(b"x");
1036 }
1037
1038 #[test]
1039 fn round_trips_all_same_bytes() {
1040 let input = vec![b'a'; 100_000];
1041 let compressed = round_trip(&input);
1042 assert!(compressed.len() < 200, "compressed to {} bytes", compressed.len());
1045 }
1046
1047 #[test]
1048 fn round_trips_random_bytes_as_stored_blocks() {
1049 let input = noise(70_000, 0xDEAD_BEEF);
1050 let compressed = round_trip(&input);
1051 assert_eq!(block_type(&compressed), 0b00);
1055 let blocks = input.len().div_ceil(BLOCK_TOKENS);
1056 assert!(compressed.len() <= input.len() + 5 * blocks, "grew to {} bytes", compressed.len());
1057 }
1058
1059 #[test]
1060 fn round_trips_repetitive_text_with_dynamic_blocks() {
1061 let input = fox_text();
1062 let compressed = round_trip(&input);
1063 assert_eq!(block_type(&compressed), 0b10, "text this size should use a dynamic block");
1064 assert!(compressed.len() < 400, "compressed to {} bytes", compressed.len());
1066 }
1067
1068 #[test]
1069 fn short_input_uses_fixed_block() {
1070 let compressed = round_trip(b"hello hello hello hello");
1072 assert_eq!(block_type(&compressed), 0b01);
1073 }
1074
1075 #[test]
1076 fn round_trips_large_input_across_several_blocks() {
1077 let mut input = Vec::new();
1080 for i in 0..1500 {
1081 input.extend_from_slice(
1082 format!("record {i} belongs to user {} in region {}\n", i % 37, i % 5).as_bytes(),
1083 );
1084 }
1085 input.extend_from_slice(&noise(70_000, 42));
1086 for i in 0..1500 {
1087 input
1088 .extend_from_slice(format!("<item id=\"{i}\"><value>{}</value></item>\n", i * 31).as_bytes());
1089 }
1090 assert!(input.len() > 64 * 1024);
1091 let compressed = round_trip(&input);
1092 assert!(compressed.len() < input.len());
1093 }
1094
1095 #[test]
1096 fn finds_match_at_maximum_distance_and_length() {
1097 let mut input = noise(WINDOW_SIZE, 7);
1100 let repeat = input[..MAX_MATCH].to_vec();
1101 input.extend_from_slice(&repeat);
1102 input.extend_from_slice(b"tail");
1103 let tokens = tokenize(&input);
1104 assert!(
1105 tokens.contains(&Token::Match { len: MAX_MATCH as u16, dist: WINDOW_SIZE as u16 }),
1106 "expected a 258-byte match at distance 32768"
1107 );
1108 round_trip(&input);
1109 }
1110
1111 #[test]
1112 fn ignores_matches_just_beyond_the_window() {
1113 let mut input = noise(WINDOW_SIZE + 1, 9);
1118 let repeat = input[..64].to_vec();
1119 input.extend_from_slice(&repeat);
1120 let tokens = tokenize(&input);
1121 let longest = tokens
1122 .iter()
1123 .map(|t| match t {
1124 Token::Match { len, .. } => *len,
1125 _ => 0,
1126 })
1127 .max();
1128 assert!(longest.unwrap_or(0) < 8, "found a match of {longest:?} bytes");
1129 round_trip(&input);
1130 }
1131
1132 #[test]
1133 fn overlapping_match_decodes() {
1134 let mut w = BitWriter::new();
1137 w.write(1, 1);
1138 w.write(0b01, 2);
1139 let litlen = assign_codes(&fixed_litlen_lengths());
1140 let dist = assign_codes(&[5u8; 32]);
1141 w.write_code(&litlen[usize::from(b'a')]);
1142 w.write_code(&litlen[257 + length_code(10)]);
1143 w.write_code(&dist[0]);
1144 w.write_code(&litlen[256]);
1145 assert_eq!(decompress(&w.finish()).unwrap(), b"aaaaaaaaaaa");
1146 }
1147
1148 #[test]
1149 fn decodes_raw_stream_from_zlib() {
1150 let stream = [0xcb, 0x48, 0xcd, 0xc9, 0xc9, 0x57, 0xc8, 0x40, 0x27, 0x01];
1152 assert_eq!(decompress(&stream).unwrap(), b"hello hello hello hello");
1153 }
1154
1155 #[test]
1156 fn decodes_dynamic_block_from_zlib() {
1157 let stream = [
1160 0x95, 0x95, 0x5b, 0x56, 0xc3, 0x30, 0x0c, 0x44, 0xff, 0x59, 0x85, 0x96, 0x60, 0x49, 0xb6, 0x63,
1161 0xb3, 0x1b, 0x1e, 0x01, 0x0a, 0xa1, 0x81, 0x96, 0xd2, 0xc2, 0xea, 0x79, 0x58, 0x93, 0xff, 0xf9,
1162 0xee, 0xb9, 0x47, 0xd1, 0xe8, 0x7a, 0xba, 0xec, 0xf6, 0xb3, 0xa4, 0x6b, 0xf9, 0x78, 0x9a, 0xe5,
1163 0xfd, 0xb4, 0xbb, 0x7b, 0x91, 0xdb, 0xc3, 0x7a, 0xde, 0xcb, 0xc3, 0x7a, 0x91, 0xe7, 0xd3, 0xeb,
1164 0xdb, 0x51, 0xd6, 0xcf, 0xf9, 0xf0, 0xff, 0xf3, 0x72, 0xf3, 0xfd, 0x25, 0xf7, 0xeb, 0xa3, 0xa4,
1165 0xab, 0xe5, 0x8f, 0x52, 0x8e, 0xd2, 0x41, 0x19, 0x47, 0xe5, 0x41, 0x39, 0x47, 0xf5, 0x41, 0x65,
1166 0xf2, 0x0b, 0xeb, 0xc0, 0x0a, 0x87, 0x59, 0x19, 0x58, 0xe5, 0x30, 0x8f, 0x69, 0x13, 0x19, 0x48,
1167 0xec, 0xd6, 0x38, 0xac, 0x46, 0x90, 0x9d, 0xc3, 0x5a, 0x5c, 0x4d, 0x49, 0x45, 0x34, 0x41, 0x12,
1168 0xd6, 0x12, 0xc3, 0x44, 0x52, 0x14, 0xcd, 0xb1, 0xa1, 0x3a, 0x7b, 0xf5, 0x48, 0x54, 0x59, 0x5d,
1169 0x7a, 0x5c, 0x50, 0x59, 0x61, 0x60, 0x8c, 0x56, 0xd6, 0x34, 0x4c, 0x24, 0xa5, 0xb1, 0x86, 0x1d,
1170 0x49, 0x6d, 0xdc, 0x90, 0x6a, 0x67, 0xed, 0xc6, 0x7b, 0x27, 0xcd, 0xc9, 0x30, 0xc7, 0x48, 0x73,
1171 0x72, 0xc6, 0x44, 0xb6, 0x62, 0x5a, 0xec, 0x68, 0xa4, 0x39, 0xc5, 0x22, 0x55, 0x23, 0xcd, 0x29,
1172 0x53, 0xdc, 0xd1, 0x48, 0x73, 0x2a, 0xcc, 0x31, 0xd2, 0x9c, 0xba, 0x4d, 0x24, 0xcd, 0x99, 0xb6,
1173 0x1d, 0x49, 0x73, 0xa6, 0x2d, 0x55, 0xb6, 0x72, 0x70, 0x47, 0x27, 0xcd, 0xe9, 0x30, 0xc7, 0x49,
1174 0x73, 0x3a, 0x5c, 0x75, 0xb6, 0x73, 0x12, 0x9e, 0x87, 0xb3, 0xa5, 0x93, 0xf0, 0x22, 0x9d, 0x6d,
1175 0x1d, 0x45, 0x09, 0x78, 0x61, 0xab, 0x15, 0xf6, 0x78, 0x65, 0x49, 0x54, 0x9d, 0x93, 0xfa, 0xa8,
1176 0xa3, 0x5d, 0xbd, 0xd1, 0x7d, 0x8e, 0x6c, 0x49, 0x81, 0xb4, 0xfc, 0xfe, 0x87, 0xfc, 0x00,
1177 ];
1178 assert_eq!(block_type(&stream), 0b10);
1179 assert_eq!(decompress(&stream).unwrap(), fox_text());
1180 }
1181
1182 #[test]
1183 fn decodes_stored_block_by_hand() {
1184 let stream = [0x01, 0x05, 0x00, 0xfa, 0xff, b'h', b'e', b'l', b'l', b'o'];
1186 assert_eq!(decompress(&stream).unwrap(), b"hello");
1187 }
1188
1189 #[test]
1190 fn rejects_truncated_streams() {
1191 let full = compress(&fox_text());
1192 for cut in [0, 1, 2, 5, full.len() / 2, full.len() - 1] {
1193 let result = decompress(&full[..cut]);
1194 assert!(matches!(result, Err(InflateError::Truncated)), "cut at {cut}: {result:?}");
1195 }
1196 assert_eq!(decompress(&[0x01, 0x05, 0x00, 0xfa, 0xff, b'h']), Err(InflateError::Truncated));
1198 }
1199
1200 #[test]
1201 fn rejects_reserved_block_type() {
1202 assert_eq!(decompress(&[0x07, 0x00]), Err(InflateError::InvalidBlockType));
1204 }
1205
1206 #[test]
1207 fn rejects_distance_before_start_of_output() {
1208 let mut w = BitWriter::new();
1210 w.write(1, 1);
1211 w.write(0b01, 2);
1212 let litlen = assign_codes(&fixed_litlen_lengths());
1213 let dist = assign_codes(&[5u8; 32]);
1214 w.write_code(&litlen[257]);
1215 w.write_code(&dist[3]);
1216 w.write_code(&litlen[256]);
1217 assert_eq!(decompress(&w.finish()), Err(InflateError::DistanceTooFar));
1218 }
1219
1220 #[test]
1221 fn rejects_stored_block_with_mismatched_lengths() {
1222 assert_eq!(
1223 decompress(&[0x01, 0x05, 0x00, 0x00, 0x00, 0, 0, 0, 0, 0]),
1224 Err(InflateError::LengthMismatch)
1225 );
1226 }
1227
1228 #[test]
1229 fn rejects_reserved_symbols_in_fixed_block() {
1230 let litlen = assign_codes(&fixed_litlen_lengths());
1233 let dist = assign_codes(&[5u8; 32]);
1234 let mut w = BitWriter::new();
1235 w.write(1, 1);
1236 w.write(0b01, 2);
1237 w.write_code(&litlen[286]);
1238 assert_eq!(decompress(&w.finish()), Err(InflateError::InvalidCode));
1239 let mut w = BitWriter::new();
1240 w.write(1, 1);
1241 w.write(0b01, 2);
1242 w.write_code(&litlen[usize::from(b'a')]);
1243 w.write_code(&litlen[257]);
1244 w.write_code(&dist[30]);
1245 assert_eq!(decompress(&w.finish()), Err(InflateError::InvalidCode));
1246 }
1247
1248 #[test]
1249 fn rejects_oversized_code_length_repeat() {
1250 let mut w = BitWriter::new();
1255 w.write(1, 1);
1256 w.write(0b10, 2);
1257 w.write(0, 5); w.write(0, 5); w.write(0, 4); w.write(0, 3); w.write(0, 3); w.write(1, 3); w.write(0, 3); w.write(0, 1); w.write(127, 7); w.write(0, 1); w.write(127, 7); assert_eq!(decompress(&w.finish()), Err(InflateError::InvalidCodeLengths));
1269 }
1270
1271 #[test]
1272 fn rejects_repeat_with_no_previous_length() {
1273 let mut w = BitWriter::new();
1275 w.write(1, 1);
1276 w.write(0b10, 2);
1277 w.write(0, 5);
1278 w.write(0, 5);
1279 w.write(0, 4);
1280 w.write(1, 3); w.write(0, 3);
1282 w.write(0, 3);
1283 w.write(0, 3);
1284 w.write(0, 1); w.write(0, 2);
1286 assert_eq!(decompress(&w.finish()), Err(InflateError::InvalidCodeLengths));
1287 }
1288
1289 #[test]
1290 fn rejects_over_subscribed_code() {
1291 assert_eq!(Decoder::new(&[1, 1, 1]).err(), Some(InflateError::InvalidCode));
1293 assert_eq!(Decoder::new(&[2, 2]).err(), Some(InflateError::InvalidCode));
1295 assert!(Decoder::new(&[0, 1]).is_ok());
1297 }
1298
1299 #[test]
1300 fn caps_output_size() {
1301 let input = vec![0u8; 1 << 20];
1302 let compressed = compress(&input);
1303 assert!(compressed.len() < 1500, "a megabyte of zeros compressed to {} bytes", compressed.len());
1306 assert_eq!(decompress_with_limit(&compressed, 4096), Err(InflateError::OutputTooLarge));
1307 assert_eq!(decompress_with_limit(&compressed, 1 << 20).unwrap().len(), 1 << 20);
1308 let stored = compress(&noise(1000, 3));
1310 assert_eq!(block_type(&stored), 0b00);
1311 assert_eq!(decompress_with_limit(&stored, 999), Err(InflateError::OutputTooLarge));
1312 }
1313
1314 #[test]
1315 fn rejects_trailing_bytes() {
1316 let mut stream = compress(b"hello").to_vec();
1317 stream.push(0);
1318 assert_eq!(decompress(&stream), Err(InflateError::TrailingData));
1319 }
1320
1321 #[test]
1322 fn garbage_never_panics() {
1323 let junk = noise(300, 0xC0FFEE);
1326 for len in 0..junk.len() {
1327 let _ = decompress(&junk[..len]);
1328 }
1329 let stream = compress(&fox_text());
1330 for i in 0..stream.len() {
1331 for bit in 0..8 {
1332 let mut corrupt = stream.clone();
1333 corrupt[i] ^= 1 << bit;
1334 let _ = decompress_with_limit(&corrupt, 1 << 16);
1335 }
1336 }
1337 }
1338
1339 #[test]
1340 fn length_and_distance_tables_agree_with_the_rfc() {
1341 assert_eq!(length_code(3), 0);
1342 assert_eq!(length_code(10), 7);
1343 assert_eq!(length_code(11), 8);
1344 assert_eq!(length_code(257), 27);
1345 assert_eq!(length_code(258), 28);
1346 assert_eq!(dist_code(1), 0);
1347 assert_eq!(dist_code(4), 3);
1348 assert_eq!(dist_code(5), 4);
1349 assert_eq!(dist_code(6), 4);
1350 assert_eq!(dist_code(24576), 28);
1351 assert_eq!(dist_code(24577), 29);
1352 assert_eq!(dist_code(32768), 29);
1353 }
1354
1355 #[test]
1356 fn code_lengths_respect_the_limit() {
1357 let mut freqs = vec![0u32; 40];
1361 let (mut a, mut b) = (1u32, 1u32);
1362 for f in freqs.iter_mut() {
1363 *f = a;
1364 let next = a.saturating_add(b);
1365 a = b;
1366 b = next;
1367 }
1368 let lengths = build_lengths(&freqs, 15);
1369 assert!(lengths.iter().all(|&l| (1..=15).contains(&l)));
1370 assert!(Decoder::new(&lengths).is_ok());
1371 let short = build_lengths(&freqs, 7);
1372 assert!(short.iter().all(|&l| (1..=7).contains(&l)));
1373 assert!(Decoder::new(&short).is_ok());
1374 }
1375
1376 #[test]
1377 fn errors_display_as_sentences() {
1378 let text = InflateError::DistanceTooFar.to_string();
1379 assert!(text.contains("before the start"));
1380 let boxed: Box<dyn std::error::Error> = Box::new(InflateError::Truncated);
1381 assert!(boxed.to_string().contains("ended before"));
1382 }
1383}