Add ICO BMP encoder

src/codecs/ico/encoder.rs

@@ -3,7 +3,10 @@ use std::borrow::Cow;
 use std::io::{self, Write};
 
 use crate::codecs::png::PngEncoder;
-use crate::error::{EncodingError, ImageError, ImageResult, ParameterError, ParameterErrorKind};
+use crate::error::{
+    EncodingError, ImageError, ImageResult, ParameterError, ParameterErrorKind, UnsupportedError,
+};
+use crate::utils::vec_try_with_capacity;
 use crate::{ExtendedColorType, ImageEncoder, ImageFormat};
 
 // Enum value indicating an ICO image (as opposed to a CUR image):
@@ -12,6 +15,8 @@ const ICO_IMAGE_TYPE: u16 = 1;
 const ICO_ICONDIR_SIZE: u32 = 6;
 // The length of an ICO file DIRENTRY structure, in bytes:
 const ICO_DIRENTRY_SIZE: u32 = 16;
+// The length of a BITMAPINFOHEADER structure, in bytes:
+const BITMAPINFOHEADER_SIZE: u32 = 40;
 
 /// ICO encoder
 pub struct IcoEncoder<W: Write> {
@@ -80,6 +85,125 @@ impl<'a> IcoFrame<'a> {
         let frame = Self::with_encoded(image_data, width, height, color_type)?;
         Ok(frame)
     }
+
+    /// Construct a new `IcoFrame` by encoding `buf` as a BMP
+    ///
+    /// The `width` and `height` must be between 1 and 256 (inclusive)
+    pub fn as_bmp(
+        buf: &[u8],
+        width: u32,
+        height: u32,
+        color_type: ExtendedColorType,
+    ) -> ImageResult<Self> {
+        if width == 0 || height == 0 {
+            return Err(ImageError::Parameter(ParameterError::from_kind(
+                ParameterErrorKind::Generic(format!(
+                    "the image width and height must be non-zero, but width {width} and height {height} were provided",
+                )),
+            )));
+        }
+        // same soft limit as BMP decoder
+        if width > u16::MAX as u32 || height > u16::MAX as u32 {
+            return Err(ImageError::Parameter(ParameterError::from_kind(
+                ParameterErrorKind::Generic("Dimensions too big.".into()),
+            )));
+        }
+
+        let mask_row_bytes = width.div_ceil(8);
+        let mask_row_bytes_padded = mask_row_bytes.next_multiple_of(4);
+
+        // Check to see that the output buffer is not going to be larger than the 4 GiB that ICO can support.
+        let output_size = BITMAPINFOHEADER_SIZE as u64 // header
+            + (width as u64 * height as u64 * 4) // XOR mask (32 bits per pixel)
+            + (height as u64 * mask_row_bytes_padded as u64); // AND mask (1 bit per pixel, padded to 32 bits per row)
+        if output_size > u32::MAX as u64 {
+            return Err(ImageError::Parameter(ParameterError::from_kind(
+                ParameterErrorKind::Generic(format!(
+                    "the encoded image data must be at most 4 GiB, but the provided image would encode to {output_size} bytes",
+                )),
+            )));
+        }
+
+        let mut image_data: Vec<u8> = vec_try_with_capacity(output_size as usize)?;
+
+        // https://learn.microsoft.com/en-us/previous-versions/ms997538(v=msdn.10)?redirectedfrom=MSDN
+        // > Only the following members are used: biSize, biWidth, biHeight, biPlanes, biBitCount,
+        // > biSizeImage. All other members must be 0.
+        // https://learn.microsoft.com/en-us/windows/win32/api/wingdi/ns-wingdi-bitmapinfoheader
+        image_data.write_u32::<LittleEndian>(BITMAPINFOHEADER_SIZE)?; // biSize
+        image_data.write_i32::<LittleEndian>(width as i32)?; // biWidth
+        image_data.write_i32::<LittleEndian>(height as i32 * 2)?; // biHeight
+        image_data.write_u16::<LittleEndian>(1)?; // biPlanes must be 1
+        image_data.write_u16::<LittleEndian>(32)?; // biBitCount (we only support 32-bit for now)
+        image_data.write_u32::<LittleEndian>(0)?; // biCompression must be 0 (BI_RGB) for uncompressed
+        image_data.write_u32::<LittleEndian>(output_size as u32 - BITMAPINFOHEADER_SIZE)?; // biSizeImage
+        image_data.extend_from_slice(&[0; 16]); // remaining unused fields
+        debug_assert_eq!(image_data.len(), BITMAPINFOHEADER_SIZE as usize);
+
+        match color_type {
+            ExtendedColorType::Rgba8 => {
+                let pixels = buf.as_chunks::<4>().0;
+                // BMP format requires the alpha channel to be stored as BGRA, so we need to convert it from RGBA.
+                for row in pixels.chunks_exact(width as usize).rev() {
+                    for [r, g, b, a] in row {
+                        image_data.extend_from_slice(&[*b, *g, *r, *a]);
+                    }
+                }
+
+                // AND mask
+                // For 32-bit BMP the AND mask is typically ignored, but it should be provided nonetheless.
+                // We use a threshold of 128 for the alpha channel to determine whether a pixel is transparent or opaque in the AND mask.
+                for row in pixels.chunks_exact(width as usize).rev() {
+                    // AND mask is 1 bpp, so 8 pixels go into 1 mask byte.
+                    for chunks in row.chunks(8) {
+                        let mut mask_byte: u8 = 0;
+                        for ([_, _, _, a], bit_pos) in chunks.iter().zip((0..8).rev()) {
+                            mask_byte |= u8::from(*a < 128) << bit_pos;
+                        }
+                        image_data.push(mask_byte);
+                    }
+
+                    // Each row of the AND mask must be padded to a multiple of 4 bytes.
+                    image_data.extend(std::iter::repeat_n(
+                        0u8,
+                        (mask_row_bytes_padded - mask_row_bytes) as usize,
+                    ));
+                }
+            }
+            ExtendedColorType::Rgb8 => {
+                let pixels = buf.as_chunks::<3>().0;
+                // ICO doesn't allow 24-bit RGB, so we have to use the 0RGB 32-bit format.
+                // This is the same as the 32-bit RGBA format but with the alpha channel set to 0 for all pixels.
+                for row in pixels.chunks_exact(width as usize).rev() {
+                    for [r, g, b] in row {
+                        image_data.extend_from_slice(&[*b, *g, *r, 0]);
+                    }
+                }
+
+                // AND mask
+                // Since there is no transparency, we can set all pixels to 0 (opaque) in the AND mask.
+                image_data.extend(std::iter::repeat_n(
+                    0u8,
+                    (height * mask_row_bytes_padded) as usize,
+                ));
+            }
+            _ => {
+                return Err(ImageError::Unsupported(
+                    UnsupportedError::from_format_and_kind(
+                        ImageFormat::Ico.into(),
+                        crate::error::UnsupportedErrorKind::Color(color_type),
+                    ),
+                ));
+            }
+        }
+
+        // verify that the output size was calculated correctly.
+        assert_eq!(image_data.len() as u64, output_size);
+
+        // color type has to be Rgba8, because the `bit_per_pixel` field in the ICO directory entry has to be set to 32.
+        let frame = Self::with_encoded(image_data, width, height, ExtendedColorType::Rgba8)?;
+        Ok(frame)
+    }
 }
 
 impl<W: Write> IcoEncoder<W> {
@@ -210,6 +334,8 @@ fn write_direntry<W: Write>(
 
 #[cfg(test)]
 mod test {
+    use crate::{DynamicImage, ImageBuffer, PixelWithColorType, Rgb, Rgba, RgbaImage};
+
     use super::*;
 
     // Test that the encoder allows image where all frames have offsets < 4GiB
@@ -235,4 +361,55 @@ mod test {
         let res = encoder.encode_images(&frames);
         assert!(res.is_err());
     }
+
+    fn encode_bmp_decode<P: PixelWithColorType<Subpixel = u8>>(
+        image: &ImageBuffer<P, Vec<P::Subpixel>>,
+    ) -> DynamicImage {
+        let frame = IcoFrame::as_bmp(
+            image.subpixels(),
+            image.width(),
+            image.height(),
+            P::COLOR_TYPE,
+        )
+        .unwrap();
+
+        let mut encoded_data: Vec<u8> = Vec::new();
+        IcoEncoder::new(&mut encoded_data)
+            .encode_images(&[frame])
+            .unwrap();
+
+        let mut reader =
+            crate::ImageReaderOptions::with_format(io::Cursor::new(encoded_data), ImageFormat::Ico);
+        reader.set_spec_compliance(crate::SpecCompliance::Strict);
+        reader.decode().unwrap()
+    }
+
+    #[test]
+    fn roundtrip() {
+        let rgba = RgbaImage::from_fn(32, 32, |x, y| {
+            Rgba([
+                x as u8 * 8,
+                y as u8 * 8,
+                0,
+                if y < 24 {
+                    255
+                } else {
+                    (x + y * 32 - 256 * 3) as u8
+                },
+            ])
+        });
+        let rgb = rgba.convert();
+        let round_rgba = encode_bmp_decode(&rgba);
+        let round_rgb = encode_bmp_decode::<Rgb<u8>>(&rgb);
+        assert_eq!(round_rgba.into_rgba8(), rgba);
+        assert_eq!(round_rgb.into_rgba8(), rgb.convert());
+    }
+
+    #[test]
+    fn roundtrip_awkward_size() {
+        // test a size that will need lots of padding in the AND mask to make sure the logic works correctly.
+        let image = RgbaImage::from_pixel(33, 5, Rgba([255, 0, 0, 255]));
+        let round_rgba = encode_bmp_decode(&image);
+        assert_eq!(round_rgba.into_rgba8(), image);
+    }
 }