Refactored code, cleaned up CZ3 implementation

Cargo.toml

@@ -4,5 +4,6 @@ version = "0.1.0"
 edition = "2021"
 
 [dependencies]
+byteorder = "1.5.0"
 image = "0.25"
 thiserror = "1.0.59"

src/cz_common.rs

@@ -1,5 +1,8 @@
 //! Shared types and traits between CZ# files
 
+use std::io::{self, Cursor, Read};
+
+use byteorder::{LittleEndian, ReadBytesExt};
 use image::Rgba;
 use thiserror::Error;
 
@@ -8,8 +11,11 @@ pub enum CzError {
     #[error("Version in header does not match expected version")]
     VersionMismatch,
 
-    #[error("Format of supplied file is incorrect")]
+    #[error("Format of supplied file is incorrect; expected {} bytes, got {}", expected, got)]
-    InvalidFormat,
+    InvalidFormat{expected: usize, got: usize},
+
+    #[error("Failed to read input")]
+    ReadError(#[from] io::Error),
 }
 
 pub trait CzHeader {
@@ -28,7 +34,7 @@ pub trait CzHeader {
 
 /// The common first part of a header of a CZ# file
 #[derive(Debug, Clone, Copy)]
-pub(crate) struct CommonHeader {
+pub struct CommonHeader {
     /// Format version from the magic bytes, (eg. CZ3, CZ4)
     pub version: u8,
 
@@ -46,14 +52,17 @@ pub(crate) struct CommonHeader {
 }
 
 impl CommonHeader {
-    pub fn new(bytes: &[u8]) -> Self {
+    pub fn new(bytes: &mut Cursor<&[u8]>) -> Result<Self, io::Error> {
-        Self {
+        let mut magic = [0u8; 4];
-            version: bytes[2] - b'0',
+        bytes.read_exact(&mut magic)?;
-            length: u32::from_le_bytes(bytes[4..8].try_into().unwrap()),
+
-            width: u16::from_le_bytes(bytes[8..10].try_into().unwrap()),
+        Ok(Self {
-            height: u16::from_le_bytes(bytes[10..12].try_into().unwrap()),
+            version: magic[2] - b'0',
-            depth: u16::from_le_bytes(bytes[12..14].try_into().unwrap()),
+            length: bytes.read_u32::<LittleEndian>()?,
-        }
+            width: bytes.read_u16::<LittleEndian>()?,
+            height: bytes.read_u16::<LittleEndian>()?,
+            depth: bytes.read_u16::<LittleEndian>()?,
+        })
     }
 }
 
@@ -73,14 +82,125 @@ pub trait CzImage {
     fn into_bitmap(self) -> Vec<u8>;
 }
 
-pub fn parse_colormap(input: &[u8], num_colors: usize) -> Vec<Rgba<u8>> {
+pub fn parse_colormap(input: &mut Cursor<&[u8]>, num_colors: usize) -> Result<Vec<Rgba<u8>>, CzError> {
     let mut colormap = Vec::with_capacity(num_colors);
+    let mut rgba_buf = [0u8; 4];
 
-    let input_iter = input.windows(4).step_by(4).take(num_colors);
+    for _ in 0..num_colors {
-
+        input.read_exact(&mut rgba_buf)?;
-    for color in input_iter {
+        colormap.push(Rgba(rgba_buf));
-        colormap.push(Rgba(color.try_into().unwrap()));
     }
 
-    colormap
+    Ok(colormap)
+}
+
+pub struct ChunkInfo {
+    pub size_compressed: usize,
+    pub size_raw: usize,
+}
+
+pub struct CompressionInfo {
+    pub chunk_count: usize,
+    pub total_size_compressed: usize,
+    pub total_size_raw: usize,
+    pub chunks: Vec<ChunkInfo>,
+
+    /// Length of the compression chunk info
+    pub length: usize,
+}
+
+/// Get info about the compression chunks
+pub fn parse_chunk_info(bytes: &mut Cursor<&[u8]>) -> Result<CompressionInfo, CzError> {
+    let parts_count = bytes.read_u32::<LittleEndian>()?;
+
+    dbg!(parts_count);
+    let mut part_sizes = vec![];
+    let mut total_size = 0;
+    let mut total_size_raw = 0;
+
+    for _ in 0..parts_count {
+        let compressed_size = bytes.read_u32::<LittleEndian>()? * 2;
+        total_size += compressed_size;
+
+        let raw_size = bytes.read_u32::<LittleEndian>()? * 4;
+        total_size_raw += raw_size;
+
+        part_sizes.push(ChunkInfo {
+            size_compressed: compressed_size as usize,
+            size_raw: raw_size as usize,
+        });
+    }
+
+    Ok(CompressionInfo {
+        chunk_count: parts_count as usize,
+        total_size_compressed: total_size as usize,
+        total_size_raw: total_size_raw as usize,
+        chunks: part_sizes,
+        length: bytes.position() as usize,
+    })
+}
+
+
+pub fn decompress(input: &mut Cursor<&[u8]>, chunk_info: CompressionInfo) -> Result<Vec<u8>, CzError> {
+    let mut m_dst = 0;
+    let mut bitmap = vec![0; chunk_info.total_size_raw];
+    for chunk in chunk_info.chunks {
+        let mut part = vec![0u8; chunk.size_compressed];
+        input.read_exact(&mut part)?;
+
+        for j in (0..part.len()).step_by(2) {
+            let ctl = part[j + 1];
+
+            if ctl == 0 {
+                bitmap[m_dst] = part[j];
+                m_dst += 1;
+            } else {
+                m_dst += copy_range(&mut bitmap, &part, get_offset(&part, j), m_dst);
+            }
+        }
+    }
+
+    Ok(bitmap)
+}
+
+fn get_offset(input: &[u8], src: usize) -> usize {
+    (((input[src] as usize) | (input[src + 1] as usize) << 8) - 0x101) * 2
+}
+
+fn copy_range(bitmap: &mut Vec<u8>, input: &[u8], src: usize, dst: usize) -> usize {
+    let mut dst = dst;
+    let start_pos = dst;
+
+    if input[src + 1] == 0 {
+        bitmap[dst] = input[src];
+        dst += 1;
+    } else if get_offset(input, src) == src {
+        bitmap[dst] = 0;
+        dst += 1;
+    } else {
+        dst += copy_range(bitmap, input, get_offset(input, src), dst);
+    }
+
+    if input[src + 3] == 0 {
+        bitmap[dst] = input[src + 2];
+        dst += 1;
+    } else if get_offset(input, src + 2) == src {
+        bitmap[dst] = bitmap[start_pos];
+        dst += 1;
+    } else {
+        bitmap[dst] = copy_one(input, get_offset(input, src + 2));
+        dst += 1;
+    }
+
+    dst - start_pos
+}
+
+fn copy_one(input: &[u8], src: usize) -> u8 {
+    if input[src + 1] == 0 {
+        input[src]
+    } else if get_offset(input, src) == src {
+        0
+    } else {
+        copy_one(input, get_offset(input, src))
+    }
 }

src/formats/cz0.rs

@@ -1,3 +1,5 @@
+use std::io::Cursor;
+
 use crate::cz_common::{CommonHeader, CzError, CzHeader, CzImage};
 
 #[derive(Debug)]
@@ -32,7 +34,8 @@ pub struct Cz0Image {
 
 impl CzHeader for Cz0Header {
     fn new(bytes: &[u8]) -> Result<Self, CzError> {
-        let common = CommonHeader::new(bytes);
+        let mut input = Cursor::new(bytes);
+        let common = CommonHeader::new(&mut input)?;
 
         if common.version != 0 {
             return Err(CzError::VersionMismatch)
@@ -40,7 +43,7 @@ impl CzHeader for Cz0Header {
 
         let mut offset_width = None;
         let mut offset_height = None;
-        if common.length < 28 {
+        if common.length > 28 {
             offset_width = Some(u16::from_le_bytes(bytes[28..30].try_into().unwrap()));
             offset_height = Some(u16::from_le_bytes(bytes[30..32].try_into().unwrap()));
         }

src/formats/cz1.rs

@@ -1,124 +1,49 @@
+use std::io::Cursor;
 use image::{ImageFormat, Rgba};
-use crate::cz_common::{parse_colormap, CommonHeader, CzError, CzHeader, CzImage};
+use crate::cz_common::{decompress, parse_chunk_info, parse_colormap, CommonHeader, CzError, CzImage};
-
-#[derive(Debug, Clone)]
-pub struct Cz1Header {
-    /// Common CZ# header
-    common: CommonHeader,
-}
 
 #[derive(Debug, Clone)]
 pub struct Cz1Image {
-    header: Cz1Header,
+    header: CommonHeader,
     bitmap: Vec<u8>,
-    palette: Option<Vec<Rgba<u8>>>,
+    palette: Vec<Rgba<u8>>,
-}
-
-impl CzHeader for Cz1Header {
-    fn new(bytes: &[u8]) -> Result<Self, CzError> {
-        let common = CommonHeader::new(bytes);
-
-        /*
-        if common.version != 1 {
-            return Err(CzError::VersionMismatch)
-        }
-        */
-
-        Ok(Self {
-            common,
-        })
-    }
-
-    fn version(&self) -> u8 {
-        self.common.version
-    }
-
-    fn header_length(&self) -> usize {
-        self.common.length as usize
-    }
-
-    fn width(&self) -> u16 {
-        self.common.width
-    }
-
-    fn height(&self) -> u16 {
-        self.common.height
-    }
-
-    fn depth(&self) -> u16 {
-        self.common.depth
-    }
 }
 
 impl CzImage for Cz1Image {
-    type Header = Cz1Header;
+    type Header = CommonHeader;
 
     fn decode(bytes: &[u8]) -> Result<Self, CzError> {
-        let mut position = 0;
+        let mut input = Cursor::new(bytes);
 
         // Get the header from the input
-        let header = Cz1Header::new(bytes)?;
+        let header = CommonHeader::new(&mut input).unwrap();
-        position += header.header_length();
 
-        // The color palette
+        // The color palette, gotten for 8 and 4 BPP images
         let mut palette = None;
-        if header.common.depth == 8 {
+        if header.depth == 8 || header.depth == 4 {
-            let temp_palette = parse_colormap(&bytes[position..], 0x100);
+            palette = Some(parse_colormap(&mut input, 1 << header.depth)?);
-            position += temp_palette.len() * 4;
-            palette = Some(temp_palette);
         }
 
-        let parts_count = u32::from_le_bytes(bytes[position..position + 4].try_into().unwrap());
+        let chunk_info = parse_chunk_info(&mut input)?;
-        position += 4;
-        dbg!(parts_count);
-        let mut part_sizes = vec![0; parts_count as usize];
-        let mut total_size = 0;
-        let mut decompressed_size = 0;
 
-        for size in &mut part_sizes {
+        if chunk_info.total_size_compressed as usize > bytes.len() {
-            let part_size = u32::from_le_bytes(bytes[position..position + 4].try_into().unwrap()) * 2;
+            return Err(CzError::InvalidFormat{
-            *size = part_size;
+                expected: chunk_info.total_size_compressed,
-            total_size += part_size;
+                got: bytes.len(),
-            position += 4;
+            })
-
-            let orig_size = u32::from_le_bytes(bytes[position..position + 4].try_into().unwrap()) * 4;
-            decompressed_size += orig_size;
-            position += 4;
-
-            dbg!(part_size, orig_size);
         }
 
-        if position + total_size as usize > bytes.len() {
+        let mut bitmap = decompress(&mut input, chunk_info).unwrap();
-            return Err(CzError::InvalidFormat)
-        }
 
-        let mut m_dst = 0;
+        // Apply the palette if it exists
-        let mut bitmap = vec![0; decompressed_size as usize];
+        if let Some(pal) = &palette {
-
+            apply_palette(&mut bitmap, pal);
-        for size in part_sizes {
-            let part = &bytes[position..position + size as usize];
-            position += size as usize;
-
-            for j in (0..part.len()).step_by(2) {
-                let ctl = part[j + 1];
-
-                if ctl == 0 {
-                    bitmap[m_dst] = part[j];
-                    m_dst += 1;
-                } else {
-                    m_dst += copy_range(&mut bitmap, part, get_offset(part, j), m_dst);
-                }
-            }
-        }
-
-        if let Some(temp_palette) = &palette {
-            apply_palette(&mut bitmap, &temp_palette);
         }
 
         let image = Self {
             header,
             bitmap,
-            palette,
+            palette: palette.unwrap(),
         };
 
         Ok(image)
@@ -126,8 +51,8 @@ impl CzImage for Cz1Image {
 
     fn save_as_png(&self, name: &str) {
         let img = image::RgbaImage::from_raw(
-            self.header.common.width as u32,
+            self.header.width as u32,
-            self.header.common.height as u32,
+            self.header.height as u32,
             self.bitmap.clone()
         ).unwrap();
 
@@ -143,11 +68,7 @@ impl CzImage for Cz1Image {
     }
 }
 
-fn get_offset(input: &[u8], src: usize) -> usize {
+fn apply_palette(input: &mut Vec<u8>, palette: &[Rgba<u8>]) {
-    (((input[src] as usize) | (input[src + 1] as usize) << 8) - 0x101) * 2
-}
-
-fn apply_palette(input: &mut Vec<u8>, palette: &Vec<Rgba<u8>>) {
     let mut output_map = Vec::new();
 
     for byte in input.iter() {
@@ -157,41 +78,3 @@ fn apply_palette(input: &mut Vec<u8>, palette: &Vec<Rgba<u8>>) {
 
     *input = output_map
 }
-
-fn copy_range(bitmap: &mut Vec<u8>, input: &[u8], src: usize, dst: usize) -> usize {
-    let mut dst = dst;
-    let start_pos = dst;
-
-    if input[src + 1] == 0 {
-        bitmap[dst] = input[src];
-        dst += 1;
-    } else if get_offset(input, src) == src {
-        bitmap[dst] = 0;
-        dst += 1;
-    } else {
-        dst += copy_range(bitmap, input, get_offset(input, src), dst);
-    }
-
-    if input[src + 3] == 0 {
-        bitmap[dst] = input[src + 2];
-        dst += 1;
-    } else if get_offset(input, src + 2) == src {
-        bitmap[dst] = bitmap[start_pos];
-        dst += 1;
-    } else {
-        bitmap[dst] = copy_one(input, get_offset(input, src + 2));
-        dst += 1;
-    }
-
-    dst - start_pos
-}
-
-fn copy_one(input: &[u8], src: usize) -> u8 {
-    if input[src + 1] == 0 {
-        input[src]
-    } else if get_offset(input, src) == src {
-        0
-    } else {
-        copy_one(input, get_offset(input, src))
-    }
-}

src/formats/cz3.rs

@@ -1,23 +1,101 @@
+use std::io::Cursor;
+
 use image::ImageFormat;
 
-use crate::cz_common::{CzError, CzHeader, CzImage};
+use crate::cz_common::{decompress, parse_chunk_info, CommonHeader, CzError, CzHeader, CzImage};
 
-use super::cz1::Cz1Header;
+#[derive(Debug, Clone, Copy)]
+pub struct Cz3Header {
+    /// Common CZ# header
+    common: CommonHeader,
 
-#[derive(Debug)]
+    /// Width of cropped image area
+    pub crop_width: u16,
+
+    /// Height of cropped image area
+    pub crop_height: u16,
+
+    /// Bounding box width
+    pub bounds_width: u16,
+
+    /// Bounding box height
+    pub bounds_height: u16,
+
+    /// Offset width
+    pub offset_width: Option<u16>,
+
+    /// Offset height
+    pub offset_height: Option<u16>,
+}
+
+impl CzHeader for Cz3Header {
+    fn new(bytes: &[u8]) -> Result<Self, CzError> where Self: Sized {
+        let mut input = Cursor::new(bytes);
+        let common = CommonHeader::new(&mut input)?;
+
+        if common.version != 3 {
+            return Err(CzError::VersionMismatch)
+        }
+
+        let mut offset_width = None;
+        let mut offset_height = None;
+        if common.length > 28 {
+            offset_width = Some(u16::from_le_bytes(bytes[28..30].try_into().unwrap()));
+            offset_height = Some(u16::from_le_bytes(bytes[30..32].try_into().unwrap()));
+        }
+
+        Ok(Self {
+            common,
+
+            crop_width: u16::from_le_bytes(bytes[20..22].try_into().unwrap()),
+            crop_height: u16::from_le_bytes(bytes[22..24].try_into().unwrap()),
+
+            bounds_width: u16::from_le_bytes(bytes[24..26].try_into().unwrap()),
+            bounds_height: u16::from_le_bytes(bytes[26..28].try_into().unwrap()),
+
+            offset_width,
+            offset_height,
+        })
+    }
+
+    fn version(&self) -> u8 {
+        self.common.version
+    }
+
+    fn header_length(&self) -> usize {
+        self.common.length as usize
+    }
+
+    fn width(&self) -> u16 {
+        self.common.width
+    }
+
+    fn height(&self) -> u16 {
+        self.common.height
+    }
+
+    fn depth(&self) -> u16 {
+        self.common.depth
+    }
+}
+
+#[derive(Debug, Clone)]
 pub struct Cz3Image {
-    header: Cz1Header,
+    header: Cz3Header,
     bitmap: Vec<u8>,
 }
 
 impl CzImage for Cz3Image {
-    type Header = Cz1Header;
+    type Header = Cz3Header;
 
     fn decode(bytes: &[u8]) -> Result<Self, CzError> {
-        let cz1_image = crate::formats::cz1::Cz1Image::decode(bytes)?;
+        let mut input = Cursor::new(bytes);
+        let header = Cz3Header::new(bytes)?;
+        input.set_position(header.header_length() as u64);
 
-        let header = cz1_image.header().clone();
+        let block_info = parse_chunk_info(&mut input)?;
-        let mut bitmap = cz1_image.into_bitmap();
+
+        let mut bitmap = decompress(&mut input, block_info)?;
 
         dbg!(bitmap.len());