image/codecs/bmp/
encoder.rs

1use byteorder_lite::{LittleEndian, WriteBytesExt};
2use std::io::{self, Write};
3
4use crate::error::{
5    EncodingError, ImageError, ImageFormatHint, ImageResult, ParameterError, ParameterErrorKind,
6    UnsupportedError, UnsupportedErrorKind,
7};
8use crate::{DynamicImage, ExtendedColorType, ImageEncoder, ImageFormat};
9
10const BITMAPFILEHEADER_SIZE: u32 = 14;
11const BITMAPINFOHEADER_SIZE: u32 = 40;
12const BITMAPV4HEADER_SIZE: u32 = 108;
13
14/// The representation of a BMP encoder.
15pub struct BmpEncoder<'a, W: 'a> {
16    writer: &'a mut W,
17}
18
19impl<'a, W: Write + 'a> BmpEncoder<'a, W> {
20    /// Create a new encoder that writes its output to ```w```.
21    pub fn new(w: &'a mut W) -> Self {
22        BmpEncoder { writer: w }
23    }
24
25    /// Encodes the image `image` that has dimensions `width` and `height` and `ExtendedColorType` `c`.
26    ///
27    /// # Panics
28    ///
29    /// Panics if `width * height * c.bytes_per_pixel() != image.len()`.
30    #[track_caller]
31    pub fn encode(
32        &mut self,
33        image: &[u8],
34        width: u32,
35        height: u32,
36        c: ExtendedColorType,
37    ) -> ImageResult<()> {
38        self.encode_with_palette(image, width, height, c, None)
39    }
40
41    /// Same as `encode`, but allow a palette to be passed in. The `palette` is ignored for color
42    /// types other than Luma/Luma-with-alpha.
43    ///
44    /// # Panics
45    ///
46    /// Panics if `width * height * c.bytes_per_pixel() != image.len()`.
47    #[track_caller]
48    pub fn encode_with_palette(
49        &mut self,
50        image: &[u8],
51        width: u32,
52        height: u32,
53        color_type: ExtendedColorType,
54        palette: Option<&[[u8; 3]]>,
55    ) -> ImageResult<()> {
56        if palette.is_some()
57            && color_type != ExtendedColorType::L8
58            && color_type != ExtendedColorType::La8
59        {
60            return Err(ImageError::Parameter(ParameterError::from_kind(
61                ParameterErrorKind::Generic(
62                    "Palette given which must only be used with L8 or La8 color types".to_string(),
63                ),
64            )));
65        }
66
67        let expected_buffer_len = color_type.buffer_size(width, height);
68        assert_eq!(
69            expected_buffer_len,
70            image.len() as u64,
71            "Invalid buffer length: expected {expected_buffer_len} got {} for {width}x{height} image",
72            image.len(),
73        );
74
75        let bmp_header_size = BITMAPFILEHEADER_SIZE;
76
77        let (dib_header_size, written_pixel_size, palette_color_count) =
78            written_pixel_info(color_type, palette)?;
79
80        let (padded_row, image_size) = width
81            .checked_mul(written_pixel_size)
82            // each row must be padded to a multiple of 4 bytes
83            .and_then(|v| v.checked_next_multiple_of(4))
84            .and_then(|v| {
85                let image_bytes = v.checked_mul(height)?;
86                Some((v, image_bytes))
87            })
88            .ok_or_else(|| {
89                ImageError::Parameter(ParameterError::from_kind(
90                    ParameterErrorKind::DimensionMismatch,
91                ))
92            })?;
93
94        let row_padding = padded_row - width * written_pixel_size;
95
96        // all palette colors are BGRA
97        let palette_size = palette_color_count.checked_mul(4).ok_or_else(|| {
98            ImageError::Encoding(EncodingError::new(
99                ImageFormatHint::Exact(ImageFormat::Bmp),
100                "calculated palette size larger than 2^32",
101            ))
102        })?;
103
104        let file_size = bmp_header_size
105            .checked_add(dib_header_size)
106            .and_then(|v| v.checked_add(palette_size))
107            .and_then(|v| v.checked_add(image_size))
108            .ok_or_else(|| {
109                ImageError::Encoding(EncodingError::new(
110                    ImageFormatHint::Exact(ImageFormat::Bmp),
111                    "calculated BMP header size larger than 2^32",
112                ))
113            })?;
114
115        let image_data_offset = bmp_header_size
116            .checked_add(dib_header_size)
117            .and_then(|v| v.checked_add(palette_size))
118            .ok_or_else(|| {
119                ImageError::Encoding(EncodingError::new(
120                    ImageFormatHint::Exact(ImageFormat::Bmp),
121                    "calculated BMP size larger than 2^32",
122                ))
123            })?;
124
125        // write BMP header
126        self.writer.write_u8(b'B')?;
127        self.writer.write_u8(b'M')?;
128        self.writer.write_u32::<LittleEndian>(file_size)?; // file size
129        self.writer.write_u16::<LittleEndian>(0)?; // reserved 1
130        self.writer.write_u16::<LittleEndian>(0)?; // reserved 2
131        self.writer.write_u32::<LittleEndian>(image_data_offset)?; // image data offset
132
133        // write DIB header
134        self.writer.write_u32::<LittleEndian>(dib_header_size)?;
135        self.writer.write_i32::<LittleEndian>(width as i32)?;
136        self.writer.write_i32::<LittleEndian>(height as i32)?;
137        self.writer.write_u16::<LittleEndian>(1)?; // color planes
138        self.writer
139            .write_u16::<LittleEndian>((written_pixel_size * 8) as u16)?; // bits per pixel
140        if dib_header_size >= BITMAPV4HEADER_SIZE {
141            // Assume BGRA32
142            self.writer.write_u32::<LittleEndian>(3)?; // compression method - bitfields
143        } else {
144            self.writer.write_u32::<LittleEndian>(0)?; // compression method - no compression
145        }
146        self.writer.write_u32::<LittleEndian>(image_size)?;
147        self.writer.write_i32::<LittleEndian>(0)?; // horizontal ppm
148        self.writer.write_i32::<LittleEndian>(0)?; // vertical ppm
149        self.writer.write_u32::<LittleEndian>(palette_color_count)?;
150        self.writer.write_u32::<LittleEndian>(0)?; // all colors are important
151        if dib_header_size >= BITMAPV4HEADER_SIZE {
152            // Assume BGRA32
153            self.writer.write_u32::<LittleEndian>(0xff << 16)?; // red mask
154            self.writer.write_u32::<LittleEndian>(0xff << 8)?; // green mask
155            self.writer.write_u32::<LittleEndian>(0xff)?; // blue mask
156            self.writer.write_u32::<LittleEndian>(0xff << 24)?; // alpha mask
157            self.writer.write_u32::<LittleEndian>(0x7352_4742)?; // colorspace - sRGB
158
159            // endpoints (3x3) and gamma (3)
160            for _ in 0..12 {
161                self.writer.write_u32::<LittleEndian>(0)?;
162            }
163        }
164
165        // write image data
166        match color_type {
167            ExtendedColorType::Rgb8 => self.encode_rgb(image, width, height, row_padding, 3)?,
168            ExtendedColorType::Rgba8 => self.encode_rgba(image, width, height, row_padding, 4)?,
169            ExtendedColorType::L8 => {
170                self.encode_gray(image, width, height, row_padding, 1, palette)?;
171            }
172            ExtendedColorType::La8 => {
173                self.encode_gray(image, width, height, row_padding, 2, palette)?;
174            }
175            _ => {
176                return Err(ImageError::Unsupported(
177                    UnsupportedError::from_format_and_kind(
178                        ImageFormat::Bmp.into(),
179                        UnsupportedErrorKind::Color(color_type),
180                    ),
181                ));
182            }
183        }
184
185        Ok(())
186    }
187
188    fn encode_rgb(
189        &mut self,
190        image: &[u8],
191        width: u32,
192        height: u32,
193        row_padding: u32,
194        bytes_per_pixel: u32,
195    ) -> io::Result<()> {
196        let width = width as usize;
197        let height = height as usize;
198        let x_stride = bytes_per_pixel as usize;
199        let y_stride = width * x_stride;
200        for row in (0..height).rev() {
201            // from the bottom up
202            let row_start = row * y_stride;
203            for px in image[row_start..][..y_stride].chunks_exact(x_stride) {
204                let r = px[0];
205                let g = px[1];
206                let b = px[2];
207                // written as BGR
208                self.writer.write_all(&[b, g, r])?;
209            }
210            self.write_row_pad(row_padding)?;
211        }
212
213        Ok(())
214    }
215
216    fn encode_rgba(
217        &mut self,
218        image: &[u8],
219        width: u32,
220        height: u32,
221        row_padding: u32,
222        bytes_per_pixel: u32,
223    ) -> io::Result<()> {
224        let width = width as usize;
225        let height = height as usize;
226        let x_stride = bytes_per_pixel as usize;
227        let y_stride = width * x_stride;
228        for row in (0..height).rev() {
229            // from the bottom up
230            let row_start = row * y_stride;
231            for px in image[row_start..][..y_stride].chunks_exact(x_stride) {
232                let r = px[0];
233                let g = px[1];
234                let b = px[2];
235                let a = px[3];
236                // written as BGRA
237                self.writer.write_all(&[b, g, r, a])?;
238            }
239            self.write_row_pad(row_padding)?;
240        }
241
242        Ok(())
243    }
244
245    fn encode_gray(
246        &mut self,
247        image: &[u8],
248        width: u32,
249        height: u32,
250        row_padding: u32,
251        bytes_per_pixel: u32,
252        palette: Option<&[[u8; 3]]>,
253    ) -> io::Result<()> {
254        // write grayscale palette
255        if let Some(palette) = palette {
256            for item in palette {
257                // each color is written as BGRA, where A is always 0
258                self.writer.write_all(&[item[2], item[1], item[0], 0])?;
259            }
260        } else {
261            for val in 0u8..=255 {
262                // each color is written as BGRA, where A is always 0 and since only grayscale is being written, B = G = R = index
263                self.writer.write_all(&[val, val, val, 0])?;
264            }
265        }
266
267        // write image data
268        let x_stride = bytes_per_pixel;
269        let y_stride = width * x_stride;
270        for row in (0..height).rev() {
271            // from the bottom up
272            let row_start = row * y_stride;
273
274            // color value is equal to the palette index
275            if x_stride == 1 {
276                // improve performance by writing the whole row at once
277                self.writer
278                    .write_all(&image[row_start as usize..][..y_stride as usize])?;
279            } else {
280                for col in 0..width {
281                    let pixel_start = (row_start + (col * x_stride)) as usize;
282                    self.writer.write_u8(image[pixel_start])?;
283                    // alpha is never written as it's not widely supported
284                }
285            }
286
287            self.write_row_pad(row_padding)?;
288        }
289
290        Ok(())
291    }
292
293    fn write_row_pad(&mut self, row_pad_size: u32) -> io::Result<()> {
294        for _ in 0..row_pad_size {
295            self.writer.write_u8(0)?;
296        }
297
298        Ok(())
299    }
300}
301
302impl<W: Write> ImageEncoder for BmpEncoder<'_, W> {
303    #[track_caller]
304    fn write_image(
305        mut self,
306        buf: &[u8],
307        width: u32,
308        height: u32,
309        color_type: ExtendedColorType,
310    ) -> ImageResult<()> {
311        self.encode(buf, width, height, color_type)
312    }
313
314    fn make_compatible_img(
315        &self,
316        _: crate::io::encoder::MethodSealedToImage,
317        img: &DynamicImage,
318    ) -> Option<DynamicImage> {
319        crate::io::encoder::dynimage_conversion_8bit(img)
320    }
321}
322
323/// Returns a tuple representing: (dib header size, written pixel size, palette color count).
324fn written_pixel_info(
325    c: ExtendedColorType,
326    palette: Option<&[[u8; 3]]>,
327) -> Result<(u32, u32, u32), ImageError> {
328    let (header, color_bytes, palette_count) = match c {
329        ExtendedColorType::Rgb8 => (BITMAPINFOHEADER_SIZE, 3, Some(0)),
330        ExtendedColorType::Rgba8 => (BITMAPV4HEADER_SIZE, 4, Some(0)),
331        ExtendedColorType::L8 => (
332            BITMAPINFOHEADER_SIZE,
333            1,
334            u32::try_from(palette.map(|p| p.len()).unwrap_or(256)).ok(),
335        ),
336        ExtendedColorType::La8 => (
337            BITMAPINFOHEADER_SIZE,
338            1,
339            u32::try_from(palette.map(|p| p.len()).unwrap_or(256)).ok(),
340        ),
341        _ => {
342            return Err(ImageError::Unsupported(
343                UnsupportedError::from_format_and_kind(
344                    ImageFormat::Bmp.into(),
345                    UnsupportedErrorKind::Color(c),
346                ),
347            ));
348        }
349    };
350
351    let palette_count = palette_count.ok_or_else(|| {
352        ImageError::Encoding(EncodingError::new(
353            ImageFormatHint::Exact(ImageFormat::Bmp),
354            "calculated palette size larger than 2^32",
355        ))
356    })?;
357
358    Ok((header, color_bytes, palette_count))
359}
360
361#[cfg(test)]
362mod tests {
363    use super::super::BmpDecoder;
364    use super::BmpEncoder;
365
366    use crate::ExtendedColorType;
367    use crate::ImageDecoder as _;
368    use std::io::Cursor;
369
370    fn round_trip_image(image: &[u8], width: u32, height: u32, c: ExtendedColorType) -> Vec<u8> {
371        let mut encoded_data = Vec::new();
372        {
373            let mut encoder = BmpEncoder::new(&mut encoded_data);
374            encoder
375                .encode(image, width, height, c)
376                .expect("could not encode image");
377        }
378
379        let decoder = BmpDecoder::new(Cursor::new(&encoded_data)).expect("failed to decode");
380
381        let mut buf = vec![0; decoder.total_bytes() as usize];
382        decoder.read_image(&mut buf).expect("failed to decode");
383        buf
384    }
385
386    #[test]
387    fn round_trip_single_pixel_rgb() {
388        let image = [255u8, 0, 0]; // single red pixel
389        let decoded = round_trip_image(&image, 1, 1, ExtendedColorType::Rgb8);
390        assert_eq!(3, decoded.len());
391        assert_eq!(255, decoded[0]);
392        assert_eq!(0, decoded[1]);
393        assert_eq!(0, decoded[2]);
394    }
395
396    #[test]
397    #[cfg(target_pointer_width = "64")]
398    fn huge_files_return_error() {
399        let mut encoded_data = Vec::new();
400        let image = vec![0u8; 3 * 40_000 * 40_000]; // 40_000x40_000 pixels, 3 bytes per pixel, allocated on the heap
401        let mut encoder = BmpEncoder::new(&mut encoded_data);
402        let result = encoder.encode(&image, 40_000, 40_000, ExtendedColorType::Rgb8);
403        assert!(result.is_err());
404    }
405
406    #[test]
407    fn round_trip_single_pixel_rgba() {
408        let image = [1, 2, 3, 4];
409        let decoded = round_trip_image(&image, 1, 1, ExtendedColorType::Rgba8);
410        assert_eq!(&decoded[..], &image[..]);
411    }
412
413    #[test]
414    fn round_trip_3px_rgb() {
415        let image = [0u8; 3 * 3 * 3]; // 3x3 pixels, 3 bytes per pixel
416        let _decoded = round_trip_image(&image, 3, 3, ExtendedColorType::Rgb8);
417    }
418
419    #[test]
420    fn round_trip_gray() {
421        let image = [0u8, 1, 2]; // 3 pixels
422        let decoded = round_trip_image(&image, 3, 1, ExtendedColorType::L8);
423        // should be read back as 3 RGB pixels
424        assert_eq!(9, decoded.len());
425        assert_eq!(0, decoded[0]);
426        assert_eq!(0, decoded[1]);
427        assert_eq!(0, decoded[2]);
428        assert_eq!(1, decoded[3]);
429        assert_eq!(1, decoded[4]);
430        assert_eq!(1, decoded[5]);
431        assert_eq!(2, decoded[6]);
432        assert_eq!(2, decoded[7]);
433        assert_eq!(2, decoded[8]);
434    }
435
436    #[test]
437    fn round_trip_graya() {
438        let image = [0u8, 0, 1, 0, 2, 0]; // 3 pixels, each with an alpha channel
439        let decoded = round_trip_image(&image, 1, 3, ExtendedColorType::La8);
440        // should be read back as 3 RGB pixels
441        assert_eq!(9, decoded.len());
442        assert_eq!(0, decoded[0]);
443        assert_eq!(0, decoded[1]);
444        assert_eq!(0, decoded[2]);
445        assert_eq!(1, decoded[3]);
446        assert_eq!(1, decoded[4]);
447        assert_eq!(1, decoded[5]);
448        assert_eq!(2, decoded[6]);
449        assert_eq!(2, decoded[7]);
450        assert_eq!(2, decoded[8]);
451    }
452
453    #[test]
454    fn regression_issue_2604() {
455        let mut image = vec![];
456        let mut encoder = BmpEncoder::new(&mut image);
457        encoder
458            .encode(&[], 1 << 31, 0, ExtendedColorType::Rgb8)
459            .unwrap_err();
460    }
461}