png/
encoder.rs

1use borrow::Cow;
2use io::{Read, Write};
3use ops::{Deref, DerefMut};
4use std::{borrow, error, fmt, io, mem, ops, result};
5
6use crc32fast::Hasher as Crc32;
7use flate2::write::ZlibEncoder;
8
9use crate::chunk::{self, ChunkType};
10use crate::common::{
11    AnimationControl, BitDepth, BlendOp, BytesPerPixel, ColorType, Compression, DisposeOp,
12    FrameControl, Info, ParameterError, ParameterErrorKind, PixelDimensions, ScaledFloat, Unit,
13};
14use crate::filter::{filter, Filter};
15use crate::text_metadata::{
16    encode_iso_8859_1, EncodableTextChunk, ITXtChunk, TEXtChunk, TextEncodingError, ZTXtChunk,
17};
18use crate::traits::WriteBytesExt;
19use crate::DeflateCompression;
20
21pub type Result<T> = result::Result<T, EncodingError>;
22
23#[derive(Debug)]
24pub enum EncodingError {
25    IoError(io::Error),
26    Format(FormatError),
27    Parameter(ParameterError),
28    LimitsExceeded,
29}
30
31#[derive(Debug)]
32pub struct FormatError {
33    inner: FormatErrorKind,
34}
35
36#[derive(Debug)]
37enum FormatErrorKind {
38    ZeroWidth,
39    ZeroHeight,
40    InvalidColorCombination(BitDepth, ColorType),
41    NoPalette,
42    // TODO: wait, what?
43    WrittenTooMuch(usize),
44    NotAnimated,
45    OutOfBounds,
46    EndReached,
47    ZeroFrames,
48    MissingFrames,
49    MissingData(usize),
50    Unrecoverable,
51    BadTextEncoding(TextEncodingError),
52}
53
54impl error::Error for EncodingError {
55    fn cause(&self) -> Option<&(dyn error::Error + 'static)> {
56        match self {
57            EncodingError::IoError(err) => Some(err),
58            _ => None,
59        }
60    }
61}
62
63impl fmt::Display for EncodingError {
64    fn fmt(&self, fmt: &mut fmt::Formatter) -> result::Result<(), fmt::Error> {
65        use self::EncodingError::*;
66        match self {
67            IoError(err) => write!(fmt, "{}", err),
68            Format(desc) => write!(fmt, "{}", desc),
69            Parameter(desc) => write!(fmt, "{}", desc),
70            LimitsExceeded => write!(fmt, "Limits are exceeded."),
71        }
72    }
73}
74
75impl fmt::Display for FormatError {
76    fn fmt(&self, fmt: &mut fmt::Formatter) -> result::Result<(), fmt::Error> {
77        use FormatErrorKind::*;
78        match self.inner {
79            ZeroWidth => write!(fmt, "Zero width not allowed"),
80            ZeroHeight => write!(fmt, "Zero height not allowed"),
81            ZeroFrames => write!(fmt, "Zero frames not allowed"),
82            InvalidColorCombination(depth, color) => write!(
83                fmt,
84                "Invalid combination of bit-depth '{:?}' and color-type '{:?}'",
85                depth, color
86            ),
87            NoPalette => write!(fmt, "can't write indexed image without palette"),
88            WrittenTooMuch(index) => write!(fmt, "wrong data size, got {} bytes too many", index),
89            NotAnimated => write!(fmt, "not an animation"),
90            OutOfBounds => write!(
91                fmt,
92                "the dimension and position go over the frame boundaries"
93            ),
94            EndReached => write!(fmt, "all the frames have been already written"),
95            MissingFrames => write!(fmt, "there are still frames to be written"),
96            MissingData(n) => write!(fmt, "there are still {} bytes to be written", n),
97            Unrecoverable => write!(
98                fmt,
99                "a previous error put the writer into an unrecoverable state"
100            ),
101            BadTextEncoding(tee) => match tee {
102                TextEncodingError::Unrepresentable => write!(
103                    fmt,
104                    "The text metadata cannot be encoded into valid ISO 8859-1"
105                ),
106                TextEncodingError::InvalidKeywordSize => write!(fmt, "Invalid keyword size"),
107                TextEncodingError::CompressionError => {
108                    write!(fmt, "Unable to compress text metadata")
109                }
110            },
111        }
112    }
113}
114
115impl From<io::Error> for EncodingError {
116    fn from(err: io::Error) -> EncodingError {
117        EncodingError::IoError(err)
118    }
119}
120
121impl From<EncodingError> for io::Error {
122    fn from(err: EncodingError) -> io::Error {
123        io::Error::new(io::ErrorKind::Other, err.to_string())
124    }
125}
126
127// Private impl.
128impl From<FormatErrorKind> for FormatError {
129    fn from(kind: FormatErrorKind) -> Self {
130        FormatError { inner: kind }
131    }
132}
133
134impl From<TextEncodingError> for EncodingError {
135    fn from(tee: TextEncodingError) -> Self {
136        EncodingError::Format(FormatError {
137            inner: FormatErrorKind::BadTextEncoding(tee),
138        })
139    }
140}
141
142/// PNG Encoder.
143///
144/// This configures the PNG format options such as animation chunks, palette use, color types,
145/// auxiliary chunks etc.
146///
147/// FIXME: Configuring APNG might be easier (less individual errors) if we had an _adapter_ which
148/// borrows this mutably but guarantees that `info.frame_control` is not `None`.
149pub struct Encoder<'a, W: Write> {
150    w: W,
151    info: Info<'a>,
152    options: Options,
153}
154
155/// Encoding options, internal type, forwarded to the Writer.
156#[derive(Default)]
157struct Options {
158    filter: Filter,
159    sep_def_img: bool,
160    validate_sequence: bool,
161    compression: DeflateCompression,
162}
163
164impl<'a, W: Write> Encoder<'a, W> {
165    pub fn new(w: W, width: u32, height: u32) -> Encoder<'static, W> {
166        Encoder {
167            w,
168            info: Info::with_size(width, height),
169            options: Options::default(),
170        }
171    }
172
173    pub fn with_info(w: W, info: Info<'a>) -> Result<Encoder<'a, W>> {
174        if info.animation_control.is_some() != info.frame_control.is_some() {
175            return Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()));
176        }
177
178        if let Some(actl) = info.animation_control {
179            if actl.num_frames == 0 {
180                return Err(EncodingError::Format(FormatErrorKind::ZeroFrames.into()));
181            }
182        }
183
184        Ok(Encoder {
185            w,
186            info,
187            options: Options::default(),
188        })
189    }
190
191    /// Specify that the image is animated.
192    ///
193    /// `num_frames` controls how many frames the animation has, while
194    /// `num_plays` controls how many times the animation should be
195    /// repeated until it stops, if it's zero then it will repeat
196    /// infinitely.
197    ///
198    /// When this method is returns successfully then the images written will be encoded as fdAT
199    /// chunks, except for the first image that is still encoded as `IDAT`. You can control if the
200    /// first frame should be treated as an animation frame with [`Encoder::set_sep_def_img()`].
201    ///
202    /// This method returns an error if `num_frames` is 0.
203    pub fn set_animated(&mut self, num_frames: u32, num_plays: u32) -> Result<()> {
204        if num_frames == 0 {
205            return Err(EncodingError::Format(FormatErrorKind::ZeroFrames.into()));
206        }
207
208        let actl = AnimationControl {
209            num_frames,
210            num_plays,
211        };
212
213        let fctl = FrameControl {
214            sequence_number: 0,
215            width: self.info.width,
216            height: self.info.height,
217            ..Default::default()
218        };
219
220        self.info.animation_control = Some(actl);
221        self.info.frame_control = Some(fctl);
222        Ok(())
223    }
224
225    /// Mark the first animated frame as a 'separate default image'.
226    ///
227    /// In APNG each animated frame is preceded by a special control chunk, `fcTL`. It's up to the
228    /// encoder to decide if the first image, the standard `IDAT` data, should be part of the
229    /// animation by emitting this chunk or by not doing so. A default image that is _not_ part of
230    /// the animation is often interpreted as a thumbnail.
231    ///
232    /// This method will return an error when animation control was not configured
233    /// (which is done by calling [`Encoder::set_animated`]).
234    pub fn set_sep_def_img(&mut self, sep_def_img: bool) -> Result<()> {
235        if self.info.animation_control.is_some() {
236            self.options.sep_def_img = sep_def_img;
237            Ok(())
238        } else {
239            Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()))
240        }
241    }
242
243    /// Sets the raw byte contents of the PLTE chunk. This method accepts
244    /// both borrowed and owned byte data.
245    pub fn set_palette<T: Into<Cow<'a, [u8]>>>(&mut self, palette: T) {
246        self.info.palette = Some(palette.into());
247    }
248
249    /// Sets the raw byte contents of the tRNS chunk. This method accepts
250    /// both borrowed and owned byte data.
251    pub fn set_trns<T: Into<Cow<'a, [u8]>>>(&mut self, trns: T) {
252        self.info.trns = Some(trns.into());
253    }
254
255    /// Set the display gamma of the source system on which the image was generated or last edited.
256    pub fn set_source_gamma(&mut self, source_gamma: ScaledFloat) {
257        self.info.source_gamma = Some(source_gamma);
258    }
259
260    /// Set the chromaticities for the source system's display channels (red, green, blue) and the whitepoint
261    /// of the source system on which the image was generated or last edited.
262    pub fn set_source_chromaticities(
263        &mut self,
264        source_chromaticities: super::SourceChromaticities,
265    ) {
266        self.info.source_chromaticities = Some(source_chromaticities);
267    }
268
269    /// Mark the image data as conforming to the SRGB color space with the specified rendering intent.
270    ///
271    /// Any ICC profiles will be ignored.
272    ///
273    /// Source gamma and chromaticities will be written only if they're set to fallback
274    /// values specified in [11.3.2.5](https://www.w3.org/TR/png-3/#sRGB-gAMA-cHRM).
275    pub fn set_source_srgb(&mut self, rendering_intent: super::SrgbRenderingIntent) {
276        self.info.set_source_srgb(rendering_intent);
277    }
278
279    /// Start encoding by writing the header data.
280    ///
281    /// The remaining data can be supplied by methods on the returned [`Writer`].
282    pub fn write_header(self) -> Result<Writer<W>> {
283        Writer::new(self.w, PartialInfo::new(&self.info), self.options).init(&self.info)
284    }
285
286    /// Set the color of the encoded image.
287    ///
288    /// These correspond to the color types in the png IHDR data that will be written. The length
289    /// of the image data that is later supplied must match the color type, otherwise an error will
290    /// be emitted.
291    pub fn set_color(&mut self, color: ColorType) {
292        self.info.color_type = color;
293    }
294
295    /// Set the indicated depth of the image data.
296    pub fn set_depth(&mut self, depth: BitDepth) {
297        self.info.bit_depth = depth;
298    }
299
300    /// Set compression parameters, see [Compression] for the available options.
301    pub fn set_compression(&mut self, compression: Compression) {
302        self.set_deflate_compression(DeflateCompression::from_simple(compression));
303        self.set_filter(Filter::from_simple(compression));
304    }
305
306    /// Provides in-depth customization of DEFLATE compression options.
307    ///
308    /// For a simpler selection of compression options see [Self::set_compression].
309    pub fn set_deflate_compression(&mut self, compression: DeflateCompression) {
310        self.options.compression = compression;
311    }
312
313    /// Set the used filter type.
314    ///
315    /// The default filter is [`Filter::Adaptive`] which automatically selects the best filter
316    /// for each row of the image.
317    ///
318    /// You should only change this if you are after very fast compression,
319    /// and either don't care about compression ratio or know exactly what works best for your images.
320    pub fn set_filter(&mut self, filter: Filter) {
321        self.options.filter = filter;
322    }
323
324    /// Set the fraction of time every frame is going to be displayed, in seconds.
325    ///
326    /// *Note that this parameter can be set for each individual frame after
327    /// [`Encoder::write_header`] is called. (see [`Writer::set_frame_delay`])*
328    ///
329    /// If the denominator is 0, it is to be treated as if it were 100
330    /// (that is, the numerator then specifies 1/100ths of a second).
331    /// If the value of the numerator is 0 the decoder should render the next frame
332    /// as quickly as possible, though viewers may impose a reasonable lower bound.
333    ///
334    /// The default value is 0 for both the numerator and denominator.
335    ///
336    /// This method will return an error if the image is not animated.
337    /// (see [`set_animated`])
338    ///
339    /// [`write_header`]: Self::write_header
340    /// [`set_animated`]: Self::set_animated
341    pub fn set_frame_delay(&mut self, numerator: u16, denominator: u16) -> Result<()> {
342        if let Some(ref mut fctl) = self.info.frame_control {
343            fctl.delay_den = denominator;
344            fctl.delay_num = numerator;
345            Ok(())
346        } else {
347            Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()))
348        }
349    }
350
351    /// Set the blend operation for every frame.
352    ///
353    /// The blend operation specifies whether the frame is to be alpha blended
354    /// into the current output buffer content, or whether it should completely
355    /// replace its region in the output buffer.
356    ///
357    /// *Note that this parameter can be set for each individual frame after
358    /// [`write_header`] is called. (see [`Writer::set_blend_op`])*
359    ///
360    /// See the [`BlendOp`] documentation for the possible values and their effects.
361    ///
362    /// *Note that for the first frame the two blend modes are functionally
363    /// equivalent due to the clearing of the output buffer at the beginning
364    /// of each play.*
365    ///
366    /// The default value is [`BlendOp::Source`].
367    ///
368    /// This method will return an error if the image is not animated.
369    /// (see [`set_animated`])
370    ///
371    /// [`write_header`]: Self::write_header
372    /// [`set_animated`]: Self::set_animated
373    pub fn set_blend_op(&mut self, op: BlendOp) -> Result<()> {
374        if let Some(ref mut fctl) = self.info.frame_control {
375            fctl.blend_op = op;
376            Ok(())
377        } else {
378            Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()))
379        }
380    }
381
382    /// Set the dispose operation for every frame.
383    ///
384    /// The dispose operation specifies how the output buffer should be changed
385    /// at the end of the delay (before rendering the next frame)
386    ///
387    /// *Note that this parameter can be set for each individual frame after
388    /// [`write_header`] is called (see [`Writer::set_dispose_op`])*
389    ///
390    /// See the [`DisposeOp`] documentation for the possible values and their effects.
391    ///
392    /// *Note that if the first frame uses [`DisposeOp::Previous`]
393    /// it will be treated as [`DisposeOp::Background`].*
394    ///
395    /// The default value is [`DisposeOp::None`].
396    ///
397    /// This method will return an error if the image is not animated.
398    /// (see [`set_animated`])
399    ///
400    /// [`set_animated`]: Self::set_animated
401    /// [`write_header`]: Self::write_header
402    pub fn set_dispose_op(&mut self, op: DisposeOp) -> Result<()> {
403        if let Some(ref mut fctl) = self.info.frame_control {
404            fctl.dispose_op = op;
405            Ok(())
406        } else {
407            Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()))
408        }
409    }
410    pub fn set_pixel_dims(&mut self, pixel_dims: Option<PixelDimensions>) {
411        self.info.pixel_dims = pixel_dims
412    }
413    /// Convenience function to add tEXt chunks to [`Info`] struct
414    pub fn add_text_chunk(&mut self, keyword: String, text: String) -> Result<()> {
415        let text_chunk = TEXtChunk::new(keyword, text);
416        self.info.uncompressed_latin1_text.push(text_chunk);
417        Ok(())
418    }
419
420    /// Convenience function to add zTXt chunks to [`Info`] struct
421    pub fn add_ztxt_chunk(&mut self, keyword: String, text: String) -> Result<()> {
422        let text_chunk = ZTXtChunk::new(keyword, text);
423        self.info.compressed_latin1_text.push(text_chunk);
424        Ok(())
425    }
426
427    /// Convenience function to add iTXt chunks to [`Info`] struct
428    ///
429    /// This function only sets the `keyword` and `text` field of the iTXt chunk.
430    /// To set the other fields, create a [`ITXtChunk`] directly, and then encode it to the output stream.
431    pub fn add_itxt_chunk(&mut self, keyword: String, text: String) -> Result<()> {
432        let text_chunk = ITXtChunk::new(keyword, text);
433        self.info.utf8_text.push(text_chunk);
434        Ok(())
435    }
436
437    /// Validate the written image sequence.
438    ///
439    /// When validation is turned on (it's turned off by default) then attempts to write more than
440    /// one `IDAT` image or images beyond the number of frames indicated in the animation control
441    /// chunk will fail and return an error result instead. Attempts to [finish][finish] the image
442    /// with missing frames will also return an error.
443    ///
444    /// [finish]: StreamWriter::finish
445    ///
446    /// (It's possible to circumvent these checks by writing raw chunks instead.)
447    pub fn validate_sequence(&mut self, validate: bool) {
448        self.options.validate_sequence = validate;
449    }
450}
451
452/// PNG writer
453///
454/// Progresses through the image by writing images, frames, or raw individual chunks. This is
455/// constructed through [`Encoder::write_header()`].
456///
457/// FIXME: Writing of animated chunks might be clearer if we had an _adapter_ that you would call
458/// to guarantee the next image to be prefaced with a fcTL-chunk, and all other chunks would be
459/// guaranteed to be `IDAT`/not affected by APNG's frame control.
460pub struct Writer<W: Write> {
461    /// The underlying writer.
462    w: W,
463    /// The local version of the `Info` struct.
464    info: PartialInfo,
465    /// Global encoding options.
466    options: Options,
467    /// The total number of image frames, counting all consecutive IDAT and fdAT chunks.
468    images_written: u64,
469    /// The total number of animation frames, that is equivalent to counting fcTL chunks.
470    animation_written: u32,
471    /// A flag to note when the IEND chunk was already added.
472    /// This is only set on code paths that drop `Self` to control the destructor.
473    iend_written: bool,
474}
475
476/// Contains the subset of attributes of [Info] needed for [Writer] to function
477struct PartialInfo {
478    width: u32,
479    height: u32,
480    bit_depth: BitDepth,
481    color_type: ColorType,
482    frame_control: Option<FrameControl>,
483    animation_control: Option<AnimationControl>,
484    has_palette: bool,
485}
486
487impl PartialInfo {
488    fn new(info: &Info) -> Self {
489        PartialInfo {
490            width: info.width,
491            height: info.height,
492            bit_depth: info.bit_depth,
493            color_type: info.color_type,
494            frame_control: info.frame_control,
495            animation_control: info.animation_control,
496            has_palette: info.palette.is_some(),
497        }
498    }
499
500    fn bpp_in_prediction(&self) -> BytesPerPixel {
501        BytesPerPixel::from_usize(self.bytes_per_pixel())
502    }
503
504    fn bytes_per_pixel(&self) -> usize {
505        self.color_type.bytes_per_pixel(self.bit_depth)
506    }
507
508    fn raw_row_length(&self) -> usize {
509        self.raw_row_length_from_width(self.width)
510    }
511
512    fn raw_row_length_from_width(&self, width: u32) -> usize {
513        self.color_type
514            .raw_row_length_from_width(self.bit_depth, width)
515    }
516}
517
518const DEFAULT_BUFFER_LENGTH: usize = 4 * 1024;
519
520pub(crate) fn write_chunk<W: Write>(mut w: W, name: chunk::ChunkType, data: &[u8]) -> Result<()> {
521    w.write_be(data.len() as u32)?;
522    w.write_all(&name.0)?;
523    w.write_all(data)?;
524    let mut crc = Crc32::new();
525    crc.update(&name.0);
526    crc.update(data);
527    w.write_be(crc.finalize())?;
528    Ok(())
529}
530
531impl<W: Write> Writer<W> {
532    fn new(w: W, info: PartialInfo, options: Options) -> Writer<W> {
533        Writer {
534            w,
535            info,
536            options,
537            images_written: 0,
538            animation_written: 0,
539            iend_written: false,
540        }
541    }
542
543    fn init(mut self, info: &Info<'_>) -> Result<Self> {
544        if self.info.width == 0 {
545            return Err(EncodingError::Format(FormatErrorKind::ZeroWidth.into()));
546        }
547
548        if self.info.height == 0 {
549            return Err(EncodingError::Format(FormatErrorKind::ZeroHeight.into()));
550        }
551
552        if self
553            .info
554            .color_type
555            .is_combination_invalid(self.info.bit_depth)
556        {
557            return Err(EncodingError::Format(
558                FormatErrorKind::InvalidColorCombination(self.info.bit_depth, self.info.color_type)
559                    .into(),
560            ));
561        }
562
563        self.encode_header(info)?;
564
565        Ok(self)
566    }
567
568    /// Encode PNG signature, IHDR, and then chunks that were added to the `Info`
569    fn encode_header(&mut self, info: &Info<'_>) -> Result<()> {
570        self.w.write_all(&[137, 80, 78, 71, 13, 10, 26, 10])?; // PNG signature
571
572        // Encode the IHDR chunk
573        let mut data = [0; 13];
574        data[..4].copy_from_slice(&info.width.to_be_bytes());
575        data[4..8].copy_from_slice(&info.height.to_be_bytes());
576        data[8] = info.bit_depth as u8;
577        data[9] = info.color_type as u8;
578        data[12] = info.interlaced as u8;
579        self.write_chunk(chunk::IHDR, &data)?;
580
581        // Encode the pHYs chunk
582        if let Some(pd) = info.pixel_dims {
583            let mut phys_data = [0; 9];
584            phys_data[0..4].copy_from_slice(&pd.xppu.to_be_bytes());
585            phys_data[4..8].copy_from_slice(&pd.yppu.to_be_bytes());
586            match pd.unit {
587                Unit::Meter => phys_data[8] = 1,
588                Unit::Unspecified => phys_data[8] = 0,
589            }
590            self.write_chunk(chunk::pHYs, &phys_data)?;
591        }
592
593        // If specified, the sRGB information overrides the source gamma and chromaticities.
594        if let Some(srgb) = &info.srgb {
595            srgb.encode(&mut self.w)?;
596
597            // gAMA and cHRM are optional, for backwards compatibility
598            let srgb_gamma = crate::srgb::substitute_gamma();
599            if Some(srgb_gamma) == info.source_gamma {
600                srgb_gamma.encode_gama(&mut self.w)?
601            }
602            let srgb_chromaticities = crate::srgb::substitute_chromaticities();
603            if Some(srgb_chromaticities) == info.source_chromaticities {
604                srgb_chromaticities.encode(&mut self.w)?;
605            }
606        } else {
607            if let Some(gma) = info.source_gamma {
608                gma.encode_gama(&mut self.w)?
609            }
610            if let Some(chrms) = info.source_chromaticities {
611                chrms.encode(&mut self.w)?;
612            }
613            if let Some(iccp) = &info.icc_profile {
614                self.write_iccp_chunk("_", iccp)?
615            }
616        }
617
618        if let Some(exif) = &info.exif_metadata {
619            self.write_chunk(chunk::eXIf, exif)?;
620        }
621
622        if let Some(actl) = info.animation_control {
623            actl.encode(&mut self.w)?;
624        }
625
626        // The position of the PLTE chunk is important, it must come before the tRNS chunk and after
627        // many of the other metadata chunks.
628        if let Some(p) = &info.palette {
629            self.write_chunk(chunk::PLTE, p)?;
630        };
631
632        if let Some(t) = &info.trns {
633            self.write_chunk(chunk::tRNS, t)?;
634        }
635
636        for text_chunk in &info.uncompressed_latin1_text {
637            self.write_text_chunk(text_chunk)?;
638        }
639
640        for text_chunk in &info.compressed_latin1_text {
641            self.write_text_chunk(text_chunk)?;
642        }
643
644        for text_chunk in &info.utf8_text {
645            self.write_text_chunk(text_chunk)?;
646        }
647
648        Ok(())
649    }
650
651    /// Write a raw chunk of PNG data.
652    ///
653    /// This function calculates the required CRC sum so this should not be included in the input
654    /// `data`, otherwise the data is not filtered in any way. This function returns an error if
655    /// the length of `data` can't be parsed as a `u32` though the length of the chunk data should
656    /// not exceed `i32::MAX` or 2,147,483,647.
657    pub fn write_chunk(&mut self, name: ChunkType, data: &[u8]) -> Result<()> {
658        use std::convert::TryFrom;
659
660        if u32::try_from(data.len()).map_or(true, |length| length > i32::MAX as u32) {
661            let kind = FormatErrorKind::WrittenTooMuch(data.len() - i32::MAX as usize);
662            return Err(EncodingError::Format(kind.into()));
663        }
664
665        write_chunk(&mut self.w, name, data)
666    }
667
668    pub fn write_text_chunk<T: EncodableTextChunk>(&mut self, text_chunk: &T) -> Result<()> {
669        text_chunk.encode(&mut self.w)
670    }
671
672    fn write_iccp_chunk(&mut self, profile_name: &str, icc_profile: &[u8]) -> Result<()> {
673        let profile_name = encode_iso_8859_1(profile_name)?;
674        if profile_name.is_empty() || profile_name.len() > 79 {
675            return Err(TextEncodingError::InvalidKeywordSize.into());
676        }
677
678        let estimated_compressed_size = icc_profile.len() * 3 / 4;
679        let chunk_size = profile_name
680            .len()
681            .checked_add(2) // string NUL + compression type. Checked add optimizes out later Vec reallocations.
682            .and_then(|s| s.checked_add(estimated_compressed_size))
683            .ok_or(EncodingError::LimitsExceeded)?;
684
685        let mut data = Vec::new();
686        data.try_reserve_exact(chunk_size)
687            .map_err(|_| EncodingError::LimitsExceeded)?;
688
689        data.extend(profile_name.into_iter().chain([0, 0]));
690
691        let mut encoder = ZlibEncoder::new(data, flate2::Compression::default());
692        encoder.write_all(icc_profile)?;
693
694        self.write_chunk(chunk::iCCP, &encoder.finish()?)
695    }
696
697    /// Check if we should allow writing another image.
698    fn validate_new_image(&self) -> Result<()> {
699        if !self.options.validate_sequence {
700            return Ok(());
701        }
702
703        match self.info.animation_control {
704            None => {
705                if self.images_written == 0 {
706                    Ok(())
707                } else {
708                    Err(EncodingError::Format(FormatErrorKind::EndReached.into()))
709                }
710            }
711            Some(_) => {
712                if self.info.frame_control.is_some() {
713                    Ok(())
714                } else {
715                    Err(EncodingError::Format(FormatErrorKind::EndReached.into()))
716                }
717            }
718        }
719    }
720
721    fn validate_sequence_done(&self) -> Result<()> {
722        if !self.options.validate_sequence {
723            return Ok(());
724        }
725
726        if (self.info.animation_control.is_some() && self.info.frame_control.is_some())
727            || self.images_written == 0
728        {
729            Err(EncodingError::Format(FormatErrorKind::MissingFrames.into()))
730        } else {
731            Ok(())
732        }
733    }
734
735    const MAX_IDAT_CHUNK_LEN: u32 = u32::MAX >> 1;
736    #[allow(non_upper_case_globals)]
737    const MAX_fdAT_CHUNK_LEN: u32 = (u32::MAX >> 1) - 4;
738
739    /// Writes the next image data.
740    pub fn write_image_data(&mut self, data: &[u8]) -> Result<()> {
741        if self.info.color_type == ColorType::Indexed && !self.info.has_palette {
742            return Err(EncodingError::Format(FormatErrorKind::NoPalette.into()));
743        }
744
745        self.validate_new_image()?;
746
747        let width: usize;
748        let height: usize;
749        if let Some(ref mut fctl) = self.info.frame_control {
750            width = fctl.width as usize;
751            height = fctl.height as usize;
752        } else {
753            width = self.info.width as usize;
754            height = self.info.height as usize;
755        }
756
757        let in_len = self.info.raw_row_length_from_width(width as u32) - 1;
758        let data_size = in_len * height;
759        if data_size != data.len() {
760            return Err(EncodingError::Parameter(
761                ParameterErrorKind::ImageBufferSize {
762                    expected: data_size,
763                    actual: data.len(),
764                }
765                .into(),
766            ));
767        }
768
769        let prev = vec![0; in_len];
770        let mut prev = prev.as_slice();
771
772        let bpp = self.info.bpp_in_prediction();
773        let filter_method = self.options.filter;
774
775        let zlib_encoded = match self.options.compression {
776            DeflateCompression::NoCompression => {
777                let mut compressor =
778                    fdeflate::StoredOnlyCompressor::new(std::io::Cursor::new(Vec::new()))?;
779                for line in data.chunks(in_len) {
780                    compressor.write_data(&[0])?;
781                    compressor.write_data(line)?;
782                }
783                compressor.finish()?.into_inner()
784            }
785            DeflateCompression::FdeflateUltraFast => {
786                let mut compressor = fdeflate::Compressor::new(std::io::Cursor::new(Vec::new()))?;
787
788                let mut current = vec![0; in_len + 1];
789                for line in data.chunks(in_len) {
790                    let filter_type = filter(filter_method, bpp, prev, line, &mut current[1..]);
791
792                    current[0] = filter_type as u8;
793                    compressor.write_data(&current)?;
794                    prev = line;
795                }
796
797                let compressed = compressor.finish()?.into_inner();
798                if compressed.len()
799                    > fdeflate::StoredOnlyCompressor::<()>::compressed_size((in_len + 1) * height)
800                {
801                    // Write uncompressed data since the result from fast compression would take
802                    // more space than that.
803                    //
804                    // This is essentially a fallback to NoCompression.
805                    let mut compressor =
806                        fdeflate::StoredOnlyCompressor::new(std::io::Cursor::new(Vec::new()))?;
807                    for line in data.chunks(in_len) {
808                        compressor.write_data(&[0])?;
809                        compressor.write_data(line)?;
810                    }
811                    compressor.finish()?.into_inner()
812                } else {
813                    compressed
814                }
815            }
816            DeflateCompression::Level(level) => {
817                let mut current = vec![0; in_len];
818
819                let mut zlib =
820                    ZlibEncoder::new(Vec::new(), flate2::Compression::new(u32::from(level)));
821                for line in data.chunks(in_len) {
822                    let filter_type = filter(filter_method, bpp, prev, line, &mut current);
823
824                    zlib.write_all(&[filter_type as u8])?;
825                    zlib.write_all(&current)?;
826                    prev = line;
827                }
828                zlib.finish()?
829            }
830        };
831
832        match self.info.frame_control {
833            None => {
834                self.write_zlib_encoded_idat(&zlib_encoded)?;
835            }
836            Some(_) if self.should_skip_frame_control_on_default_image() => {
837                self.write_zlib_encoded_idat(&zlib_encoded)?;
838            }
839            Some(ref mut fctl) => {
840                fctl.encode(&mut self.w)?;
841                fctl.sequence_number = fctl.sequence_number.wrapping_add(1);
842                self.animation_written += 1;
843
844                // If the default image is the first frame of an animation, it's still an IDAT.
845                if self.images_written == 0 {
846                    self.write_zlib_encoded_idat(&zlib_encoded)?;
847                } else {
848                    let buff_size = zlib_encoded.len().min(Self::MAX_fdAT_CHUNK_LEN as usize);
849                    let mut alldata = vec![0u8; 4 + buff_size];
850                    for chunk in zlib_encoded.chunks(Self::MAX_fdAT_CHUNK_LEN as usize) {
851                        alldata[..4].copy_from_slice(&fctl.sequence_number.to_be_bytes());
852                        alldata[4..][..chunk.len()].copy_from_slice(chunk);
853                        write_chunk(&mut self.w, chunk::fdAT, &alldata[..4 + chunk.len()])?;
854                        fctl.sequence_number = fctl.sequence_number.wrapping_add(1);
855                    }
856                }
857            }
858        }
859
860        self.increment_images_written();
861
862        Ok(())
863    }
864
865    fn increment_images_written(&mut self) {
866        self.images_written = self.images_written.saturating_add(1);
867
868        if let Some(actl) = self.info.animation_control {
869            if actl.num_frames <= self.animation_written {
870                // If we've written all animation frames, all following will be normal image chunks.
871                self.info.frame_control = None;
872            }
873        }
874    }
875
876    fn write_iend(&mut self) -> Result<()> {
877        self.iend_written = true;
878        self.write_chunk(chunk::IEND, &[])
879    }
880
881    fn should_skip_frame_control_on_default_image(&self) -> bool {
882        self.options.sep_def_img && self.images_written == 0
883    }
884
885    fn write_zlib_encoded_idat(&mut self, zlib_encoded: &[u8]) -> Result<()> {
886        for chunk in zlib_encoded.chunks(Self::MAX_IDAT_CHUNK_LEN as usize) {
887            self.write_chunk(chunk::IDAT, chunk)?;
888        }
889        Ok(())
890    }
891
892    /// Set the used filter type.
893    ///
894    /// The default filter is [`Filter::Adaptive`] which automatically selects the best filter
895    /// for each row of the image.
896    ///
897    /// You should only change this if you are after very fast compression,
898    /// and either don't care about compression ratio or know exactly what works best for your images.
899    pub fn set_filter(&mut self, filter: Filter) {
900        self.options.filter = filter;
901    }
902
903    /// Set the fraction of time the following frames are going to be displayed,
904    /// in seconds
905    ///
906    /// If the denominator is 0, it is to be treated as if it were 100
907    /// (that is, the numerator then specifies 1/100ths of a second).
908    /// If the value of the numerator is 0 the decoder should render the next frame
909    /// as quickly as possible, though viewers may impose a reasonable lower bound.
910    ///
911    /// This method will return an error if the image is not animated.
912    pub fn set_frame_delay(&mut self, numerator: u16, denominator: u16) -> Result<()> {
913        if let Some(ref mut fctl) = self.info.frame_control {
914            fctl.delay_den = denominator;
915            fctl.delay_num = numerator;
916            Ok(())
917        } else {
918            Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()))
919        }
920    }
921
922    /// Set the dimension of the following frames.
923    ///
924    /// This function will return an error when:
925    /// - The image is not an animated;
926    ///
927    /// - The selected dimension, considering also the current frame position,
928    ///   goes outside the image boundaries;
929    ///
930    /// - One or both the width and height are 0;
931    ///
932    // ??? TODO ???
933    // - The next frame is the default image
934    pub fn set_frame_dimension(&mut self, width: u32, height: u32) -> Result<()> {
935        if let Some(ref mut fctl) = self.info.frame_control {
936            if Some(width) > self.info.width.checked_sub(fctl.x_offset)
937                || Some(height) > self.info.height.checked_sub(fctl.y_offset)
938            {
939                return Err(EncodingError::Format(FormatErrorKind::OutOfBounds.into()));
940            } else if width == 0 {
941                return Err(EncodingError::Format(FormatErrorKind::ZeroWidth.into()));
942            } else if height == 0 {
943                return Err(EncodingError::Format(FormatErrorKind::ZeroHeight.into()));
944            }
945            fctl.width = width;
946            fctl.height = height;
947            Ok(())
948        } else {
949            Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()))
950        }
951    }
952
953    /// Set the position of the following frames.
954    ///
955    /// An error will be returned if:
956    /// - The image is not animated;
957    ///
958    /// - The selected position, considering also the current frame dimension,
959    ///   goes outside the image boundaries;
960    ///
961    // ??? TODO ???
962    // - The next frame is the default image
963    pub fn set_frame_position(&mut self, x: u32, y: u32) -> Result<()> {
964        if let Some(ref mut fctl) = self.info.frame_control {
965            if Some(x) > self.info.width.checked_sub(fctl.width)
966                || Some(y) > self.info.height.checked_sub(fctl.height)
967            {
968                return Err(EncodingError::Format(FormatErrorKind::OutOfBounds.into()));
969            }
970            fctl.x_offset = x;
971            fctl.y_offset = y;
972            Ok(())
973        } else {
974            Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()))
975        }
976    }
977
978    /// Set the frame dimension to occupy all the image, starting from
979    /// the current position.
980    ///
981    /// To reset the frame to the full image size [`reset_frame_position`]
982    /// should be called first.
983    ///
984    /// This method will return an error if the image is not animated.
985    ///
986    /// [`reset_frame_position`]: Writer::reset_frame_position
987    pub fn reset_frame_dimension(&mut self) -> Result<()> {
988        if let Some(ref mut fctl) = self.info.frame_control {
989            fctl.width = self.info.width - fctl.x_offset;
990            fctl.height = self.info.height - fctl.y_offset;
991            Ok(())
992        } else {
993            Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()))
994        }
995    }
996
997    /// Set the frame position to (0, 0).
998    ///
999    /// Equivalent to calling [`set_frame_position(0, 0)`].
1000    ///
1001    /// This method will return an error if the image is not animated.
1002    ///
1003    /// [`set_frame_position(0, 0)`]: Writer::set_frame_position
1004    pub fn reset_frame_position(&mut self) -> Result<()> {
1005        if let Some(ref mut fctl) = self.info.frame_control {
1006            fctl.x_offset = 0;
1007            fctl.y_offset = 0;
1008            Ok(())
1009        } else {
1010            Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()))
1011        }
1012    }
1013
1014    /// Set the blend operation for the following frames.
1015    ///
1016    /// The blend operation specifies whether the frame is to be alpha blended
1017    /// into the current output buffer content, or whether it should completely
1018    /// replace its region in the output buffer.
1019    ///
1020    /// See the [`BlendOp`] documentation for the possible values and their effects.
1021    ///
1022    /// *Note that for the first frame the two blend modes are functionally
1023    /// equivalent due to the clearing of the output buffer at the beginning
1024    /// of each play.*
1025    ///
1026    /// This method will return an error if the image is not animated.
1027    pub fn set_blend_op(&mut self, op: BlendOp) -> Result<()> {
1028        if let Some(ref mut fctl) = self.info.frame_control {
1029            fctl.blend_op = op;
1030            Ok(())
1031        } else {
1032            Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()))
1033        }
1034    }
1035
1036    /// Set the dispose operation for the following frames.
1037    ///
1038    /// The dispose operation specifies how the output buffer should be changed
1039    /// at the end of the delay (before rendering the next frame)
1040    ///
1041    /// See the [`DisposeOp`] documentation for the possible values and their effects.
1042    ///
1043    /// *Note that if the first frame uses [`DisposeOp::Previous`]
1044    /// it will be treated as [`DisposeOp::Background`].*
1045    ///
1046    /// This method will return an error if the image is not animated.
1047    pub fn set_dispose_op(&mut self, op: DisposeOp) -> Result<()> {
1048        if let Some(ref mut fctl) = self.info.frame_control {
1049            fctl.dispose_op = op;
1050            Ok(())
1051        } else {
1052            Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()))
1053        }
1054    }
1055
1056    /// Create a stream writer.
1057    ///
1058    /// This allows you to create images that do not fit in memory. The default
1059    /// chunk size is 4K, use `stream_writer_with_size` to set another chunk
1060    /// size.
1061    ///
1062    /// This borrows the writer which allows for manually appending additional
1063    /// chunks after the image data has been written.
1064    pub fn stream_writer(&mut self) -> Result<StreamWriter<'_, W>> {
1065        self.stream_writer_with_size(DEFAULT_BUFFER_LENGTH)
1066    }
1067
1068    /// Create a stream writer with custom buffer size.
1069    ///
1070    /// See [`stream_writer`].
1071    ///
1072    /// [`stream_writer`]: Self::stream_writer
1073    pub fn stream_writer_with_size(&mut self, size: usize) -> Result<StreamWriter<'_, W>> {
1074        StreamWriter::new(ChunkOutput::Borrowed(self), size)
1075    }
1076
1077    /// Turn this into a stream writer for image data.
1078    ///
1079    /// This allows you to create images that do not fit in memory. The default
1080    /// chunk size is 4K, use [`stream_writer_with_size`] to set another chunk
1081    /// size.
1082    ///
1083    /// [`stream_writer_with_size`]: Self::stream_writer_with_size
1084    pub fn into_stream_writer(self) -> Result<StreamWriter<'static, W>> {
1085        self.into_stream_writer_with_size(DEFAULT_BUFFER_LENGTH)
1086    }
1087
1088    /// Turn this into a stream writer with custom buffer size.
1089    ///
1090    /// See [`into_stream_writer`].
1091    ///
1092    /// [`into_stream_writer`]: Self::into_stream_writer
1093    pub fn into_stream_writer_with_size(self, size: usize) -> Result<StreamWriter<'static, W>> {
1094        StreamWriter::new(ChunkOutput::Owned(self), size)
1095    }
1096
1097    /// Consume the stream writer with validation.
1098    ///
1099    /// Unlike a simple drop this ensures that the final chunk was written correctly. When other
1100    /// validation options (chunk sequencing) had been turned on in the configuration then it will
1101    /// also do a check on their correctness _before_ writing the final chunk.
1102    pub fn finish(mut self) -> Result<()> {
1103        self.validate_sequence_done()?;
1104        self.write_iend()?;
1105        self.w.flush()?;
1106
1107        // Explicitly drop `self` just for clarity.
1108        drop(self);
1109        Ok(())
1110    }
1111}
1112
1113impl<W: Write> Drop for Writer<W> {
1114    fn drop(&mut self) {
1115        if !self.iend_written {
1116            let _ = self.write_iend();
1117        }
1118    }
1119}
1120
1121enum ChunkOutput<'a, W: Write> {
1122    Borrowed(&'a mut Writer<W>),
1123    Owned(Writer<W>),
1124}
1125
1126// opted for deref for practical reasons
1127impl<'a, W: Write> Deref for ChunkOutput<'a, W> {
1128    type Target = Writer<W>;
1129
1130    fn deref(&self) -> &Self::Target {
1131        match self {
1132            ChunkOutput::Borrowed(writer) => writer,
1133            ChunkOutput::Owned(writer) => writer,
1134        }
1135    }
1136}
1137
1138impl<'a, W: Write> DerefMut for ChunkOutput<'a, W> {
1139    fn deref_mut(&mut self) -> &mut Self::Target {
1140        match self {
1141            ChunkOutput::Borrowed(writer) => writer,
1142            ChunkOutput::Owned(writer) => writer,
1143        }
1144    }
1145}
1146
1147/// This writer is used between the actual writer and the
1148/// ZlibEncoder and has the job of packaging the compressed
1149/// data into a PNG chunk, based on the image metadata
1150///
1151/// Currently the way it works is that the specified buffer
1152/// will hold one chunk at the time and buffer the incoming
1153/// data until `flush` is called or the maximum chunk size
1154/// is reached.
1155///
1156/// The maximum chunk is the smallest between the selected buffer size
1157/// and `u32::MAX >> 1` (`0x7fffffff` or `2147483647` dec)
1158///
1159/// When a chunk has to be flushed the length (that is now known)
1160/// and the CRC will be written at the correct locations in the chunk.
1161struct ChunkWriter<'a, W: Write> {
1162    writer: ChunkOutput<'a, W>,
1163    buffer: Vec<u8>,
1164    /// keeps track of where the last byte was written
1165    index: usize,
1166    curr_chunk: ChunkType,
1167}
1168
1169impl<'a, W: Write> ChunkWriter<'a, W> {
1170    fn new(writer: ChunkOutput<'a, W>, buf_len: usize) -> ChunkWriter<'a, W> {
1171        // currently buf_len will determine the size of each chunk
1172        // the len is capped to the maximum size every chunk can hold
1173        // (this wont ever overflow an u32)
1174        //
1175        // TODO (maybe): find a way to hold two chunks at a time if `usize`
1176        //               is 64 bits.
1177        const CAP: usize = u32::MAX as usize >> 1;
1178        let curr_chunk = if writer.images_written == 0 {
1179            chunk::IDAT
1180        } else {
1181            chunk::fdAT
1182        };
1183        ChunkWriter {
1184            writer,
1185            buffer: vec![0; CAP.min(buf_len)],
1186            index: 0,
1187            curr_chunk,
1188        }
1189    }
1190
1191    /// Returns the size of each scanline for the next frame
1192    /// paired with the size of the whole frame
1193    ///
1194    /// This is used by the `StreamWriter` to know when the scanline ends
1195    /// so it can filter compress it and also to know when to start
1196    /// the next one
1197    fn next_frame_info(&self) -> (usize, usize) {
1198        let wrt = self.writer.deref();
1199
1200        let width: usize;
1201        let height: usize;
1202        if let Some(fctl) = wrt.info.frame_control {
1203            width = fctl.width as usize;
1204            height = fctl.height as usize;
1205        } else {
1206            width = wrt.info.width as usize;
1207            height = wrt.info.height as usize;
1208        }
1209
1210        let in_len = wrt.info.raw_row_length_from_width(width as u32) - 1;
1211        let data_size = in_len * height;
1212
1213        (in_len, data_size)
1214    }
1215
1216    /// NOTE: this bypasses the internal buffer so the flush method should be called before this
1217    ///       in the case there is some data left in the buffer when this is called, it will panic
1218    fn write_header(&mut self) -> Result<()> {
1219        assert_eq!(self.index, 0, "Called when not flushed");
1220        let wrt = self.writer.deref_mut();
1221
1222        self.curr_chunk = if wrt.images_written == 0 {
1223            chunk::IDAT
1224        } else {
1225            chunk::fdAT
1226        };
1227
1228        match wrt.info.frame_control {
1229            Some(_) if wrt.should_skip_frame_control_on_default_image() => {}
1230            Some(ref mut fctl) => {
1231                fctl.encode(&mut wrt.w)?;
1232                fctl.sequence_number += 1;
1233            }
1234            _ => {}
1235        }
1236
1237        Ok(())
1238    }
1239
1240    /// Set the [`FrameControl`] for the following frame
1241    ///
1242    /// It will ignore the `sequence_number` of the parameter
1243    /// as it is updated internally.
1244    fn set_fctl(&mut self, f: FrameControl) {
1245        if let Some(ref mut fctl) = self.writer.info.frame_control {
1246            // Ignore the sequence number
1247            *fctl = FrameControl {
1248                sequence_number: fctl.sequence_number,
1249                ..f
1250            };
1251        } else {
1252            panic!("This function must be called on an animated PNG")
1253        }
1254    }
1255
1256    /// Flushes the current chunk
1257    fn flush_inner(&mut self) -> io::Result<()> {
1258        if self.index > 0 {
1259            // flush the chunk and reset everything
1260            write_chunk(
1261                &mut self.writer.w,
1262                self.curr_chunk,
1263                &self.buffer[..self.index],
1264            )?;
1265
1266            self.index = 0;
1267        }
1268        Ok(())
1269    }
1270}
1271
1272impl<'a, W: Write> Write for ChunkWriter<'a, W> {
1273    fn write(&mut self, mut data: &[u8]) -> io::Result<usize> {
1274        if data.is_empty() {
1275            return Ok(0);
1276        }
1277
1278        // index == 0 means a chunk has been flushed out
1279        if self.index == 0 {
1280            let wrt = self.writer.deref_mut();
1281
1282            // Prepare the next animated frame, if any.
1283            let no_fctl = wrt.should_skip_frame_control_on_default_image();
1284            if wrt.info.frame_control.is_some() && !no_fctl {
1285                let fctl = wrt.info.frame_control.as_mut().unwrap();
1286                self.buffer[0..4].copy_from_slice(&fctl.sequence_number.to_be_bytes());
1287                fctl.sequence_number += 1;
1288                self.index = 4;
1289            }
1290        }
1291
1292        // Cap the buffer length to the maximum number of bytes that can't still
1293        // be added to the current chunk
1294        let written = data.len().min(self.buffer.len() - self.index);
1295        data = &data[..written];
1296
1297        self.buffer[self.index..][..written].copy_from_slice(data);
1298        self.index += written;
1299
1300        // if the maximum data for this chunk as been reached it needs to be flushed
1301        if self.index == self.buffer.len() {
1302            self.flush_inner()?;
1303        }
1304
1305        Ok(written)
1306    }
1307
1308    fn flush(&mut self) -> io::Result<()> {
1309        self.flush_inner()
1310    }
1311}
1312
1313impl<W: Write> Drop for ChunkWriter<'_, W> {
1314    fn drop(&mut self) {
1315        let _ = self.flush();
1316    }
1317}
1318
1319// TODO: find a better name
1320//
1321/// This enum is used to be allow the `StreamWriter` to keep
1322/// its inner `ChunkWriter` without wrapping it inside a
1323/// `ZlibEncoder`. This is used in the case that between the
1324/// change of state that happens when the last write of a frame
1325/// is performed an error occurs, which obviously has to be returned.
1326/// This creates the problem of where to store the writer before
1327/// exiting the function, and this is where `Wrapper` comes in.
1328///
1329/// Unfortunately the `ZlibWriter` can't be used because on the
1330/// write following the error, `finish` would be called and that
1331/// would write some data even if 0 bytes where compressed.
1332///
1333/// If the `finish` function fails then there is nothing much to
1334/// do as the `ChunkWriter` would get lost so the `Unrecoverable`
1335/// variant is used to signal that.
1336enum Wrapper<'a, W: Write> {
1337    Chunk(ChunkWriter<'a, W>),
1338    Flate2(ZlibEncoder<ChunkWriter<'a, W>>),
1339    FDeflate(fdeflate::Compressor<ChunkWriter<'a, W>>),
1340    Unrecoverable,
1341    /// This is used in-between, should never be matched
1342    None,
1343}
1344
1345impl<'a, W: Write> Wrapper<'a, W> {
1346    fn from_level(writer: ChunkWriter<'a, W>, compression: DeflateCompression) -> io::Result<Self> {
1347        Ok(match compression {
1348            DeflateCompression::NoCompression => {
1349                Wrapper::Flate2(ZlibEncoder::new(writer, flate2::Compression::none()))
1350            }
1351            DeflateCompression::FdeflateUltraFast => {
1352                Wrapper::FDeflate(fdeflate::Compressor::new(writer)?)
1353            }
1354            DeflateCompression::Level(level) => Wrapper::Flate2(ZlibEncoder::new(
1355                writer,
1356                flate2::Compression::new(u32::from(level)),
1357            )),
1358        })
1359    }
1360
1361    /// Like `Option::take` this returns the `Wrapper` contained
1362    /// in `self` and replaces it with `Wrapper::None`
1363    fn take(&mut self) -> Wrapper<'a, W> {
1364        let mut swap = Wrapper::None;
1365        mem::swap(self, &mut swap);
1366        swap
1367    }
1368}
1369
1370/// Streaming PNG writer
1371///
1372/// This may silently fail in the destructor, so it is a good idea to call
1373/// [`finish`] or [`flush`] before dropping.
1374///
1375/// [`finish`]: Self::finish
1376/// [`flush`]: Write::flush
1377pub struct StreamWriter<'a, W: Write> {
1378    /// The option here is needed in order to access the inner `ChunkWriter` in-between
1379    /// each frame, which is needed for writing the fcTL chunks between each frame
1380    writer: Wrapper<'a, W>,
1381    prev_buf: Vec<u8>,
1382    curr_buf: Vec<u8>,
1383    filtered_buf: Vec<u8>,
1384    /// Amount of data already written
1385    index: usize,
1386    /// length of the current scanline
1387    line_len: usize,
1388    /// size of the frame (width * height * sample_size)
1389    to_write: usize,
1390
1391    width: u32,
1392    height: u32,
1393
1394    bpp: BytesPerPixel,
1395    filter: Filter,
1396    fctl: Option<FrameControl>,
1397    compression: DeflateCompression,
1398}
1399
1400impl<'a, W: Write> StreamWriter<'a, W> {
1401    fn new(writer: ChunkOutput<'a, W>, buf_len: usize) -> Result<StreamWriter<'a, W>> {
1402        let PartialInfo {
1403            width,
1404            height,
1405            frame_control: fctl,
1406            ..
1407        } = writer.info;
1408
1409        let bpp = writer.info.bpp_in_prediction();
1410        let in_len = writer.info.raw_row_length() - 1;
1411        let filter = writer.options.filter;
1412        let compression = writer.options.compression;
1413        let prev_buf = vec![0; in_len];
1414        let curr_buf = vec![0; in_len];
1415        let filtered_buf = vec![0; in_len];
1416
1417        let mut chunk_writer = ChunkWriter::new(writer, buf_len);
1418        let (line_len, to_write) = chunk_writer.next_frame_info();
1419        chunk_writer.write_header()?;
1420
1421        Ok(StreamWriter {
1422            writer: Wrapper::from_level(chunk_writer, compression)?,
1423            index: 0,
1424            prev_buf,
1425            curr_buf,
1426            filtered_buf,
1427            bpp,
1428            filter,
1429            width,
1430            height,
1431            line_len,
1432            to_write,
1433            fctl,
1434            compression,
1435        })
1436    }
1437
1438    /// Set the used filter type.
1439    ///
1440    /// The default filter is [`Filter::Adaptive`] which automatically selects the best filter
1441    /// for each row of the image.
1442    ///
1443    /// You should only change this if you are after very fast compression,
1444    /// and either don't care about compression ratio or know exactly what works best for your images.
1445    pub fn set_filter(&mut self, filter: Filter) {
1446        self.filter = filter;
1447    }
1448
1449    /// Set the fraction of time the following frames are going to be displayed,
1450    /// in seconds
1451    ///
1452    /// If the denominator is 0, it is to be treated as if it were 100
1453    /// (that is, the numerator then specifies 1/100ths of a second).
1454    /// If the value of the numerator is 0 the decoder should render the next frame
1455    /// as quickly as possible, though viewers may impose a reasonable lower bound.
1456    ///
1457    /// This method will return an error if the image is not animated.
1458    pub fn set_frame_delay(&mut self, numerator: u16, denominator: u16) -> Result<()> {
1459        if let Some(ref mut fctl) = self.fctl {
1460            fctl.delay_den = denominator;
1461            fctl.delay_num = numerator;
1462            Ok(())
1463        } else {
1464            Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()))
1465        }
1466    }
1467
1468    /// Set the dimension of the following frames.
1469    ///
1470    /// This function will return an error when:
1471    /// - The image is not an animated;
1472    ///
1473    /// - The selected dimension, considering also the current frame position,
1474    ///   goes outside the image boundaries;
1475    ///
1476    /// - One or both the width and height are 0;
1477    ///
1478    pub fn set_frame_dimension(&mut self, width: u32, height: u32) -> Result<()> {
1479        if let Some(ref mut fctl) = self.fctl {
1480            if Some(width) > self.width.checked_sub(fctl.x_offset)
1481                || Some(height) > self.height.checked_sub(fctl.y_offset)
1482            {
1483                return Err(EncodingError::Format(FormatErrorKind::OutOfBounds.into()));
1484            } else if width == 0 {
1485                return Err(EncodingError::Format(FormatErrorKind::ZeroWidth.into()));
1486            } else if height == 0 {
1487                return Err(EncodingError::Format(FormatErrorKind::ZeroHeight.into()));
1488            }
1489            fctl.width = width;
1490            fctl.height = height;
1491            Ok(())
1492        } else {
1493            Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()))
1494        }
1495    }
1496
1497    /// Set the position of the following frames.
1498    ///
1499    /// An error will be returned if:
1500    /// - The image is not animated;
1501    ///
1502    /// - The selected position, considering also the current frame dimension,
1503    ///   goes outside the image boundaries;
1504    ///
1505    pub fn set_frame_position(&mut self, x: u32, y: u32) -> Result<()> {
1506        if let Some(ref mut fctl) = self.fctl {
1507            if Some(x) > self.width.checked_sub(fctl.width)
1508                || Some(y) > self.height.checked_sub(fctl.height)
1509            {
1510                return Err(EncodingError::Format(FormatErrorKind::OutOfBounds.into()));
1511            }
1512            fctl.x_offset = x;
1513            fctl.y_offset = y;
1514            Ok(())
1515        } else {
1516            Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()))
1517        }
1518    }
1519
1520    /// Set the frame dimension to occupy all the image, starting from
1521    /// the current position.
1522    ///
1523    /// To reset the frame to the full image size [`reset_frame_position`]
1524    /// should be called first.
1525    ///
1526    /// This method will return an error if the image is not animated.
1527    ///
1528    /// [`reset_frame_position`]: Writer::reset_frame_position
1529    pub fn reset_frame_dimension(&mut self) -> Result<()> {
1530        if let Some(ref mut fctl) = self.fctl {
1531            fctl.width = self.width - fctl.x_offset;
1532            fctl.height = self.height - fctl.y_offset;
1533            Ok(())
1534        } else {
1535            Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()))
1536        }
1537    }
1538
1539    /// Set the frame position to (0, 0).
1540    ///
1541    /// Equivalent to calling [`set_frame_position(0, 0)`].
1542    ///
1543    /// This method will return an error if the image is not animated.
1544    ///
1545    /// [`set_frame_position(0, 0)`]: Writer::set_frame_position
1546    pub fn reset_frame_position(&mut self) -> Result<()> {
1547        if let Some(ref mut fctl) = self.fctl {
1548            fctl.x_offset = 0;
1549            fctl.y_offset = 0;
1550            Ok(())
1551        } else {
1552            Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()))
1553        }
1554    }
1555
1556    /// Set the blend operation for the following frames.
1557    ///
1558    /// The blend operation specifies whether the frame is to be alpha blended
1559    /// into the current output buffer content, or whether it should completely
1560    /// replace its region in the output buffer.
1561    ///
1562    /// See the [`BlendOp`] documentation for the possible values and their effects.
1563    ///
1564    /// *Note that for the first frame the two blend modes are functionally
1565    /// equivalent due to the clearing of the output buffer at the beginning
1566    /// of each play.*
1567    ///
1568    /// This method will return an error if the image is not animated.
1569    pub fn set_blend_op(&mut self, op: BlendOp) -> Result<()> {
1570        if let Some(ref mut fctl) = self.fctl {
1571            fctl.blend_op = op;
1572            Ok(())
1573        } else {
1574            Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()))
1575        }
1576    }
1577
1578    /// Set the dispose operation for the following frames.
1579    ///
1580    /// The dispose operation specifies how the output buffer should be changed
1581    /// at the end of the delay (before rendering the next frame)
1582    ///
1583    /// See the [`DisposeOp`] documentation for the possible values and their effects.
1584    ///
1585    /// *Note that if the first frame uses [`DisposeOp::Previous`]
1586    /// it will be treated as [`DisposeOp::Background`].*
1587    ///
1588    /// This method will return an error if the image is not animated.
1589    pub fn set_dispose_op(&mut self, op: DisposeOp) -> Result<()> {
1590        if let Some(ref mut fctl) = self.fctl {
1591            fctl.dispose_op = op;
1592            Ok(())
1593        } else {
1594            Err(EncodingError::Format(FormatErrorKind::NotAnimated.into()))
1595        }
1596    }
1597
1598    /// Consume the stream writer with validation.
1599    ///
1600    /// Unlike a simple drop this ensures that the all data was written correctly. When other
1601    /// validation options (chunk sequencing) had been turned on in the configuration of inner
1602    /// [`Writer`], then it will also do a check on their correctness. Differently from
1603    /// [`Writer::finish`], this just `flush`es, returns error if some data is abandoned.
1604    pub fn finish(mut self) -> Result<()> {
1605        self.finish_mut()
1606    }
1607
1608    /// Internal helper that can be called both from `fn finish(mut self)`
1609    /// and from `fn drop(&mut self)`.
1610    fn finish_mut(&mut self) -> Result<()> {
1611        if self.to_write > 0 {
1612            let err = FormatErrorKind::MissingData(self.to_write).into();
1613            return Err(EncodingError::Format(err));
1614        }
1615
1616        self.flush()?;
1617        match self.writer.take() {
1618            Wrapper::Chunk(wrt) => {
1619                wrt.writer.validate_sequence_done()?;
1620            }
1621            Wrapper::FDeflate(wrt) => {
1622                wrt.finish()?;
1623            }
1624            Wrapper::Flate2(wrt) => {
1625                wrt.finish()?;
1626            }
1627            Wrapper::None => unreachable!(),
1628            Wrapper::Unrecoverable => {
1629                let err = FormatErrorKind::Unrecoverable.into();
1630                return Err(EncodingError::Format(err));
1631            }
1632        }
1633
1634        Ok(())
1635    }
1636
1637    /// Flushes the buffered chunk, checks if it was the last frame,
1638    /// writes the next frame header and gets the next frame scanline size
1639    /// and image size.
1640    /// NOTE: This method must only be called when the writer is the variant Chunk(_)
1641    fn new_frame(&mut self) -> Result<()> {
1642        let wrt = match &mut self.writer {
1643            Wrapper::Chunk(wrt) => wrt,
1644            Wrapper::Unrecoverable => {
1645                let err = FormatErrorKind::Unrecoverable.into();
1646                return Err(EncodingError::Format(err));
1647            }
1648            Wrapper::Flate2(_) | Wrapper::FDeflate(_) => {
1649                unreachable!("never called on a half-finished frame")
1650            }
1651            Wrapper::None => unreachable!(),
1652        };
1653        wrt.flush()?;
1654        wrt.writer.validate_new_image()?;
1655
1656        if let Some(fctl) = self.fctl {
1657            wrt.set_fctl(fctl);
1658        }
1659        let (scansize, size) = wrt.next_frame_info();
1660        self.line_len = scansize;
1661        self.to_write = size;
1662
1663        wrt.write_header()?;
1664        wrt.writer.increment_images_written();
1665
1666        // now it can be taken because the next statements cannot cause any errors
1667        match self.writer.take() {
1668            Wrapper::Chunk(wrt) => match Wrapper::from_level(wrt, self.compression) {
1669                Ok(writer) => self.writer = writer,
1670                Err(err) => {
1671                    self.writer = Wrapper::Unrecoverable;
1672                    return Err(err.into());
1673                }
1674            },
1675            _ => unreachable!(),
1676        };
1677
1678        Ok(())
1679    }
1680}
1681
1682impl<'a, W: Write> Write for StreamWriter<'a, W> {
1683    fn write(&mut self, mut data: &[u8]) -> io::Result<usize> {
1684        if let Wrapper::Unrecoverable = self.writer {
1685            let err = FormatErrorKind::Unrecoverable.into();
1686            return Err(EncodingError::Format(err).into());
1687        }
1688
1689        if data.is_empty() {
1690            return Ok(0);
1691        }
1692
1693        if self.to_write == 0 {
1694            match self.writer.take() {
1695                Wrapper::Flate2(wrt) => match wrt.finish() {
1696                    Ok(chunk) => self.writer = Wrapper::Chunk(chunk),
1697                    Err(err) => {
1698                        self.writer = Wrapper::Unrecoverable;
1699                        return Err(err);
1700                    }
1701                },
1702                Wrapper::FDeflate(wrt) => match wrt.finish() {
1703                    Ok(chunk) => self.writer = Wrapper::Chunk(chunk),
1704                    Err(err) => {
1705                        self.writer = Wrapper::Unrecoverable;
1706                        return Err(err);
1707                    }
1708                },
1709                chunk @ Wrapper::Chunk(_) => self.writer = chunk,
1710                Wrapper::Unrecoverable => unreachable!(),
1711                Wrapper::None => unreachable!(),
1712            };
1713
1714            // Transition Wrapper::Chunk to Wrapper::Zlib.
1715            self.new_frame()?;
1716        }
1717
1718        let written = data.read(&mut self.curr_buf[..self.line_len][self.index..])?;
1719        self.index += written;
1720        self.to_write -= written;
1721
1722        if self.index == self.line_len {
1723            let filter_type = filter(
1724                self.filter,
1725                self.bpp,
1726                &self.prev_buf,
1727                &self.curr_buf,
1728                &mut self.filtered_buf,
1729            );
1730            // This can't fail as the other variant is used only to allow the zlib encoder to finish
1731            match &mut self.writer {
1732                Wrapper::Flate2(wrt) => {
1733                    wrt.write_all(&[filter_type as u8])?;
1734                    wrt.write_all(&self.filtered_buf)?;
1735                }
1736                Wrapper::FDeflate(wrt) => {
1737                    wrt.write_data(&[filter_type as u8])?;
1738                    wrt.write_data(&self.filtered_buf)?;
1739                }
1740                _ => unreachable!(),
1741            };
1742
1743            mem::swap(&mut self.prev_buf, &mut self.curr_buf);
1744            self.index = 0;
1745        }
1746
1747        Ok(written)
1748    }
1749
1750    fn flush(&mut self) -> io::Result<()> {
1751        match &mut self.writer {
1752            Wrapper::Flate2(wrt) => wrt.flush()?,
1753            Wrapper::Chunk(wrt) => wrt.flush()?,
1754            Wrapper::FDeflate(_) => (), // TODO: Add `flush()` to `fdeflate::Compressor`?
1755            // This handles both the case where we entered an unrecoverable state after zlib
1756            // decoding failure and after a panic while we had taken the chunk/zlib reader.
1757            Wrapper::Unrecoverable | Wrapper::None => {
1758                let err = FormatErrorKind::Unrecoverable.into();
1759                return Err(EncodingError::Format(err).into());
1760            }
1761        }
1762
1763        if self.index > 0 {
1764            let err = FormatErrorKind::WrittenTooMuch(self.index).into();
1765            return Err(EncodingError::Format(err).into());
1766        }
1767
1768        Ok(())
1769    }
1770}
1771
1772impl<W: Write> Drop for StreamWriter<'_, W> {
1773    fn drop(&mut self) {
1774        let _ = self.finish_mut();
1775    }
1776}
1777
1778#[cfg(test)]
1779mod tests {
1780    use super::*;
1781    use crate::Decoder;
1782
1783    use io::BufReader;
1784    use rand::{rng, Rng};
1785    use std::cmp;
1786    use std::fs::File;
1787    use std::io::Cursor;
1788
1789    #[test]
1790    fn roundtrip1() {
1791        roundtrip_inner();
1792    }
1793
1794    #[test]
1795    fn roundtrip2() {
1796        roundtrip_inner();
1797    }
1798
1799    fn roundtrip_inner() {
1800        // More loops = more random testing, but also more test wait time
1801        for _ in 0..5 {
1802            for path in glob::glob("tests/pngsuite/*.png")
1803                .unwrap()
1804                .map(|r| r.unwrap())
1805            {
1806                if path.file_name().unwrap().to_str().unwrap().starts_with('x') {
1807                    // x* files are expected to fail to decode
1808                    continue;
1809                }
1810                eprintln!("{}", path.display());
1811                // Decode image
1812                let decoder = Decoder::new(BufReader::new(File::open(path).unwrap()));
1813                let mut reader = decoder.read_info().unwrap();
1814                let mut buf = vec![0; reader.output_buffer_size().unwrap()];
1815                let info = reader.next_frame(&mut buf).unwrap();
1816                use DeflateCompression::*;
1817                for compression in [NoCompression, FdeflateUltraFast, Level(4)] {
1818                    // Encode decoded image
1819                    let mut out = Vec::new();
1820                    {
1821                        let mut wrapper = RandomChunkWriter {
1822                            rng: rng(),
1823                            w: &mut out,
1824                        };
1825
1826                        let mut encoder = Encoder::new(&mut wrapper, info.width, info.height);
1827                        encoder.set_color(info.color_type);
1828                        encoder.set_depth(info.bit_depth);
1829                        encoder.set_deflate_compression(compression);
1830                        if let Some(palette) = &reader.info().palette {
1831                            encoder.set_palette(palette.clone());
1832                        }
1833                        let mut encoder = encoder.write_header().unwrap();
1834                        encoder.write_image_data(&buf).unwrap();
1835                    }
1836                    // Decode encoded decoded image
1837                    let decoder = Decoder::new(Cursor::new(&*out));
1838                    let mut reader = decoder.read_info().unwrap();
1839                    let mut buf2 = vec![0; reader.output_buffer_size().unwrap()];
1840                    reader.next_frame(&mut buf2).unwrap();
1841                    // check if the encoded image is ok:
1842                    assert_eq!(buf, buf2);
1843                }
1844            }
1845        }
1846    }
1847
1848    #[test]
1849    fn roundtrip_stream1() {
1850        roundtrip_stream_inner();
1851    }
1852
1853    #[test]
1854    fn roundtrip_stream2() {
1855        roundtrip_stream_inner();
1856    }
1857
1858    fn roundtrip_stream_inner() {
1859        // More loops = more random testing, but also more test wait time
1860        for _ in 0..5 {
1861            for path in glob::glob("tests/pngsuite/*.png")
1862                .unwrap()
1863                .map(|r| r.unwrap())
1864            {
1865                if path.file_name().unwrap().to_str().unwrap().starts_with('x') {
1866                    // x* files are expected to fail to decode
1867                    continue;
1868                }
1869                // Decode image
1870                let decoder = Decoder::new(BufReader::new(File::open(path).unwrap()));
1871                let mut reader = decoder.read_info().unwrap();
1872                let mut buf = vec![0; reader.output_buffer_size().unwrap()];
1873                let info = reader.next_frame(&mut buf).unwrap();
1874                use DeflateCompression::*;
1875                for compression in [NoCompression, FdeflateUltraFast, Level(4)] {
1876                    // Encode decoded image
1877                    let mut out = Vec::new();
1878                    {
1879                        let mut wrapper = RandomChunkWriter {
1880                            rng: rng(),
1881                            w: &mut out,
1882                        };
1883
1884                        let mut encoder = Encoder::new(&mut wrapper, info.width, info.height);
1885                        encoder.set_color(info.color_type);
1886                        encoder.set_depth(info.bit_depth);
1887                        encoder.set_deflate_compression(compression);
1888                        if let Some(palette) = &reader.info().palette {
1889                            encoder.set_palette(palette.clone());
1890                        }
1891                        let mut encoder = encoder.write_header().unwrap();
1892                        let mut stream_writer = encoder.stream_writer().unwrap();
1893
1894                        let mut outer_wrapper = RandomChunkWriter {
1895                            rng: rng(),
1896                            w: &mut stream_writer,
1897                        };
1898
1899                        outer_wrapper.write_all(&buf).unwrap();
1900                    }
1901                    // Decode encoded decoded image
1902                    let decoder = Decoder::new(Cursor::new(&*out));
1903                    let mut reader = decoder.read_info().unwrap();
1904                    let mut buf2 = vec![0; reader.output_buffer_size().unwrap()];
1905                    reader.next_frame(&mut buf2).unwrap();
1906                    // check if the encoded image is ok:
1907                    assert_eq!(buf, buf2);
1908                }
1909            }
1910        }
1911    }
1912
1913    #[test]
1914    fn image_palette() -> Result<()> {
1915        for &bit_depth in &[1u8, 2, 4, 8] {
1916            // Do a reference decoding, choose a fitting palette image from pngsuite
1917            let path = format!("tests/pngsuite/basn3p0{}.png", bit_depth);
1918            let decoder = Decoder::new(BufReader::new(File::open(&path).unwrap()));
1919            let mut reader = decoder.read_info().unwrap();
1920
1921            let mut decoded_pixels = vec![0; reader.output_buffer_size().unwrap()];
1922            let info = reader.info();
1923            assert_eq!(
1924                info.width as usize * info.height as usize * usize::from(bit_depth),
1925                decoded_pixels.len() * 8
1926            );
1927            let info = reader.next_frame(&mut decoded_pixels).unwrap();
1928            let indexed_data = decoded_pixels;
1929
1930            let palette = reader.info().palette.as_ref().unwrap();
1931            let mut out = Vec::new();
1932            {
1933                let mut encoder = Encoder::new(&mut out, info.width, info.height);
1934                encoder.set_depth(BitDepth::from_u8(bit_depth).unwrap());
1935                encoder.set_color(ColorType::Indexed);
1936                encoder.set_palette(palette.as_ref());
1937
1938                let mut writer = encoder.write_header().unwrap();
1939                writer.write_image_data(&indexed_data).unwrap();
1940            }
1941
1942            // Decode re-encoded image
1943            let decoder = Decoder::new(Cursor::new(&*out));
1944            let mut reader = decoder.read_info().unwrap();
1945            let mut redecoded = vec![0; reader.output_buffer_size().unwrap()];
1946            reader.next_frame(&mut redecoded).unwrap();
1947            // check if the encoded image is ok:
1948            assert_eq!(indexed_data, redecoded);
1949        }
1950        Ok(())
1951    }
1952
1953    #[test]
1954    fn expect_error_on_wrong_image_len() -> Result<()> {
1955        let width = 10;
1956        let height = 10;
1957
1958        let output = vec![0u8; 1024];
1959        let writer = Cursor::new(output);
1960        let mut encoder = Encoder::new(writer, width as u32, height as u32);
1961        encoder.set_depth(BitDepth::Eight);
1962        encoder.set_color(ColorType::Rgb);
1963        let mut png_writer = encoder.write_header()?;
1964
1965        let correct_image_size = width * height * 3;
1966        let image = vec![0u8; correct_image_size + 1];
1967        let result = png_writer.write_image_data(image.as_ref());
1968        assert!(result.is_err());
1969
1970        Ok(())
1971    }
1972
1973    #[test]
1974    fn expect_error_on_empty_image() -> Result<()> {
1975        let output = vec![0u8; 1024];
1976        let mut writer = Cursor::new(output);
1977
1978        let encoder = Encoder::new(&mut writer, 0, 0);
1979        assert!(encoder.write_header().is_err());
1980
1981        let encoder = Encoder::new(&mut writer, 100, 0);
1982        assert!(encoder.write_header().is_err());
1983
1984        let encoder = Encoder::new(&mut writer, 0, 100);
1985        assert!(encoder.write_header().is_err());
1986
1987        Ok(())
1988    }
1989
1990    #[test]
1991    fn expect_error_on_invalid_bit_depth_color_type_combination() -> Result<()> {
1992        let output = vec![0u8; 1024];
1993        let mut writer = Cursor::new(output);
1994
1995        let mut encoder = Encoder::new(&mut writer, 1, 1);
1996        encoder.set_depth(BitDepth::One);
1997        encoder.set_color(ColorType::Rgb);
1998        assert!(encoder.write_header().is_err());
1999
2000        let mut encoder = Encoder::new(&mut writer, 1, 1);
2001        encoder.set_depth(BitDepth::One);
2002        encoder.set_color(ColorType::GrayscaleAlpha);
2003        assert!(encoder.write_header().is_err());
2004
2005        let mut encoder = Encoder::new(&mut writer, 1, 1);
2006        encoder.set_depth(BitDepth::One);
2007        encoder.set_color(ColorType::Rgba);
2008        assert!(encoder.write_header().is_err());
2009
2010        let mut encoder = Encoder::new(&mut writer, 1, 1);
2011        encoder.set_depth(BitDepth::Two);
2012        encoder.set_color(ColorType::Rgb);
2013        assert!(encoder.write_header().is_err());
2014
2015        let mut encoder = Encoder::new(&mut writer, 1, 1);
2016        encoder.set_depth(BitDepth::Two);
2017        encoder.set_color(ColorType::GrayscaleAlpha);
2018        assert!(encoder.write_header().is_err());
2019
2020        let mut encoder = Encoder::new(&mut writer, 1, 1);
2021        encoder.set_depth(BitDepth::Two);
2022        encoder.set_color(ColorType::Rgba);
2023        assert!(encoder.write_header().is_err());
2024
2025        let mut encoder = Encoder::new(&mut writer, 1, 1);
2026        encoder.set_depth(BitDepth::Four);
2027        encoder.set_color(ColorType::Rgb);
2028        assert!(encoder.write_header().is_err());
2029
2030        let mut encoder = Encoder::new(&mut writer, 1, 1);
2031        encoder.set_depth(BitDepth::Four);
2032        encoder.set_color(ColorType::GrayscaleAlpha);
2033        assert!(encoder.write_header().is_err());
2034
2035        let mut encoder = Encoder::new(&mut writer, 1, 1);
2036        encoder.set_depth(BitDepth::Four);
2037        encoder.set_color(ColorType::Rgba);
2038        assert!(encoder.write_header().is_err());
2039
2040        let mut encoder = Encoder::new(&mut writer, 1, 1);
2041        encoder.set_depth(BitDepth::Sixteen);
2042        encoder.set_color(ColorType::Indexed);
2043        assert!(encoder.write_header().is_err());
2044
2045        Ok(())
2046    }
2047
2048    #[test]
2049    fn can_write_header_with_valid_bit_depth_color_type_combination() -> Result<()> {
2050        let output = vec![0u8; 1024];
2051        let mut writer = Cursor::new(output);
2052
2053        let mut encoder = Encoder::new(&mut writer, 1, 1);
2054        encoder.set_depth(BitDepth::One);
2055        encoder.set_color(ColorType::Grayscale);
2056        assert!(encoder.write_header().is_ok());
2057
2058        let mut encoder = Encoder::new(&mut writer, 1, 1);
2059        encoder.set_depth(BitDepth::One);
2060        encoder.set_color(ColorType::Indexed);
2061        assert!(encoder.write_header().is_ok());
2062
2063        let mut encoder = Encoder::new(&mut writer, 1, 1);
2064        encoder.set_depth(BitDepth::Two);
2065        encoder.set_color(ColorType::Grayscale);
2066        assert!(encoder.write_header().is_ok());
2067
2068        let mut encoder = Encoder::new(&mut writer, 1, 1);
2069        encoder.set_depth(BitDepth::Two);
2070        encoder.set_color(ColorType::Indexed);
2071        assert!(encoder.write_header().is_ok());
2072
2073        let mut encoder = Encoder::new(&mut writer, 1, 1);
2074        encoder.set_depth(BitDepth::Four);
2075        encoder.set_color(ColorType::Grayscale);
2076        assert!(encoder.write_header().is_ok());
2077
2078        let mut encoder = Encoder::new(&mut writer, 1, 1);
2079        encoder.set_depth(BitDepth::Four);
2080        encoder.set_color(ColorType::Indexed);
2081        assert!(encoder.write_header().is_ok());
2082
2083        let mut encoder = Encoder::new(&mut writer, 1, 1);
2084        encoder.set_depth(BitDepth::Eight);
2085        encoder.set_color(ColorType::Grayscale);
2086        assert!(encoder.write_header().is_ok());
2087
2088        let mut encoder = Encoder::new(&mut writer, 1, 1);
2089        encoder.set_depth(BitDepth::Eight);
2090        encoder.set_color(ColorType::Rgb);
2091        assert!(encoder.write_header().is_ok());
2092
2093        let mut encoder = Encoder::new(&mut writer, 1, 1);
2094        encoder.set_depth(BitDepth::Eight);
2095        encoder.set_color(ColorType::Indexed);
2096        assert!(encoder.write_header().is_ok());
2097
2098        let mut encoder = Encoder::new(&mut writer, 1, 1);
2099        encoder.set_depth(BitDepth::Eight);
2100        encoder.set_color(ColorType::GrayscaleAlpha);
2101        assert!(encoder.write_header().is_ok());
2102
2103        let mut encoder = Encoder::new(&mut writer, 1, 1);
2104        encoder.set_depth(BitDepth::Eight);
2105        encoder.set_color(ColorType::Rgba);
2106        assert!(encoder.write_header().is_ok());
2107
2108        let mut encoder = Encoder::new(&mut writer, 1, 1);
2109        encoder.set_depth(BitDepth::Sixteen);
2110        encoder.set_color(ColorType::Grayscale);
2111        assert!(encoder.write_header().is_ok());
2112
2113        let mut encoder = Encoder::new(&mut writer, 1, 1);
2114        encoder.set_depth(BitDepth::Sixteen);
2115        encoder.set_color(ColorType::Rgb);
2116        assert!(encoder.write_header().is_ok());
2117
2118        let mut encoder = Encoder::new(&mut writer, 1, 1);
2119        encoder.set_depth(BitDepth::Sixteen);
2120        encoder.set_color(ColorType::GrayscaleAlpha);
2121        assert!(encoder.write_header().is_ok());
2122
2123        let mut encoder = Encoder::new(&mut writer, 1, 1);
2124        encoder.set_depth(BitDepth::Sixteen);
2125        encoder.set_color(ColorType::Rgba);
2126        assert!(encoder.write_header().is_ok());
2127
2128        Ok(())
2129    }
2130
2131    #[test]
2132    fn all_filters_roundtrip() -> io::Result<()> {
2133        let pixel: Vec<_> = (0..48).collect();
2134
2135        let roundtrip = |filter: Filter| -> io::Result<()> {
2136            let mut buffer = vec![];
2137            let mut encoder = Encoder::new(&mut buffer, 4, 4);
2138            encoder.set_depth(BitDepth::Eight);
2139            encoder.set_color(ColorType::Rgb);
2140            encoder.set_filter(filter);
2141            encoder.write_header()?.write_image_data(&pixel)?;
2142
2143            let decoder = crate::Decoder::new(Cursor::new(buffer));
2144            let mut reader = decoder.read_info()?;
2145            let info = reader.info();
2146            assert_eq!(info.width, 4);
2147            assert_eq!(info.height, 4);
2148            let mut dest = vec![0; pixel.len()];
2149            reader.next_frame(&mut dest)?;
2150            assert_eq!(dest, pixel, "Deviation with filter type {:?}", filter);
2151
2152            Ok(())
2153        };
2154
2155        roundtrip(Filter::NoFilter)?;
2156        roundtrip(Filter::Sub)?;
2157        roundtrip(Filter::Up)?;
2158        roundtrip(Filter::Avg)?;
2159        roundtrip(Filter::Paeth)?;
2160
2161        Ok(())
2162    }
2163
2164    #[test]
2165    fn some_gamma_roundtrip() -> io::Result<()> {
2166        let pixel: Vec<_> = (0..48).collect();
2167
2168        let roundtrip = |gamma: Option<ScaledFloat>| -> io::Result<()> {
2169            let mut buffer = vec![];
2170            let mut encoder = Encoder::new(&mut buffer, 4, 4);
2171            encoder.set_depth(BitDepth::Eight);
2172            encoder.set_color(ColorType::Rgb);
2173            encoder.set_filter(Filter::Avg);
2174            if let Some(gamma) = gamma {
2175                encoder.set_source_gamma(gamma);
2176            }
2177            encoder.write_header()?.write_image_data(&pixel)?;
2178
2179            let decoder = crate::Decoder::new(Cursor::new(buffer));
2180            let mut reader = decoder.read_info()?;
2181            assert_eq!(
2182                reader.info().gamma(),
2183                gamma,
2184                "Deviation with gamma {:?}",
2185                gamma
2186            );
2187            let mut dest = vec![0; pixel.len()];
2188            let info = reader.next_frame(&mut dest)?;
2189            assert_eq!(info.width, 4);
2190            assert_eq!(info.height, 4);
2191
2192            Ok(())
2193        };
2194
2195        roundtrip(None)?;
2196        roundtrip(Some(ScaledFloat::new(0.35)))?;
2197        roundtrip(Some(ScaledFloat::new(0.45)))?;
2198        roundtrip(Some(ScaledFloat::new(0.55)))?;
2199        roundtrip(Some(ScaledFloat::new(0.7)))?;
2200        roundtrip(Some(ScaledFloat::new(1.0)))?;
2201        roundtrip(Some(ScaledFloat::new(2.5)))?;
2202
2203        Ok(())
2204    }
2205
2206    #[test]
2207    fn write_image_chunks_beyond_first() -> Result<()> {
2208        let width = 10;
2209        let height = 10;
2210
2211        let output = vec![0u8; 1024];
2212        let writer = Cursor::new(output);
2213
2214        // Not an animation but we should still be able to write multiple images
2215        // See issue: <https://github.com/image-rs/image-png/issues/301>
2216        // This is technically all valid png so there is no issue with correctness.
2217        let mut encoder = Encoder::new(writer, width, height);
2218        encoder.set_depth(BitDepth::Eight);
2219        encoder.set_color(ColorType::Grayscale);
2220        let mut png_writer = encoder.write_header()?;
2221
2222        for _ in 0..3 {
2223            let correct_image_size = (width * height) as usize;
2224            let image = vec![0u8; correct_image_size];
2225            png_writer.write_image_data(image.as_ref())?;
2226        }
2227
2228        Ok(())
2229    }
2230
2231    #[test]
2232    fn image_validate_sequence_without_animation() -> Result<()> {
2233        let width = 10;
2234        let height = 10;
2235
2236        let output = vec![0u8; 1024];
2237        let writer = Cursor::new(output);
2238
2239        let mut encoder = Encoder::new(writer, width, height);
2240        encoder.set_depth(BitDepth::Eight);
2241        encoder.set_color(ColorType::Grayscale);
2242        encoder.validate_sequence(true);
2243        let mut png_writer = encoder.write_header()?;
2244
2245        let correct_image_size = (width * height) as usize;
2246        let image = vec![0u8; correct_image_size];
2247        png_writer.write_image_data(image.as_ref())?;
2248
2249        assert!(png_writer.write_image_data(image.as_ref()).is_err());
2250        Ok(())
2251    }
2252
2253    #[test]
2254    fn image_validate_animation() -> Result<()> {
2255        let width = 10;
2256        let height = 10;
2257
2258        let output = vec![0u8; 1024];
2259        let writer = Cursor::new(output);
2260        let correct_image_size = (width * height) as usize;
2261        let image = vec![0u8; correct_image_size];
2262
2263        let mut encoder = Encoder::new(writer, width, height);
2264        encoder.set_depth(BitDepth::Eight);
2265        encoder.set_color(ColorType::Grayscale);
2266        encoder.set_animated(1, 0)?;
2267        encoder.validate_sequence(true);
2268        let mut png_writer = encoder.write_header()?;
2269
2270        png_writer.write_image_data(image.as_ref())?;
2271
2272        Ok(())
2273    }
2274
2275    #[test]
2276    fn image_validate_animation2() -> Result<()> {
2277        let width = 10;
2278        let height = 10;
2279
2280        let output = vec![0u8; 1024];
2281        let writer = Cursor::new(output);
2282        let correct_image_size = (width * height) as usize;
2283        let image = vec![0u8; correct_image_size];
2284
2285        let mut encoder = Encoder::new(writer, width, height);
2286        encoder.set_depth(BitDepth::Eight);
2287        encoder.set_color(ColorType::Grayscale);
2288        encoder.set_animated(2, 0)?;
2289        encoder.validate_sequence(true);
2290        let mut png_writer = encoder.write_header()?;
2291
2292        png_writer.write_image_data(image.as_ref())?;
2293        png_writer.write_image_data(image.as_ref())?;
2294        png_writer.finish()?;
2295
2296        Ok(())
2297    }
2298
2299    #[test]
2300    fn image_validate_animation_sep_def_image() -> Result<()> {
2301        let width = 10;
2302        let height = 10;
2303
2304        let output = vec![0u8; 1024];
2305        let writer = Cursor::new(output);
2306        let correct_image_size = (width * height) as usize;
2307        let image = vec![0u8; correct_image_size];
2308
2309        let mut encoder = Encoder::new(writer, width, height);
2310        encoder.set_depth(BitDepth::Eight);
2311        encoder.set_color(ColorType::Grayscale);
2312        encoder.set_animated(1, 0)?;
2313        encoder.set_sep_def_img(true)?;
2314        encoder.validate_sequence(true);
2315        let mut png_writer = encoder.write_header()?;
2316
2317        png_writer.write_image_data(image.as_ref())?;
2318        png_writer.write_image_data(image.as_ref())?;
2319        png_writer.finish()?;
2320
2321        Ok(())
2322    }
2323
2324    #[test]
2325    fn image_validate_missing_image() -> Result<()> {
2326        let width = 10;
2327        let height = 10;
2328
2329        let output = vec![0u8; 1024];
2330        let writer = Cursor::new(output);
2331
2332        let mut encoder = Encoder::new(writer, width, height);
2333        encoder.set_depth(BitDepth::Eight);
2334        encoder.set_color(ColorType::Grayscale);
2335        encoder.validate_sequence(true);
2336        let png_writer = encoder.write_header()?;
2337
2338        assert!(png_writer.finish().is_err());
2339        Ok(())
2340    }
2341
2342    #[test]
2343    fn image_validate_missing_animated_frame() -> Result<()> {
2344        let width = 10;
2345        let height = 10;
2346
2347        let output = vec![0u8; 1024];
2348        let writer = Cursor::new(output);
2349        let correct_image_size = (width * height) as usize;
2350        let image = vec![0u8; correct_image_size];
2351
2352        let mut encoder = Encoder::new(writer, width, height);
2353        encoder.set_depth(BitDepth::Eight);
2354        encoder.set_color(ColorType::Grayscale);
2355        encoder.set_animated(2, 0)?;
2356        encoder.validate_sequence(true);
2357        let mut png_writer = encoder.write_header()?;
2358
2359        png_writer.write_image_data(image.as_ref())?;
2360        assert!(png_writer.finish().is_err());
2361
2362        Ok(())
2363    }
2364
2365    #[test]
2366    fn issue_307_stream_validation() -> Result<()> {
2367        let output = vec![0u8; 1024];
2368        let mut cursor = Cursor::new(output);
2369
2370        let encoder = Encoder::new(&mut cursor, 1, 1); // Create a 1-pixel image
2371        let mut writer = encoder.write_header()?;
2372        let mut stream = writer.stream_writer()?;
2373
2374        let written = stream.write(&[1, 2, 3, 4])?;
2375        assert_eq!(written, 1);
2376        stream.finish()?;
2377        drop(writer);
2378
2379        {
2380            cursor.set_position(0);
2381            let mut decoder = Decoder::new(cursor).read_info().expect("A valid image");
2382            let mut buffer = [0u8; 1];
2383            decoder.next_frame(&mut buffer[..]).expect("Valid read");
2384            assert_eq!(buffer, [1]);
2385        }
2386
2387        Ok(())
2388    }
2389
2390    #[test]
2391    fn stream_filtering() -> Result<()> {
2392        let output = vec![0u8; 1024];
2393        let mut cursor = Cursor::new(output);
2394
2395        let mut encoder = Encoder::new(&mut cursor, 8, 8);
2396        encoder.set_color(ColorType::Rgba);
2397        encoder.set_filter(Filter::Paeth);
2398        let mut writer = encoder.write_header()?;
2399        let mut stream = writer.stream_writer()?;
2400
2401        for _ in 0..8 {
2402            let written = stream.write(&[1; 32])?;
2403            assert_eq!(written, 32);
2404        }
2405        stream.finish()?;
2406        drop(writer);
2407
2408        {
2409            cursor.set_position(0);
2410            let mut decoder = Decoder::new(cursor).read_info().expect("A valid image");
2411            let mut buffer = [0u8; 256];
2412            decoder.next_frame(&mut buffer[..]).expect("Valid read");
2413            assert_eq!(buffer, [1; 256]);
2414        }
2415
2416        Ok(())
2417    }
2418
2419    fn test_stream_flushing(compression: Compression) -> Result<()> {
2420        let output = vec![0u8; 1024];
2421        let mut cursor = Cursor::new(output);
2422
2423        let mut encoder = Encoder::new(&mut cursor, 8, 8);
2424        encoder.set_color(ColorType::Rgba);
2425        encoder.set_compression(compression);
2426        let mut writer = encoder.write_header()?;
2427        let mut stream = writer.stream_writer()?;
2428
2429        for _ in 0..8 {
2430            let written = stream.write(&[1; 32])?;
2431            assert_eq!(written, 32);
2432            stream.flush()?;
2433        }
2434        stream.finish()?;
2435        drop(writer);
2436
2437        {
2438            cursor.set_position(0);
2439            let mut decoder = Decoder::new(cursor).read_info().expect("A valid image");
2440            let mut buffer = [0u8; 256];
2441            decoder.next_frame(&mut buffer[..]).expect("Valid read");
2442            assert_eq!(buffer, [1; 256]);
2443        }
2444
2445        Ok(())
2446    }
2447
2448    #[test]
2449    fn stream_flushing_with_high_compression() -> Result<()> {
2450        test_stream_flushing(Compression::High)
2451    }
2452
2453    #[test]
2454    fn stream_flushing_with_balanced_compression() -> Result<()> {
2455        test_stream_flushing(Compression::Balanced)
2456    }
2457
2458    #[test]
2459    fn stream_flushing_with_fast_compression() -> Result<()> {
2460        test_stream_flushing(Compression::Fast)
2461    }
2462
2463    #[test]
2464    fn stream_flushing_with_fastest_compression() -> Result<()> {
2465        test_stream_flushing(Compression::Fastest)
2466    }
2467
2468    #[test]
2469    fn stream_flushing_with_no_compression() -> Result<()> {
2470        test_stream_flushing(Compression::NoCompression)
2471    }
2472
2473    #[test]
2474    #[cfg(all(unix, not(target_pointer_width = "32")))]
2475    fn exper_error_on_huge_chunk() -> Result<()> {
2476        // Okay, so we want a proper 4 GB chunk but not actually spend the memory for reserving it.
2477        // Let's rely on overcommit? Otherwise we got the rather dumb option of mmap-ing /dev/zero.
2478        let empty = vec![0; 1usize << 31];
2479        let writer = Cursor::new(vec![0u8; 1024]);
2480
2481        let mut encoder = Encoder::new(writer, 10, 10);
2482        encoder.set_depth(BitDepth::Eight);
2483        encoder.set_color(ColorType::Grayscale);
2484        let mut png_writer = encoder.write_header()?;
2485
2486        assert!(png_writer.write_chunk(chunk::fdAT, &empty).is_err());
2487        Ok(())
2488    }
2489
2490    #[test]
2491    #[cfg(all(unix, not(target_pointer_width = "32")))]
2492    fn exper_error_on_non_u32_chunk() -> Result<()> {
2493        // Okay, so we want a proper 4 GB chunk but not actually spend the memory for reserving it.
2494        // Let's rely on overcommit? Otherwise we got the rather dumb option of mmap-ing /dev/zero.
2495        let empty = vec![0; 1usize << 32];
2496        let writer = Cursor::new(vec![0u8; 1024]);
2497
2498        let mut encoder = Encoder::new(writer, 10, 10);
2499        encoder.set_depth(BitDepth::Eight);
2500        encoder.set_color(ColorType::Grayscale);
2501        let mut png_writer = encoder.write_header()?;
2502
2503        assert!(png_writer.write_chunk(chunk::fdAT, &empty).is_err());
2504        Ok(())
2505    }
2506
2507    #[test]
2508    fn finish_drops_inner_writer() -> Result<()> {
2509        struct NoWriter<'flag>(&'flag mut bool);
2510
2511        impl Write for NoWriter<'_> {
2512            fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
2513                Ok(buf.len())
2514            }
2515            fn flush(&mut self) -> io::Result<()> {
2516                Ok(())
2517            }
2518        }
2519        impl Drop for NoWriter<'_> {
2520            fn drop(&mut self) {
2521                *self.0 = true;
2522            }
2523        }
2524
2525        let mut flag = false;
2526
2527        {
2528            let mut encoder = Encoder::new(NoWriter(&mut flag), 10, 10);
2529            encoder.set_depth(BitDepth::Eight);
2530            encoder.set_color(ColorType::Grayscale);
2531
2532            let mut writer = encoder.write_header()?;
2533            writer.write_image_data(&[0; 100])?;
2534            writer.finish()?;
2535        }
2536
2537        assert!(flag, "PNG finished but writer was not dropped");
2538        Ok(())
2539    }
2540
2541    /// A Writer that only writes a few bytes at a time
2542    struct RandomChunkWriter<R: Rng, W: Write> {
2543        rng: R,
2544        w: W,
2545    }
2546
2547    impl<R: Rng, W: Write> Write for RandomChunkWriter<R, W> {
2548        fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
2549            // choose a random length to write
2550            let len = cmp::min(self.rng.random_range(1..50), buf.len());
2551
2552            self.w.write(&buf[0..len])
2553        }
2554
2555        fn flush(&mut self) -> io::Result<()> {
2556            self.w.flush()
2557        }
2558    }
2559}