image/codecs/bmp/
encoder.rs

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