Improved functions, added error type

Cargo.toml

@@ -6,3 +6,4 @@ edition = "2021"
 [dependencies]
 byteorder = "1.5.0"
 image = { version = "0.25.2", default-features = false, features = ["png", "jpeg"] }
+thiserror = "1.0.63"

src/compression.rs

@@ -162,9 +162,7 @@ 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 data = Cursor::new(input_data);
     let mut dictionary = HashMap::new();
     for i in 0..256 {
         dictionary.insert(i as u64, vec![i as u8]);
@@ -179,6 +177,10 @@ fn decompress_lzw(input_data: &[u8], size: usize) -> Vec<u8> {
 
     let mut element;
     loop {
+        if bit_io.byte_offset() >= data_size - 1 {
+            break;
+        }
+
         let flag = bit_io.read_bit(1);
         if flag == 0 {
             element = bit_io.read_bit(15);
@@ -186,10 +188,6 @@ fn decompress_lzw(input_data: &[u8], size: usize) -> Vec<u8> {
             element = bit_io.read_bit(18);
         }
 
-        if bit_io.byte_offset() > data_size {
-            break;
-        }
-
         let mut entry;
         if let Some(x) = dictionary.get(&element) {
             // If the element was already in the dict, get it

src/header.rs

@@ -4,8 +4,15 @@ use byteorder::{WriteBytesExt, LE};
 pub struct Header {
     pub magic: [u8; 8],
 
+    /// Width of the image in pixels
     pub width: u32,
+    /// Height of the image in pixels
     pub height: u32,
+
+    /// Bit depth in bits per pixel
+    pub depth: u16,
+
+    pub encoding: ImageEncoding,
 }
 
 impl Default for Header {
@@ -14,6 +21,8 @@ impl Default for Header {
             magic: *b"dangoimg",
             width: 0,
             height: 0,
+            depth: 32,
+            encoding: ImageEncoding::LosslessCompressed,
         }
     }
 }
@@ -29,3 +38,15 @@ impl Header {
         buf.into_inner().try_into().unwrap()
     }
 }
+
+#[repr(u16)]
+pub enum ImageEncoding {
+    /// Uncompressed raw bitmap
+    Bitmap = 0,
+
+    /// Losslessly compressed
+    LosslessCompressed = 1,
+
+    /// Lossily compressed
+    LossyCompressed = 2,
+}

src/main.rs

@@ -4,33 +4,34 @@ mod operations;
 mod binio;
 mod picture;
 
-use std::fs::File;
+use std::{fs::File, io::{BufReader, BufWriter}, time::Instant};
 use header::Header;
 use picture::DangoPicture;
 
 use image::RgbaImage;
 
 fn main() {
-    let image_data = image::open("clouds_dark.png").unwrap().to_rgba8();
-    let encoded_dpf = DangoPicture {
-        header: Header {
-            width: image_data.width(),
-            height: image_data.height(),
+    let image_data = image::open("small_transparency.png").unwrap().to_rgba8();
+    let encoded_dpf = DangoPicture::from_raw(
+        image_data.width(),
+        image_data.height(),
+        &image_data
+    );
 
-            ..Default::default()
-        },
-        bitmap: image_data.into_vec(),
-    };
-
-    let mut outfile = File::create("test.dpf").unwrap();
+    let timer = Instant::now();
+    let mut outfile = BufWriter::new(File::create("test.dpf").unwrap());
     encoded_dpf.encode(&mut outfile);
+    println!("Encoding took {}ms", timer.elapsed().as_millis());
+
+    let timer = Instant::now();
+    let mut infile = BufReader::new(File::open("test.dpf").unwrap());
+    let decoded_dpf = DangoPicture::decode(&mut infile).unwrap();
+    println!("Decoding took {}ms", timer.elapsed().as_millis());
 
-    let mut infile = File::open("test.dpf").unwrap();
-    let decoded_dpf = DangoPicture::decode(&mut infile);
     let out_image = RgbaImage::from_raw(
         decoded_dpf.header.width,
         decoded_dpf.header.height,
-        decoded_dpf.bitmap
+        decoded_dpf.bitmap.into()
     ).unwrap();
     out_image.save("test.png").unwrap();
 }

src/operations.rs

@@ -1,7 +1,9 @@
+use image::{DynamicImage, GrayImage};
+
 pub fn line_diff(width: u32, height: u32, data: &[u8]) -> Vec<u8> {
     let mut output_buf = Vec::with_capacity((width * height * 4) as usize);
 
-    let block_height = (f32::ceil(height as f32 / 3.0) as u16) as u32;
+    let block_height = f32::ceil(height as f32 / 3.0) as u32;
 
     let mut curr_line;
     let mut prev_line = Vec::with_capacity(width as usize * 3);
@@ -53,7 +55,7 @@ pub fn diff_line(width: u32, height: u32, input: &[u8]) -> Vec<u8> {
     let mut data = Vec::with_capacity(width as usize * 3);
     let mut alpha_data = Vec::with_capacity(width as usize);
 
-    let block_height = (f32::ceil(height as f32 / 3.0) as u16) as usize;
+    let block_height = f32::ceil(height as f32 / 3.0) as u32;
     let pixel_byte_count = 4;
     let line_byte_count = (width * pixel_byte_count as u32) as usize;
 
@@ -68,8 +70,7 @@ pub fn diff_line(width: u32, height: u32, input: &[u8]) -> Vec<u8> {
         curr_line = input[i..i + line_byte_count]
             .windows(4)
             .step_by(4)
-            .flat_map(|r| &r[0..3])
-            .copied()
+            .flat_map(|r| [r[0], r[1], r[2]])
             .collect();
         curr_alpha = input[i..i + line_byte_count]
             .iter()

src/picture.rs

@@ -1,6 +1,7 @@
 use std::io::{Read, Write};
 
 use byteorder::{ReadBytesExt, WriteBytesExt, LE};
+use thiserror::Error;
 
 use crate::{compression::{compress, decompress, ChunkInfo, CompressionInfo}, header::Header, operations::{diff_line, line_diff}};
 
@@ -35,10 +36,14 @@ impl DangoPicture {
     }
 
     /// Decode the image from anything that implements [Read]
-    pub fn decode<I: Read + ReadBytesExt>(mut input: I) -> DangoPicture {
+    pub fn decode<I: Read + ReadBytesExt>(mut input: I) -> Result<DangoPicture, Error> {
         let mut magic = [0u8; 8];
         input.read_exact(&mut magic).unwrap();
 
+        if magic != *b"dangoimg" {
+            return Err(Error::InvalidIdentifier(magic))
+        }
+
         let header = Header {
             magic,
             width: input.read_u32::<LE>().unwrap(),
@@ -58,11 +63,32 @@ impl DangoPicture {
         }
 
         let preprocessed_bitmap = decompress(&mut input, &compression_info);
+
         let bitmap = line_diff(header.width, header.height, &preprocessed_bitmap);
 
+        Ok(DangoPicture {
+            header,
+            bitmap
+        })
+    }
+
+    pub fn from_raw(width: u32, height: u32, bitmap: &[u8]) -> Self {
+        let header = Header {
+            width,
+            height,
+
+            ..Default::default()
+        };
+
         DangoPicture {
             header,
-            bitmap
+            bitmap: bitmap.into(),
         }
     }
 }
+
+#[derive(Error, Debug)]
+pub enum Error {
+    #[error("incorrect identifier, got {}", 0)]
+    InvalidIdentifier([u8; 8]),
+}