This commit is contained in:
nora 2026-06-07 23:12:43 +02:00
parent 00a411a728
commit 8e1d90e0ce
2 changed files with 223 additions and 46 deletions

View file

@ -4,30 +4,60 @@ struct Bitstream<'a> {
}
impl Bitstream<'_> {
fn read_bits_normal(&mut self, mut len: usize) -> Option<u64> {
assert!((len as u32) < u64::BITS);
fn read_bits(&mut self, mut len: usize) -> Result<u16, Error> {
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<u16, Error> {
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<u16, Error> {
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<u8>) {
pub fn inflate(data: &[u8], out: &mut Vec<u8>) -> 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<u16, Error> {
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<u16, Error> {
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<u8>) {
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);
}
}

View file

@ -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();
}