From abedabd06fa078f75182e838e55603f508d6bcee Mon Sep 17 00:00:00 2001 From: G2-Games Date: Tue, 23 Jul 2024 23:19:39 -0500 Subject: [PATCH] Improved bit reader/writer --- Cargo.toml | 2 +- src/binio.rs | 172 +++++++++++++++++++++++++++++---------------- src/compression.rs | 24 ++++--- src/main.rs | 3 +- 4 files changed, 130 insertions(+), 71 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index e40e1df..3310476 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -5,4 +5,4 @@ edition = "2021" [dependencies] byteorder = "1.5.0" -image = { version = "0.25.2", default-features = false, features = ["png"] } +image = { version = "0.25.2", default-features = false, features = ["png", "jpeg"] } diff --git a/src/binio.rs b/src/binio.rs index 334d0e4..243b3d7 100644 --- a/src/binio.rs +++ b/src/binio.rs @@ -1,35 +1,123 @@ -pub struct BitIo { - data: Vec, +use std::io::{Read, Write}; + +use byteorder::{ReadBytesExt, WriteBytesExt}; + +pub struct BitWriter<'a, O: Write + WriteBytesExt> { + output: &'a mut O, + + current_byte: u8, + byte_offset: usize, bit_offset: usize, byte_size: usize, } -impl BitIo { +impl<'a, O: Write + WriteBytesExt> BitWriter<'a, O> { /// Create a new BitIO reader and writer over some data - pub fn new(data: Vec) -> Self { + pub fn new(output: &'a mut O) -> Self { Self { - data, + output, + + current_byte: 0, + byte_offset: 0, bit_offset: 0, + byte_size: 0, } } - /// Get the byte offset of the reader - pub fn byte_offset(&self) -> usize { - self.byte_offset - } - /// Get the byte size of the reader pub fn byte_size(&self) -> usize { self.byte_size } - /// Get the current bytes up to `byte_size` in the reader - pub fn bytes(&self) -> Vec { - self.data[..self.byte_size].to_vec() + /// Write some bits to the buffer + pub fn write_bit(&mut self, data: u64, bit_len: usize) { + if bit_len > 8 * 8 { + panic!("Cannot write more than 64 bits at once"); + } + + if bit_len % 8 == 0 && self.bit_offset == 0 { + self.write(data, bit_len / 8); + return; + } + + for i in 0..bit_len { + let bit_value = (data >> i) & 1; + + self.current_byte &= !(1 << self.bit_offset); + + self.current_byte |= (bit_value << self.bit_offset) as u8; + + self.bit_offset += 1; + if self.bit_offset >= 8 { + self.byte_offset += 1; + self.bit_offset = 0; + + self.output.write_u8(self.current_byte).unwrap(); + self.current_byte = 0; + } + } + + self.byte_size = self.byte_offset + (self.bit_offset + 7) / 8; + } + + pub fn write(&mut self, data: u64, byte_len: usize) { + if byte_len > 8 { + panic!("Cannot write more than 8 bytes at once") + } + + self.output.write_all(&data.to_le_bytes()[..byte_len]).unwrap(); + self.byte_offset += byte_len; + + self.byte_size = self.byte_offset + (self.bit_offset + 7) / 8; + } +} + +impl<'a, O: Write + WriteBytesExt> Drop for BitWriter<'_, O> { + fn drop(&mut self) { + let _ = self.output.write_u8(self.current_byte); + } +} + +pub struct BitReader<'a, I: Read + ReadBytesExt> { + input: &'a mut I, + + current_byte: Option, + + byte_offset: usize, + bit_offset: usize, + + byte_size: usize, +} + + +impl<'a, I: Read + ReadBytesExt> BitReader<'a, I> { + /// Create a new BitIO reader and writer over some data + pub fn new(input: &'a mut I) -> Self { + let first = input.read_u8().unwrap(); + Self { + input, + + current_byte: Some(first), + + byte_offset: 0, + bit_offset: 0, + + byte_size: 0, + } + } + + /// Get the byte size of the reader + pub fn byte_offset(&self) -> usize { + self.byte_offset + } + + /// Get the byte size of the reader + pub fn byte_size(&self) -> usize { + self.byte_size } /// Read some bits from the buffer @@ -44,12 +132,14 @@ impl BitIo { let mut result = 0; for i in 0..bit_len { - let bit_value = ((self.data[self.byte_offset] as usize >> self.bit_offset) & 1) as u64; + let bit_value = ((self.current_byte.unwrap() as usize >> self.bit_offset) & 1) as u64; self.bit_offset += 1; if self.bit_offset == 8 { self.byte_offset += 1; self.bit_offset = 0; + + self.current_byte = Some(self.input.read_u8().unwrap()); } result |= bit_value << i; @@ -64,53 +154,17 @@ impl BitIo { panic!("Cannot read more than 8 bytes") } - let mut padded_slice = [0u8; 8]; - padded_slice.copy_from_slice(&self.data[self.byte_offset..self.byte_offset + byte_len]); + if self.current_byte.is_none() { + self.current_byte = Some(self.input.read_u8().unwrap()); + } + + let mut padded_slice = vec![0u8; byte_len]; + self.input.read_exact(&mut padded_slice).unwrap(); self.byte_offset += byte_len; - u64::from_le_bytes(padded_slice) - } + let extra_length = padded_slice.len() - byte_len; + padded_slice.extend_from_slice(&vec![0u8; extra_length]); - /// Write some bits to the buffer - pub fn write_bit(&mut self, data: u64, bit_len: usize) { - if bit_len > 8 * 8 { - panic!("Cannot write more than 64 bits"); - } - - if bit_len % 8 == 0 && self.bit_offset == 0 { - self.write(data, bit_len / 8); - return; - } - - for i in 0..bit_len { - let bit_value = (data >> i) & 1; - - self.data[self.byte_offset] &= !(1 << self.bit_offset); - - self.data[self.byte_offset] |= (bit_value << self.bit_offset) as u8; - - self.bit_offset += 1; - if self.bit_offset == 8 { - self.byte_offset += 1; - self.bit_offset = 0; - } - } - - self.byte_size = self.byte_offset + (self.bit_offset + 7) / 8; - } - - pub fn write(&mut self, data: u64, byte_len: usize) { - if byte_len > 8 { - panic!("Cannot write more than 8 bytes") - } - - let mut padded_slice = [0u8; 8]; - padded_slice.copy_from_slice(&data.to_le_bytes()); - - self.data[self.byte_offset..self.byte_offset + byte_len] - .copy_from_slice(&padded_slice[..byte_len]); - self.byte_offset += byte_len; - - self.byte_size = self.byte_offset + (self.bit_offset + 7) / 8; + u64::from_le_bytes(padded_slice.try_into().unwrap()) } } diff --git a/src/compression.rs b/src/compression.rs index bbe16e6..82ca7e0 100644 --- a/src/compression.rs +++ b/src/compression.rs @@ -1,8 +1,8 @@ -use std::{collections::HashMap, io::{Read, Write}}; +use std::{collections::HashMap, io::{Cursor, Read, Write}}; use byteorder::{ReadBytesExt, WriteBytesExt, LE}; -use crate::binio::BitIo; +use crate::binio::{BitReader, BitWriter}; /// The size of compressed data in each chunk #[derive(Debug, Clone, Copy)] @@ -89,8 +89,9 @@ fn compress_lzw(data: &[u8], last: Vec) -> (usize, Vec, Vec) { element = last } - let mut bit_io = BitIo::new(vec![0u8; 0xF0000]); - let write_bit = |bit_io: &mut BitIo, code: u64| { + let mut output_buf = Vec::new(); + let mut bit_io = BitWriter::new(&mut output_buf); + let write_bit = |bit_io: &mut BitWriter>, code: u64| { if code > 0x7FFF { bit_io.write_bit(1, 1); bit_io.write_bit(code, 18); @@ -128,15 +129,18 @@ fn compress_lzw(data: &[u8], last: Vec) -> (usize, Vec, Vec) { write_bit(&mut bit_io, *dictionary.get(&vec![c]).unwrap()); } } - return (count, bit_io.bytes(), Vec::new()); + drop(bit_io); + return (count, output_buf, Vec::new()); } else if dictionary_count < 0x3FFFE { if !last_element.is_empty() { write_bit(&mut bit_io, *dictionary.get(&last_element).unwrap()); } - return (count, bit_io.bytes(), Vec::new()); + drop(bit_io); + return (count, output_buf, Vec::new()); } - (count, bit_io.bytes(), last_element) + drop(bit_io); + (count, output_buf, last_element) } pub fn decompress( @@ -159,6 +163,8 @@ pub fn decompress( fn decompress_lzw(input_data: &[u8], size: usize) -> Vec { let mut data = input_data.to_vec(); + data.extend_from_slice(&[0, 0]); + let mut data = Cursor::new(data); let mut dictionary = HashMap::new(); for i in 0..256 { dictionary.insert(i as u64, vec![i as u8]); @@ -167,8 +173,8 @@ fn decompress_lzw(input_data: &[u8], size: usize) -> Vec { let mut result = Vec::with_capacity(size); let data_size = input_data.len(); - data.extend_from_slice(&[0, 0]); - let mut bit_io = BitIo::new(data); + + let mut bit_io = BitReader::new(&mut data); let mut w = dictionary.get(&0).unwrap().clone(); let mut element; diff --git a/src/main.rs b/src/main.rs index fc82cb7..2231a27 100644 --- a/src/main.rs +++ b/src/main.rs @@ -11,7 +11,7 @@ use picture::DangoPicture; use image::RgbaImage; fn main() { - let image_data = image::open("test.png").unwrap().to_rgba8(); + let image_data = image::open("clouds_dark.png").unwrap().to_rgba8(); let encoded_dpf = DangoPicture { header: Header { width: image_data.width(), @@ -25,7 +25,6 @@ fn main() { let mut outfile = File::create("test.dpf").unwrap(); encoded_dpf.encode(&mut outfile); - let mut infile = File::open("test.dpf").unwrap(); let decoded_dpf = DangoPicture::decode(&mut infile); let out_image = RgbaImage::from_raw(