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