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(image:buf, width, height, c: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: (u32, u32, u32) = 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(default:256) as u32,
304 ),
305 ExtendedColorType::La8 => (
306 BITMAPINFOHEADER_SIZE,
307 1,
308 palette.map(|p| p.len()).unwrap_or(default:256) as u32,
309 ),
310 _ => {
311 return Err(io::Error::new(
312 kind: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}
413