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(image:buf, width, height, c: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: (u32, u32, u32) = 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(default:256) as u32,
299 ),
300 color::ColorType::La8 => (
301 BITMAPINFOHEADER_SIZE,
302 1,
303 palette.map(|p| p.len()).unwrap_or(default:256) as u32,
304 ),
305 _ => {
306 return Err(io::Error::new(
307 kind: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}
407