diff options
Diffstat (limited to 'src/zlib.rs')
| -rw-r--r-- | src/zlib.rs | 258 |
1 files changed, 258 insertions, 0 deletions
diff --git a/src/zlib.rs b/src/zlib.rs new file mode 100644 index 0000000..dcbdb55 --- /dev/null +++ b/src/zlib.rs @@ -0,0 +1,258 @@ +const MAX_BITS: usize = 15; +const CODE_TREE_LEN: usize = 288; +const DIST_TREE_LEN: usize = 32; + +const LEN_START: [usize; 29] = [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, 227, 258]; +const LEN_EXTRA: [usize; 29] = [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]; + +const DIST_START: [usize; 30] = [1, 2, 3, 4, 5, 7, 9, 13, 17, 25, 33, 49, 65, 97, 129, 193, 257, 385, 513, 769, 1025, 1537, 2049, 3073, 4097, 6145, 8193, 12289, 16385, 24577]; +const DIST_EXTRA: [usize; 30] = [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]; + +#[derive(Clone)] +struct HuffmanTree<const LENGTH: usize> +{ + pub len_counts: [u8; MAX_BITS + 1], + pub ordered_codes: [u16; LENGTH] +} + +impl<const LENGTH: usize> HuffmanTree<LENGTH> +{ + const fn new(code_lens: [u8; LENGTH]) -> Self + { + let mut len_counts = [0u8; MAX_BITS + 1]; + let mut ordered_codes = [0u16; LENGTH]; + + let mut code: usize = 0; + while code < LENGTH + { + len_counts[code_lens[code] as usize] += 1; + code += 1 + } + + let mut offsets = [0usize; MAX_BITS + 1]; + let mut len: usize = 1; + while len < MAX_BITS + { + offsets[len + 1] = offsets[len] + len_counts[len] as usize; + len += 1; + } + + let mut code = 0; + while code < LENGTH + { + if code_lens[code] != 0 + { + ordered_codes[offsets[code_lens[code] as usize]] = code as u16; + offsets[code_lens[code] as usize] += 1; + } + code += 1; + } + + Self + { + len_counts, + ordered_codes + } + } + + + + fn decode<T: Bitstream>(&self, bitstream: &mut T) -> u16 + { + let mut code: u16 = 0; + let mut first: u16 = 0; + let mut index: usize = 0; + + for len in 1..MAX_BITS + { + code |= bitstream.next_bit() as u16; + let count: u16 = self.len_counts[len] as u16; + if code - count < first + { + return self.ordered_codes[index + ((code - first) as usize)]; + } + index += count as usize; + first += count; + first <<= 1; + code <<= 1; + } + + panic!("No codes left") + } + + const fn get_static_code_lens() -> [u8; CODE_TREE_LEN] + { + let mut code_lens = [0u8; CODE_TREE_LEN]; + let mut i = 0; + while i < CODE_TREE_LEN + { + code_lens[i] = match i + { + 0..144 => 8, + 144..256 => 9, + 256..280 => 7, + 280..CODE_TREE_LEN => 8, + _ => panic!() + }; + i += 1 + } + code_lens + } + const fn get_static_dist_code_lens() -> [u8; DIST_TREE_LEN] + { + [5u8; DIST_TREE_LEN] + } +} + +static STATIC_HUFFMAN_TREE: HuffmanTree<CODE_TREE_LEN> = HuffmanTree::new(HuffmanTree::<CODE_TREE_LEN>::get_static_code_lens()); +static STATIC_HUFFMAN_DIST_TREE: HuffmanTree<DIST_TREE_LEN> = HuffmanTree::new(HuffmanTree::<DIST_TREE_LEN>::get_static_dist_code_lens()); + +#[derive(PartialEq)] +enum CompressionType +{ + None = 0, + StaticCodes = 1, + DynamicCodes = 2, + Reserved = 3, +} + +impl CompressionType +{ + fn new(x: u8) -> Self + { + match x + { + 0 => Self::None, + 1 => Self::StaticCodes, + 2 => Self::DynamicCodes, + 3 => Self::Reserved, + _ => panic!("Compression type should not be greater than 4") + } + } +} + +pub struct ByteBuffer<'a> +{ + buffer: &'a [u8], + pos: usize, +} + +impl<'a> ByteBuffer<'a> +{ + pub fn new(buffer: &'a [u8]) -> Self + { + Self + { + buffer, + pos: 0 + } + } +} + +pub trait Bitstream { + fn next_bit(&mut self) -> bool; + + fn next_bits(&mut self, c: usize) -> u16 + { + if c > 16 + { + panic!("Cannot fetch more than 16 bits at a time"); + } + let mut bits = 0u16; + + for i in 0..c + { + bits |= (self.next_bit() as u16) << i; + } + + bits + } + + fn skip_to_byte_boundary(&mut self); +} + +impl<'a> Bitstream for ByteBuffer<'a> +{ + fn next_bit(&mut self) -> bool + { + let bit = (self.buffer[self.pos/8 as usize] & 1 << self.pos%8) != 0; + self.pos += 1; + bit + } + + fn skip_to_byte_boundary(&mut self) + { + // println!("{}", (8 - self.pos%8)%8); + self.pos += (8 - self.pos%8)%8; + } +} + +pub fn zlib_decode<T: Bitstream>(mut instream: T) -> Vec<u8> +{ + let mut output_buffer: Vec<u8> = Vec::new(); + // let mut output_pos: usize = 0; + + let mut final_chunk = false; + while !final_chunk + { + final_chunk = instream.next_bit(); + let compression_type = CompressionType::new(instream.next_bits(2) as u8); + if compression_type == CompressionType::Reserved + { + panic!("Reserved compression type should not be used"); + } + else if compression_type == CompressionType::None + { + instream.skip_to_byte_boundary(); + let len = instream.next_bits(16); + let nlen = instream.next_bits(16); + if len != !nlen + { + panic!("len and nlen don't match"); + } + for _ in 0..len + { + output_buffer.push(instream.next_bits(8) as u8); + } + // todo!("No compression"); + } + else + { + let (code_tree, dist_tree) = + if compression_type == CompressionType::StaticCodes + { + (STATIC_HUFFMAN_TREE.clone(), STATIC_HUFFMAN_DIST_TREE.clone()) + } + else // Dynamic codes + { + todo!("Dynamic codes") + }; + + + let mut code = 0u16; + while code != 256 + { + code = code_tree.decode(&mut instream); + if code < 256 + { + output_buffer.push(code as u8); + } + else if code > 256 + { + let len = LEN_START[code as usize - 257] + instream.next_bits(LEN_EXTRA[code as usize - 257]) as usize; + + let dist_code = dist_tree.decode(&mut instream); + let dist = DIST_START[dist_code as usize] + instream.next_bits(DIST_EXTRA[dist_code as usize]) as usize; + + let out_pos = output_buffer.len() - 1; + for _ in 0..len + { + output_buffer.push(output_buffer[out_pos - dist]); + } + } + } + } + } + + output_buffer +} |