fix bug where decoding was failing due to 16 bit math in C++ version (#…

…38)

* Revert "Revert "fix bug where decoding was failing due to 16 bit math in C++ version""

This reverts commit d9bfc0f.

* added unit tests

* added tests

src/enabled_features.rs

@@ -11,6 +11,9 @@ pub struct EnabledFeatures {
 
     // maximum jpeg height
     pub max_jpeg_height: i32,
+
+    // Sadly C++ version has a bug where it uses 16 bit math in the SIMD path and 32 bit math in the scalar path
+    pub use_16bit_dc_estimate: bool,
 }
 
 impl Default for EnabledFeatures {
@@ -20,18 +23,21 @@ impl Default for EnabledFeatures {
             reject_dqts_with_zeros: true,
             max_jpeg_width: 16386,
             max_jpeg_height: 16386,
+            use_16bit_dc_estimate: false,
         }
     }
 }
 
 impl EnabledFeatures {
     /// parameters that allow everything
+    #[allow(dead_code)]
     pub fn all() -> Self {
         Self {
             progressive: true,
             reject_dqts_with_zeros: true,
             max_jpeg_height: i32::MAX,
             max_jpeg_width: i32::MAX,
+            use_16bit_dc_estimate: false,
         }
     }
 }

src/lib.rs

@@ -134,32 +134,44 @@ pub unsafe extern "C" fn WrapperDecompressImage(
     result_size: *mut u64,
 ) -> i32 {
     match catch_unwind(|| {
-        let input = std::slice::from_raw_parts(input_buffer, input_buffer_size as usize);
-
-        let output = std::slice::from_raw_parts_mut(output_buffer, output_buffer_size as usize);
-
-        let mut reader = Cursor::new(input);
-        let mut writer = Cursor::new(output);
-
         // For back-compat with C++ version we allow decompression of images with zeros in DQT tables
         let mut enabled_features = EnabledFeatures::all();
         enabled_features.reject_dqts_with_zeros = false;
 
-        match decode_lepton_wrapper(
-            &mut reader,
-            &mut writer,
-            number_of_threads as usize,
-            &enabled_features,
-        ) {
-            Ok(_) => {}
-            Err(e) => {
-                return translate_error(e).exit_code as i32;
+        loop {
+            let input = std::slice::from_raw_parts(input_buffer, input_buffer_size as usize);
+            let output = std::slice::from_raw_parts_mut(output_buffer, output_buffer_size as usize);
+
+            let mut reader = Cursor::new(input);
+            let mut writer = Cursor::new(output);
+
+            match decode_lepton_wrapper(
+                &mut reader,
+                &mut writer,
+                number_of_threads as usize,
+                &enabled_features,
+            ) {
+                Ok(_) => {
+                    *result_size = writer.position().into();
+                    return 0;
+                }
+                Err(e) => {
+                    let exit_code = translate_error(e).exit_code;
+
+                    // there's a bug in the C++ version where it uses 16 bit math in the SIMD path and 32 bit math in the scalar path depending on the compiler options.
+                    // unfortunately there's no way to tell ahead of time other than the fact that the image will be decoded with an error.
+                    if exit_code == ExitCode::StreamInconsistent
+                        && !enabled_features.use_16bit_dc_estimate
+                    {
+                        // try again with the flag set
+                        enabled_features.use_16bit_dc_estimate = true;
+                        continue;
+                    }
+
+                    return exit_code as i32;
+                }
             }
         }
-
-        *result_size = writer.position().into();
-
-        return 0;
     }) {
         Ok(code) => {
             return code;

src/main.rs

@@ -111,7 +111,12 @@ fn main_with_result() -> anyhow::Result<()> {
             let _metrics;
 
             (block_image, _metrics) = lh
-                .decode_as_single_image(&mut reader, filelen, num_threads as usize)
+                .decode_as_single_image(
+                    &mut reader,
+                    filelen,
+                    num_threads as usize,
+                    &enabled_features,
+                )
                 .context(here!())?;
 
             loop {

src/structs/lepton_decoder.rs

@@ -12,6 +12,7 @@ use std::cmp;
 use std::io::Read;
 
 use crate::consts::{LOG_TABLE_256, RASTER_TO_ALIGNED, UNZIGZAG_49};
+use crate::enabled_features::EnabledFeatures;
 use crate::helpers::{err_exit_code, here, u16_bit_length};
 use crate::lepton_error::ExitCode;
 
@@ -38,6 +39,7 @@ pub fn lepton_decode_row_range<R: Read>(
     max_y: i32,
     is_last_thread: bool,
     full_file_compression: bool,
+    features: &EnabledFeatures,
 ) -> Result<Metrics> {
     let component_size_in_blocks = trunc.get_component_sizes_in_blocks();
     let max_coded_heights = trunc.get_max_coded_heights();
@@ -98,6 +100,7 @@ pub fn lepton_decode_row_range<R: Read>(
             &component_size_in_blocks[..],
             cur_row.component,
             cur_row.curr_y,
+            features,
         )
         .context(here!())?;
     }
@@ -116,6 +119,7 @@ fn decode_row_wrapper<R: Read>(
     component_size_in_blocks: &[i32],
     component: usize,
     curr_y: i32,
+    features: &EnabledFeatures,
 ) -> Result<()> {
     let mut context = image_data.off_y(curr_y);
 
@@ -133,6 +137,7 @@ fn decode_row_wrapper<R: Read>(
             &mut context,
             num_non_zeros,
             component_size_in_blocks[component],
+            features,
         )
         .context(here!())?;
     } else if block_width > 1 {
@@ -148,6 +153,7 @@ fn decode_row_wrapper<R: Read>(
             &mut context,
             num_non_zeros,
             component_size_in_blocks[component],
+            features,
         )
         .context(here!())?;
     } else {
@@ -163,6 +169,7 @@ fn decode_row_wrapper<R: Read>(
             &mut context,
             num_non_zeros,
             component_size_in_blocks[component],
+            features,
         )
         .context(here!())?;
     }
@@ -181,6 +188,7 @@ fn decode_row<R: Read>(
     block_context: &mut BlockContext,
     num_non_zeros: &mut [NeighborSummary],
     component_size_in_blocks: i32,
+    features: &EnabledFeatures,
 ) -> Result<()> {
     let block_width = image_data.get_block_width();
     if block_width > 0 {
@@ -192,6 +200,7 @@ fn decode_row<R: Read>(
             num_non_zeros,
             qt,
             left_model,
+            features,
         )
         .context(here!())?;
         let offset = block_context.next(true);
@@ -211,6 +220,7 @@ fn decode_row<R: Read>(
                 num_non_zeros,
                 qt,
                 middle_model,
+                features,
             )
             .context(here!())?;
         } else {
@@ -222,6 +232,7 @@ fn decode_row<R: Read>(
                 num_non_zeros,
                 qt,
                 middle_model,
+                features,
             )
             .context(here!())?;
         }
@@ -243,6 +254,7 @@ fn decode_row<R: Read>(
                 num_non_zeros,
                 qt,
                 right_model,
+                features,
             )
             .context(here!())?;
         } else {
@@ -254,6 +266,7 @@ fn decode_row<R: Read>(
                 num_non_zeros,
                 qt,
                 right_model,
+                features,
             )
             .context(here!())?;
         }
@@ -272,6 +285,7 @@ fn parse_token<R: Read, const ALL_PRESENT: bool>(
     num_non_zeros: &mut [NeighborSummary],
     qt: &QuantizationTables,
     pt: &ProbabilityTables,
+    features: &EnabledFeatures,
 ) -> Result<()> {
     debug_assert!(pt.is_all_present() == ALL_PRESENT);
 
@@ -370,12 +384,14 @@ fn parse_token<R: Read, const ALL_PRESENT: bool>(
         &predicted_dc.advanced_predict_dc_pixels_sans_dc,
         qt.get_quantization_table(),
         output.get_dc(),
+        features,
     );
 
     here.set_vertical(
         &predicted_dc.advanced_predict_dc_pixels_sans_dc,
         qt.get_quantization_table(),
         output.get_dc(),
+        features,
     );
 
     image_data.append_block(output);

src/structs/lepton_encoder.rs

@@ -10,6 +10,7 @@ use std::cmp;
 use std::io::Write;
 
 use crate::consts::*;
+use crate::enabled_features::EnabledFeatures;
 use crate::helpers::*;
 use crate::lepton_error::ExitCode;
 
@@ -36,6 +37,7 @@ pub fn lepton_encode_row_range<W: Write>(
     max_y: i32,
     is_last_thread: bool,
     full_file_compression: bool,
+    features: &EnabledFeatures,
 ) -> Result<Metrics> {
     let mut model = Model::default_boxed();
     let mut bool_writer = VPXBoolWriter::new(writer)?;
@@ -106,6 +108,7 @@ pub fn lepton_encode_row_range<W: Write>(
                 &mut num_non_zeros[bt][..],
                 block_width,
                 component_size_in_blocks[bt],
+                features,
             )
             .context(here!())?;
         } else if block_width > 1 {
@@ -122,6 +125,7 @@ pub fn lepton_encode_row_range<W: Write>(
                 &mut num_non_zeros[bt][..],
                 block_width,
                 component_size_in_blocks[bt],
+                features,
             )
             .context(here!())?;
         } else {
@@ -139,6 +143,7 @@ pub fn lepton_encode_row_range<W: Write>(
                 &mut num_non_zeros[bt][..],
                 block_width,
                 component_size_in_blocks[bt],
+                features,
             )
             .context(here!())?;
         }
@@ -182,6 +187,7 @@ fn process_row<W: Write>(
     num_non_zeros: &mut [NeighborSummary],
     block_width: i32,
     component_size_in_block: i32,
+    features: &EnabledFeatures,
 ) -> Result<()> {
     if block_width > 0 {
         state
@@ -196,6 +202,7 @@ fn process_row<W: Write>(
             image_data,
             num_non_zeros,
             bool_writer,
+            features,
         )
         .context(here!())?;
         let offset = state.next(true);
@@ -220,6 +227,7 @@ fn process_row<W: Write>(
                 image_data,
                 num_non_zeros,
                 bool_writer,
+                features,
             )
             .context(here!())?;
         } else {
@@ -231,6 +239,7 @@ fn process_row<W: Write>(
                 image_data,
                 num_non_zeros,
                 bool_writer,
+                features,
             )
             .context(here!())?;
         }
@@ -256,6 +265,7 @@ fn process_row<W: Write>(
                 image_data,
                 num_non_zeros,
                 bool_writer,
+                features,
             )
             .context(here!())?;
         } else {
@@ -267,6 +277,7 @@ fn process_row<W: Write>(
                 image_data,
                 num_non_zeros,
                 bool_writer,
+                features,
             )
             .context(here!())?;
         }
@@ -285,6 +296,7 @@ fn serialize_tokens<W: Write, const ALL_PRESENT: bool>(
     image_data: &BlockBasedImage,
     num_non_zeros: &mut [NeighborSummary],
     bool_writer: &mut VPXBoolWriter<W>,
+    features: &EnabledFeatures,
 ) -> Result<()> {
     debug_assert!(ALL_PRESENT == pt.is_all_present());
 
@@ -401,12 +413,14 @@ fn serialize_tokens<W: Write, const ALL_PRESENT: bool>(
         &predicted_val.advanced_predict_dc_pixels_sans_dc,
         qt.get_quantization_table(),
         block.get_dc(),
+        features,
     );
 
     here.set_vertical(
         &predicted_val.advanced_predict_dc_pixels_sans_dc,
         qt.get_quantization_table(),
         block.get_dc(),
+        features,
     );
 
     Ok(())