png/
encoder.rs

Help
1extern crate deflate;
2
3use std::borrow::Cow;
4use std::error;
5use std::fmt;
6use std::io::{self, Write};
7use std::mem;
8use std::result;
9
10use chunk;
11use crc::Crc32;
12use common::{Info, ColorType, BitDepth, Compression};
13use filter::{FilterType, filter};
14use traits::{WriteBytesExt, HasParameters, Parameter};
15
16pub type Result<T> = result::Result<T, EncodingError>;
17
18#[derive(Debug)]
19pub enum EncodingError {
20    IoError(io::Error),
21    Format(Cow<'static, str>),
22}
23
24impl error::Error for EncodingError {
25    fn description(&self) -> &str {
26        use self::EncodingError::*;
27        match *self {
28            IoError(ref err) => err.description(),
29            Format(ref desc) => &desc,
30        }
31    }
32}
33
34impl fmt::Display for EncodingError {
35    fn fmt(&self, fmt: &mut fmt::Formatter) -> result::Result<(), fmt::Error> {
36        write!(fmt, "{}", (self as &error::Error).description())
37    }
38}
39
40impl From<io::Error> for EncodingError {
41    fn from(err: io::Error) -> EncodingError {
42        EncodingError::IoError(err)
43    }
44}
45impl From<EncodingError> for io::Error {
46    fn from(err: EncodingError) -> io::Error {
47        io::Error::new(io::ErrorKind::Other, (&err as &error::Error).description())
48    }
49}
50
51/// PNG Encoder
52pub struct Encoder<W: Write> {
53    w: W,
54    info: Info,
55}
56
57impl<W: Write> Encoder<W> {
58    pub fn new(w: W, width: u32, height: u32) -> Encoder<W> {
59        let mut info = Info::default();
60        info.width = width;
61        info.height = height;
62        Encoder { w: w, info: info }
63    }
64
65    pub fn write_header(self) -> Result<Writer<W>> {
66        Writer::new(self.w, self.info).init()
67    }
68}
69
70impl<W: Write> HasParameters for Encoder<W> {}
71
72impl<W: Write> Parameter<Encoder<W>> for ColorType {
73    fn set_param(self, this: &mut Encoder<W>) {
74        this.info.color_type = self
75    }
76}
77
78impl<W: Write> Parameter<Encoder<W>> for BitDepth {
79    fn set_param(self, this: &mut Encoder<W>) {
80        this.info.bit_depth = self
81    }
82}
83
84/// Set compression param for a ```Compression``` or any type that can transform 
85/// into a ```Compression```, notably ```deflate::Compression``` and 
86/// ```deflate::CompressionOptions``` which "just work".
87impl<W: Write, C: Into<Compression>> Parameter<Encoder<W>> for C {
88    fn set_param(self, this: &mut Encoder<W>) {
89        this.info.compression = self.into()
90    }
91}
92
93impl <W: Write> Parameter<Encoder<W>> for FilterType {
94    fn set_param(self, this: &mut Encoder<W>) {
95        this.info.filter = self
96    }
97}
98
99/// PNG writer
100pub struct Writer<W: Write> {
101    w: W,
102    info: Info,
103}
104
105impl<W: Write> Writer<W> {
106    fn new(w: W, info: Info) -> Writer<W> {
107        let w = Writer { w: w, info: info };
108        w
109    }
110
111    fn init(mut self) -> Result<Self> {
112        try!(self.w.write_all(&[137, 80, 78, 71, 13, 10, 26, 10]));
113        let mut data = [0; 13];
114        try!((&mut data[..]).write_be(self.info.width));
115        try!((&mut data[4..]).write_be(self.info.height));
116        data[8] = self.info.bit_depth as u8;
117        data[9] = self.info.color_type as u8;
118        data[12] = if self.info.interlaced { 1 } else { 0 };
119        try!(self.write_chunk(chunk::IHDR, &data));
120        Ok(self)
121    }
122
123    pub fn write_chunk(&mut self, name: [u8; 4], data: &[u8]) -> Result<()> {
124        try!(self.w.write_be(data.len() as u32));
125        try!(self.w.write_all(&name));
126        try!(self.w.write_all(data));
127        let mut crc = Crc32::new();
128        crc.update(&name);
129        crc.update(data);
130        try!(self.w.write_be(crc.checksum()));
131        Ok(())
132    }
133
134    /// Writes the image data.
135    pub fn write_image_data(&mut self, data: &[u8]) -> Result<()> {
136        let bpp = self.info.bytes_per_pixel();
137        let in_len = self.info.raw_row_length() - 1;
138        let mut prev = vec![0; in_len];
139        let mut current = vec![0; in_len];
140        let data_size = in_len * self.info.height as usize;
141        if data_size != data.len() {
142            let message = format!("wrong data size, expected {} got {}", data_size, data.len());
143            return Err(EncodingError::Format(message.into()));
144        }
145        let mut zlib = deflate::write::ZlibEncoder::new(Vec::new(), self.info.compression.clone());
146        let filter_method = self.info.filter;
147        for line in data.chunks(in_len) {
148            current.copy_from_slice(&line);
149            try!(zlib.write_all(&[filter_method as u8]));
150            filter(filter_method, bpp, &prev, &mut current);
151            try!(zlib.write_all(&current));
152            mem::swap(&mut prev, &mut current);
153        }
154        self.write_chunk(chunk::IDAT, &try!(zlib.finish()))
155    }
156}
157
158impl<W: Write> Drop for Writer<W> {
159    fn drop(&mut self) {
160        let _ = self.write_chunk(chunk::IEND, &[]);
161    }
162}
163
164#[cfg(test)]
165mod tests {
166    use super::*;
167
168    extern crate rand;
169    extern crate glob;
170
171    use self::rand::Rng;
172    use std::{io, cmp};
173    use std::io::Write;
174    use std::fs::File;
175
176    #[test]
177    fn roundtrip() {
178        // More loops = more random testing, but also more test wait time
179        for _ in 0..10 {
180            for path in glob::glob("tests/pngsuite/*.png").unwrap().map(|r| r.unwrap()) {
181                if path.file_name().unwrap().to_str().unwrap().starts_with("x") {
182                    // x* files are expected to fail to decode
183                    continue;
184                }
185                // Decode image
186                let decoder = ::Decoder::new(File::open(path).unwrap());
187                let (info, mut reader) = decoder.read_info().unwrap();
188                if info.line_size != 32 {
189                    // TODO encoding only works with line size 32?
190                    continue;
191                }
192                let mut buf = vec![0; info.buffer_size()];
193                reader.next_frame(&mut buf).unwrap();
194                // Encode decoded image
195                let mut out = Vec::new();
196                {
197                    let mut wrapper = RandomChunkWriter {
198                        rng: self::rand::thread_rng(),
199                        w: &mut out
200                    };
201
202                    let mut encoder = Encoder::new(&mut wrapper, info.width, info.height).write_header().unwrap();
203                    encoder.write_image_data(&buf).unwrap();
204                }
205                // Decode encoded decoded image
206                let decoder = ::Decoder::new(&*out);
207                let (info, mut reader) = decoder.read_info().unwrap();
208                let mut buf2 = vec![0; info.buffer_size()];
209                reader.next_frame(&mut buf2).unwrap();
210                // check if the encoded image is ok:
211                assert_eq!(buf, buf2);
212            }
213        }
214    }
215
216    #[test]
217    fn expect_error_on_wrong_image_len() -> Result<()> {
218        use std::io::Cursor;
219
220        let width = 10;
221        let height = 10;
222
223        let output = vec![0u8; 1024];
224        let writer = Cursor::new(output);
225        let mut encoder = Encoder::new(writer, width as u32, height as u32);
226        encoder.set(BitDepth::Eight);
227        encoder.set(ColorType::RGB);
228        let mut png_writer = encoder.write_header()?;
229
230        let correct_image_size = width * height * 3;
231        let image = vec![0u8; correct_image_size + 1];
232        let result = png_writer.write_image_data(image.as_ref());
233        assert!(result.is_err());
234
235        Ok(())
236    }
237
238    /// A Writer that only writes a few bytes at a time
239    struct RandomChunkWriter<'a, R: Rng, W: Write + 'a> {
240        rng: R,
241        w: &'a mut W
242    }
243
244    impl<'a, R: Rng, W: Write + 'a> Write for RandomChunkWriter<'a, R, W> {
245        fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
246            // choose a random length to write
247            let len = cmp::min(self.rng.gen_range(1, 50), buf.len());
248
249            self.w.write(&buf[0..len])
250        }
251
252        fn flush(&mut self) -> io::Result<()> {
253            self.w.flush()
254        }
255    }
256
257}
png/encoder.rs

png/
encoder.rs
