Improved bit reader/writer

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"] }

src/binio.rs

@@ -1,35 +1,123 @@
-pub struct BitIo {
-    data: Vec<u8>,
+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<u8>) -> 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<u8> {
-        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<u8>,
+
+    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())
     }
 }

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<u8>) -> (usize, Vec<u8>, Vec<u8>) {
         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<Vec<u8>>, 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<u8>) -> (usize, Vec<u8>, Vec<u8>) {
                 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<T: ReadBytesExt + Read>(
@@ -159,6 +163,8 @@ pub fn decompress<T: ReadBytesExt + Read>(
 
 fn decompress_lzw(input_data: &[u8], size: usize) -> Vec<u8> {
     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<u8> {
     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;

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(