Fixed issues with changing color format

src/header.rs

@@ -93,6 +93,12 @@ pub enum ColorFormat {
 
     /// RGB, 8 bits per channel
     Rgb8 = 1,
+
+    /// Grayscale with alpha, 8 bits per channel
+    GrayA8 = 2,
+
+    /// Grayscale, 8 bits per channel
+    Gray8 = 3,
 }
 
 impl ColorFormat {
@@ -101,8 +107,10 @@ impl ColorFormat {
     /// Ex. `Rgba8` has `8bpc`
     pub fn bpc(&self) -> u8 {
         match self {
-            ColorFormat::Rgba8 => 8,
-            ColorFormat::Rgb8 => 8,
+            Self::Rgba8 => 8,
+            Self::Rgb8 => 8,
+            Self::GrayA8 => 8,
+            Self::Gray8 => 8,
         }
     }
 
@@ -111,8 +119,10 @@ impl ColorFormat {
     /// Ex. `Rgba8` has `32bpp`
     pub fn bpp(&self) -> u16 {
         match self {
-            ColorFormat::Rgba8 => 32,
-            ColorFormat::Rgb8 => 24,
+            Self::Rgba8 => 32,
+            Self::Rgb8 => 24,
+            Self::GrayA8 => 16,
+            Self::Gray8 => 8,
         }
     }
 
@@ -121,23 +131,30 @@ impl ColorFormat {
     /// Ex. `Rgba8` has `4` channels
     pub fn channels(&self) -> u16 {
         match self {
-            ColorFormat::Rgba8 => 4,
-            ColorFormat::Rgb8 => 3,
+            Self::Rgba8 => 4,
+            Self::Rgb8 => 3,
+            Self::GrayA8 => 2,
+            Self::Gray8 => 1,
         }
     }
 
     /// The channel in which alpha is contained, or [`None`] if there is none.
     ///
     /// Ex. `Rgba8`'s 3rd channel is alpha
-    pub fn alpha_channel(&self) -> Option<u8> {
+    pub fn alpha_channel(&self) -> Option<usize> {
         match self {
-            ColorFormat::Rgba8 => Some(4),
-            ColorFormat::Rgb8 => None,
+            Self::Rgba8 => Some(3),
+            Self::Rgb8 => None,
+            Self::GrayA8 => Some(1),
+            Self::Gray8 => None,
         }
     }
 
-    pub fn pixel_byte_count(&self) -> u16 {
-        self.bpp() / 8
+    /// Pixel Byte Count, The number of bytes per pixel.
+    ///
+    /// Convenience method over [`Self::bpp`]
+    pub fn pbc(&self) -> usize {
+        (self.bpp() / 8).into()
     }
 }
 
@@ -148,6 +165,8 @@ impl TryFrom<u8> for ColorFormat {
         Ok(match value {
             0 => Self::Rgba8,
             1 => Self::Rgb8,
+            2 => Self::GrayA8,
+            3 => Self::Gray8,
             v => return Err(format!("invalid color format {v}")),
         })
     }

src/operations.rs

@@ -2,10 +2,10 @@ use crate::ColorFormat;
 use rayon::prelude::*;
 
 pub fn sub_rows(width: u32, height: u32, color_format: ColorFormat, input: &[u8]) -> Vec<u8> {
-    let mut data = Vec::with_capacity(width as usize * (color_format.bpp() / 8) as usize);
+    let mut data = Vec::with_capacity(width as usize * color_format.pbc());
 
     let block_height = f32::ceil(height as f32 / 3.0) as u32;
-    let line_byte_count = (width * color_format.pixel_byte_count() as u32) as usize;
+    let line_byte_count = (width * color_format.pbc() as u32) as usize;
 
     let mut curr_line: Vec<u8>;
     let mut prev_line: Vec<u8> = Vec::new();
@@ -30,22 +30,22 @@ pub fn sub_rows(width: u32, height: u32, color_format: ColorFormat, input: &[u8]
     }
 
     if color_format.alpha_channel().is_some() {
-        let (pixels, alpha): (Vec<[u8; 3]>, Vec<u8>) =
-            data.chunks(4)
+        let (pixels, alpha): (Vec<&[u8]>, Vec<u8>) =
+            data.chunks(color_format.pbc())
                 .map(|i| (
-                    [i[0], i[1], i[2]],
-                    i[3]
+                    &i[..color_format.pbc() - 1],
+                    i[color_format.alpha_channel().unwrap()]
                 ))
                 .unzip();
 
-        pixels.into_iter().flatten().chain(alpha).collect()
+        pixels.into_iter().flatten().copied().chain(alpha).collect()
     } else {
         data
     }
 }
 
 pub fn add_rows(width: u32, height: u32, color_format: ColorFormat, data: &[u8]) -> Vec<u8> {
-    let mut output_buf = Vec::with_capacity((width * height * color_format.channels() as u32) as usize);
+    let mut output_buf = Vec::with_capacity((width * height * color_format.pbc() as u32) as usize);
 
     let block_height = f32::ceil(height as f32 / 3.0) as u32;
 
@@ -53,11 +53,12 @@ pub fn add_rows(width: u32, height: u32, color_format: ColorFormat, data: &[u8])
     let mut prev_line = Vec::new();
 
     let mut rgb_index = 0;
-    let mut alpha_index = (width * height * (color_format.channels() as u32 - 1)) as usize;
+    let mut alpha_index = (width * height * (color_format.pbc() - 1) as u32) as usize;
     for y in 0..height {
         curr_line = if color_format.alpha_channel().is_some() {
-            data[rgb_index..rgb_index + width as usize * 3]
-                .chunks(3)
+            // Interleave the offset alpha into the RGB bytes
+            data[rgb_index..rgb_index + width as usize * (color_format.pbc() - 1)]
+                .chunks(color_format.pbc() - 1)
                 .zip(data[alpha_index..alpha_index + width as usize].into_iter())
                 .flat_map(|(a, b)| {
                     a.into_iter().chain(vec![b])
@@ -65,7 +66,7 @@ pub fn add_rows(width: u32, height: u32, color_format: ColorFormat, data: &[u8])
                 .copied()
                 .collect()
         } else {
-            data[rgb_index..rgb_index + width as usize * 3].to_vec()
+            data[rgb_index..rgb_index + width as usize * color_format.pbc()].to_vec()
         };
 
         if y % block_height != 0 {
@@ -81,7 +82,11 @@ pub fn add_rows(width: u32, height: u32, color_format: ColorFormat, data: &[u8])
         output_buf.extend_from_slice(&curr_line);
 
         prev_line.clone_from(&curr_line);
-        rgb_index += width as usize * 3;
+        rgb_index += if color_format.alpha_channel().is_some() {
+            width as usize * (color_format.pbc() - 1)
+        } else {
+            width as usize * color_format.pbc()
+        };
         alpha_index += width as usize;
     }
 

src/picture.rs

@@ -156,9 +156,6 @@ impl SquishyPicture {
             },
         };
 
-        let mut inspection_file = File::create("raw_data").unwrap();
-        inspection_file.write_all(&modified_data).unwrap();
-
         // Compress the final image data using the basic LZW scheme
         let (compressed_data, compression_info) = compress(modified_data)?;
 
@@ -192,7 +189,12 @@ impl SquishyPicture {
         let bitmap = match header.compression_type {
             CompressionType::None => pre_bitmap,
             CompressionType::Lossless => {
-                add_rows(header.width, header.height, header.color_format, &pre_bitmap)
+                add_rows(
+                    header.width,
+                    header.height,
+                    header.color_format,
+                    &pre_bitmap
+                )
             },
             CompressionType::LossyDct => {
                 dct_decompress(