jpeg_encoder/
writer.rs

1use crate::EncodingError;
2use crate::encoder::Component;
3use crate::huffman::{CodingClass, HuffmanTable};
4use crate::marker::{Marker, SOFType};
5use crate::quantization::QuantizationTable;
6
7/// Represents the pixel density of an image
8///
9/// For example, a 300 DPI image is represented by:
10///
11/// ```rust
12/// # use jpeg_encoder::{PixelDensity, PixelDensityUnit};
13/// let hdpi = PixelDensity::dpi(300);
14/// assert_eq!(hdpi, PixelDensity {density: (300,300), unit: PixelDensityUnit::Inches})
15/// ```
16#[derive(Clone, Copy, Debug, Eq, PartialEq)]
17pub struct PixelDensity {
18    /// A couple of values for (Xdensity, Ydensity)
19    pub density: (u16, u16),
20    /// The unit in which the density is measured
21    pub unit: PixelDensityUnit,
22}
23
24impl PixelDensity {
25    /// Creates the most common pixel density type:
26    /// the horizontal and the vertical density are equal,
27    /// and measured in pixels per inch.
28    #[must_use]
29    pub fn dpi(density: u16) -> Self {
30        PixelDensity {
31            density: (density, density),
32            unit: PixelDensityUnit::Inches,
33        }
34    }
35}
36
37impl Default for PixelDensity {
38    /// Returns a pixel density with a pixel aspect ratio of 1
39    fn default() -> Self {
40        PixelDensity {
41            density: (1, 1),
42            unit: PixelDensityUnit::PixelAspectRatio,
43        }
44    }
45}
46
47/// Represents a unit in which the density of an image is measured
48#[derive(Clone, Copy, Debug, Eq, PartialEq)]
49pub enum PixelDensityUnit {
50    /// Represents the absence of a unit, the values indicate only a
51    /// [pixel aspect ratio](https://en.wikipedia.org/wiki/Pixel_aspect_ratio)
52    PixelAspectRatio,
53
54    /// Pixels per inch (2.54 cm)
55    Inches,
56
57    /// Pixels per centimeter
58    Centimeters,
59}
60
61/// Zig-zag sequence of quantized DCT coefficients
62///
63/// Figure A.6
64pub static ZIGZAG: [u8; 64] = [
65    0, 1, 8, 16, 9, 2, 3, 10, 17, 24, 32, 25, 18, 11, 4, 5, 12, 19, 26, 33, 40, 48, 41, 34, 27, 20,
66    13, 6, 7, 14, 21, 28, 35, 42, 49, 56, 57, 50, 43, 36, 29, 22, 15, 23, 30, 37, 44, 51, 58, 59,
67    52, 45, 38, 31, 39, 46, 53, 60, 61, 54, 47, 55, 62, 63,
68];
69
70const BUFFER_SIZE: usize = core::mem::size_of::<usize>() * 8;
71
72/// A no_std alternative for `std::io::Write`
73///
74/// An implementation of a subset of `std::io::Write` necessary to use the encoder without `std`.
75/// This trait is implemented for `std::io::Write` if the `std` feature is enabled.
76pub trait JfifWrite {
77    /// Writes the whole buffer. The behavior must be identical to std::io::Write::write_all
78    /// # Errors
79    ///
80    /// Return an error if the data can't be written
81    fn write_all(&mut self, buf: &[u8]) -> Result<(), EncodingError>;
82}
83
84#[cfg(not(feature = "std"))]
85impl<W: JfifWrite + ?Sized> JfifWrite for &mut W {
86    fn write_all(&mut self, buf: &[u8]) -> Result<(), EncodingError> {
87        (**self).write_all(buf)
88    }
89}
90
91#[cfg(not(feature = "std"))]
92impl JfifWrite for alloc::vec::Vec<u8> {
93    fn write_all(&mut self, buf: &[u8]) -> Result<(), EncodingError> {
94        self.extend_from_slice(buf);
95        Ok(())
96    }
97}
98
99#[cfg(feature = "std")]
100impl<W: std::io::Write + ?Sized> JfifWrite for W {
101    #[inline(always)]
102    fn write_all(&mut self, buf: &[u8]) -> Result<(), EncodingError> {
103        self.write_all(buf)?;
104        Ok(())
105    }
106}
107
108pub(crate) struct JfifWriter<W: JfifWrite> {
109    w: W,
110    bit_buffer: usize,
111    free_bits: i8,
112}
113
114impl<W: JfifWrite> JfifWriter<W> {
115    pub fn new(w: W) -> Self {
116        JfifWriter {
117            w,
118            bit_buffer: 0,
119            free_bits: BUFFER_SIZE as i8,
120        }
121    }
122
123    #[inline(always)]
124    pub fn write(&mut self, buf: &[u8]) -> Result<(), EncodingError> {
125        self.w.write_all(buf)
126    }
127
128    #[inline(always)]
129    pub fn write_u8(&mut self, value: u8) -> Result<(), EncodingError> {
130        self.w.write_all(&[value])
131    }
132
133    #[inline(always)]
134    pub fn write_u16(&mut self, value: u16) -> Result<(), EncodingError> {
135        self.w.write_all(&value.to_be_bytes())
136    }
137
138    pub fn finalize_bit_buffer(&mut self) -> Result<(), EncodingError> {
139        self.write_bits(0x7F, 7)?;
140        self.flush_bit_buffer()?;
141        self.bit_buffer = 0;
142        self.free_bits = BUFFER_SIZE as i8;
143
144        Ok(())
145    }
146
147    pub fn flush_bit_buffer(&mut self) -> Result<(), EncodingError> {
148        while self.free_bits <= (BUFFER_SIZE as i8 - 8) {
149            self.flush_byte_from_bit_buffer(self.free_bits)?;
150            self.free_bits += 8;
151        }
152
153        Ok(())
154    }
155
156    #[inline(always)]
157    fn flush_byte_from_bit_buffer(&mut self, free_bits: i8) -> Result<(), EncodingError> {
158        let value = (self.bit_buffer >> (BUFFER_SIZE as i8 - 8 - free_bits)) & 0xFF;
159
160        self.write_u8(value as u8)?;
161
162        if value == 0xFF {
163            self.write_u8(0x00)?;
164        }
165
166        Ok(())
167    }
168
169    #[inline(always)]
170    #[allow(overflowing_literals)]
171    fn write_bit_buffer(&mut self) -> Result<(), EncodingError> {
172        if (self.bit_buffer
173            & 0x8080808080808080
174            & !(self.bit_buffer.wrapping_add(0x0101010101010101)))
175            != 0
176        {
177            for i in 0..(BUFFER_SIZE / 8) {
178                self.flush_byte_from_bit_buffer((i * 8) as i8)?;
179            }
180            Ok(())
181        } else {
182            self.w.write_all(&self.bit_buffer.to_be_bytes())
183        }
184    }
185
186    pub fn write_bits(&mut self, value: u32, size: u8) -> Result<(), EncodingError> {
187        let size = size as i8;
188        let value = value as usize;
189
190        let free_bits = self.free_bits - size;
191
192        if free_bits < 0 {
193            self.bit_buffer = (self.bit_buffer << (size + free_bits)) | (value >> -free_bits);
194            self.write_bit_buffer()?;
195            self.bit_buffer = value;
196            self.free_bits = free_bits + BUFFER_SIZE as i8;
197        } else {
198            self.free_bits = free_bits;
199            self.bit_buffer = (self.bit_buffer << size) | value;
200        }
201        Ok(())
202    }
203
204    pub fn write_marker(&mut self, marker: Marker) -> Result<(), EncodingError> {
205        self.write(&[0xFF, marker.into()])
206    }
207
208    pub fn write_segment(&mut self, marker: Marker, data: &[u8]) -> Result<(), EncodingError> {
209        self.write_marker(marker)?;
210        self.write_u16(data.len() as u16 + 2)?;
211        self.write(data)?;
212
213        Ok(())
214    }
215
216    pub fn write_header(&mut self, density: &PixelDensity) -> Result<(), EncodingError> {
217        self.write_marker(Marker::APP(0))?;
218        self.write_u16(16)?;
219
220        self.write(b"JFIF\0")?;
221        self.write(&[0x01, 0x02])?;
222
223        match density.unit {
224            PixelDensityUnit::PixelAspectRatio => {
225                self.write_u8(0x00)?;
226            }
227            PixelDensityUnit::Inches => {
228                self.write_u8(0x01)?;
229            }
230            PixelDensityUnit::Centimeters => {
231                self.write_u8(0x02)?;
232            }
233        }
234        let (x, y) = density.density;
235        self.write_u16(x)?;
236        self.write_u16(y)?;
237
238        self.write(&[0x00, 0x00])
239    }
240
241    /// Append huffman table segment
242    ///
243    /// - `class`: 0 for DC or 1 for AC
244    /// - `dest`: 0 for luma or 1 for chroma tables
245    ///
246    /// Layout:
247    /// ```txt
248    /// |--------|---------------|--------------------------|--------------------|--------|
249    /// | 0xFFC4 | 16 bit length | 4 bit class / 4 bit dest |  16 byte num codes | values |
250    /// |--------|---------------|--------------------------|--------------------|--------|
251    /// ```
252    ///
253    pub fn write_huffman_segment(
254        &mut self,
255        class: CodingClass,
256        destination: u8,
257        table: &HuffmanTable,
258    ) -> Result<(), EncodingError> {
259        assert!(destination < 4, "Bad destination: {}", destination);
260
261        self.write_marker(Marker::DHT)?;
262        self.write_u16(2 + 1 + 16 + table.values().len() as u16)?;
263
264        self.write_u8(((class as u8) << 4) | destination)?;
265        self.write(table.length())?;
266        self.write(table.values())?;
267
268        Ok(())
269    }
270
271    /// Append a quantization table
272    ///
273    /// - `precision`: 0 which means 1 byte per value.
274    /// - `dest`: 0 for luma or 1 for chroma tables
275    ///
276    /// Layout:
277    /// ```txt
278    /// |--------|---------------|------------------------------|--------|--------|-----|--------|
279    /// | 0xFFDB | 16 bit length | 4 bit precision / 4 bit dest | V(0,0) | V(0,1) | ... | V(7,7) |
280    /// |--------|---------------|------------------------------|--------|--------|-----|--------|
281    /// ```
282    ///
283    pub fn write_quantization_segment(
284        &mut self,
285        destination: u8,
286        table: &QuantizationTable,
287    ) -> Result<(), EncodingError> {
288        assert!(destination < 4, "Bad destination: {}", destination);
289
290        self.write_marker(Marker::DQT)?;
291        self.write_u16(2 + 1 + 64)?;
292
293        self.write_u8(destination)?;
294
295        for &v in ZIGZAG.iter() {
296            self.write_u8(table.get(v as usize))?;
297        }
298
299        Ok(())
300    }
301
302    pub fn write_dri(&mut self, restart_interval: u16) -> Result<(), EncodingError> {
303        self.write_marker(Marker::DRI)?;
304        self.write_u16(4)?;
305        self.write_u16(restart_interval)
306    }
307
308    #[inline]
309    pub fn huffman_encode(&mut self, val: u8, table: &HuffmanTable) -> Result<(), EncodingError> {
310        let &(size, code) = table.get_for_value(val);
311        self.write_bits(code as u32, size)
312    }
313
314    #[inline]
315    pub fn huffman_encode_value(
316        &mut self,
317        size: u8,
318        symbol: u8,
319        value: u16,
320        table: &HuffmanTable,
321    ) -> Result<(), EncodingError> {
322        let &(num_bits, code) = table.get_for_value(symbol);
323
324        let mut temp = value as u32;
325        temp |= (code as u32) << size;
326        let size = size + num_bits;
327
328        self.write_bits(temp, size)
329    }
330
331    pub fn write_block(
332        &mut self,
333        block: &[i16; 64],
334        prev_dc: i16,
335        dc_table: &HuffmanTable,
336        ac_table: &HuffmanTable,
337    ) -> Result<(), EncodingError> {
338        self.write_dc(block[0], prev_dc, dc_table)?;
339        self.write_ac_block(block, 1, 64, ac_table)
340    }
341
342    pub fn write_dc(
343        &mut self,
344        value: i16,
345        prev_dc: i16,
346        dc_table: &HuffmanTable,
347    ) -> Result<(), EncodingError> {
348        let diff = value - prev_dc;
349        let (size, value) = get_code(diff);
350
351        self.huffman_encode_value(size, size, value, dc_table)?;
352
353        Ok(())
354    }
355
356    pub fn write_ac_block(
357        &mut self,
358        block: &[i16; 64],
359        start: usize,
360        end: usize,
361        ac_table: &HuffmanTable,
362    ) -> Result<(), EncodingError> {
363        let mut zero_run = 0;
364
365        for &value in &block[start..end] {
366            if value == 0 {
367                zero_run += 1;
368            } else {
369                while zero_run > 15 {
370                    self.huffman_encode(0xF0, ac_table)?;
371                    zero_run -= 16;
372                }
373
374                let (size, value) = get_code(value);
375                let symbol = (zero_run << 4) | size;
376
377                self.huffman_encode_value(size, symbol, value, ac_table)?;
378
379                zero_run = 0;
380            }
381        }
382
383        if zero_run > 0 {
384            self.huffman_encode(0x00, ac_table)?;
385        }
386
387        Ok(())
388    }
389
390    pub fn write_frame_header(
391        &mut self,
392        width: u16,
393        height: u16,
394        components: &[Component],
395        progressive: bool,
396    ) -> Result<(), EncodingError> {
397        if progressive {
398            self.write_marker(Marker::SOF(SOFType::ProgressiveDCT))?;
399        } else {
400            self.write_marker(Marker::SOF(SOFType::BaselineDCT))?;
401        }
402
403        self.write_u16(2 + 1 + 2 + 2 + 1 + (components.len() as u16) * 3)?;
404
405        // Precision
406        self.write_u8(8)?;
407
408        self.write_u16(height)?;
409        self.write_u16(width)?;
410
411        self.write_u8(components.len() as u8)?;
412
413        for component in components.iter() {
414            self.write_u8(component.id)?;
415            self.write_u8(
416                (component.horizontal_sampling_factor << 4) | component.vertical_sampling_factor,
417            )?;
418            self.write_u8(component.quantization_table)?;
419        }
420
421        Ok(())
422    }
423
424    pub fn write_scan_header(
425        &mut self,
426        components: &[&Component],
427        spectral: Option<(u8, u8)>,
428    ) -> Result<(), EncodingError> {
429        self.write_marker(Marker::SOS)?;
430
431        self.write_u16(2 + 1 + (components.len() as u16) * 2 + 3)?;
432
433        self.write_u8(components.len() as u8)?;
434
435        for component in components.iter() {
436            self.write_u8(component.id)?;
437            self.write_u8((component.dc_huffman_table << 4) | component.ac_huffman_table)?;
438        }
439
440        let (spectral_start, spectral_end) = spectral.unwrap_or((0, 63));
441
442        // Start of spectral or predictor selection
443        self.write_u8(spectral_start)?;
444
445        // End of spectral selection
446        self.write_u8(spectral_end)?;
447
448        // Successive approximation bit position high and low
449        self.write_u8(0)?;
450
451        Ok(())
452    }
453}
454
455#[inline]
456pub(crate) fn get_code(value: i16) -> (u8, u16) {
457    let temp = value - (value.is_negative() as i16);
458    let temp2 = value.abs();
459
460    /*
461     * Doing this instead of 16 - temp2.leading_zeros()
462     * Gives the compiler the information that leadings_zeros
463     * is always called on a non zero value, which removes a branch on x86
464     */
465    let num_bits = 15 - (temp2 << 1 | 1).leading_zeros() as u16;
466
467    let coefficient = temp & ((1 << num_bits as usize) - 1);
468
469    (num_bits as u8, coefficient as u16)
470}
jpeg_encoder/writer.rs

jpeg_encoder/
writer.rs
