1use crate::utils::vec_try_with_capacity;
3use std::fmt;
4use std::io;
5use std::io::Write;
6
7use super::AutoBreak;
8use super::{ArbitraryHeader, ArbitraryTuplType, BitmapHeader, GraymapHeader, PixmapHeader};
9use super::{HeaderRecord, PnmHeader, PnmSubtype, SampleEncoding};
10
11use crate::color::ExtendedColorType;
12use crate::error::{
13 ImageError, ImageResult, ParameterError, ParameterErrorKind, UnsupportedError,
14 UnsupportedErrorKind,
15};
16use crate::{ImageEncoder, ImageFormat};
17
18use byteorder_lite::{BigEndian, WriteBytesExt};
19
20enum HeaderStrategy {
21 Dynamic,
22 Subtype(PnmSubtype),
23 Chosen(PnmHeader),
24}
25
26#[derive(Clone, Copy)]
27pub enum FlatSamples<'a> {
28 U8(&'a [u8]),
29 U16(&'a [u16]),
30}
31
32pub struct PnmEncoder<W: Write> {
34 writer: W,
35 header: HeaderStrategy,
36}
37
38struct CheckedImageBuffer<'a> {
41 _image: FlatSamples<'a>,
42 _width: u32,
43 _height: u32,
44 _color: ExtendedColorType,
45}
46
47struct UncheckedHeader<'a> {
49 header: &'a PnmHeader,
50}
51
52struct CheckedDimensions<'a> {
53 unchecked: UncheckedHeader<'a>,
54 width: u32,
55 height: u32,
56}
57
58struct CheckedHeaderColor<'a> {
59 dimensions: CheckedDimensions<'a>,
60 color: ExtendedColorType,
61}
62
63struct CheckedHeader<'a> {
64 color: CheckedHeaderColor<'a>,
65 encoding: TupleEncoding<'a>,
66 _image: CheckedImageBuffer<'a>,
67}
68
69enum TupleEncoding<'a> {
70 PbmBits {
71 samples: FlatSamples<'a>,
72 width: u32,
73 },
74 Ascii {
75 samples: FlatSamples<'a>,
76 },
77 Bytes {
78 samples: FlatSamples<'a>,
79 },
80}
81
82impl<W: Write> PnmEncoder<W> {
83 pub fn new(writer: W) -> Self {
89 PnmEncoder {
90 writer,
91 header: HeaderStrategy::Dynamic,
92 }
93 }
94
95 pub fn with_subtype(self, subtype: PnmSubtype) -> Self {
103 PnmEncoder {
104 writer: self.writer,
105 header: HeaderStrategy::Subtype(subtype),
106 }
107 }
108
109 pub fn with_header(self, header: PnmHeader) -> Self {
119 PnmEncoder {
120 writer: self.writer,
121 header: HeaderStrategy::Chosen(header),
122 }
123 }
124
125 pub fn with_dynamic_header(self) -> Self {
133 PnmEncoder {
134 writer: self.writer,
135 header: HeaderStrategy::Dynamic,
136 }
137 }
138
139 pub fn encode<'s, S>(
149 &mut self,
150 image: S,
151 width: u32,
152 height: u32,
153 color: ExtendedColorType,
154 ) -> ImageResult<()>
155 where
156 S: Into<FlatSamples<'s>>,
157 {
158 let image = image.into();
159
160 let image = match (image, color) {
164 (
165 FlatSamples::U8(samples),
166 ExtendedColorType::L16
167 | ExtendedColorType::La16
168 | ExtendedColorType::Rgb16
169 | ExtendedColorType::Rgba16,
170 ) => {
171 match bytemuck::try_cast_slice(samples) {
172 Ok(samples) => FlatSamples::U16(samples),
174 Err(_e) => {
175 let new_samples: Vec<u16> = samples
177 .chunks(2)
178 .map(|chunk| u16::from_ne_bytes([chunk[0], chunk[1]]))
179 .collect();
180
181 let image = FlatSamples::U16(&new_samples);
182
183 return self.encode_impl(image, width, height, color);
186 }
187 }
188 }
189 _ => image,
191 };
192
193 self.encode_impl(image, width, height, color)
194 }
195
196 fn encode_impl(
198 &mut self,
199 samples: FlatSamples<'_>,
200 width: u32,
201 height: u32,
202 color: ExtendedColorType,
203 ) -> ImageResult<()> {
204 match self.header {
205 HeaderStrategy::Dynamic => self.write_dynamic_header(samples, width, height, color),
206 HeaderStrategy::Subtype(subtype) => {
207 self.write_subtyped_header(subtype, samples, width, height, color)
208 }
209 HeaderStrategy::Chosen(ref header) => {
210 Self::write_with_header(&mut self.writer, header, samples, width, height, color)
211 }
212 }
213 }
214
215 fn write_dynamic_header(
219 &mut self,
220 image: FlatSamples,
221 width: u32,
222 height: u32,
223 color: ExtendedColorType,
224 ) -> ImageResult<()> {
225 let depth = u32::from(color.channel_count());
226 let (maxval, tupltype) = match color {
227 ExtendedColorType::L1 => (1, ArbitraryTuplType::BlackAndWhite),
228 ExtendedColorType::L8 => (0xff, ArbitraryTuplType::Grayscale),
229 ExtendedColorType::L16 => (0xffff, ArbitraryTuplType::Grayscale),
230 ExtendedColorType::La1 => (1, ArbitraryTuplType::BlackAndWhiteAlpha),
231 ExtendedColorType::La8 => (0xff, ArbitraryTuplType::GrayscaleAlpha),
232 ExtendedColorType::La16 => (0xffff, ArbitraryTuplType::GrayscaleAlpha),
233 ExtendedColorType::Rgb8 => (0xff, ArbitraryTuplType::RGB),
234 ExtendedColorType::Rgb16 => (0xffff, ArbitraryTuplType::RGB),
235 ExtendedColorType::Rgba8 => (0xff, ArbitraryTuplType::RGBAlpha),
236 ExtendedColorType::Rgba16 => (0xffff, ArbitraryTuplType::RGBAlpha),
237 _ => {
238 return Err(ImageError::Unsupported(
239 UnsupportedError::from_format_and_kind(
240 ImageFormat::Pnm.into(),
241 UnsupportedErrorKind::Color(color),
242 ),
243 ))
244 }
245 };
246
247 let header = PnmHeader {
248 decoded: HeaderRecord::Arbitrary(ArbitraryHeader {
249 width,
250 height,
251 depth,
252 maxval,
253 tupltype: Some(tupltype),
254 }),
255 encoded: None,
256 };
257
258 Self::write_with_header(&mut self.writer, &header, image, width, height, color)
259 }
260
261 fn write_subtyped_header(
263 &mut self,
264 subtype: PnmSubtype,
265 image: FlatSamples,
266 width: u32,
267 height: u32,
268 color: ExtendedColorType,
269 ) -> ImageResult<()> {
270 let header = match (subtype, color) {
271 (PnmSubtype::ArbitraryMap, color) => {
272 return self.write_dynamic_header(image, width, height, color)
273 }
274 (PnmSubtype::Pixmap(encoding), ExtendedColorType::Rgb8) => PnmHeader {
275 decoded: HeaderRecord::Pixmap(PixmapHeader {
276 encoding,
277 width,
278 height,
279 maxval: 255,
280 }),
281 encoded: None,
282 },
283 (PnmSubtype::Graymap(encoding), ExtendedColorType::L8) => PnmHeader {
284 decoded: HeaderRecord::Graymap(GraymapHeader {
285 encoding,
286 width,
287 height,
288 maxwhite: 255,
289 }),
290 encoded: None,
291 },
292 (PnmSubtype::Bitmap(encoding), ExtendedColorType::L8 | ExtendedColorType::L1) => {
293 PnmHeader {
294 decoded: HeaderRecord::Bitmap(BitmapHeader {
295 encoding,
296 height,
297 width,
298 }),
299 encoded: None,
300 }
301 }
302 (_, _) => {
303 return Err(ImageError::Unsupported(
304 UnsupportedError::from_format_and_kind(
305 ImageFormat::Pnm.into(),
306 UnsupportedErrorKind::Color(color),
307 ),
308 ))
309 }
310 };
311
312 Self::write_with_header(&mut self.writer, &header, image, width, height, color)
313 }
314
315 fn write_with_header(
319 writer: &mut dyn Write,
320 header: &PnmHeader,
321 image: FlatSamples,
322 width: u32,
323 height: u32,
324 color: ExtendedColorType,
325 ) -> ImageResult<()> {
326 let unchecked = UncheckedHeader { header };
327
328 unchecked
329 .check_header_dimensions(width, height)?
330 .check_header_color(color)?
331 .check_sample_values(image)?
332 .write_header(writer)?
333 .write_image(writer)
334 }
335}
336
337impl<W: Write> ImageEncoder for PnmEncoder<W> {
338 #[track_caller]
339 fn write_image(
340 mut self,
341 buf: &[u8],
342 width: u32,
343 height: u32,
344 color_type: ExtendedColorType,
345 ) -> ImageResult<()> {
346 let expected_buffer_len = color_type.buffer_size(width, height);
347 assert_eq!(
348 expected_buffer_len,
349 buf.len() as u64,
350 "Invalid buffer length: expected {expected_buffer_len} got {} for {width}x{height} image",
351 buf.len(),
352 );
353
354 self.encode(buf, width, height, color_type)
355 }
356}
357
358impl<'a> CheckedImageBuffer<'a> {
359 fn check(
360 image: FlatSamples<'a>,
361 width: u32,
362 height: u32,
363 color: ExtendedColorType,
364 ) -> ImageResult<CheckedImageBuffer<'a>> {
365 let components = color.channel_count() as usize;
366 let uwidth = width as usize;
367 let uheight = height as usize;
368 let expected_len = components
369 .checked_mul(uwidth)
370 .and_then(|v| v.checked_mul(uheight));
371 if Some(image.len()) != expected_len {
372 return Err(ImageError::Parameter(ParameterError::from_kind(
374 ParameterErrorKind::DimensionMismatch,
375 )));
376 }
377 Ok(CheckedImageBuffer {
378 _image: image,
379 _width: width,
380 _height: height,
381 _color: color,
382 })
383 }
384}
385
386impl<'a> UncheckedHeader<'a> {
387 fn check_header_dimensions(
388 self,
389 width: u32,
390 height: u32,
391 ) -> ImageResult<CheckedDimensions<'a>> {
392 if self.header.width() != width || self.header.height() != height {
393 return Err(ImageError::Parameter(ParameterError::from_kind(
395 ParameterErrorKind::DimensionMismatch,
396 )));
397 }
398
399 Ok(CheckedDimensions {
400 unchecked: self,
401 width,
402 height,
403 })
404 }
405}
406
407impl<'a> CheckedDimensions<'a> {
408 fn check_header_color(self, color: ExtendedColorType) -> ImageResult<CheckedHeaderColor<'a>> {
412 let components = u32::from(color.channel_count());
413
414 match *self.unchecked.header {
415 PnmHeader {
416 decoded: HeaderRecord::Bitmap(_),
417 ..
418 } => match color {
419 ExtendedColorType::L1 | ExtendedColorType::L8 | ExtendedColorType::L16 => (),
420 _ => {
421 return Err(ImageError::Parameter(ParameterError::from_kind(
422 ParameterErrorKind::Generic(
423 "PBM format only support luma color types".to_owned(),
424 ),
425 )))
426 }
427 },
428 PnmHeader {
429 decoded: HeaderRecord::Graymap(_),
430 ..
431 } => match color {
432 ExtendedColorType::L1 | ExtendedColorType::L8 | ExtendedColorType::L16 => (),
433 _ => {
434 return Err(ImageError::Parameter(ParameterError::from_kind(
435 ParameterErrorKind::Generic(
436 "PGM format only support luma color types".to_owned(),
437 ),
438 )))
439 }
440 },
441 PnmHeader {
442 decoded: HeaderRecord::Pixmap(_),
443 ..
444 } => match color {
445 ExtendedColorType::Rgb8 => (),
446 _ => {
447 return Err(ImageError::Parameter(ParameterError::from_kind(
448 ParameterErrorKind::Generic(
449 "PPM format only support ExtendedColorType::Rgb8".to_owned(),
450 ),
451 )))
452 }
453 },
454 PnmHeader {
455 decoded:
456 HeaderRecord::Arbitrary(ArbitraryHeader {
457 depth,
458 ref tupltype,
459 ..
460 }),
461 ..
462 } => match (tupltype, color) {
463 (&Some(ArbitraryTuplType::BlackAndWhite), ExtendedColorType::L1) => (),
464 (&Some(ArbitraryTuplType::BlackAndWhiteAlpha), ExtendedColorType::La8) => (),
465
466 (&Some(ArbitraryTuplType::Grayscale), ExtendedColorType::L1) => (),
467 (&Some(ArbitraryTuplType::Grayscale), ExtendedColorType::L8) => (),
468 (&Some(ArbitraryTuplType::Grayscale), ExtendedColorType::L16) => (),
469 (&Some(ArbitraryTuplType::GrayscaleAlpha), ExtendedColorType::La8) => (),
470
471 (&Some(ArbitraryTuplType::RGB), ExtendedColorType::Rgb8) => (),
472 (&Some(ArbitraryTuplType::RGB), ExtendedColorType::Rgb16) => (),
473 (&Some(ArbitraryTuplType::RGBAlpha), ExtendedColorType::Rgba8) => (),
474 (&Some(ArbitraryTuplType::RGBAlpha), ExtendedColorType::Rgba16) => (),
475
476 (&None, _) if depth == components => (),
477 (&Some(ArbitraryTuplType::Custom(_)), _) if depth == components => (),
478 _ if depth != components => {
479 return Err(ImageError::Parameter(ParameterError::from_kind(
480 ParameterErrorKind::Generic(format!(
481 "Depth mismatch: header {depth} vs. color {components}"
482 )),
483 )))
484 }
485 _ => {
486 return Err(ImageError::Parameter(ParameterError::from_kind(
487 ParameterErrorKind::Generic(
488 "Invalid color type for selected PAM color type".to_owned(),
489 ),
490 )))
491 }
492 },
493 }
494
495 Ok(CheckedHeaderColor {
496 dimensions: self,
497 color,
498 })
499 }
500}
501
502impl<'a> CheckedHeaderColor<'a> {
503 fn check_sample_values(self, image: FlatSamples<'a>) -> ImageResult<CheckedHeader<'a>> {
504 let header_maxval = match self.dimensions.unchecked.header.decoded {
505 HeaderRecord::Bitmap(_) => 1,
506 HeaderRecord::Graymap(GraymapHeader { maxwhite, .. }) => maxwhite,
507 HeaderRecord::Pixmap(PixmapHeader { maxval, .. }) => maxval,
508 HeaderRecord::Arbitrary(ArbitraryHeader { maxval, .. }) => maxval,
509 };
510
511 let max_sample = match self.color {
513 ExtendedColorType::Unknown(n) if n <= 16 => (1 << n) - 1,
514 ExtendedColorType::L1 => 1,
515 ExtendedColorType::L8
516 | ExtendedColorType::La8
517 | ExtendedColorType::Rgb8
518 | ExtendedColorType::Rgba8
519 | ExtendedColorType::Bgr8
520 | ExtendedColorType::Bgra8 => 0xff,
521 ExtendedColorType::L16
522 | ExtendedColorType::La16
523 | ExtendedColorType::Rgb16
524 | ExtendedColorType::Rgba16 => 0xffff,
525 _ => {
526 return Err(ImageError::Unsupported(
528 UnsupportedError::from_format_and_kind(
529 ImageFormat::Pnm.into(),
530 UnsupportedErrorKind::Color(self.color),
531 ),
532 ));
533 }
534 };
535
536 if header_maxval < max_sample && !image.all_smaller(header_maxval) {
538 return Err(ImageError::Unsupported(
540 UnsupportedError::from_format_and_kind(
541 ImageFormat::Pnm.into(),
542 UnsupportedErrorKind::GenericFeature(
543 "Sample value greater than allowed for chosen header".to_owned(),
544 ),
545 ),
546 ));
547 }
548
549 let encoding = image.encoding_for(&self.dimensions.unchecked.header.decoded);
550
551 let image = CheckedImageBuffer::check(
552 image,
553 self.dimensions.width,
554 self.dimensions.height,
555 self.color,
556 )?;
557
558 Ok(CheckedHeader {
559 color: self,
560 encoding,
561 _image: image,
562 })
563 }
564}
565
566impl<'a> CheckedHeader<'a> {
567 fn write_header(self, writer: &mut dyn Write) -> ImageResult<TupleEncoding<'a>> {
568 self.header().write(writer)?;
569 Ok(self.encoding)
570 }
571
572 fn header(&self) -> &PnmHeader {
573 self.color.dimensions.unchecked.header
574 }
575}
576
577struct SampleWriter<'a>(&'a mut dyn Write);
578
579impl SampleWriter<'_> {
580 fn write_samples_ascii<V>(self, samples: V) -> io::Result<()>
581 where
582 V: Iterator,
583 V::Item: fmt::Display,
584 {
585 let mut auto_break_writer = AutoBreak::new(self.0, 70)?;
586 for value in samples {
587 write!(auto_break_writer, "{value} ")?;
588 }
589 auto_break_writer.flush()
590 }
591
592 fn write_pbm_bits<V>(self, samples: &[V], width: u32) -> io::Result<()>
593 where
595 V: Default + Eq + Copy,
596 {
597 let line_width = (width - 1) / 8 + 1;
599
600 let mut line_buffer = vec_try_with_capacity(line_width as usize)?;
602
603 for line in samples.chunks(width as usize) {
604 for byte_bits in line.chunks(8) {
605 let mut byte = 0u8;
606 for i in 0..8 {
607 if let Some(&v) = byte_bits.get(i) {
609 if v == V::default() {
610 byte |= 1u8 << (7 - i);
611 }
612 }
613 }
614 line_buffer.push(byte);
615 }
616 self.0.write_all(line_buffer.as_slice())?;
617 line_buffer.clear();
618 }
619
620 self.0.flush()
621 }
622}
623
624impl<'a> FlatSamples<'a> {
625 fn len(&self) -> usize {
626 match *self {
627 FlatSamples::U8(arr) => arr.len(),
628 FlatSamples::U16(arr) => arr.len(),
629 }
630 }
631
632 fn all_smaller(&self, max_val: u32) -> bool {
633 match *self {
634 FlatSamples::U8(arr) => arr.iter().all(|&val| u32::from(val) <= max_val),
635 FlatSamples::U16(arr) => arr.iter().all(|&val| u32::from(val) <= max_val),
636 }
637 }
638
639 fn encoding_for(&self, header: &HeaderRecord) -> TupleEncoding<'a> {
640 match *header {
641 HeaderRecord::Bitmap(BitmapHeader {
642 encoding: SampleEncoding::Binary,
643 width,
644 ..
645 }) => TupleEncoding::PbmBits {
646 samples: *self,
647 width,
648 },
649
650 HeaderRecord::Bitmap(BitmapHeader {
651 encoding: SampleEncoding::Ascii,
652 ..
653 }) => TupleEncoding::Ascii { samples: *self },
654
655 HeaderRecord::Arbitrary(_) => TupleEncoding::Bytes { samples: *self },
656
657 HeaderRecord::Graymap(GraymapHeader {
658 encoding: SampleEncoding::Ascii,
659 ..
660 })
661 | HeaderRecord::Pixmap(PixmapHeader {
662 encoding: SampleEncoding::Ascii,
663 ..
664 }) => TupleEncoding::Ascii { samples: *self },
665
666 HeaderRecord::Graymap(GraymapHeader {
667 encoding: SampleEncoding::Binary,
668 ..
669 })
670 | HeaderRecord::Pixmap(PixmapHeader {
671 encoding: SampleEncoding::Binary,
672 ..
673 }) => TupleEncoding::Bytes { samples: *self },
674 }
675 }
676}
677
678impl<'a> From<&'a [u8]> for FlatSamples<'a> {
679 fn from(samples: &'a [u8]) -> Self {
680 FlatSamples::U8(samples)
681 }
682}
683
684impl<'a> From<&'a [u16]> for FlatSamples<'a> {
685 fn from(samples: &'a [u16]) -> Self {
686 FlatSamples::U16(samples)
687 }
688}
689
690impl TupleEncoding<'_> {
691 fn write_image(&self, writer: &mut dyn Write) -> ImageResult<()> {
692 match *self {
693 TupleEncoding::PbmBits {
694 samples: FlatSamples::U8(samples),
695 width,
696 } => SampleWriter(writer)
697 .write_pbm_bits(samples, width)
698 .map_err(ImageError::IoError),
699 TupleEncoding::PbmBits {
700 samples: FlatSamples::U16(samples),
701 width,
702 } => SampleWriter(writer)
703 .write_pbm_bits(samples, width)
704 .map_err(ImageError::IoError),
705
706 TupleEncoding::Bytes {
707 samples: FlatSamples::U8(samples),
708 } => writer.write_all(samples).map_err(ImageError::IoError),
709 TupleEncoding::Bytes {
710 samples: FlatSamples::U16(samples),
711 } => samples.iter().try_for_each(|&sample| {
712 writer
713 .write_u16::<BigEndian>(sample)
714 .map_err(ImageError::IoError)
715 }),
716
717 TupleEncoding::Ascii {
718 samples: FlatSamples::U8(samples),
719 } => SampleWriter(writer)
720 .write_samples_ascii(samples.iter())
721 .map_err(ImageError::IoError),
722 TupleEncoding::Ascii {
723 samples: FlatSamples::U16(samples),
724 } => SampleWriter(writer)
725 .write_samples_ascii(samples.iter())
726 .map_err(ImageError::IoError),
727 }
728 }
729}
730
731#[test]
732fn pbm_allows_black() {
733 let imgbuf = crate::DynamicImage::new_luma8(50, 50);
734
735 let mut buffer = vec![];
736 let encoder =
737 PnmEncoder::new(&mut buffer).with_subtype(PnmSubtype::Bitmap(SampleEncoding::Ascii));
738
739 imgbuf
740 .write_with_encoder(encoder)
741 .expect("all-zeroes is a black image");
742}
743
744#[test]
745fn pbm_allows_white() {
746 let imgbuf =
747 crate::DynamicImage::ImageLuma8(crate::ImageBuffer::from_pixel(50, 50, crate::Luma([1])));
748
749 let mut buffer = vec![];
750 let encoder =
751 PnmEncoder::new(&mut buffer).with_subtype(PnmSubtype::Bitmap(SampleEncoding::Ascii));
752
753 imgbuf
754 .write_with_encoder(encoder)
755 .expect("all-zeroes is a white image");
756}
757
758#[test]
759fn pbm_verifies_pixels() {
760 let imgbuf =
761 crate::DynamicImage::ImageLuma8(crate::ImageBuffer::from_pixel(50, 50, crate::Luma([255])));
762
763 let mut buffer = vec![];
764 let encoder =
765 PnmEncoder::new(&mut buffer).with_subtype(PnmSubtype::Bitmap(SampleEncoding::Ascii));
766
767 imgbuf
768 .write_with_encoder(encoder)
769 .expect_err("failed to catch violating samples");
770}