diff --git a/src/lib.rs b/src/lib.rs index 00e524c..8633c3c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -4,30 +4,60 @@ struct Bitstream<'a> { } impl Bitstream<'_> { - fn read_bits_normal(&mut self, mut len: usize) -> Option { - assert!((len as u32) < u64::BITS); + fn read_bits(&mut self, mut len: usize) -> Result { + assert!((len as u32) < u16::BITS); let mut result = 0; while len > 0 { if self.data.is_empty() { - return None; + return Err(Error::Incomplete); } let to_read_from_current_byte = std::cmp::min(len % 8, 8 - self.pos_bit); result <<= to_read_from_current_byte; result |= - ((self.data[0] >> self.pos_bit) & ((1 << to_read_from_current_byte) - 1)) as u64; + ((self.data[0] >> self.pos_bit) & ((1 << to_read_from_current_byte) - 1)) as u16; len -= to_read_from_current_byte; self.pos_bit += to_read_from_current_byte; if self.pos_bit == 8 { - self.data = &self.data[1..]; + self.data = &self.data.get(1..).ok_or(Error::Incomplete)?; self.pos_bit = 0; } } - Some(result) + Ok(result) + } + + fn skip_current_byte(&mut self) -> Result<(), Error> { + if self.pos_bit > 0 { + self.data = &self.data.get(1..).ok_or(Error::Incomplete)?; + self.pos_bit = 0; + } + Ok(()) + } + + fn read_aligned_u16(&mut self) -> Result { + assert_eq!(self.pos_bit, 0); + let value = u16::from_le_bytes( + self.data + .get(..2) + .ok_or(Error::Incomplete)? + .try_into() + .unwrap(), + ); + self.data = self.data.get(2..).ok_or(Error::Incomplete)?; + Ok(value) + } + + fn read_aligned_bytes(&mut self, len: usize) -> Result<&[u8], Error> { + if self.data.len() < len { + return Err(Error::Incomplete); + } + let (data, rest) = self.data.split_at(len); + self.data = rest; + Ok(data) } } @@ -38,6 +68,7 @@ struct HuffmanTree { #[derive(Debug)] enum HuffmanNode { Leaf(u16), + // todo optimize this into an array index lol Cont { zero: usize, one: usize }, } @@ -47,7 +78,7 @@ enum HuffmanLookupResult { } impl HuffmanTree { - fn fixed_code_lengths() -> [u8; 288] { + fn fixed_lit_len_code_lengths() -> [u8; 288] { let mut codes = [0; 288]; codes[0..144].fill(8); codes[144..256].fill(9); @@ -56,6 +87,12 @@ impl HuffmanTree { codes } + fn fixed_dist_code_lengths() -> [u8; 32] { + let mut codes = [0; 32]; + codes.fill(5); + codes + } + fn from_lengths(lengths: &[u8]) -> Self { const MAX_BITS: usize = 10; @@ -114,7 +151,7 @@ impl HuffmanTree { Self { nodes } } - fn lookup_with_state(&self, state: usize, bit: u64) -> HuffmanLookupResult { + fn lookup_with_state(&self, state: usize, bit: u16) -> HuffmanLookupResult { let HuffmanNode::Cont { zero, one } = self.nodes[state] else { unreachable!("invalid state, should point at continuation node"); }; @@ -125,6 +162,20 @@ impl HuffmanTree { } } + fn read_value(&self, data: &mut Bitstream) -> Result { + let mut node_state = 0; + loop { + let bit = data.read_bits(1)?; + let result = self.lookup_with_state(node_state, bit); + match result { + HuffmanLookupResult::Done(value) => { + return Ok(value); + } + HuffmanLookupResult::Incomplete { next_state } => node_state = next_state, + } + } + } + fn to_dot(&self) -> String { use std::fmt::Write; @@ -184,49 +235,174 @@ impl HuffmanTree { } } +#[derive(Debug)] +pub enum Error { + ReservedBlockType, + Incomplete, + InvalidCode, +} + // https://datatracker.ietf.org/doc/html/rfc1951 -pub fn inflate(data: &[u8], out: &mut Vec) { +pub fn inflate(data: &[u8], out: &mut Vec) -> Result<(), Error> { std::fs::write( "output.dot", - HuffmanTree::from_lengths(&HuffmanTree::fixed_code_lengths()).to_dot(), + HuffmanTree::from_lengths(&HuffmanTree::fixed_lit_len_code_lengths()).to_dot(), ) .unwrap(); let mut data = Bitstream { data, pos_bit: 0 }; loop { - let bfinal = data.read_bits_normal(1).unwrap(); + // 3.2.3. Details of block format + // Each block of compressed data begins with 3 header bits containing the following data: + // first bit BFINAL + // next 2 bits BTYPE + let bfinal = data.read_bits(1)?; + let btype = data.read_bits(2)?; - let btype = data.read_bits_normal(2).unwrap(); + 'end_block: { + // BTYPE specifies how the data are compressed, as follows: + // 00 - no compression + // 01 - compressed with fixed Huffman codes + // 10 - compressed with dynamic Huffman codes + // 11 - reserved (error) + let (code_lengths_lit_len, code_lengths_dist) = match btype { + 0b00 => { + // Any bits of input up to the next byte boundary are ignored. + // The rest of the block consists of the following information: - assert_eq!(btype, 1, "not a static huffman tree construction"); + // 0 1 2 3 4... + // +---+---+---+---+================================+ + // | LEN | NLEN |... LEN bytes of literal data...| + // +---+---+---+---+================================+ - let tree = HuffmanTree::from_lengths(&HuffmanTree::fixed_code_lengths()); - - let mut node_state = 0; - loop { - let bit = data.read_bits_normal(1).unwrap(); - let result = tree.lookup_with_state(node_state, bit); - match result { - HuffmanLookupResult::Done(value) => { - dbg!(value); - node_state = 0; - - match value { - 0..256 => { - out.push(value as u8); - } - 256 => break, - 257..286 => { - let length = match value { - 257..265 => value - (257 - 3), - _ => todo!("lz77 more"), - }; - } - 286.. => unreachable!("invalid byte"), - } + // LEN is the number of data bytes in the block. NLEN is the + // one's complement of LEN. + data.skip_current_byte()?; + let len = data.read_aligned_u16()?; + let _nlen = data.read_aligned_u16()?; + out.extend_from_slice(data.read_aligned_bytes(len as usize)?); + break 'end_block; + } + 0b01 => ( + &HuffmanTree::fixed_lit_len_code_lengths(), + &HuffmanTree::fixed_dist_code_lengths(), + ), + 0b10 => todo!("dynamic huffman codes"), + 0b11 => return Err(Error::ReservedBlockType), + _ => unreachable!("only 2 bits"), + }; + + let lit_len_tree = HuffmanTree::from_lengths(code_lengths_lit_len); + let dist_tree = HuffmanTree::from_lengths(code_lengths_dist); + + loop { + let lit_len = lit_len_tree.read_value(&mut data)?; + match lit_len { + 0..=255 => { + out.push(lit_len as /* literal */ u8); + } + 256 => break 'end_block, + 257..286 => { + let mut compute_value = + |first_code, extra_bits_amount, value_start| -> Result { + let extra_value = if extra_bits_amount == 0 { + 0 + } else { + data.read_bits(extra_bits_amount)? as u16 + }; + + let extra_values_amount = 2_u16.pow(extra_bits_amount as u32); + + Ok(((lit_len - first_code) * extra_values_amount) + + value_start + + extra_value) + }; + // Extra Extra Extra + // Code Bits Length(s) Code Bits Lengths Code Bits Length(s) + // ---- ---- ------ ---- ---- ------- ---- ---- ------- + // 257 0 3 267 1 15,16 277 4 67-82 + // 258 0 4 268 1 17,18 278 4 83-98 + // 259 0 5 269 2 19-22 279 4 99-114 + // 260 0 6 270 2 23-26 280 4 115-130 + // 261 0 7 271 2 27-30 281 5 131-162 + // 262 0 8 272 2 31-34 282 5 163-194 + // 263 0 9 273 3 35-42 283 5 195-226 + // 264 0 10 274 3 43-50 284 5 227-257 + // 265 1 11,12 275 3 51-58 285 0 258 + // 266 1 13,14 276 3 59-66 + let length = match lit_len { + 257..=264 => compute_value(257, 0, 3), + 265..=268 => compute_value(265, 1, 11), + 269..=272 => compute_value(269, 2, 19), + 273..=276 => compute_value(273, 3, 35), + 277..=280 => compute_value(277, 4, 67), + 281..=284 => compute_value(281, 5, 131), + 285 => Ok(258), + _ => return Err(Error::InvalidCode), + }?; + + let dist_value = dist_tree.read_value(&mut data)?; + + let mut compute_value = + |first_code, extra_bits_amount, value_start| -> Result { + assert_eq!(extra_bits_amount, 3); + dbg!(data.read_bits(1).unwrap()); + dbg!(data.read_bits(1).unwrap()); + dbg!(data.read_bits(1).unwrap()); + std::process::exit(1); + + // TODO: read_bits is broken (i suspect the byte parts are reversed) + + let extra_value = if extra_bits_amount == 0 { + 0 + } else { + data.read_bits(extra_bits_amount)? as u16 + }; + + dbg!(extra_value); + + let extra_values_amount = 2_u16.pow(extra_bits_amount as u32); + + Ok(((dist_value - first_code) * extra_values_amount) + + value_start + + extra_value) + }; + // Extra Extra Extra + // Code Bits Dist Code Bits Dist Code Bits Distance + // ---- ---- ---- ---- ---- ------ ---- ---- -------- + // 0 0 1 10 4 33-48 20 9 1025-1536 + // 1 0 2 11 4 49-64 21 9 1537-2048 + // 2 0 3 12 5 65-96 22 10 2049-3072 + // 3 0 4 13 5 97-128 23 10 3073-4096 + // 4 1 5,6 14 6 129-192 24 11 4097-6144 + // 5 1 7,8 15 6 193-256 25 11 6145-8192 + // 6 2 9-12 16 7 257-384 26 12 8193-12288 + // 7 2 13-16 17 7 385-512 27 12 12289-16384 + // 8 3 17-24 18 8 513-768 28 13 16385-24576 + // 9 3 25-32 19 8 769-1024 29 13 24577-32768 + let distance = match dist_value { + 0..=3 => compute_value(0, 0, 1), + 4..=5 => compute_value(4, 1, 5), + 6..=7 => compute_value(6, 2, 9), + 8..=9 => compute_value(8, 3, 17), + 10..=11 => compute_value(10, 4, 33), + 12..=13 => compute_value(12, 5, 65), + 14..=15 => compute_value(14, 6, 129), + 16..=17 => compute_value(16, 7, 257), + 18..=19 => compute_value(18, 8, 513), + 20..=21 => compute_value(20, 9, 1025), + 22..=23 => compute_value(22, 10, 2049), + 24..=25 => compute_value(24, 11, 4097), + 26..=27 => compute_value(26, 12, 8193), + 28..=29 => compute_value(28, 13, 16385), + _ => return Err(Error::InvalidCode), + }?; + + dbg!((length, dist_value, distance)); + } + 286.. => return Err(Error::InvalidCode), } - HuffmanLookupResult::Incomplete { next_state } => node_state = next_state, } } @@ -234,6 +410,8 @@ pub fn inflate(data: &[u8], out: &mut Vec) { break; } } + + Ok(()) } #[cfg(test)] @@ -247,16 +425,16 @@ mod tests { data: &bytes, pos_bit: 0, }; - assert_eq!(stream.read_bits_normal(2).unwrap(), 0b01); - assert_eq!(stream.read_bits_normal(3).unwrap(), 0b010); - assert_eq!(stream.read_bits_normal(5).unwrap(), 0b11001); - assert_eq!(stream.read_bits_normal(3).unwrap(), 0b111); - assert_eq!(stream.read_bits_normal(3).unwrap(), 0b010); + assert_eq!(stream.read_bits(2).unwrap(), 0b01); + assert_eq!(stream.read_bits(3).unwrap(), 0b010); + assert_eq!(stream.read_bits(5).unwrap(), 0b11001); + assert_eq!(stream.read_bits(3).unwrap(), 0b111); + assert_eq!(stream.read_bits(3).unwrap(), 0b010); } #[test] fn decode() { - let lengths = HuffmanTree::fixed_code_lengths(); + let lengths = HuffmanTree::fixed_lit_len_code_lengths(); HuffmanTree::from_lengths(&lengths); } } diff --git a/src/main.rs b/src/main.rs index de35c24..a518d7f 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,7 +1,7 @@ use std::ffi::CStr; fn main() { - let gz = std::env::args().nth(1).unwrap(); + let gz = std::env::args().nth(1).expect("must provide a cli argument"); let gz = std::fs::read(gz).unwrap(); assert_eq!(gz[0], 31, "ID"); @@ -27,6 +27,5 @@ fn main() { zwergli::inflate(blocks, &mut out); - dbg!(&out); dbg!(String::from_utf8(out)).ok(); }