1 #![allow(clippy::too_many_arguments)]
2 
3 use byteorder::{BigEndian, WriteBytesExt};
4 use math::utils::clamp;
5 use num_iter::range_step;
6 use std::io::{self, Write};
7 
8 use color;
9 
10 use super::entropy::build_huff_lut;
11 use super::transform;
12 
13 // Markers
14 // Baseline DCT
15 static SOF0: u8 = 0xC0;
16 // Huffman Tables
17 static DHT: u8 = 0xC4;
18 // Start of Image (standalone)
19 static SOI: u8 = 0xD8;
20 // End of image (standalone)
21 static EOI: u8 = 0xD9;
22 // Start of Scan
23 static SOS: u8 = 0xDA;
24 // Quantization Tables
25 static DQT: u8 = 0xDB;
26 // Application segments start and end
27 static APP0: u8 = 0xE0;
28 
29 // section K.1
30 // table K.1
31 #[rustfmt::skip]
32 static STD_LUMA_QTABLE: [u8; 64] = [
33     16, 11, 10, 16,  24,  40,  51,  61,
34     12, 12, 14, 19,  26,  58,  60,  55,
35     14, 13, 16, 24,  40,  57,  69,  56,
36     14, 17, 22, 29,  51,  87,  80,  62,
37     18, 22, 37, 56,  68, 109, 103,  77,
38     24, 35, 55, 64,  81, 104, 113,  92,
39     49, 64, 78, 87, 103, 121, 120, 101,
40     72, 92, 95, 98, 112, 100, 103,  99,
41 ];
42 
43 // table K.2
44 #[rustfmt::skip]
45 static STD_CHROMA_QTABLE: [u8; 64] = [
46     17, 18, 24, 47, 99, 99, 99, 99,
47     18, 21, 26, 66, 99, 99, 99, 99,
48     24, 26, 56, 99, 99, 99, 99, 99,
49     47, 66, 99, 99, 99, 99, 99, 99,
50     99, 99, 99, 99, 99, 99, 99, 99,
51     99, 99, 99, 99, 99, 99, 99, 99,
52     99, 99, 99, 99, 99, 99, 99, 99,
53     99, 99, 99, 99, 99, 99, 99, 99,
54 ];
55 
56 // section K.3
57 // Code lengths and values for table K.3
58 static STD_LUMA_DC_CODE_LENGTHS: [u8; 16] = [
59     0x00, 0x01, 0x05, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
60 ];
61 
62 static STD_LUMA_DC_VALUES: [u8; 12] = [
63     0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A, 0x0B,
64 ];
65 
66 // Code lengths and values for table K.4
67 static STD_CHROMA_DC_CODE_LENGTHS: [u8; 16] = [
68     0x00, 0x03, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00,
69 ];
70 
71 static STD_CHROMA_DC_VALUES: [u8; 12] = [
72     0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A, 0x0B,
73 ];
74 
75 // Code lengths and values for table k.5
76 static STD_LUMA_AC_CODE_LENGTHS: [u8; 16] = [
77     0x00, 0x02, 0x01, 0x03, 0x03, 0x02, 0x04, 0x03, 0x05, 0x05, 0x04, 0x04, 0x00, 0x00, 0x01, 0x7D,
78 ];
79 
80 static STD_LUMA_AC_VALUES: [u8; 162] = [
81     0x01, 0x02, 0x03, 0x00, 0x04, 0x11, 0x05, 0x12, 0x21, 0x31, 0x41, 0x06, 0x13, 0x51, 0x61, 0x07,
82     0x22, 0x71, 0x14, 0x32, 0x81, 0x91, 0xA1, 0x08, 0x23, 0x42, 0xB1, 0xC1, 0x15, 0x52, 0xD1, 0xF0,
83     0x24, 0x33, 0x62, 0x72, 0x82, 0x09, 0x0A, 0x16, 0x17, 0x18, 0x19, 0x1A, 0x25, 0x26, 0x27, 0x28,
84     0x29, 0x2A, 0x34, 0x35, 0x36, 0x37, 0x38, 0x39, 0x3A, 0x43, 0x44, 0x45, 0x46, 0x47, 0x48, 0x49,
85     0x4A, 0x53, 0x54, 0x55, 0x56, 0x57, 0x58, 0x59, 0x5A, 0x63, 0x64, 0x65, 0x66, 0x67, 0x68, 0x69,
86     0x6A, 0x73, 0x74, 0x75, 0x76, 0x77, 0x78, 0x79, 0x7A, 0x83, 0x84, 0x85, 0x86, 0x87, 0x88, 0x89,
87     0x8A, 0x92, 0x93, 0x94, 0x95, 0x96, 0x97, 0x98, 0x99, 0x9A, 0xA2, 0xA3, 0xA4, 0xA5, 0xA6, 0xA7,
88     0xA8, 0xA9, 0xAA, 0xB2, 0xB3, 0xB4, 0xB5, 0xB6, 0xB7, 0xB8, 0xB9, 0xBA, 0xC2, 0xC3, 0xC4, 0xC5,
89     0xC6, 0xC7, 0xC8, 0xC9, 0xCA, 0xD2, 0xD3, 0xD4, 0xD5, 0xD6, 0xD7, 0xD8, 0xD9, 0xDA, 0xE1, 0xE2,
90     0xE3, 0xE4, 0xE5, 0xE6, 0xE7, 0xE8, 0xE9, 0xEA, 0xF1, 0xF2, 0xF3, 0xF4, 0xF5, 0xF6, 0xF7, 0xF8,
91     0xF9, 0xFA,
92 ];
93 
94 // Code lengths and values for table k.6
95 static STD_CHROMA_AC_CODE_LENGTHS: [u8; 16] = [
96     0x00, 0x02, 0x01, 0x02, 0x04, 0x04, 0x03, 0x04, 0x07, 0x05, 0x04, 0x04, 0x00, 0x01, 0x02, 0x77,
97 ];
98 static STD_CHROMA_AC_VALUES: [u8; 162] = [
99     0x00, 0x01, 0x02, 0x03, 0x11, 0x04, 0x05, 0x21, 0x31, 0x06, 0x12, 0x41, 0x51, 0x07, 0x61, 0x71,
100     0x13, 0x22, 0x32, 0x81, 0x08, 0x14, 0x42, 0x91, 0xA1, 0xB1, 0xC1, 0x09, 0x23, 0x33, 0x52, 0xF0,
101     0x15, 0x62, 0x72, 0xD1, 0x0A, 0x16, 0x24, 0x34, 0xE1, 0x25, 0xF1, 0x17, 0x18, 0x19, 0x1A, 0x26,
102     0x27, 0x28, 0x29, 0x2A, 0x35, 0x36, 0x37, 0x38, 0x39, 0x3A, 0x43, 0x44, 0x45, 0x46, 0x47, 0x48,
103     0x49, 0x4A, 0x53, 0x54, 0x55, 0x56, 0x57, 0x58, 0x59, 0x5A, 0x63, 0x64, 0x65, 0x66, 0x67, 0x68,
104     0x69, 0x6A, 0x73, 0x74, 0x75, 0x76, 0x77, 0x78, 0x79, 0x7A, 0x82, 0x83, 0x84, 0x85, 0x86, 0x87,
105     0x88, 0x89, 0x8A, 0x92, 0x93, 0x94, 0x95, 0x96, 0x97, 0x98, 0x99, 0x9A, 0xA2, 0xA3, 0xA4, 0xA5,
106     0xA6, 0xA7, 0xA8, 0xA9, 0xAA, 0xB2, 0xB3, 0xB4, 0xB5, 0xB6, 0xB7, 0xB8, 0xB9, 0xBA, 0xC2, 0xC3,
107     0xC4, 0xC5, 0xC6, 0xC7, 0xC8, 0xC9, 0xCA, 0xD2, 0xD3, 0xD4, 0xD5, 0xD6, 0xD7, 0xD8, 0xD9, 0xDA,
108     0xE2, 0xE3, 0xE4, 0xE5, 0xE6, 0xE7, 0xE8, 0xE9, 0xEA, 0xF2, 0xF3, 0xF4, 0xF5, 0xF6, 0xF7, 0xF8,
109     0xF9, 0xFA,
110 ];
111 
112 static DCCLASS: u8 = 0;
113 static ACCLASS: u8 = 1;
114 
115 static LUMADESTINATION: u8 = 0;
116 static CHROMADESTINATION: u8 = 1;
117 
118 static LUMAID: u8 = 1;
119 static CHROMABLUEID: u8 = 2;
120 static CHROMAREDID: u8 = 3;
121 
122 /// The permutation of dct coefficients.
123 #[rustfmt::skip]
124 static UNZIGZAG: [u8; 64] = [
125      0,  1,  8, 16,  9,  2,  3, 10,
126     17, 24, 32, 25, 18, 11,  4,  5,
127     12, 19, 26, 33, 40, 48, 41, 34,
128     27, 20, 13,  6,  7, 14, 21, 28,
129     35, 42, 49, 56, 57, 50, 43, 36,
130     29, 22, 15, 23, 30, 37, 44, 51,
131     58, 59, 52, 45, 38, 31, 39, 46,
132     53, 60, 61, 54, 47, 55, 62, 63,
133 ];
134 
135 /// A representation of a JPEG component
136 #[derive(Copy, Clone)]
137 struct Component {
138     /// The Component's identifier
139     id: u8,
140 
141     /// Horizontal sampling factor
142     h: u8,
143 
144     /// Vertical sampling factor
145     v: u8,
146 
147     /// The quantization table selector
148     tq: u8,
149 
150     /// Index to the Huffman DC Table
151     dc_table: u8,
152 
153     /// Index to the AC Huffman Table
154     ac_table: u8,
155 
156     /// The dc prediction of the component
157     _dc_pred: i32,
158 }
159 
160 pub(crate) struct BitWriter<'a, W: 'a> {
161     w: &'a mut W,
162     accumulator: u32,
163     nbits: u8,
164 }
165 
166 impl<'a, W: Write + 'a> BitWriter<'a, W> {
new(w: &'a mut W) -> Self167     fn new(w: &'a mut W) -> Self {
168         BitWriter {
169             w,
170             accumulator: 0,
171             nbits: 0,
172         }
173     }
174 
write_bits(&mut self, bits: u16, size: u8) -> io::Result<()>175     fn write_bits(&mut self, bits: u16, size: u8) -> io::Result<()> {
176         if size == 0 {
177             return Ok(());
178         }
179 
180         self.accumulator |= u32::from(bits) << (32 - (self.nbits + size)) as usize;
181         self.nbits += size;
182 
183         while self.nbits >= 8 {
184             let byte = (self.accumulator & (0xFFFF_FFFFu32 << 24)) >> 24;
185             self.w.write_all(&[byte as u8])?;
186 
187             if byte == 0xFF {
188                 self.w.write_all(&[0x00])?;
189             }
190 
191             self.nbits -= 8;
192             self.accumulator <<= 8;
193         }
194 
195         Ok(())
196     }
197 
pad_byte(&mut self) -> io::Result<()>198     fn pad_byte(&mut self) -> io::Result<()> {
199         self.write_bits(0x7F, 7)
200     }
201 
huffman_encode(&mut self, val: u8, table: &[(u8, u16)]) -> io::Result<()>202     fn huffman_encode(&mut self, val: u8, table: &[(u8, u16)]) -> io::Result<()> {
203         let (size, code) = table[val as usize];
204 
205         if size > 16 {
206             panic!("bad huffman value");
207         }
208 
209         self.write_bits(code, size)
210     }
211 
write_block( &mut self, block: &[i32], prevdc: i32, dctable: &[(u8, u16)], actable: &[(u8, u16)], ) -> io::Result<i32>212     fn write_block(
213         &mut self,
214         block: &[i32],
215         prevdc: i32,
216         dctable: &[(u8, u16)],
217         actable: &[(u8, u16)],
218     ) -> io::Result<i32> {
219         // Differential DC encoding
220         let dcval = block[0];
221         let diff = dcval - prevdc;
222         let (size, value) = encode_coefficient(diff);
223 
224         self.huffman_encode(size, dctable)?;
225         self.write_bits(value, size)?;
226 
227         // Figure F.2
228         let mut zero_run = 0;
229         let mut k = 0usize;
230 
231         loop {
232             k += 1;
233 
234             if block[UNZIGZAG[k] as usize] == 0 {
235                 if k == 63 {
236                     self.huffman_encode(0x00, actable)?;
237                     break;
238                 }
239 
240                 zero_run += 1;
241             } else {
242                 while zero_run > 15 {
243                     self.huffman_encode(0xF0, actable)?;
244                     zero_run -= 16;
245                 }
246 
247                 let (size, value) = encode_coefficient(block[UNZIGZAG[k] as usize]);
248                 let symbol = (zero_run << 4) | size;
249 
250                 self.huffman_encode(symbol, actable)?;
251                 self.write_bits(value, size)?;
252 
253                 zero_run = 0;
254 
255                 if k == 63 {
256                     break;
257                 }
258             }
259         }
260 
261         Ok(dcval)
262     }
263 
write_segment(&mut self, marker: u8, data: Option<&[u8]>) -> io::Result<()>264     fn write_segment(&mut self, marker: u8, data: Option<&[u8]>) -> io::Result<()> {
265         self.w.write_all(&[0xFF])?;
266         self.w.write_all(&[marker])?;
267 
268         if let Some(b) = data {
269             self.w.write_u16::<BigEndian>(b.len() as u16 + 2)?;
270             self.w.write_all(b)?;
271         }
272         Ok(())
273     }
274 }
275 
276 /// The representation of a JPEG encoder
277 pub struct JPEGEncoder<'a, W: 'a> {
278     writer: BitWriter<'a, W>,
279 
280     components: Vec<Component>,
281     tables: Vec<u8>,
282 
283     luma_dctable: Vec<(u8, u16)>,
284     luma_actable: Vec<(u8, u16)>,
285     chroma_dctable: Vec<(u8, u16)>,
286     chroma_actable: Vec<(u8, u16)>,
287 }
288 
289 impl<'a, W: Write> JPEGEncoder<'a, W> {
290     /// Create a new encoder that writes its output to ```w```
new(w: &mut W) -> JPEGEncoder<W>291     pub fn new(w: &mut W) -> JPEGEncoder<W> {
292         JPEGEncoder::new_with_quality(w, 75)
293     }
294 
295     /// Create a new encoder that writes its output to ```w```, and has
296     /// the quality parameter ```quality``` with a value in the range 1-100
297     /// where 1 is the worst and 100 is the best.
new_with_quality(w: &mut W, quality: u8) -> JPEGEncoder<W>298     pub fn new_with_quality(w: &mut W, quality: u8) -> JPEGEncoder<W> {
299         let ld = build_huff_lut(&STD_LUMA_DC_CODE_LENGTHS, &STD_LUMA_DC_VALUES);
300         let la = build_huff_lut(&STD_LUMA_AC_CODE_LENGTHS, &STD_LUMA_AC_VALUES);
301 
302         let cd = build_huff_lut(&STD_CHROMA_DC_CODE_LENGTHS, &STD_CHROMA_DC_VALUES);
303         let ca = build_huff_lut(&STD_CHROMA_AC_CODE_LENGTHS, &STD_CHROMA_AC_VALUES);
304 
305         let components = vec![
306             Component {
307                 id: LUMAID,
308                 h: 1,
309                 v: 1,
310                 tq: LUMADESTINATION,
311                 dc_table: LUMADESTINATION,
312                 ac_table: LUMADESTINATION,
313                 _dc_pred: 0,
314             },
315             Component {
316                 id: CHROMABLUEID,
317                 h: 1,
318                 v: 1,
319                 tq: CHROMADESTINATION,
320                 dc_table: CHROMADESTINATION,
321                 ac_table: CHROMADESTINATION,
322                 _dc_pred: 0,
323             },
324             Component {
325                 id: CHROMAREDID,
326                 h: 1,
327                 v: 1,
328                 tq: CHROMADESTINATION,
329                 dc_table: CHROMADESTINATION,
330                 ac_table: CHROMADESTINATION,
331                 _dc_pred: 0,
332             },
333         ];
334 
335         // Derive our quantization table scaling value using the libjpeg algorithm
336         let scale = u32::from(clamp(quality, 1, 100));
337         let scale = if scale < 50 {
338             5000 / scale
339         } else {
340             200 - scale * 2
341         };
342 
343         let mut tables = Vec::new();
344         let scale_value = |&v: &u8| {
345             let value = (u32::from(v) * scale + 50) / 100;
346 
347             clamp(value, 1, u32::from(u8::max_value())) as u8
348         };
349         tables.extend(STD_LUMA_QTABLE.iter().map(&scale_value));
350         tables.extend(STD_CHROMA_QTABLE.iter().map(&scale_value));
351 
352         JPEGEncoder {
353             writer: BitWriter::new(w),
354 
355             components,
356             tables,
357 
358             luma_dctable: ld,
359             luma_actable: la,
360             chroma_dctable: cd,
361             chroma_actable: ca,
362         }
363     }
364 
365     /// Encodes the image ```image```
366     /// that has dimensions ```width``` and ```height```
367     /// and ```ColorType``` ```c```
368     ///
369     /// The Image in encoded with subsampling ratio 4:2:2
encode( &mut self, image: &[u8], width: u32, height: u32, c: color::ColorType, ) -> io::Result<()>370     pub fn encode(
371         &mut self,
372         image: &[u8],
373         width: u32,
374         height: u32,
375         c: color::ColorType,
376     ) -> io::Result<()> {
377         let n = color::channel_count(c);
378         let num_components = if n == 1 || n == 2 { 1 } else { 3 };
379 
380         self.writer.write_segment(SOI, None)?;
381 
382         let mut buf = Vec::new();
383 
384         build_jfif_header(&mut buf);
385         self.writer.write_segment(APP0, Some(&buf))?;
386 
387         build_frame_header(
388             &mut buf,
389             8,
390             width as u16,
391             height as u16,
392             &self.components[..num_components],
393         );
394         self.writer.write_segment(SOF0, Some(&buf))?;
395 
396         assert_eq!(self.tables.len() / 64, 2);
397         let numtables = if num_components == 1 { 1 } else { 2 };
398 
399         for (i, table) in self.tables.chunks(64).enumerate().take(numtables) {
400             build_quantization_segment(&mut buf, 8, i as u8, table);
401             self.writer.write_segment(DQT, Some(&buf))?;
402         }
403 
404         build_huffman_segment(
405             &mut buf,
406             DCCLASS,
407             LUMADESTINATION,
408             &STD_LUMA_DC_CODE_LENGTHS,
409             &STD_LUMA_DC_VALUES,
410         );
411         self.writer.write_segment(DHT, Some(&buf))?;
412 
413         build_huffman_segment(
414             &mut buf,
415             ACCLASS,
416             LUMADESTINATION,
417             &STD_LUMA_AC_CODE_LENGTHS,
418             &STD_LUMA_AC_VALUES,
419         );
420         self.writer.write_segment(DHT, Some(&buf))?;
421 
422         if num_components == 3 {
423             build_huffman_segment(
424                 &mut buf,
425                 DCCLASS,
426                 CHROMADESTINATION,
427                 &STD_CHROMA_DC_CODE_LENGTHS,
428                 &STD_CHROMA_DC_VALUES,
429             );
430             self.writer.write_segment(DHT, Some(&buf))?;
431 
432             build_huffman_segment(
433                 &mut buf,
434                 ACCLASS,
435                 CHROMADESTINATION,
436                 &STD_CHROMA_AC_CODE_LENGTHS,
437                 &STD_CHROMA_AC_VALUES,
438             );
439             self.writer.write_segment(DHT, Some(&buf))?;
440         }
441 
442         build_scan_header(&mut buf, &self.components[..num_components]);
443         self.writer.write_segment(SOS, Some(&buf))?;
444 
445         match c {
446             color::ColorType::RGB(8) => {
447                 self.encode_rgb(image, width as usize, height as usize, 3)?
448             }
449             color::ColorType::RGBA(8) => {
450                 self.encode_rgb(image, width as usize, height as usize, 4)?
451             }
452             color::ColorType::Gray(8) => {
453                 self.encode_gray(image, width as usize, height as usize, 1)?
454             }
455             color::ColorType::GrayA(8) => {
456                 self.encode_gray(image, width as usize, height as usize, 2)?
457             }
458             _ => {
459                 return Err(io::Error::new(
460                     io::ErrorKind::InvalidInput,
461                     &format!(
462                     "Unsupported color type {:?}. Use 8 bit per channel RGB(A) or Gray(A) instead.",
463                     c
464                 )[..],
465                 ))
466             }
467         };
468 
469         self.writer.pad_byte()?;
470         self.writer.write_segment(EOI, None)?;
471         Ok(())
472     }
473 
encode_gray( &mut self, image: &[u8], width: usize, height: usize, bpp: usize, ) -> io::Result<()>474     fn encode_gray(
475         &mut self,
476         image: &[u8],
477         width: usize,
478         height: usize,
479         bpp: usize,
480     ) -> io::Result<()> {
481         let mut yblock = [0u8; 64];
482         let mut y_dcprev = 0;
483         let mut dct_yblock = [0i32; 64];
484 
485         for y in range_step(0, height, 8) {
486             for x in range_step(0, width, 8) {
487                 // RGB -> YCbCr
488                 copy_blocks_gray(image, x, y, width, bpp, &mut yblock);
489 
490                 // Level shift and fdct
491                 // Coeffs are scaled by 8
492                 transform::fdct(&yblock, &mut dct_yblock);
493 
494                 // Quantization
495                 for (i, dct) in dct_yblock.iter_mut().enumerate().take(64) {
496                     *dct = ((*dct / 8) as f32 / f32::from(self.tables[i])).round() as i32;
497                 }
498 
499                 let la = &*self.luma_actable;
500                 let ld = &*self.luma_dctable;
501 
502                 y_dcprev = self.writer.write_block(&dct_yblock, y_dcprev, ld, la)?;
503             }
504         }
505 
506         Ok(())
507     }
508 
encode_rgb( &mut self, image: &[u8], width: usize, height: usize, bpp: usize, ) -> io::Result<()>509     fn encode_rgb(
510         &mut self,
511         image: &[u8],
512         width: usize,
513         height: usize,
514         bpp: usize,
515     ) -> io::Result<()> {
516         let mut y_dcprev = 0;
517         let mut cb_dcprev = 0;
518         let mut cr_dcprev = 0;
519 
520         let mut dct_yblock = [0i32; 64];
521         let mut dct_cb_block = [0i32; 64];
522         let mut dct_cr_block = [0i32; 64];
523 
524         let mut yblock = [0u8; 64];
525         let mut cb_block = [0u8; 64];
526         let mut cr_block = [0u8; 64];
527 
528         for y in range_step(0, height, 8) {
529             for x in range_step(0, width, 8) {
530                 // RGB -> YCbCr
531                 copy_blocks_ycbcr(
532                     image,
533                     x,
534                     y,
535                     width,
536                     bpp,
537                     &mut yblock,
538                     &mut cb_block,
539                     &mut cr_block,
540                 );
541 
542                 // Level shift and fdct
543                 // Coeffs are scaled by 8
544                 transform::fdct(&yblock, &mut dct_yblock);
545                 transform::fdct(&cb_block, &mut dct_cb_block);
546                 transform::fdct(&cr_block, &mut dct_cr_block);
547 
548                 // Quantization
549                 for i in 0usize..64 {
550                     dct_yblock[i] =
551                         ((dct_yblock[i] / 8) as f32 / f32::from(self.tables[i])).round() as i32;
552                     dct_cb_block[i] = ((dct_cb_block[i] / 8) as f32
553                         / f32::from(self.tables[64..][i]))
554                         .round() as i32;
555                     dct_cr_block[i] = ((dct_cr_block[i] / 8) as f32
556                         / f32::from(self.tables[64..][i]))
557                         .round() as i32;
558                 }
559 
560                 let la = &*self.luma_actable;
561                 let ld = &*self.luma_dctable;
562                 let cd = &*self.chroma_dctable;
563                 let ca = &*self.chroma_actable;
564 
565                 y_dcprev = self.writer.write_block(&dct_yblock, y_dcprev, ld, la)?;
566                 cb_dcprev = self.writer.write_block(&dct_cb_block, cb_dcprev, cd, ca)?;
567                 cr_dcprev = self.writer.write_block(&dct_cr_block, cr_dcprev, cd, ca)?;
568             }
569         }
570 
571         Ok(())
572     }
573 }
574 
build_jfif_header(m: &mut Vec<u8>)575 fn build_jfif_header(m: &mut Vec<u8>) {
576     m.clear();
577 
578     let _ = write!(m, "JFIF");
579     let _ = m.write_all(&[0]);
580     let _ = m.write_all(&[0x01]);
581     let _ = m.write_all(&[0x02]);
582     let _ = m.write_all(&[0]);
583     let _ = m.write_u16::<BigEndian>(1);
584     let _ = m.write_u16::<BigEndian>(1);
585     let _ = m.write_all(&[0]);
586     let _ = m.write_all(&[0]);
587 }
588 
build_frame_header( m: &mut Vec<u8>, precision: u8, width: u16, height: u16, components: &[Component], )589 fn build_frame_header(
590     m: &mut Vec<u8>,
591     precision: u8,
592     width: u16,
593     height: u16,
594     components: &[Component],
595 ) {
596     m.clear();
597 
598     let _ = m.write_all(&[precision]);
599     let _ = m.write_u16::<BigEndian>(height);
600     let _ = m.write_u16::<BigEndian>(width);
601     let _ = m.write_all(&[components.len() as u8]);
602 
603     for &comp in components.iter() {
604         let _ = m.write_all(&[comp.id]);
605         let hv = (comp.h << 4) | comp.v;
606         let _ = m.write_all(&[hv]);
607         let _ = m.write_all(&[comp.tq]);
608     }
609 }
610 
build_scan_header(m: &mut Vec<u8>, components: &[Component])611 fn build_scan_header(m: &mut Vec<u8>, components: &[Component]) {
612     m.clear();
613 
614     let _ = m.write_all(&[components.len() as u8]);
615 
616     for &comp in components.iter() {
617         let _ = m.write_all(&[comp.id]);
618         let tables = (comp.dc_table << 4) | comp.ac_table;
619         let _ = m.write_all(&[tables]);
620     }
621 
622     // spectral start and end, approx. high and low
623     let _ = m.write_all(&[0]);
624     let _ = m.write_all(&[63]);
625     let _ = m.write_all(&[0]);
626 }
627 
build_huffman_segment( m: &mut Vec<u8>, class: u8, destination: u8, numcodes: &[u8], values: &[u8], )628 fn build_huffman_segment(
629     m: &mut Vec<u8>,
630     class: u8,
631     destination: u8,
632     numcodes: &[u8],
633     values: &[u8],
634 ) {
635     m.clear();
636 
637     let tcth = (class << 4) | destination;
638     let _ = m.write_all(&[tcth]);
639 
640     assert_eq!(numcodes.len(), 16);
641 
642     let mut sum = 0usize;
643 
644     for &i in numcodes.iter() {
645         let _ = m.write_all(&[i]);
646         sum += i as usize;
647     }
648 
649     assert_eq!(sum, values.len());
650 
651     for &i in values.iter() {
652         let _ = m.write_all(&[i]);
653     }
654 }
655 
build_quantization_segment(m: &mut Vec<u8>, precision: u8, identifier: u8, qtable: &[u8])656 fn build_quantization_segment(m: &mut Vec<u8>, precision: u8, identifier: u8, qtable: &[u8]) {
657     assert_eq!(qtable.len() % 64, 0);
658     m.clear();
659 
660     let p = if precision == 8 { 0 } else { 1 };
661 
662     let pqtq = (p << 4) | identifier;
663     let _ = m.write_all(&[pqtq]);
664 
665     for i in 0usize..64 {
666         let _ = m.write_all(&[qtable[UNZIGZAG[i] as usize]]);
667     }
668 }
669 
encode_coefficient(coefficient: i32) -> (u8, u16)670 fn encode_coefficient(coefficient: i32) -> (u8, u16) {
671     let mut magnitude = coefficient.abs() as u16;
672     let mut num_bits = 0u8;
673 
674     while magnitude > 0 {
675         magnitude >>= 1;
676         num_bits += 1;
677     }
678 
679     let mask = (1 << num_bits as usize) - 1;
680 
681     let val = if coefficient < 0 {
682         (coefficient - 1) as u16 & mask
683     } else {
684         coefficient as u16 & mask
685     };
686 
687     (num_bits, val)
688 }
689 
rgb_to_ycbcr(r: u8, g: u8, b: u8) -> (u8, u8, u8)690 fn rgb_to_ycbcr(r: u8, g: u8, b: u8) -> (u8, u8, u8) {
691     let r = f32::from(r);
692     let g = f32::from(g);
693     let b = f32::from(b);
694 
695     let y = 0.299f32 * r + 0.587f32 * g + 0.114f32 * b;
696     let cb = -0.1687f32 * r - 0.3313f32 * g + 0.5f32 * b + 128f32;
697     let cr = 0.5f32 * r - 0.4187f32 * g - 0.0813f32 * b + 128f32;
698 
699     (y as u8, cb as u8, cr as u8)
700 }
701 
value_at(s: &[u8], index: usize) -> u8702 fn value_at(s: &[u8], index: usize) -> u8 {
703     if index < s.len() {
704         s[index]
705     } else {
706         s[s.len() - 1]
707     }
708 }
709 
copy_blocks_ycbcr( source: &[u8], x0: usize, y0: usize, width: usize, bpp: usize, yb: &mut [u8; 64], cbb: &mut [u8; 64], crb: &mut [u8; 64], )710 fn copy_blocks_ycbcr(
711     source: &[u8],
712     x0: usize,
713     y0: usize,
714     width: usize,
715     bpp: usize,
716     yb: &mut [u8; 64],
717     cbb: &mut [u8; 64],
718     crb: &mut [u8; 64],
719 ) {
720     for y in 0usize..8 {
721         let ystride = (y0 + y) * bpp * width;
722 
723         for x in 0usize..8 {
724             let xstride = x0 * bpp + x * bpp;
725 
726             let r = value_at(source, ystride + xstride);
727             let g = value_at(source, ystride + xstride + 1);
728             let b = value_at(source, ystride + xstride + 2);
729 
730             let (yc, cb, cr) = rgb_to_ycbcr(r, g, b);
731 
732             yb[y * 8 + x] = yc;
733             cbb[y * 8 + x] = cb;
734             crb[y * 8 + x] = cr;
735         }
736     }
737 }
738 
copy_blocks_gray( source: &[u8], x0: usize, y0: usize, width: usize, bpp: usize, gb: &mut [u8; 64], )739 fn copy_blocks_gray(
740     source: &[u8],
741     x0: usize,
742     y0: usize,
743     width: usize,
744     bpp: usize,
745     gb: &mut [u8; 64],
746 ) {
747     for y in 0usize..8 {
748         let ystride = (y0 + y) * bpp * width;
749 
750         for x in 0usize..8 {
751             let xstride = x0 * bpp + x * bpp;
752             gb[y * 8 + x] = value_at(source, ystride + xstride);
753         }
754     }
755 }
756 
757 #[cfg(test)]
758 mod tests {
759     use super::super::JPEGDecoder;
760     use super::JPEGEncoder;
761     use color::ColorType;
762     use image::ImageDecoder;
763     use std::io::Cursor;
764 
765     #[test]
roundtrip_sanity_check()766     fn roundtrip_sanity_check() {
767         // create a 1x1 8-bit image buffer containing a single red pixel
768         let img = [255u8, 0, 0];
769 
770         // encode it into a memory buffer
771         let mut encoded_img = Vec::new();
772         {
773             let mut encoder = JPEGEncoder::new_with_quality(&mut encoded_img, 100);
774             encoder
775                 .encode(&img, 1, 1, ColorType::RGB(8))
776                 .expect("Could not encode image");
777         }
778 
779         // decode it from the memory buffer
780         {
781             let decoder = JPEGDecoder::new(Cursor::new(&encoded_img))
782                 .expect("Could not decode image");
783             let decoded = decoder.read_image().expect("Could not decode image");
784             // note that, even with the encode quality set to 100, we do not get the same image
785             // back. Therefore, we're going to assert that it's at least red-ish:
786             assert_eq!(3, decoded.len());
787             assert!(decoded[0] > 0x80);
788             assert!(decoded[1] < 0x80);
789             assert!(decoded[2] < 0x80);
790         }
791     }
792 
793     #[test]
grayscale_roundtrip_sanity_check()794     fn grayscale_roundtrip_sanity_check() {
795         // create a 2x2 8-bit image buffer containing a white diagonal
796         let img = [255u8, 0, 0, 255];
797 
798         // encode it into a memory buffer
799         let mut encoded_img = Vec::new();
800         {
801             let mut encoder = JPEGEncoder::new_with_quality(&mut encoded_img, 100);
802             encoder
803                 .encode(&img, 2, 2, ColorType::Gray(8))
804                 .expect("Could not encode image");
805         }
806 
807         // decode it from the memory buffer
808         {
809             let decoder = JPEGDecoder::new(Cursor::new(&encoded_img))
810                 .expect("Could not decode image");
811             let decoded = decoder.read_image().expect("Could not decode image");
812             // note that, even with the encode quality set to 100, we do not get the same image
813             // back. Therefore, we're going to assert that the diagonal is at least white-ish:
814             assert_eq!(4, decoded.len());
815             assert!(decoded[0] > 0x80);
816             assert!(decoded[1] < 0x80);
817             assert!(decoded[2] < 0x80);
818             assert!(decoded[3] > 0x80);
819         }
820     }
821 }
822