Skip to main content

otf_pixels_codec_png/
encoder.rs

1//! The PNG encoder.
2//!
3//! # Memory
4//!
5//! Encoding is **internally buffered**: DEFLATE needs the whole scanline
6//! stream before it can emit its Huffman-coded blocks, so filtered rows
7//! accumulate until [`Encoder::finish`]. That is ADR-0005's stated allowance
8//! for codecs that cannot work incrementally, and it is symmetric with the
9//! decoder — see the module docs there for why an incremental inflate and
10//! deflate are deferred rather than absent.
11//!
12//! # What is written
13//!
14//! Non-interlaced PNG: `IHDR`, one `IDAT`, `IEND`, in the image's own colour
15//! type. Adam7 would cost bytes rather than save them, and every decoder must
16//! accept the plain form.
17//!
18//! Asked for a palette (sharp's `palette`), the encoder instead quantizes to
19//! at most 256 colours with the median cut GIF uses, and writes colour type
20//! 3: `PLTE`, a `tRNS` when any entry is translucent, and indices at the
21//! smallest bit depth that holds them, unfiltered, as libpng recommends for
22//! palette images. That is lossy unless the image has no more colours than
23//! the palette, and usually much smaller.
24
25use otf_pixels_core::{
26    Density, EncodeOptions, Encoder, ImageDescriptor, PixelFormat, PixelsError, ResolutionUnit,
27    Result, Sink,
28};
29
30use crate::format::{ColorType, Filter, SIGNATURE, apply_filter, write_chunk};
31use otf_pixels_compress::{Level, build_palette, build_palette_rgba, quantize_with, zlib_compress};
32use otf_pixels_core::{ChannelLayout, SampleKind};
33
34/// Encodes a PNG stream.
35#[derive(Debug)]
36pub struct PngEncoder {
37    level: Level,
38    /// Palette output: the palette size and the dithering level.
39    palette: Option<(usize, f32)>,
40    /// Set by `write_header`; its presence means the header was written.
41    state: Option<State>,
42    /// The ICC profile to write as `iCCP`, if any.
43    icc: Option<Vec<u8>>,
44    /// The EXIF block to write as `eXIf`, if any.
45    exif: Option<Vec<u8>>,
46    /// The XMP packet to write as an `iTXt` chunk, if any.
47    xmp: Option<Vec<u8>>,
48    /// The resolution to write as `pHYs`, if any.
49    density: Option<Density>,
50}
51
52/// The ICC profile, EXIF block, XMP packet and density to write.
53type Ancillary<'a> = (
54    Option<&'a [u8]>,
55    Option<&'a [u8]>,
56    Option<&'a [u8]>,
57    Option<Density>,
58);
59
60/// The ancillary chunks that go before the image data: `iCCP`, `pHYs`,
61/// `eXIf` and the XMP `iTXt`.
62fn write_metadata(
63    out: &mut Vec<u8>,
64    level: Level,
65    icc: Option<&[u8]>,
66    exif: Option<&[u8]>,
67    xmp: Option<&[u8]>,
68    density: Option<Density>,
69) -> Result<()> {
70    if let Some(profile) = icc {
71        // iCCP (§11.3.3.3): a name, NUL, method 0, the zlib stream.
72        let mut iccp = b"ICC Profile\0\0".to_vec();
73        iccp.extend(zlib_compress(profile, level).map_err(crate::compress_error)?);
74        write_chunk(out, b"iCCP", &iccp);
75    }
76    if let Some(d) = density {
77        // pHYs (§11.3.5.3) counts pixels per metre.
78        let per_metre = |v: f64| {
79            (v * if d.unit == ResolutionUnit::Inch {
80                100.0 / 2.54
81            } else {
82                100.0
83            })
84            .round()
85            .clamp(0.0, f64::from(u32::MAX)) as u32
86        };
87        let mut phys = per_metre(d.x).to_be_bytes().to_vec();
88        phys.extend_from_slice(&per_metre(d.y).to_be_bytes());
89        phys.push(1);
90        write_chunk(out, b"pHYs", &phys);
91    }
92    if let Some(exif) = exif {
93        // eXIf goes before the image data (§11.3.6.1).
94        write_chunk(out, b"eXIf", exif);
95    }
96    if let Some(xmp) = xmp {
97        // iTXt keyed for XMP, uncompressed, no language or translation, as
98        // the XMP specification (part 3, §1.1.5) asks.
99        let mut itxt = b"XML:com.adobe.xmp\0\0\0\0\0".to_vec();
100        itxt.extend_from_slice(xmp);
101        write_chunk(out, b"iTXt", &itxt);
102    }
103    Ok(())
104}
105
106/// Everything fixed once the descriptor is known.
107#[derive(Debug)]
108struct State {
109    descriptor: ImageDescriptor,
110    bit_depth: u8,
111    /// Bytes per pixel, rounded up — the filter's left-neighbour offset.
112    stride: usize,
113    /// Filtered scanlines, each prefixed by its filter byte.
114    filtered: Vec<u8>,
115    /// The previous *unfiltered* row, which filters predict from.
116    previous: Vec<u8>,
117    /// Reusable big-endian staging buffer; empty unless `bit_depth` is 16.
118    swapped: Vec<u8>,
119    /// For palette output, the image as 8-bit RGB or RGBA, gathered until
120    /// `finish` because the palette depends on every pixel.
121    true_colour: Option<Vec<u8>>,
122    rows_written: u32,
123}
124
125impl PngEncoder {
126    /// An encoder at the default compression level.
127    #[must_use]
128    pub const fn new() -> Self {
129        Self::with_level(Level::DEFAULT)
130    }
131
132    /// An encoder at an explicit DEFLATE level.
133    #[must_use]
134    pub const fn with_level(level: Level) -> Self {
135        Self {
136            level,
137            palette: None,
138            state: None,
139            icc: None,
140            exif: None,
141            xmp: None,
142            density: None,
143        }
144    }
145
146    /// This encoder writing a palette image of at most `colours` entries,
147    /// diffusing `dither` (0..=1) of the quantization error.
148    ///
149    /// # Errors
150    ///
151    /// Returns [`PixelsError::InvalidArgument`] for `colours` outside
152    /// 2..=256 or `dither` outside 0..=1.
153    pub fn with_palette(mut self, colours: usize, dither: f32) -> Result<Self> {
154        if !(2..=256).contains(&colours) {
155            return Err(PixelsError::invalid_argument(
156                "colours",
157                format!("palette size must be in 2..=256, got {colours}"),
158            ));
159        }
160        if !(0.0..=1.0).contains(&dither) {
161            return Err(PixelsError::invalid_argument(
162                "dither",
163                format!("must be in 0..=1, got {dither}"),
164            ));
165        }
166        self.palette = Some((colours, dither));
167        Ok(self)
168    }
169
170    /// An encoder configured from generic encode options.
171    ///
172    /// [`EncodeOptions::compression_level`] is the DEFLATE level. Without
173    /// one, [`EncodeOptions::quality`], which cannot trade fidelity for size
174    /// in a lossless format, is read as compression *effort*, mapping 1..=100
175    /// onto levels 1..=9.
176    ///
177    /// [`EncodeOptions::wants_palette`] selects palette output, of
178    /// [`EncodeOptions::colours`] entries (256 by default) dithered by
179    /// [`EncodeOptions::dither`] (1 by default).
180    #[must_use]
181    pub fn from_options(options: &EncodeOptions) -> Self {
182        let level = options.compression_level.unwrap_or_else(|| {
183            let quality = u32::from(options.quality.clamp(1, 100));
184            // 1..=100 onto 1..=9, so the default quality of 80 lands on
185            // level 7, just above zlib's own default of 6.
186            ((quality - 1) * 8 / 99 + 1) as u8
187        });
188        let mut encoder = Self::with_level(Level::new(level).unwrap_or(Level::DEFAULT));
189        if options.wants_palette() {
190            let colours = options.colours.map_or(256, usize::from).clamp(2, 256);
191            let dither = options.dither.unwrap_or(1.0).clamp(0.0, 1.0);
192            encoder.palette = Some((colours, dither));
193        }
194        encoder
195    }
196
197    /// The colour type and bit depth for a pixel format.
198    fn png_type_of(format: PixelFormat) -> Result<(ColorType, u8)> {
199        match format {
200            PixelFormat::Gray8 => Ok((ColorType::Grayscale, 8)),
201            PixelFormat::Gray16 => Ok((ColorType::Grayscale, 16)),
202            PixelFormat::GrayA8 => Ok((ColorType::GrayscaleAlpha, 8)),
203            PixelFormat::Rgb8 => Ok((ColorType::Rgb, 8)),
204            PixelFormat::Rgb16 => Ok((ColorType::Rgb, 16)),
205            PixelFormat::Rgba8 => Ok((ColorType::Rgba, 8)),
206            PixelFormat::Rgba16 => Ok((ColorType::Rgba, 16)),
207            // PNG has no float sample type at any bit depth (§11.2.2).
208            other => Err(PixelsError::unsupported(format!(
209                "PNG cannot represent {other}; convert to an integer format first"
210            ))),
211        }
212    }
213}
214
215impl Default for PngEncoder {
216    fn default() -> Self {
217        Self::new()
218    }
219}
220
221impl State {
222    /// Choose a filter for `row` and append it, filter byte first.
223    ///
224    /// Uses the minimum-sum-of-absolute-differences heuristic from the PNG
225    /// specification's §12.8: filtered bytes near zero compress best, and the
226    /// sum of their signed magnitudes estimates that without running DEFLATE
227    /// five times.
228    fn append_filtered(&mut self, row: &[u8]) {
229        let candidates: &[Filter] = &[
230            Filter::None,
231            Filter::Sub,
232            Filter::Up,
233            Filter::Average,
234            Filter::Paeth,
235        ];
236
237        let mut best = Filter::None;
238        let mut best_score = u64::MAX;
239        let mut scratch = Vec::with_capacity(row.len());
240        for &filter in candidates {
241            scratch.clear();
242            apply_filter(filter, row, &self.previous, self.stride, &mut scratch);
243            let score: u64 = scratch
244                .iter()
245                .map(|&byte| u64::from((byte as i8).unsigned_abs()))
246                .sum();
247            if score < best_score {
248                best_score = score;
249                best = filter;
250            }
251        }
252
253        self.filtered.push(best.to_byte());
254        apply_filter(best, row, &self.previous, self.stride, &mut self.filtered);
255        self.previous.clear();
256        self.previous.extend_from_slice(row);
257    }
258}
259
260/// Rewrite native-endian 16-bit samples as the big-endian ones PNG stores.
261///
262/// On a big-endian host this is a copy; the swap is expressed in terms of
263/// `to_be_bytes` rather than a `cfg`, so there is one code path to be wrong.
264fn to_big_endian_16(row: &[u8], out: &mut Vec<u8>) {
265    out.clear();
266    for pair in row.chunks_exact(2) {
267        let value = u16::from_ne_bytes([
268            pair.first().copied().unwrap_or(0),
269            pair.get(1).copied().unwrap_or(0),
270        ]);
271        out.extend_from_slice(&value.to_be_bytes());
272    }
273}
274
275/// Append one row as 8-bit RGB, or RGBA if `format` has alpha.
276fn append_8_bit(row: &[u8], format: PixelFormat, out: &mut Vec<u8>) {
277    let wide = format.sample_kind() == SampleKind::U16;
278    let sample = |pixel: &[u8], index: usize| -> u8 {
279        if wide {
280            let at = index * 2;
281            let v = u16::from_ne_bytes([
282                pixel.get(at).copied().unwrap_or(0),
283                pixel.get(at + 1).copied().unwrap_or(0),
284            ]);
285            // Rounded to nearest, as every narrowing to 8 bits is.
286            ((u32::from(v) * 255 + 32_767) / 65_535) as u8
287        } else {
288            pixel.get(index).copied().unwrap_or(0)
289        }
290    };
291    for pixel in row.chunks_exact(format.bytes_per_pixel()) {
292        match format.layout() {
293            ChannelLayout::Gray => {
294                let g = sample(pixel, 0);
295                out.extend_from_slice(&[g, g, g]);
296            }
297            ChannelLayout::GrayAlpha => {
298                let g = sample(pixel, 0);
299                out.extend_from_slice(&[g, g, g, sample(pixel, 1)]);
300            }
301            ChannelLayout::Rgb => out.extend((0..3).map(|i| sample(pixel, i))),
302            ChannelLayout::Rgba => out.extend((0..4).map(|i| sample(pixel, i))),
303        }
304    }
305}
306
307/// A whole palette PNG for `pixels`, 8-bit RGB or (when `descriptor` has
308/// alpha) RGBA.
309fn encode_palette(
310    descriptor: &ImageDescriptor,
311    pixels: &[u8],
312    colours: usize,
313    dither: f32,
314    level: Level,
315    metadata: Ancillary<'_>,
316) -> Result<Vec<u8>> {
317    let width = descriptor.width as usize;
318    let (entries, indices): (Vec<[u8; 4]>, Vec<u8>) = if descriptor.pixel.layout().has_alpha() {
319        let palette = build_palette_rgba(pixels, colours);
320        let indices = quantize_with(pixels, width, &palette, dither);
321        (palette.entries().to_vec(), indices)
322    } else {
323        let palette = build_palette(pixels, colours);
324        let indices = quantize_with(pixels, width, &palette, dither);
325        let entries = palette
326            .entries()
327            .iter()
328            .map(|&[r, g, b]| [r, g, b, 255])
329            .collect();
330        (entries, indices)
331    };
332
333    // The smallest depth whose indices reach every entry.
334    let bit_depth: u8 = match entries.len() {
335        0..=2 => 1,
336        3..=4 => 2,
337        5..=16 => 4,
338        _ => 8,
339    };
340    let per_byte = 8 / usize::from(bit_depth);
341    let row_bytes = width.div_ceil(per_byte);
342    let mut scanlines = Vec::with_capacity((row_bytes + 1) * descriptor.height as usize);
343    for row in indices.chunks_exact(width) {
344        scanlines.push(Filter::None.to_byte());
345        for group in row.chunks(per_byte) {
346            let mut byte = 0_u8;
347            for (i, &index) in group.iter().enumerate() {
348                // Leftmost pixel in the highest bits (§7.2).
349                byte |= index << (8 - usize::from(bit_depth) * (i + 1));
350            }
351            scanlines.push(byte);
352        }
353    }
354
355    let mut out = SIGNATURE.to_vec();
356    let mut ihdr = Vec::with_capacity(13);
357    ihdr.extend_from_slice(&descriptor.width.to_be_bytes());
358    ihdr.extend_from_slice(&descriptor.height.to_be_bytes());
359    ihdr.extend_from_slice(&[bit_depth, ColorType::Palette.to_byte(), 0, 0, 0]);
360    write_chunk(&mut out, b"IHDR", &ihdr);
361    let (icc, exif, xmp, density) = metadata;
362    write_metadata(&mut out, level, icc, exif, xmp, density)?;
363    let plte: Vec<u8> = entries.iter().flat_map(|e| [e[0], e[1], e[2]]).collect();
364    write_chunk(&mut out, b"PLTE", &plte);
365    // tRNS lists alphas up to the last translucent entry; the rest are
366    // opaque by default (§11.3.2.1).
367    if let Some(last) = entries.iter().rposition(|e| e[3] != 255) {
368        let alphas: Vec<u8> = entries.iter().take(last + 1).map(|e| e[3]).collect();
369        write_chunk(&mut out, b"tRNS", &alphas);
370    }
371    let compressed = zlib_compress(&scanlines, level).map_err(crate::compress_error)?;
372    write_chunk(&mut out, b"IDAT", &compressed);
373    write_chunk(&mut out, b"IEND", &[]);
374    Ok(out)
375}
376
377impl Encoder for PngEncoder {
378    fn set_icc_profile(&mut self, profile: Option<&[u8]>) -> Result<()> {
379        if self.state.is_some() {
380            return Err(PixelsError::invalid_argument(
381                "profile",
382                "the ICC profile must be set before write_header",
383            ));
384        }
385        self.icc = profile.map(<[u8]>::to_vec);
386        Ok(())
387    }
388
389    fn set_exif(&mut self, exif: Option<&[u8]>) -> Result<()> {
390        if self.state.is_some() {
391            return Err(PixelsError::invalid_argument(
392                "exif",
393                "the EXIF block must be set before write_header",
394            ));
395        }
396        self.exif = exif.map(<[u8]>::to_vec);
397        Ok(())
398    }
399
400    fn set_xmp(&mut self, xmp: Option<&[u8]>) -> Result<()> {
401        if self.state.is_some() {
402            return Err(PixelsError::invalid_argument(
403                "xmp",
404                "the XMP packet must be set before write_header",
405            ));
406        }
407        self.xmp = xmp.map(<[u8]>::to_vec);
408        Ok(())
409    }
410
411    fn set_density(&mut self, density: Option<Density>) -> Result<()> {
412        if self.state.is_some() {
413            return Err(PixelsError::invalid_argument(
414                "density",
415                "the density must be set before write_header",
416            ));
417        }
418        self.density = density;
419        Ok(())
420    }
421
422    fn write_header(&mut self, desc: &ImageDescriptor, sink: &mut dyn Sink) -> Result<()> {
423        if self.state.is_some() {
424            return Err(PixelsError::invalid_argument(
425                "descriptor",
426                "write_header called more than once",
427            ));
428        }
429        if desc.width == 0 || desc.height == 0 {
430            return Err(PixelsError::invalid_argument(
431                "descriptor",
432                format!(
433                    "PNG dimensions must be non-zero, got {}x{}",
434                    desc.width, desc.height
435                ),
436            ));
437        }
438        let (color_type, bit_depth) = Self::png_type_of(desc.pixel)?;
439
440        if self.palette.is_some() {
441            // Nothing can be written yet: the header's bit depth depends on
442            // how many colours the palette ends up with.
443            let channels = if desc.pixel.layout().has_alpha() {
444                4
445            } else {
446                3
447            };
448            self.state = Some(State {
449                descriptor: *desc,
450                bit_depth: 8,
451                stride: 1,
452                filtered: Vec::new(),
453                previous: Vec::new(),
454                swapped: Vec::new(),
455                true_colour: Some(Vec::with_capacity(
456                    desc.width as usize * desc.height as usize * channels,
457                )),
458                rows_written: 0,
459            });
460            return Ok(());
461        }
462
463        sink.write_all(&SIGNATURE)?;
464        let mut ihdr = Vec::with_capacity(13);
465        ihdr.extend_from_slice(&desc.width.to_be_bytes());
466        ihdr.extend_from_slice(&desc.height.to_be_bytes());
467        ihdr.push(bit_depth);
468        ihdr.push(color_type.to_byte());
469        // Compression 0 (DEFLATE), filter 0 (adaptive), interlace 0 (none) —
470        // the only values the specification defines.
471        ihdr.extend_from_slice(&[0, 0, 0]);
472        let mut chunk = Vec::new();
473        write_chunk(&mut chunk, b"IHDR", &ihdr);
474        write_metadata(
475            &mut chunk,
476            self.level,
477            self.icc.as_deref(),
478            self.exif.as_deref(),
479            self.xmp.as_deref(),
480            self.density,
481        )?;
482        sink.write_all(&chunk)?;
483
484        let row_bytes = desc.row_bytes();
485        let stride = (color_type.channels() * bit_depth as usize).div_ceil(8);
486        self.state = Some(State {
487            descriptor: *desc,
488            bit_depth,
489            stride,
490            // Each row costs its bytes plus one filter byte.
491            filtered: Vec::with_capacity((row_bytes + 1) * desc.height as usize),
492            previous: vec![0_u8; row_bytes],
493            swapped: Vec::new(),
494            true_colour: None,
495            rows_written: 0,
496        });
497        Ok(())
498    }
499
500    fn write_row(&mut self, row: &[u8], _sink: &mut dyn Sink) -> Result<()> {
501        let Some(state) = self.state.as_mut() else {
502            return Err(PixelsError::invalid_argument(
503                "row",
504                "write_row called before write_header",
505            ));
506        };
507        let expected = state.descriptor.row_bytes();
508        if row.len() != expected {
509            return Err(PixelsError::invalid_argument(
510                "row",
511                format!("row is {} bytes, expected {expected}", row.len()),
512            ));
513        }
514        if state.rows_written >= state.descriptor.height {
515            return Err(PixelsError::invalid_argument(
516                "row",
517                format!("more than {} rows written", state.descriptor.height),
518            ));
519        }
520
521        if let Some(pixels) = state.true_colour.as_mut() {
522            append_8_bit(row, state.descriptor.pixel, pixels);
523        } else if state.bit_depth == 16 {
524            let mut swapped = std::mem::take(&mut state.swapped);
525            to_big_endian_16(row, &mut swapped);
526            state.append_filtered(&swapped);
527            state.swapped = swapped;
528        } else {
529            state.append_filtered(row);
530        }
531        state.rows_written += 1;
532        Ok(())
533    }
534
535    fn finish(&mut self, sink: &mut dyn Sink) -> Result<()> {
536        let Some(state) = self.state.as_mut() else {
537            return Err(PixelsError::malformed(
538                "png",
539                "finish called before write_header",
540            ));
541        };
542        if state.rows_written != state.descriptor.height {
543            return Err(PixelsError::malformed(
544                "png",
545                format!(
546                    "{} of {} rows written; a partial image is never emitted",
547                    state.rows_written, state.descriptor.height
548                ),
549            ));
550        }
551
552        if let (Some(pixels), Some((colours, dither))) = (state.true_colour.take(), self.palette) {
553            let file = encode_palette(
554                &state.descriptor,
555                &pixels,
556                colours,
557                dither,
558                self.level,
559                (
560                    self.icc.as_deref(),
561                    self.exif.as_deref(),
562                    self.xmp.as_deref(),
563                    self.density,
564                ),
565            )?;
566            sink.write_all(&file)?;
567            return sink.flush();
568        }
569
570        let compressed =
571            zlib_compress(&state.filtered, self.level).map_err(crate::compress_error)?;
572        let mut chunk = Vec::with_capacity(compressed.len() + 12);
573        write_chunk(&mut chunk, b"IDAT", &compressed);
574        write_chunk(&mut chunk, b"IEND", &[]);
575        sink.write_all(&chunk)?;
576        sink.flush()?;
577
578        // Release the raster now rather than at drop; a caller that keeps the
579        // encoder around to inspect it should not keep the image too.
580        state.filtered = Vec::new();
581        state.previous = Vec::new();
582        state.swapped = Vec::new();
583        Ok(())
584    }
585}
586
587#[cfg(test)]
588#[allow(
589    clippy::unwrap_used,
590    clippy::expect_used,
591    clippy::indexing_slicing,
592    clippy::panic,
593    reason = "tests operate on known-good values and assert shapes directly"
594)]
595mod tests {
596    use super::*;
597    use crate::decoder::PngDecoder;
598    use otf_pixels_core::{Decoder, Limits};
599
600    /// Encode a full raster, returning the PNG bytes.
601    fn encode(desc: &ImageDescriptor, raster: &[u8], level: Level) -> Result<Vec<u8>> {
602        let mut encoder = PngEncoder::with_level(level);
603        let mut out: Vec<u8> = Vec::new();
604        encoder.write_header(desc, &mut out)?;
605        for row in raster.chunks_exact(desc.row_bytes()) {
606            encoder.write_row(row, &mut out)?;
607        }
608        encoder.finish(&mut out)?;
609        Ok(out)
610    }
611
612    /// Decode a full raster back.
613    fn decode(bytes: &[u8]) -> Result<(ImageDescriptor, Vec<u8>)> {
614        let mut decoder = PngDecoder::new(bytes, Limits::default())?;
615        let mut row = vec![0_u8; decoder.descriptor().row_bytes()];
616        let mut raster = Vec::new();
617        for _ in 0..decoder.descriptor().height {
618            decoder.read_row(&mut row)?;
619            raster.extend_from_slice(&row);
620        }
621        Ok((decoder.descriptor(), raster))
622    }
623
624    /// A deterministic non-trivial raster: gradients plus a hard edge, so the
625    /// filter heuristic has something to choose between.
626    fn sample(desc: &ImageDescriptor) -> Vec<u8> {
627        let mut raster = vec![0_u8; desc.byte_len().unwrap()];
628        for (index, byte) in raster.iter_mut().enumerate() {
629            let value = if index % 97 < 40 {
630                (index % 251) as u8
631            } else {
632                ((index * 7) % 13) as u8
633            };
634            *byte = value;
635        }
636        raster
637    }
638
639    #[test]
640    fn every_supported_format_round_trips() {
641        for format in [
642            PixelFormat::Gray8,
643            PixelFormat::Gray16,
644            PixelFormat::GrayA8,
645            PixelFormat::Rgb8,
646            PixelFormat::Rgb16,
647            PixelFormat::Rgba8,
648            PixelFormat::Rgba16,
649        ] {
650            let desc = ImageDescriptor::new(23, 17, format).unwrap();
651            let raster = sample(&desc);
652            let bytes = encode(&desc, &raster, Level::DEFAULT)
653                .unwrap_or_else(|e| panic!("encoding {format}: {e}"));
654            let (decoded_desc, decoded) =
655                decode(&bytes).unwrap_or_else(|e| panic!("decoding {format}: {e}"));
656            assert_eq!(decoded_desc.pixel, format, "{format} changed format");
657            assert_eq!(
658                (decoded_desc.width, decoded_desc.height),
659                (23, 17),
660                "{format}"
661            );
662            assert_eq!(decoded, raster, "{format} did not round-trip");
663        }
664    }
665
666    #[test]
667    fn every_level_round_trips_to_the_same_pixels() {
668        let desc = ImageDescriptor::new(19, 11, PixelFormat::Rgba8).unwrap();
669        let raster = sample(&desc);
670        for level in 0..=9 {
671            let level = Level::new(level).unwrap();
672            let bytes = encode(&desc, &raster, level).unwrap();
673            let (_, decoded) = decode(&bytes).unwrap();
674            assert_eq!(decoded, raster, "level {} did not round-trip", level.get());
675        }
676    }
677
678    #[test]
679    fn the_output_starts_with_the_signature_and_ihdr() {
680        let desc = ImageDescriptor::new(4, 4, PixelFormat::Rgb8).unwrap();
681        let bytes = encode(&desc, &sample(&desc), Level::FAST).unwrap();
682        assert_eq!(&bytes[..8], &SIGNATURE, "signature");
683        assert_eq!(&bytes[12..16], b"IHDR", "first chunk");
684        assert_eq!(
685            &bytes[bytes.len() - 8..bytes.len() - 4],
686            b"IEND",
687            "last chunk"
688        );
689    }
690
691    #[test]
692    fn sixteen_bit_samples_are_written_big_endian() {
693        // The one place a host-endianness bug would hide: encode a known
694        // sample and read the two bytes straight out of the decompressed
695        // scanline via a decode, which is byte-order-symmetric, plus assert
696        // the swap helper directly.
697        let mut out = Vec::new();
698        to_big_endian_16(&0x1234_u16.to_ne_bytes(), &mut out);
699        assert_eq!(out, vec![0x12, 0x34]);
700    }
701
702    #[test]
703    fn compression_shrinks_a_compressible_image() {
704        // A flat image is the clearest case: level 0 stores it, level 9 must
705        // do dramatically better, which proves the level reaches DEFLATE.
706        let desc = ImageDescriptor::new(64, 64, PixelFormat::Rgb8).unwrap();
707        let raster = vec![7_u8; desc.byte_len().unwrap()];
708        let stored = encode(&desc, &raster, Level::NONE).unwrap();
709        let packed = encode(&desc, &raster, Level::BEST).unwrap();
710        assert!(
711            packed.len() * 10 < stored.len(),
712            "level 9 produced {} bytes against level 0's {}",
713            packed.len(),
714            stored.len()
715        );
716        assert_eq!(decode(&packed).unwrap().1, raster);
717    }
718
719    #[test]
720    fn a_float_format_is_unsupported_not_a_panic() {
721        let desc = ImageDescriptor::new(2, 2, PixelFormat::RgbaF32).unwrap();
722        let mut encoder = PngEncoder::new();
723        let mut out = Vec::new();
724        let error = encoder.write_header(&desc, &mut out).unwrap_err();
725        assert!(matches!(error, PixelsError::Unsupported { .. }), "{error}");
726        assert!(
727            out.is_empty(),
728            "nothing should be written for a rejected format"
729        );
730    }
731
732    #[test]
733    fn a_short_image_is_an_error_not_a_truncated_png() {
734        let desc = ImageDescriptor::new(4, 4, PixelFormat::Gray8).unwrap();
735        let mut encoder = PngEncoder::new();
736        let mut out = Vec::new();
737        encoder.write_header(&desc, &mut out).unwrap();
738        encoder.write_row(&[0; 4], &mut out).unwrap();
739        let error = encoder.finish(&mut out).unwrap_err();
740        assert!(matches!(error, PixelsError::Malformed { .. }), "{error}");
741    }
742
743    #[test]
744    fn extra_rows_and_wrong_row_lengths_are_errors() {
745        let desc = ImageDescriptor::new(4, 1, PixelFormat::Gray8).unwrap();
746        let mut encoder = PngEncoder::new();
747        let mut out = Vec::new();
748        encoder.write_header(&desc, &mut out).unwrap();
749        assert!(encoder.write_row(&[0; 3], &mut out).is_err(), "short row");
750        assert!(encoder.write_row(&[0; 5], &mut out).is_err(), "long row");
751        encoder.write_row(&[0; 4], &mut out).unwrap();
752        assert!(encoder.write_row(&[0; 4], &mut out).is_err(), "extra row");
753    }
754
755    #[test]
756    fn rows_before_the_header_are_an_error() {
757        let mut encoder = PngEncoder::new();
758        let mut out = Vec::new();
759        assert!(encoder.write_row(&[0; 4], &mut out).is_err());
760        assert!(encoder.finish(&mut out).is_err());
761    }
762
763    #[test]
764    fn a_second_header_is_an_error() {
765        let desc = ImageDescriptor::new(2, 2, PixelFormat::Gray8).unwrap();
766        let mut encoder = PngEncoder::new();
767        let mut out = Vec::new();
768        encoder.write_header(&desc, &mut out).unwrap();
769        assert!(encoder.write_header(&desc, &mut out).is_err());
770    }
771
772    #[test]
773    fn zero_dimensions_never_reach_the_encoder() {
774        // `ImageDescriptor` refuses to hold a zero dimension, so the guard in
775        // `write_header` is defence in depth rather than the only check. Both
776        // matter: PNG's IHDR has no representation for an empty image, and a
777        // future descriptor change must not quietly start emitting one.
778        assert!(ImageDescriptor::new(0, 4, PixelFormat::Gray8).is_err());
779        assert!(ImageDescriptor::new(4, 0, PixelFormat::Gray8).is_err());
780    }
781
782    #[test]
783    fn options_map_quality_onto_compression_effort() {
784        assert_eq!(
785            PngEncoder::from_options(&EncodeOptions::default())
786                .level
787                .get(),
788            7
789        );
790        let lowest = EncodeOptions::with_quality(1).unwrap();
791        assert_eq!(PngEncoder::from_options(&lowest).level.get(), 1);
792        let highest = EncodeOptions::with_quality(100).unwrap();
793        assert_eq!(PngEncoder::from_options(&highest).level.get(), 9);
794    }
795
796    #[test]
797    fn a_single_pixel_image_round_trips() {
798        // The degenerate case where every filter's neighbours are all zero.
799        let desc = ImageDescriptor::new(1, 1, PixelFormat::Rgba8).unwrap();
800        let raster = vec![1, 2, 3, 4];
801        let bytes = encode(&desc, &raster, Level::BEST).unwrap();
802        assert_eq!(decode(&bytes).unwrap().1, raster);
803    }
804
805    /// Encode with `encoder`, returning the PNG bytes.
806    fn encode_with(desc: &ImageDescriptor, raster: &[u8], mut encoder: PngEncoder) -> Vec<u8> {
807        let mut out: Vec<u8> = Vec::new();
808        encoder.write_header(desc, &mut out).unwrap();
809        for row in raster.chunks_exact(desc.row_bytes()) {
810            encoder.write_row(row, &mut out).unwrap();
811        }
812        encoder.finish(&mut out).unwrap();
813        out
814    }
815
816    fn palette(colours: usize, dither: f32) -> PngEncoder {
817        PngEncoder::new().with_palette(colours, dither).unwrap()
818    }
819
820    /// The IHDR bit depth and colour type.
821    fn header(bytes: &[u8]) -> (u8, u8) {
822        (bytes[24], bytes[25])
823    }
824
825    fn has_chunk(bytes: &[u8], kind: &[u8; 4]) -> bool {
826        bytes.windows(4).any(|w| w == kind)
827    }
828
829    #[test]
830    fn a_palette_image_with_few_colours_round_trips_exactly() {
831        let desc = ImageDescriptor::new(13, 7, PixelFormat::Rgba8).unwrap();
832        let colours = [
833            [255, 0, 0, 255],
834            [0, 255, 0, 128],
835            [0, 0, 255, 0],
836            [9, 9, 9, 255],
837        ];
838        let raster: Vec<u8> = (0..13 * 7).flat_map(|i| colours[i % 4]).collect();
839        let bytes = encode_with(&desc, &raster, palette(256, 1.0));
840        assert_eq!(header(&bytes), (2, 3), "four colours fit two bits");
841        assert!(has_chunk(&bytes, b"PLTE") && has_chunk(&bytes, b"tRNS"));
842        let (out, decoded) = decode(&bytes).unwrap();
843        assert_eq!(out.pixel, PixelFormat::Rgba8);
844        // Every invisible pixel is the same entry, whatever colour it hid.
845        let visible: Vec<u8> = raster
846            .chunks_exact(4)
847            .flat_map(|p| {
848                if p[3] == 0 {
849                    [0; 4]
850                } else {
851                    [p[0], p[1], p[2], p[3]]
852                }
853            })
854            .collect();
855        assert_eq!(decoded, visible);
856    }
857
858    #[test]
859    fn the_bit_depth_is_the_smallest_that_holds_the_palette() {
860        for (count, depth) in [
861            (2_usize, 1_u8),
862            (3, 2),
863            (4, 2),
864            (5, 4),
865            (16, 4),
866            (17, 8),
867            (200, 8),
868        ] {
869            // Widths that leave a partial last byte at every depth.
870            let desc = ImageDescriptor::new(11, 3, PixelFormat::Rgb8).unwrap();
871            let raster: Vec<u8> = (0..33)
872                .flat_map(|i| {
873                    let v = (i % count) as u8;
874                    [v, v.wrapping_mul(3), 255 - v]
875                })
876                .collect();
877            let bytes = encode_with(&desc, &raster, palette(256, 1.0));
878            assert_eq!(header(&bytes), (depth, 3), "{count} colours");
879            assert!(!has_chunk(&bytes, b"tRNS"), "an opaque image needs no tRNS");
880            let (out, decoded) = decode(&bytes).unwrap();
881            assert_eq!(out.pixel, PixelFormat::Rgb8);
882            assert_eq!(decoded, raster, "{count} colours");
883        }
884    }
885
886    #[test]
887    fn a_photograph_is_reduced_to_the_palette_size() {
888        let desc = ImageDescriptor::new(64, 64, PixelFormat::Rgb8).unwrap();
889        let raster: Vec<u8> = (0..64 * 64 * 3)
890            .map(|i| ((i * 37 + i / 191) % 256) as u8)
891            .collect();
892        for dither in [0.0, 0.5, 1.0] {
893            let bytes = encode_with(&desc, &raster, palette(16, dither));
894            let plte = bytes.windows(4).position(|w| w == b"PLTE").unwrap();
895            let length = u32::from_be_bytes(bytes[plte - 4..plte].try_into().unwrap());
896            assert!(length <= 16 * 3, "{length} bytes of palette");
897            let (_, decoded) = decode(&bytes).unwrap();
898            let mut distinct: Vec<&[u8]> = decoded.chunks_exact(3).collect();
899            distinct.sort_unstable();
900            distinct.dedup();
901            assert!(distinct.len() <= 16);
902        }
903        // And on photographic noise, which true colour cannot compress, a
904        // palette is much smaller.
905        let mut seed = 12_345_u32;
906        let noisy: Vec<u8> = (0..64 * 64 * 3)
907            .map(|_| {
908                seed = seed.wrapping_mul(1_103_515_245).wrapping_add(12_345);
909                (seed >> 16) as u8
910            })
911            .collect();
912        let true_colour = encode(&desc, &noisy, Level::DEFAULT).unwrap();
913        let reduced = encode_with(&desc, &noisy, palette(16, 1.0));
914        assert!(
915            reduced.len() * 2 < true_colour.len(),
916            "{} against {}",
917            reduced.len(),
918            true_colour.len()
919        );
920    }
921
922    #[test]
923    fn grey_and_wide_images_take_a_palette_too() {
924        let desc = ImageDescriptor::new(4, 1, PixelFormat::Gray16).unwrap();
925        let raster: Vec<u8> = [0_u16, 65535, 0, 65535]
926            .iter()
927            .flat_map(|v| v.to_ne_bytes())
928            .collect();
929        let bytes = encode_with(&desc, &raster, palette(256, 1.0));
930        assert_eq!(header(&bytes), (1, 3));
931        assert_eq!(
932            decode(&bytes).unwrap().1,
933            [0, 0, 0, 255, 255, 255, 0, 0, 0, 255, 255, 255]
934        );
935    }
936
937    #[test]
938    fn options_choose_the_level_and_the_palette() {
939        let desc = ImageDescriptor::new(32, 32, PixelFormat::Rgb8).unwrap();
940        let raster: Vec<u8> = (0..32 * 32 * 3).map(|i| (i / 7 % 256) as u8).collect();
941        let with = |options: EncodeOptions| {
942            encode_with(&desc, &raster, PngEncoder::from_options(&options))
943        };
944        let stored = with(EncodeOptions::default().with_compression_level(0).unwrap());
945        let best = with(EncodeOptions::default().with_compression_level(9).unwrap());
946        assert!(best.len() < stored.len());
947        // Unset, the level still comes from quality, as it always did.
948        assert_eq!(
949            with(EncodeOptions::default()),
950            encode(&desc, &raster, Level::new(7).unwrap()).unwrap()
951        );
952        for options in [
953            EncodeOptions::default().with_palette(true),
954            EncodeOptions::default().with_colours(64).unwrap(),
955            EncodeOptions::default().with_dither(0.0).unwrap(),
956        ] {
957            assert_eq!(header(&with(options)).1, 3, "{options:?}");
958        }
959        assert_eq!(header(&with(EncodeOptions::default())).1, 2);
960    }
961
962    #[test]
963    fn palette_parameters_are_validated() {
964        assert!(PngEncoder::new().with_palette(1, 1.0).is_err());
965        assert!(PngEncoder::new().with_palette(257, 1.0).is_err());
966        assert!(PngEncoder::new().with_palette(2, 1.5).is_err());
967    }
968}