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