Improve API for IDCT scaling

Author: kevinmehall

src/decoder.rs

@@ -50,7 +50,7 @@ pub struct ImageInfo {
 /// JPEG decoder
 pub struct Decoder<R> {
     reader: R,
-    scale: usize,
+    requested_size: Option<Dimensions>,
 
     frame: Option<FrameInfo>,
     dc_huffman_tables: Vec<Option<HuffmanTable>>,
@@ -71,15 +71,26 @@ pub struct Decoder<R> {
 impl<R: Read> Decoder<R> {
     /// Creates a new `Decoder` using the reader `reader`.
     pub fn new(reader: R) -> Decoder<R> {
-        Decoder::new_scaled(reader, 8)
+        Decoder::init(reader, None)
     }
 
-    /// Creates a new `Decoder` using the reader `reader` that scales
-    /// down the image by a factor of `scale/8`.
-    pub fn new_scaled(reader: R, scale: usize) -> Decoder<R> {
+    /// Creates a new `Decoder` using the reader `reader` that returns a
+    /// scaled image that is equal to or larger than the requested size in at
+    /// least one axis, or the full size of the image if the requested size is
+    /// larger.
+    /// 
+    /// This efficiently scales down the image by one of a fixed set of
+    /// available ratios during decoding. To generate a thumbnail of an
+    /// exact size, pass the desired size or larger and then scale to the
+    /// final size using a traditional resampling algorithm.
+    pub fn scaled(reader: R, requested_width: u16, requested_height: u16) -> Decoder<R> {
+        Decoder::init(reader, Some(Dimensions{ width: requested_width, height: requested_height }))
+    }
+    
+    fn init(reader: R, requested_size: Option<Dimensions>) -> Decoder<R> {
         Decoder {
             reader: reader,
-            scale: scale,
+            requested_size,
             frame: None,
             dc_huffman_tables: vec![None, None, None, None],
             ac_huffman_tables: vec![None, None, None, None],
@@ -161,7 +172,7 @@ impl<R: Read> Decoder<R> {
                         return Err(Error::Unsupported(UnsupportedFeature::Hierarchical));
                     }
 
-                    let frame = parse_sof(&mut self.reader, marker, self.scale)?;
+                    let frame = parse_sof(&mut self.reader, marker, self.requested_size)?;
                     let component_count = frame.components.len();
 
                     if frame.is_differential {

src/idct.rs

@@ -1,9 +1,51 @@
 // Malicious JPEG files can cause operations in the idct to overflow.
 // One example is tests/crashtest/images/imagetestsuite/b0b8914cc5f7a6eff409f16d8cc236c5.jpg
 // That's why wrapping operators are needed.
+use crate::parser::Dimensions;
+
+pub fn choose_idct_size(full_size: Dimensions, requested_size: Dimensions) -> usize {
+    fn scaled(len: u16, scale: usize) -> u16 { ((len as u32 * scale as u32 - 1) / 8 + 1) as u16 }
+
+    for &scale in &[1, 2, 4] {
+        if scaled(full_size.width, scale) >= requested_size.width || scaled(full_size.height, scale) >= requested_size.height {
+            return scale;
+        }
+    }
+
+    return 8;
+}
+
+#[test]
+fn test_choose_idct_size() {
+    assert_eq!(choose_idct_size(Dimensions{width: 5472, height: 3648}, Dimensions{width: 200, height: 200}), 1);
+    assert_eq!(choose_idct_size(Dimensions{width: 5472, height: 3648}, Dimensions{width: 500, height: 500}), 1);
+    assert_eq!(choose_idct_size(Dimensions{width: 5472, height: 3648}, Dimensions{width: 684, height: 456}), 1);
+    assert_eq!(choose_idct_size(Dimensions{width: 5472, height: 3648}, Dimensions{width: 999, height: 456}), 1);
+    assert_eq!(choose_idct_size(Dimensions{width: 5472, height: 3648}, Dimensions{width: 684, height: 999}), 1);
+    assert_eq!(choose_idct_size(Dimensions{width: 500, height: 333}, Dimensions{width: 63, height: 42}), 1);
+
+    assert_eq!(choose_idct_size(Dimensions{width: 5472, height: 3648}, Dimensions{width: 685, height: 999}), 2);
+    assert_eq!(choose_idct_size(Dimensions{width: 5472, height: 3648}, Dimensions{width: 1000, height: 1000}), 2);
+    assert_eq!(choose_idct_size(Dimensions{width: 5472, height: 3648}, Dimensions{width: 1400, height: 1400}), 4);
+    
+    assert_eq!(choose_idct_size(Dimensions{width: 5472, height: 3648}, Dimensions{width: 5472, height: 3648}), 8);
+    assert_eq!(choose_idct_size(Dimensions{width: 5472, height: 3648}, Dimensions{width: 16384, height: 16384}), 8);
+    assert_eq!(choose_idct_size(Dimensions{width: 1, height: 1}, Dimensions{width: 65535, height: 65535}), 8);
+    assert_eq!(choose_idct_size(Dimensions{width: 5472, height: 3648}, Dimensions{width: 16384, height: 16384}), 8);
+}
+
+pub fn dequantize_and_idct_block(scale: usize, coefficients: &[i16], quantization_table: &[u16; 64], output_linestride: usize, output: &mut [u8]) {
+    match scale {
+        8 => dequantize_and_idct_block_8x8(coefficients, quantization_table, output_linestride, output),
+        4 => dequantize_and_idct_block_4x4(coefficients, quantization_table, output_linestride, output),
+        2 => dequantize_and_idct_block_2x2(coefficients, quantization_table, output_linestride, output),
+        1 => dequantize_and_idct_block_1x1(coefficients, quantization_table, output_linestride, output),
+        _ => panic!("Unsupported IDCT scale {}/8", scale),
+    }
+}
 
 // This is based on stb_image's 'stbi__idct_block'.
-pub fn dequantize_and_idct_block_8x8(coefficients: &[i16], quantization_table: &[u16; 64], output_linestride: usize, output: &mut [u8]) {
+fn dequantize_and_idct_block_8x8(coefficients: &[i16], quantization_table: &[u16; 64], output_linestride: usize, output: &mut [u8]) {
     debug_assert_eq!(coefficients.len(), 64);
 
     let mut temp = [0i32; 64];
@@ -157,7 +199,7 @@ pub fn dequantize_and_idct_block_8x8(coefficients: &[i16], quantization_table: &
 
 // 4x4 and 2x2 IDCT based on Rakesh Dugad and Narendra Ahuja: "A Fast Scheme for Image Size Change in the Compressed Domain" (2001).
 // http://sylvana.net/jpegcrop/jidctred/
-pub fn dequantize_and_idct_block_4x4(coefficients: &[i16], quantization_table: &[u16; 64], output_linestride: usize, output: &mut [u8]) {
+fn dequantize_and_idct_block_4x4(coefficients: &[i16], quantization_table: &[u16; 64], output_linestride: usize, output: &mut [u8]) {
     debug_assert_eq!(coefficients.len(), 64);
     let mut temp = [0i32; 4*4];
 
@@ -214,7 +256,7 @@ pub fn dequantize_and_idct_block_4x4(coefficients: &[i16], quantization_table: &
     }
 }
 
-pub fn dequantize_and_idct_block_2x2(coefficients: &[i16], quantization_table: &[u16; 64], output_linestride: usize, output: &mut [u8]) {
+fn dequantize_and_idct_block_2x2(coefficients: &[i16], quantization_table: &[u16; 64], output_linestride: usize, output: &mut [u8]) {
     debug_assert_eq!(coefficients.len(), 64);
 
     const SCALE_BITS: u32 = 3;
@@ -245,7 +287,7 @@ pub fn dequantize_and_idct_block_2x2(coefficients: &[i16], quantization_table: &
     output[output_linestride + 1] = stbi_clamp(x2.wrapping_sub(x3).wrapping_shr(SCALE_BITS));
 }
 
-pub fn dequantize_and_idct_block_1x1(coefficients: &[i16], quantization_table: &[u16; 64], _output_linestride: usize, output: &mut [u8]) {
+fn dequantize_and_idct_block_1x1(coefficients: &[i16], quantization_table: &[u16; 64], _output_linestride: usize, output: &mut [u8]) {
     debug_assert_eq!(coefficients.len(), 64);
 
     let s0 = (coefficients[0] as i32 * quantization_table[0] as i32).wrapping_add(128 * 8) / 8;

src/parser.rs

@@ -107,7 +107,7 @@ fn skip_bytes<R: Read>(reader: &mut R, length: usize) -> Result<()> {
 }
 
 // Section B.2.2
-pub fn parse_sof<R: Read>(reader: &mut R, marker: Marker, scale: usize) -> Result<FrameInfo> {
+pub fn parse_sof<R: Read>(reader: &mut R, marker: Marker, requested_size: Option<Dimensions>) -> Result<FrameInfo> {
     let length = read_length(reader, marker)?;
 
     if length <= 6 {
@@ -159,6 +159,10 @@ pub fn parse_sof<R: Read>(reader: &mut R, marker: Marker, scale: usize) -> Resul
         return Err(Error::Format("zero width in frame header".to_owned()));
     }
 
+    let scale = if let Some(req) = requested_size {
+        crate::idct::choose_idct_size(Dimensions { width, height }, req)
+    } else { 8 };
+
     let component_count = reader.read_u8()?;
 
     if component_count == 0 {

src/worker/immediate.rs

@@ -1,6 +1,6 @@
 use decoder::MAX_COMPONENTS;
 use error::Result;
-use idct::{ dequantize_and_idct_block_8x8, dequantize_and_idct_block_4x4, dequantize_and_idct_block_2x2, dequantize_and_idct_block_1x1 };
+use idct::dequantize_and_idct_block;
 use std::mem;
 use std::sync::Arc;
 use parser::Component;
@@ -47,13 +47,7 @@ impl ImmediateWorker {
             let coefficients = &data[i * 64..(i + 1) * 64];
             let output = &mut self.results[index][self.offsets[index] + y * line_stride + x..];
 
-            match component.dct_scale {
-                8 => dequantize_and_idct_block_8x8(coefficients, quantization_table, line_stride, output),
-                4 => dequantize_and_idct_block_4x4(coefficients, quantization_table, line_stride, output),
-                2 => dequantize_and_idct_block_2x2(coefficients, quantization_table, line_stride, output),
-                1 => dequantize_and_idct_block_1x1(coefficients, quantization_table, line_stride, output),
-                _ => unimplemented!(),
-            }
+            dequantize_and_idct_block(component.dct_scale, coefficients, quantization_table, line_stride, output);
         }
 
         self.offsets[index] += block_count * component.dct_scale * component.dct_scale;