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