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};
7use crate::image::ImageEncoder;
8use crate::{ExtendedColorType, 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        c: ExtendedColorType,
54        palette: Option<&[[u8; 3]]>,
55    ) -> ImageResult<()> {
56        if palette.is_some() && c != ExtendedColorType::L8 && c != ExtendedColorType::La8 {
57            return Err(ImageError::IoError(io::Error::new(
58                io::ErrorKind::InvalidInput,
59                format!(
60                    "Unsupported color type {c:?} when using a non-empty palette. Supported types: Gray(8), GrayA(8)."
61                ),
62            )));
63        }
64
65        let expected_buffer_len = c.buffer_size(width, height);
66        assert_eq!(
67            expected_buffer_len,
68            image.len() as u64,
69            "Invalid buffer length: expected {expected_buffer_len} got {} for {width}x{height} image",
70            image.len(),
71        );
72
73        let bmp_header_size = BITMAPFILEHEADER_SIZE;
74
75        let (dib_header_size, written_pixel_size, palette_color_count) =
76            get_pixel_info(c, palette)?;
77        let row_pad_size = (4 - (width * written_pixel_size) % 4) % 4; // each row must be padded to a multiple of 4 bytes
78        let image_size = width
79            .checked_mul(height)
80            .and_then(|v| v.checked_mul(written_pixel_size))
81            .and_then(|v| v.checked_add(height * row_pad_size))
82            .ok_or_else(|| {
83                ImageError::Parameter(ParameterError::from_kind(
84                    ParameterErrorKind::DimensionMismatch,
85                ))
86            })?;
87        let palette_size = palette_color_count * 4; // all palette colors are BGRA
88        let file_size = bmp_header_size
89            .checked_add(dib_header_size)
90            .and_then(|v| v.checked_add(palette_size))
91            .and_then(|v| v.checked_add(image_size))
92            .ok_or_else(|| {
93                ImageError::Encoding(EncodingError::new(
94                    ImageFormatHint::Exact(ImageFormat::Bmp),
95                    "calculated BMP header size larger than 2^32",
96                ))
97            })?;
98
99        // write BMP header
100        self.writer.write_u8(b'B')?;
101        self.writer.write_u8(b'M')?;
102        self.writer.write_u32::<LittleEndian>(file_size)?; // file size
103        self.writer.write_u16::<LittleEndian>(0)?; // reserved 1
104        self.writer.write_u16::<LittleEndian>(0)?; // reserved 2
105        self.writer
106            .write_u32::<LittleEndian>(bmp_header_size + dib_header_size + palette_size)?; // image data offset
107
108        // write DIB header
109        self.writer.write_u32::<LittleEndian>(dib_header_size)?;
110        self.writer.write_i32::<LittleEndian>(width as i32)?;
111        self.writer.write_i32::<LittleEndian>(height as i32)?;
112        self.writer.write_u16::<LittleEndian>(1)?; // color planes
113        self.writer
114            .write_u16::<LittleEndian>((written_pixel_size * 8) as u16)?; // bits per pixel
115        if dib_header_size >= BITMAPV4HEADER_SIZE {
116            // Assume BGRA32
117            self.writer.write_u32::<LittleEndian>(3)?; // compression method - bitfields
118        } else {
119            self.writer.write_u32::<LittleEndian>(0)?; // compression method - no compression
120        }
121        self.writer.write_u32::<LittleEndian>(image_size)?;
122        self.writer.write_i32::<LittleEndian>(0)?; // horizontal ppm
123        self.writer.write_i32::<LittleEndian>(0)?; // vertical ppm
124        self.writer.write_u32::<LittleEndian>(palette_color_count)?;
125        self.writer.write_u32::<LittleEndian>(0)?; // all colors are important
126        if dib_header_size >= BITMAPV4HEADER_SIZE {
127            // Assume BGRA32
128            self.writer.write_u32::<LittleEndian>(0xff << 16)?; // red mask
129            self.writer.write_u32::<LittleEndian>(0xff << 8)?; // green mask
130            self.writer.write_u32::<LittleEndian>(0xff)?; // blue mask
131            self.writer.write_u32::<LittleEndian>(0xff << 24)?; // alpha mask
132            self.writer.write_u32::<LittleEndian>(0x7352_4742)?; // colorspace - sRGB
133
134            // endpoints (3x3) and gamma (3)
135            for _ in 0..12 {
136                self.writer.write_u32::<LittleEndian>(0)?;
137            }
138        }
139
140        // write image data
141        match c {
142            ExtendedColorType::Rgb8 => self.encode_rgb(image, width, height, row_pad_size, 3)?,
143            ExtendedColorType::Rgba8 => self.encode_rgba(image, width, height, row_pad_size, 4)?,
144            ExtendedColorType::L8 => {
145                self.encode_gray(image, width, height, row_pad_size, 1, palette)?;
146            }
147            ExtendedColorType::La8 => {
148                self.encode_gray(image, width, height, row_pad_size, 2, palette)?;
149            }
150            _ => {
151                return Err(ImageError::IoError(io::Error::new(
152                    io::ErrorKind::InvalidInput,
153                    &get_unsupported_error_message(c)[..],
154                )))
155            }
156        }
157
158        Ok(())
159    }
160
161    fn encode_rgb(
162        &mut self,
163        image: &[u8],
164        width: u32,
165        height: u32,
166        row_pad_size: u32,
167        bytes_per_pixel: u32,
168    ) -> io::Result<()> {
169        let width = width as usize;
170        let height = height as usize;
171        let x_stride = bytes_per_pixel as usize;
172        let y_stride = width * x_stride;
173        for row in (0..height).rev() {
174            // from the bottom up
175            let row_start = row * y_stride;
176            for px in image[row_start..][..y_stride].chunks_exact(x_stride) {
177                let r = px[0];
178                let g = px[1];
179                let b = px[2];
180                // written as BGR
181                self.writer.write_all(&[b, g, r])?;
182            }
183            self.write_row_pad(row_pad_size)?;
184        }
185
186        Ok(())
187    }
188
189    fn encode_rgba(
190        &mut self,
191        image: &[u8],
192        width: u32,
193        height: u32,
194        row_pad_size: u32,
195        bytes_per_pixel: u32,
196    ) -> io::Result<()> {
197        let width = width as usize;
198        let height = height as usize;
199        let x_stride = bytes_per_pixel as usize;
200        let y_stride = width * x_stride;
201        for row in (0..height).rev() {
202            // from the bottom up
203            let row_start = row * y_stride;
204            for px in image[row_start..][..y_stride].chunks_exact(x_stride) {
205                let r = px[0];
206                let g = px[1];
207                let b = px[2];
208                let a = px[3];
209                // written as BGRA
210                self.writer.write_all(&[b, g, r, a])?;
211            }
212            self.write_row_pad(row_pad_size)?;
213        }
214
215        Ok(())
216    }
217
218    fn encode_gray(
219        &mut self,
220        image: &[u8],
221        width: u32,
222        height: u32,
223        row_pad_size: u32,
224        bytes_per_pixel: u32,
225        palette: Option<&[[u8; 3]]>,
226    ) -> io::Result<()> {
227        // write grayscale palette
228        if let Some(palette) = palette {
229            for item in palette {
230                // each color is written as BGRA, where A is always 0
231                self.writer.write_all(&[item[2], item[1], item[0], 0])?;
232            }
233        } else {
234            for val in 0u8..=255 {
235                // each color is written as BGRA, where A is always 0 and since only grayscale is being written, B = G = R = index
236                self.writer.write_all(&[val, val, val, 0])?;
237            }
238        }
239
240        // write image data
241        let x_stride = bytes_per_pixel;
242        let y_stride = width * x_stride;
243        for row in (0..height).rev() {
244            // from the bottom up
245            let row_start = row * y_stride;
246
247            // color value is equal to the palette index
248            if x_stride == 1 {
249                // improve performance by writing the whole row at once
250                self.writer
251                    .write_all(&image[row_start as usize..][..y_stride as usize])?;
252            } else {
253                for col in 0..width {
254                    let pixel_start = (row_start + (col * x_stride)) as usize;
255                    self.writer.write_u8(image[pixel_start])?;
256                    // alpha is never written as it's not widely supported
257                }
258            }
259
260            self.write_row_pad(row_pad_size)?;
261        }
262
263        Ok(())
264    }
265
266    fn write_row_pad(&mut self, row_pad_size: u32) -> io::Result<()> {
267        for _ in 0..row_pad_size {
268            self.writer.write_u8(0)?;
269        }
270
271        Ok(())
272    }
273}
274
275impl<W: Write> ImageEncoder for BmpEncoder<'_, W> {
276    #[track_caller]
277    fn write_image(
278        mut self,
279        buf: &[u8],
280        width: u32,
281        height: u32,
282        color_type: ExtendedColorType,
283    ) -> ImageResult<()> {
284        self.encode(buf, width, height, color_type)
285    }
286}
287
288fn get_unsupported_error_message(c: ExtendedColorType) -> String {
289    format!("Unsupported color type {c:?}.  Supported types: RGB(8), RGBA(8), Gray(8), GrayA(8).")
290}
291
292/// Returns a tuple representing: (dib header size, written pixel size, palette color count).
293fn get_pixel_info(
294    c: ExtendedColorType,
295    palette: Option<&[[u8; 3]]>,
296) -> io::Result<(u32, u32, u32)> {
297    let sizes = match c {
298        ExtendedColorType::Rgb8 => (BITMAPINFOHEADER_SIZE, 3, 0),
299        ExtendedColorType::Rgba8 => (BITMAPV4HEADER_SIZE, 4, 0),
300        ExtendedColorType::L8 => (
301            BITMAPINFOHEADER_SIZE,
302            1,
303            palette.map(|p| p.len()).unwrap_or(256) as u32,
304        ),
305        ExtendedColorType::La8 => (
306            BITMAPINFOHEADER_SIZE,
307            1,
308            palette.map(|p| p.len()).unwrap_or(256) as u32,
309        ),
310        _ => {
311            return Err(io::Error::new(
312                io::ErrorKind::InvalidInput,
313                &get_unsupported_error_message(c)[..],
314            ))
315        }
316    };
317
318    Ok(sizes)
319}
320
321#[cfg(test)]
322mod tests {
323    use super::super::BmpDecoder;
324    use super::BmpEncoder;
325
326    use crate::image::ImageDecoder;
327    use crate::ExtendedColorType;
328    use std::io::Cursor;
329
330    fn round_trip_image(image: &[u8], width: u32, height: u32, c: ExtendedColorType) -> Vec<u8> {
331        let mut encoded_data = Vec::new();
332        {
333            let mut encoder = BmpEncoder::new(&mut encoded_data);
334            encoder
335                .encode(image, width, height, c)
336                .expect("could not encode image");
337        }
338
339        let decoder = BmpDecoder::new(Cursor::new(&encoded_data)).expect("failed to decode");
340
341        let mut buf = vec![0; decoder.total_bytes() as usize];
342        decoder.read_image(&mut buf).expect("failed to decode");
343        buf
344    }
345
346    #[test]
347    fn round_trip_single_pixel_rgb() {
348        let image = [255u8, 0, 0]; // single red pixel
349        let decoded = round_trip_image(&image, 1, 1, ExtendedColorType::Rgb8);
350        assert_eq!(3, decoded.len());
351        assert_eq!(255, decoded[0]);
352        assert_eq!(0, decoded[1]);
353        assert_eq!(0, decoded[2]);
354    }
355
356    #[test]
357    #[cfg(target_pointer_width = "64")]
358    fn huge_files_return_error() {
359        let mut encoded_data = Vec::new();
360        let image = vec![0u8; 3 * 40_000 * 40_000]; // 40_000x40_000 pixels, 3 bytes per pixel, allocated on the heap
361        let mut encoder = BmpEncoder::new(&mut encoded_data);
362        let result = encoder.encode(&image, 40_000, 40_000, ExtendedColorType::Rgb8);
363        assert!(result.is_err());
364    }
365
366    #[test]
367    fn round_trip_single_pixel_rgba() {
368        let image = [1, 2, 3, 4];
369        let decoded = round_trip_image(&image, 1, 1, ExtendedColorType::Rgba8);
370        assert_eq!(&decoded[..], &image[..]);
371    }
372
373    #[test]
374    fn round_trip_3px_rgb() {
375        let image = [0u8; 3 * 3 * 3]; // 3x3 pixels, 3 bytes per pixel
376        let _decoded = round_trip_image(&image, 3, 3, ExtendedColorType::Rgb8);
377    }
378
379    #[test]
380    fn round_trip_gray() {
381        let image = [0u8, 1, 2]; // 3 pixels
382        let decoded = round_trip_image(&image, 3, 1, ExtendedColorType::L8);
383        // should be read back as 3 RGB pixels
384        assert_eq!(9, decoded.len());
385        assert_eq!(0, decoded[0]);
386        assert_eq!(0, decoded[1]);
387        assert_eq!(0, decoded[2]);
388        assert_eq!(1, decoded[3]);
389        assert_eq!(1, decoded[4]);
390        assert_eq!(1, decoded[5]);
391        assert_eq!(2, decoded[6]);
392        assert_eq!(2, decoded[7]);
393        assert_eq!(2, decoded[8]);
394    }
395
396    #[test]
397    fn round_trip_graya() {
398        let image = [0u8, 0, 1, 0, 2, 0]; // 3 pixels, each with an alpha channel
399        let decoded = round_trip_image(&image, 1, 3, ExtendedColorType::La8);
400        // should be read back as 3 RGB pixels
401        assert_eq!(9, decoded.len());
402        assert_eq!(0, decoded[0]);
403        assert_eq!(0, decoded[1]);
404        assert_eq!(0, decoded[2]);
405        assert_eq!(1, decoded[3]);
406        assert_eq!(1, decoded[4]);
407        assert_eq!(1, decoded[5]);
408        assert_eq!(2, decoded[6]);
409        assert_eq!(2, decoded[7]);
410        assert_eq!(2, decoded[8]);
411    }
412}
image/codecs/bmp/encoder.rs

image/codecs/bmp/
encoder.rs
