Skip to main content

png/
encoder.rs

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