957 lines
40 KiB
Rust
957 lines
40 KiB
Rust
/*
|
|
* Copyright (c) 2023.
|
|
*
|
|
* This software is free software;
|
|
*
|
|
* You can redistribute it or modify it under terms of the MIT, Apache License or Zlib license
|
|
*/
|
|
|
|
use alloc::vec::Vec;
|
|
use alloc::{format, vec};
|
|
use core::cmp::min;
|
|
|
|
use makepad_zune_core::bytestream::ZByteReaderTrait;
|
|
use makepad_zune_core::colorspace::ColorSpace;
|
|
use makepad_zune_core::colorspace::ColorSpace::Luma;
|
|
use makepad_zune_core::log::{error, trace, warn};
|
|
|
|
use crate::bitstream::BitStream;
|
|
use crate::components::SampleRatios;
|
|
use crate::decoder::MAX_COMPONENTS;
|
|
use crate::errors::DecodeErrors;
|
|
use crate::marker::Marker;
|
|
use crate::mcu_prog::get_marker;
|
|
use crate::misc::{calculate_padded_width, setup_component_params};
|
|
use crate::worker::{color_convert, upsample};
|
|
use crate::JpegDecoder;
|
|
|
|
/// The size of a DC block for a MCU.
|
|
|
|
pub const DCT_BLOCK: usize = 64;
|
|
|
|
impl<T: ZByteReaderTrait> JpegDecoder<T> {
|
|
/// Check for existence of DC and AC Huffman Tables
|
|
pub(crate) fn check_tables(&self) -> Result<(), DecodeErrors> {
|
|
// check that dc and AC tables exist outside the hot path
|
|
for component in &self.components {
|
|
let _ = &self
|
|
.dc_huffman_tables
|
|
.get(component.dc_huff_table)
|
|
.as_ref()
|
|
.ok_or_else(|| {
|
|
DecodeErrors::HuffmanDecode(format!(
|
|
"No Huffman DC table for component {:?} ",
|
|
component.component_id
|
|
))
|
|
})?
|
|
.as_ref()
|
|
.ok_or_else(|| {
|
|
DecodeErrors::HuffmanDecode(format!(
|
|
"No DC table for component {:?}",
|
|
component.component_id
|
|
))
|
|
})?;
|
|
|
|
let _ = &self
|
|
.ac_huffman_tables
|
|
.get(component.ac_huff_table)
|
|
.as_ref()
|
|
.ok_or_else(|| {
|
|
DecodeErrors::HuffmanDecode(format!(
|
|
"No Huffman AC table for component {:?} ",
|
|
component.component_id
|
|
))
|
|
})?
|
|
.as_ref()
|
|
.ok_or_else(|| {
|
|
DecodeErrors::HuffmanDecode(format!(
|
|
"No AC table for component {:?}",
|
|
component.component_id
|
|
))
|
|
})?;
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
/// Decode MCUs and carry out post processing.
|
|
///
|
|
/// This is the main decoder loop for the library, the hot path.
|
|
///
|
|
/// Because of this, we pull in some very crazy optimization tricks hence readability is a pinch
|
|
/// here.
|
|
#[allow(
|
|
clippy::similar_names,
|
|
clippy::too_many_lines,
|
|
clippy::cast_possible_truncation
|
|
)]
|
|
#[inline(never)]
|
|
pub(crate) fn decode_mcu_ycbcr_baseline(
|
|
&mut self, pixels: &mut [u8]
|
|
) -> Result<(), DecodeErrors> {
|
|
setup_component_params(self)?;
|
|
|
|
// check dc and AC tables
|
|
self.check_tables()?;
|
|
|
|
let (mut mcu_width, mut mcu_height);
|
|
|
|
if self.is_interleaved {
|
|
// set upsampling functions
|
|
self.set_upsampling()?;
|
|
|
|
mcu_width = self.mcu_x;
|
|
mcu_height = self.mcu_y;
|
|
} else {
|
|
// For non-interleaved images( (1*1) subsampling)
|
|
// number of MCU's are the widths (+7 to account for paddings) divided bu 8.
|
|
mcu_width = (self.info.width as usize + 7) / 8;
|
|
mcu_height = (self.info.height as usize + 7) / 8;
|
|
}
|
|
if self.is_interleaved
|
|
&& self.input_colorspace.num_components() > 1
|
|
&& self.options.jpeg_get_out_colorspace().num_components() == 1
|
|
&& (self.info.sample_ratio == SampleRatios::V
|
|
|| self.info.sample_ratio == SampleRatios::HV)
|
|
{
|
|
// For a specific set of images, e.g interleaved,
|
|
// when converting from YcbCr to grayscale, we need to
|
|
// take into account mcu height since the MCU decoding needs to take
|
|
// it into account for padding purposes and the post processor
|
|
// parses two rows per mcu width.
|
|
//
|
|
// set coeff to be 2 to ensure that we increment two rows
|
|
// for every mcu processed also
|
|
mcu_height *= self.v_max;
|
|
mcu_height /= self.h_max;
|
|
self.coeff = 2;
|
|
}
|
|
|
|
if self.input_colorspace == ColorSpace::Luma && self.is_interleaved {
|
|
warn!("Grayscale image with down-sampled component, resetting component details");
|
|
|
|
self.reset_params();
|
|
|
|
mcu_width = ((self.info.width + 7) / 8) as usize;
|
|
mcu_height = ((self.info.height + 7) / 8) as usize;
|
|
}
|
|
let width = usize::from(self.info.width);
|
|
|
|
let padded_width = calculate_padded_width(width, self.info.sample_ratio);
|
|
|
|
let mut stream = BitStream::new();
|
|
let mut tmp = [0_i32; DCT_BLOCK];
|
|
|
|
let comp_len = self.components.len();
|
|
|
|
for (pos, comp) in self.components.iter_mut().enumerate() {
|
|
// Allocate only needed components.
|
|
//
|
|
// For special colorspaces i.e YCCK and CMYK, just allocate all of the needed
|
|
// components.
|
|
if min(
|
|
self.options.jpeg_get_out_colorspace().num_components() - 1,
|
|
pos
|
|
) == pos
|
|
|| comp_len == 4
|
|
// Special colorspace
|
|
{
|
|
// allocate enough space to hold a whole MCU width
|
|
// this means we should take into account sampling ratios
|
|
// `*8` is because each MCU spans 8 widths.
|
|
let len = comp.width_stride * comp.vertical_sample * 8;
|
|
|
|
comp.needed = true;
|
|
comp.raw_coeff = vec![0; len];
|
|
} else {
|
|
comp.needed = false;
|
|
}
|
|
}
|
|
|
|
// If all components are contained in the first scan of MCUs, then we can process into
|
|
// (upsampled) pixels immediately after each MCU, for convenience we use each row of MCUS.
|
|
// Otherwise, we must first wait until following SOS provide the remaining components.
|
|
let all_components_in_first_scan = usize::from(self.num_scans) == self.components.len();
|
|
let mut progressive_mcus: [Vec<i16>; 4] = core::array::from_fn(|_| vec![]);
|
|
|
|
if !all_components_in_first_scan {
|
|
for (component, mcu) in self.components.iter().zip(&mut progressive_mcus) {
|
|
let len = mcu_width
|
|
* component.vertical_sample
|
|
* component.horizontal_sample
|
|
* mcu_height
|
|
* 64;
|
|
*mcu = vec![0; len];
|
|
}
|
|
}
|
|
|
|
let mut pixels_written = 0;
|
|
|
|
let is_hv = usize::from(self.is_interleaved);
|
|
let upsampler_scratch_size = is_hv * self.components.iter().map(|x| x.width_stride).max().unwrap_or(0) * 8;
|
|
let mut upsampler_scratch_space = vec![0; upsampler_scratch_size];
|
|
|
|
'sos: loop {
|
|
trace!(
|
|
"Baseline decoding of components: {:?}",
|
|
&self.z_order[..usize::from(self.num_scans)]
|
|
);
|
|
|
|
trace!("Decoding MCU width: {mcu_width}, height: {mcu_height}");
|
|
|
|
for i in 0..mcu_height {
|
|
if stream.overread_by > 0 {
|
|
pixels.get_mut(pixels_written..).map(|v| v.fill(128));
|
|
if self.options.strict_mode() {
|
|
return Err(DecodeErrors::FormatStatic("Premature end of buffer"));
|
|
};
|
|
|
|
error!("Premature end of buffer");
|
|
break;
|
|
}
|
|
|
|
// decode a whole MCU width,
|
|
// this takes into account interleaved components.
|
|
let terminate = if all_components_in_first_scan {
|
|
self.decode_mcu_width::<false>(
|
|
mcu_width,
|
|
i,
|
|
&mut tmp,
|
|
&mut stream,
|
|
&mut progressive_mcus
|
|
)?
|
|
} else {
|
|
/* NB: (cae). This code was added due to the issue at https://github.com/etemesi254/zune-image/issues/277
|
|
*
|
|
* There is a particular set of images that interleave the start of scan (SOS) with the MCU,
|
|
* E.g if it's a three component image, we have SOS->MCU ->SOS->MCU ->SOS->MCU
|
|
* which presents a problem on decoding, we need to buffer the whole image before continuing since
|
|
* we won't have a row containing all the component data which will be needed e.g for color conversion.
|
|
*
|
|
* The mechanisms is that we decode the whole image upfront, which goes against the normal
|
|
* routine of decoding MCU width , so this requires more memory upfront than initial routines
|
|
* but it is a single image out of the many corpuses that exist, so its fine.
|
|
* (image in test-images/jpeg/sos_news.jpeg)
|
|
|
|
* Code contributed by Aurelia Molzer (https://github.com/197g)
|
|
|
|
*
|
|
*/
|
|
|
|
self.decode_mcu_width::<true>(
|
|
mcu_width,
|
|
i,
|
|
&mut tmp,
|
|
&mut stream,
|
|
&mut progressive_mcus
|
|
)?
|
|
};
|
|
|
|
// process that width up until it's impossible. This is faster than allocation the
|
|
// full components, which we skipped earlier.
|
|
if all_components_in_first_scan {
|
|
self.post_process(
|
|
pixels,
|
|
i,
|
|
mcu_height,
|
|
width,
|
|
padded_width,
|
|
&mut pixels_written,
|
|
&mut upsampler_scratch_space
|
|
)?;
|
|
}
|
|
|
|
match terminate {
|
|
McuContinuation::Ok => {}
|
|
McuContinuation::AnotherSos if all_components_in_first_scan => {
|
|
warn!("More than one SOS despite already having all components");
|
|
return Ok(());
|
|
}
|
|
McuContinuation::AnotherSos => continue 'sos,
|
|
McuContinuation::InterScanMarker(marker) => {
|
|
// Handle inter-scan markers (DHT/DQT/etc) uniformly here.
|
|
// This keeps all marker handling in the outer loop.
|
|
if self.advance_to_next_sos(marker, &mut stream)? {
|
|
continue 'sos;
|
|
} else {
|
|
// Hit EOI
|
|
break;
|
|
}
|
|
}
|
|
McuContinuation::Terminate => {
|
|
warn!("Got terminate signal, will not process further");
|
|
pixels.get_mut(pixels_written..).map(|v| v.fill(128));
|
|
return Ok(());
|
|
}
|
|
}
|
|
}
|
|
|
|
// Breaks if we get here, looping only if we have restarted, i.e. found another SOS and
|
|
// continued at `'sos'.
|
|
break;
|
|
}
|
|
|
|
if !all_components_in_first_scan {
|
|
self.finish_baseline_decoding(&progressive_mcus, mcu_width, pixels)?;
|
|
}
|
|
|
|
// it may happen that some images don't have the whole buffer
|
|
// so we can't panic in case of that
|
|
// assert_eq!(pixels_written, pixels.len());
|
|
|
|
// For UHD usecases that tie two images separating them with EOI and
|
|
// SOI markers, it may happen that we do not reach this image end of image
|
|
// So this ensures we reach it
|
|
// Ensure we read EOI
|
|
if !stream.seen_eoi {
|
|
let marker = get_marker(&mut self.stream, &mut stream);
|
|
match marker {
|
|
Ok(_m) => {
|
|
trace!("Found marker {:?}", _m);
|
|
}
|
|
Err(_) => {
|
|
// ignore error
|
|
}
|
|
}
|
|
}
|
|
|
|
trace!("Finished decoding image");
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Process all MCUs when baseline decoding has been processing them component-after-component.
|
|
/// For simplicity this assembles the dequantized blocks in the order that the post processing
|
|
/// of an interleaved baseline decoding would use.
|
|
#[allow(clippy::too_many_lines)]
|
|
#[allow(clippy::cast_sign_loss)]
|
|
pub(crate) fn finish_baseline_decoding(
|
|
&mut self, block: &[Vec<i16>; MAX_COMPONENTS], _mcu_width: usize, pixels: &mut [u8]
|
|
) -> Result<(), DecodeErrors> {
|
|
let mcu_height = self.mcu_y;
|
|
|
|
// Size of our output image(width*height)
|
|
let is_hv = usize::from(self.is_interleaved);
|
|
let upsampler_scratch_size = is_hv * self.components[0].width_stride;
|
|
let width = usize::from(self.info.width);
|
|
let padded_width = calculate_padded_width(width, self.info.sample_ratio);
|
|
|
|
let mut upsampler_scratch_space = vec![0; upsampler_scratch_size];
|
|
|
|
for (pos, comp) in self.components.iter_mut().enumerate() {
|
|
// Mark only needed components for computing output colors.
|
|
if min(
|
|
self.options.jpeg_get_out_colorspace().num_components() - 1,
|
|
pos
|
|
) == pos
|
|
|| self.input_colorspace == ColorSpace::YCCK
|
|
|| self.input_colorspace == ColorSpace::CMYK
|
|
{
|
|
comp.needed = true;
|
|
} else {
|
|
comp.needed = false;
|
|
}
|
|
}
|
|
|
|
let mut pixels_written = 0;
|
|
|
|
// dequantize and idct have been performed, only color convert.
|
|
for i in 0..mcu_height {
|
|
// All the data is already in the right order, we just need to be able to pass it to
|
|
// the post_process & upsample method. That expects all the data to be stored as one
|
|
// row of MCUs in each component's `raw_coeff`.
|
|
'component: for (position, component) in &mut self.components.iter_mut().enumerate() {
|
|
if !component.needed {
|
|
continue 'component;
|
|
}
|
|
|
|
// step is the number of pixels this iteration wil be handling
|
|
// Given by the number of mcu's height and the length of the component block
|
|
// Since the component block contains the whole channel as raw pixels
|
|
// we this evenly divides the pixels into MCU blocks
|
|
//
|
|
// For interleaved images, this gives us the exact pixels comprising a whole MCU
|
|
// block
|
|
let step = block[position].len() / mcu_height;
|
|
|
|
// where we will be reading our pixels from.
|
|
let slice = &block[position][i * step..][..step];
|
|
let temp_channel = &mut component.raw_coeff;
|
|
temp_channel[..step].copy_from_slice(slice);
|
|
}
|
|
|
|
// process that whole stripe of MCUs
|
|
self.post_process(
|
|
pixels,
|
|
i,
|
|
mcu_height,
|
|
width,
|
|
padded_width,
|
|
&mut pixels_written,
|
|
&mut upsampler_scratch_space
|
|
)?;
|
|
}
|
|
|
|
return Ok(());
|
|
}
|
|
|
|
fn decode_mcu_width<const PROGRESSIVE: bool>(
|
|
&mut self, mcu_width: usize, mcu_height: usize, tmp: &mut [i32; 64],
|
|
stream: &mut BitStream, progressive: &mut [Vec<i16>; 4]
|
|
) -> Result<McuContinuation, DecodeErrors> {
|
|
let is_one_by_one = !self.scan_subsampled;
|
|
|
|
// The definition of MCU depends on the sampling factor of involved scans. When components
|
|
// have different factors then each Minimal-Coding-Unit is the least common multiple such
|
|
// that we have an integer number of blocks from each component. But the decoding of these
|
|
// components differs from it otherwise, we need an inner loop with a dynamic amount of
|
|
// coefficients per component, whereas otherwise we have exactly one block of coefficients
|
|
// encoded for each component in the bitstream order.
|
|
//
|
|
// We statically specialize on this to improve code generation of the common case a little
|
|
// bit. We could also special case common sub-sampling cases but be mindful of code bloat.
|
|
if is_one_by_one {
|
|
self.inner_decode_mcu_width::<PROGRESSIVE, false>(
|
|
mcu_width,
|
|
mcu_height,
|
|
tmp,
|
|
stream,
|
|
progressive
|
|
)
|
|
} else {
|
|
self.inner_decode_mcu_width::<PROGRESSIVE, true>(
|
|
mcu_width,
|
|
mcu_height,
|
|
tmp,
|
|
stream,
|
|
progressive
|
|
)
|
|
}
|
|
}
|
|
|
|
// Inline-never ensures we do get this function optimize on its own, into two different
|
|
// versions, without the optimizer tripping up over the complexity that comes with the
|
|
// constant folding. And constant folding is quite important for performance here as
|
|
// when `not SAMPLED` then the inner loop has exactly one iteration per component in
|
|
// the scan. The difference was ~1% or a bit more.
|
|
fn inner_decode_mcu_width<const PROGRESSIVE: bool, const SAMPLED: bool>(
|
|
&mut self, mcu_width: usize, mcu_height: usize, tmp: &mut [i32; 64],
|
|
stream: &mut BitStream, progressive: &mut [Vec<i16>; 4]
|
|
) -> Result<McuContinuation, DecodeErrors> {
|
|
let z_order = self.z_order;
|
|
let z_scans = &z_order[..usize::from(self.num_scans)];
|
|
|
|
// How much of the head of `tmp` was written by the last MCU decoding? We only check for
|
|
// two different cases and not all possible outcomes as this is only used to optimize the
|
|
// bytes written in `fill`. Since the clobber happens in UNZIGZAG order we'd be straddling
|
|
// most cache lines anyways even if we did a partial write with the exact length of the
|
|
// coefficient data which was written into `tmp`.
|
|
let mut clobber_more_than_4x4 = true;
|
|
|
|
// For non-interleaved scans (PROGRESSIVE=true), each scan contains a single component
|
|
// and we iterate over that component's actual data unit count, not the interleaved MCU
|
|
// width multiplied by sampling factor.
|
|
let mut scan_du_width = if PROGRESSIVE {
|
|
let k = z_scans[0];
|
|
let comp = &self.components[k];
|
|
// Calculate actual data units for this component: ceil(width / (8 * subsampling_ratio))
|
|
(self.info.width as usize * comp.horizontal_sample + self.h_max * 8 - 1)
|
|
/ (self.h_max * 8)
|
|
} else {
|
|
mcu_width
|
|
};
|
|
// In malformed scans that list multiple components, clamp to the smallest row capacity
|
|
// to avoid writing past the row buffer.
|
|
if PROGRESSIVE && z_scans.len() > 1 {
|
|
let min_du = z_scans
|
|
.iter()
|
|
.map(|&k| self.components[k].width_stride / 8)
|
|
.min()
|
|
.unwrap_or(0);
|
|
scan_du_width = scan_du_width.min(min_du);
|
|
}
|
|
|
|
for j in 0..scan_du_width {
|
|
// iterate over components
|
|
for &k in z_scans {
|
|
// we made this loop body massive due to several different paths that depend on
|
|
// static conditions. Note we (potentially) call into other functions so the
|
|
// compiler will not unroll anything here anyways. The gains from separating
|
|
// differently optimized loop bodies are much greater than a single additional jump
|
|
// here.
|
|
let component = &mut self.components[k];
|
|
|
|
let dc_table = self.dc_huffman_tables[component.dc_huff_table % MAX_COMPONENTS]
|
|
.as_ref()
|
|
.ok_or(DecodeErrors::FormatStatic("DC table not found"))?;
|
|
|
|
let ac_table = self.ac_huffman_tables[component.ac_huff_table % MAX_COMPONENTS]
|
|
.as_ref()
|
|
.ok_or(DecodeErrors::FormatStatic("AC table not found"))?;
|
|
|
|
let qt_table = &component.quantization_table;
|
|
let channel = if PROGRESSIVE {
|
|
let offset =
|
|
mcu_height * component.width_stride * 8 * component.vertical_sample;
|
|
// Small stopgap for https://github.com/etemesi254/zune-image/issues/362
|
|
if offset >= progressive[k].len(){
|
|
return Err(DecodeErrors::FormatStatic("Would panic on slice iteration"))
|
|
}
|
|
&mut progressive[k][offset..]
|
|
} else {
|
|
&mut component.raw_coeff
|
|
};
|
|
|
|
let component_samples_needed = component.needed;
|
|
|
|
// If image is interleaved iterate over scan components,
|
|
// otherwise if it-s non-interleaved, these routines iterate in
|
|
// trivial scanline order(Y,Cb,Cr)
|
|
//
|
|
// Turn the bounds into a compile time constant for a common special case. This
|
|
// allows the compiler to unroll the loop and then do a bunch of interleaving.
|
|
//
|
|
// For PROGRESSIVE (non-interleaved), we iterate data units directly so
|
|
// h_samp/v_samp loops run exactly once.
|
|
let v_step =
|
|
if SAMPLED && !PROGRESSIVE { 0..component.vertical_sample } else { 0..1 };
|
|
|
|
for v_samp in v_step {
|
|
let h_step =
|
|
if SAMPLED && !PROGRESSIVE { 0..component.horizontal_sample } else { 0..1 };
|
|
|
|
for h_samp in h_step {
|
|
let result = if component_samples_needed {
|
|
// Fill the array with zeroes, decode_mcu_block expects
|
|
// a zero based array. Clobber is in zig-zag order though.
|
|
// Writing consecutive entries is basically free in terms
|
|
// of memory throughput so we opt for a larger power of
|
|
// two which lets the compiler turn this into a repeated
|
|
// write of a zeroed vector register, which does not have
|
|
// any branches, instead of a more difficult pattern where
|
|
// we attempt to overwrite exactly one coefficient.
|
|
let clobber_len = if !clobber_more_than_4x4 { 32 } else { 64 };
|
|
|
|
tmp[..clobber_len].fill(0);
|
|
|
|
stream.decode_mcu_block(
|
|
&mut self.stream,
|
|
dc_table,
|
|
ac_table,
|
|
qt_table,
|
|
tmp,
|
|
&mut component.dc_pred
|
|
)
|
|
} else {
|
|
// We do not touch tmp so there is no need to reset it.
|
|
stream.discard_mcu_block(&mut self.stream, dc_table, ac_table)
|
|
};
|
|
|
|
// If an error occurs we can either propagate it
|
|
// as an error or print it and call terminate.
|
|
//
|
|
// This allows even corrupt images to render something,
|
|
// even if its bad, matching browsers.
|
|
//
|
|
// See example in https://github.com/etemesi254/zune-image/issues/293
|
|
let len = if let Ok(len) = result {
|
|
len
|
|
} else {
|
|
// result.is_err()
|
|
return if self.options.strict_mode() {
|
|
Err(result.err().unwrap())
|
|
} else {
|
|
error!("{}", result.err().unwrap());
|
|
Ok(McuContinuation::Terminate)
|
|
};
|
|
};
|
|
|
|
if component_samples_needed {
|
|
// tmp was only written partially, note that len is in ZigZag order.
|
|
clobber_more_than_4x4 = len > 10;
|
|
|
|
let idct_position = if PROGRESSIVE {
|
|
// For non-interleaved, j indexes data units directly
|
|
j * 8
|
|
} else {
|
|
// derived from stb and rewritten for my tastes
|
|
let c2 = v_samp * 8;
|
|
let c3 = ((j * component.horizontal_sample) + h_samp) * 8;
|
|
|
|
component.width_stride * c2 + c3
|
|
};
|
|
|
|
let idct_pos = channel.get_mut(idct_position..).unwrap();
|
|
|
|
if len <= 1 {
|
|
(self.idct_1x1_func)(tmp, idct_pos, component.width_stride);
|
|
} else if len <= 10 {
|
|
(self.idct_4x4_func)(tmp, idct_pos, component.width_stride);
|
|
} else {
|
|
// call idct.
|
|
(self.idct_func)(tmp, idct_pos, component.width_stride);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
self.todo = self.todo.wrapping_sub(1);
|
|
|
|
if self.todo == 0 {
|
|
self.handle_rst_main(stream)?;
|
|
continue;
|
|
}
|
|
|
|
if stream.marker.is_some() && stream.bits_left == 0 {
|
|
break;
|
|
}
|
|
}
|
|
|
|
self.check_stream_marker_after_mcu_width(stream)
|
|
}
|
|
|
|
fn check_stream_marker_after_mcu_width(
|
|
&mut self, stream: &mut BitStream
|
|
) -> Result<McuContinuation, DecodeErrors> {
|
|
// After all interleaved components, that's an MCU
|
|
// handle stream markers
|
|
//
|
|
// In some corrupt images, it may occur that header markers occur in the stream.
|
|
// The spec EXPLICITLY FORBIDS this, specifically, in
|
|
// routine F.2.2.5 it says
|
|
// `The only valid marker which may occur within the Huffman coded data is the RSTm marker.`
|
|
//
|
|
// But libjpeg-turbo allows it because of some weird reason. so I'll also
|
|
// allow it because of some weird reason.
|
|
if let Some(m) = stream.marker {
|
|
if m == Marker::EOI {
|
|
// acknowledge and ignore EOI marker.
|
|
stream.marker.take();
|
|
trace!("Found EOI marker");
|
|
// Google Introduced the Ultra-HD image format which is basically
|
|
// stitching two images into one container.
|
|
// They basically separate two images via a EOI and SOI marker
|
|
// so let's just ensure if we ever see EOI, we never read past that
|
|
// ever.
|
|
// https://github.com/google/libultrahdr
|
|
stream.seen_eoi = true;
|
|
} else if let Marker::RST(_) = m {
|
|
//debug_assert_eq!(self.todo, 0);
|
|
if self.todo == 0 {
|
|
self.handle_rst(stream)?;
|
|
}
|
|
} else if let Marker::SOS = m {
|
|
self.parse_marker_inner(m)?;
|
|
stream.marker.take();
|
|
stream.reset();
|
|
trace!("Found SOS marker");
|
|
return Ok(McuContinuation::AnotherSos);
|
|
} else if matches!(m, Marker::DHT | Marker::DQT | Marker::DRI | Marker::COM)
|
|
|| matches!(m, Marker::APP(_))
|
|
{
|
|
// For non-interleaved images, setup markers can appear between scans.
|
|
// Signal the caller to handle this marker and find the next SOS.
|
|
// This keeps all marker parsing in the caller's loop.
|
|
stream.marker.take();
|
|
trace!("Found inter-scan marker {:?}", m);
|
|
return Ok(McuContinuation::InterScanMarker(m));
|
|
} else {
|
|
if self.options.strict_mode() {
|
|
return Err(DecodeErrors::Format(format!(
|
|
"Marker {m:?} found where not expected"
|
|
)));
|
|
}
|
|
error!(
|
|
"Marker `{:?}` Found within Huffman Stream, possibly corrupt jpeg",
|
|
m
|
|
);
|
|
|
|
self.parse_marker_inner(m)?;
|
|
stream.marker.take();
|
|
stream.reset();
|
|
return Ok(McuContinuation::Terminate);
|
|
}
|
|
}
|
|
|
|
Ok(McuContinuation::Ok)
|
|
}
|
|
|
|
/// Scan for the next SOS marker, parsing setup markers along the way.
|
|
///
|
|
/// This is the unified marker scanning function used after encountering an
|
|
/// inter-scan marker. It handles DHT, DQT, DRI, COM, and APP markers that
|
|
/// can appear between scans in non-interleaved images.
|
|
///
|
|
/// # Arguments
|
|
/// * `first_marker` - The first marker that was already detected (not yet parsed)
|
|
/// * `stream` - The bitstream state
|
|
///
|
|
/// # Returns
|
|
/// * `Ok(true)` - Found SOS, ready to continue decoding
|
|
/// * `Ok(false)` - Found EOI, decoding complete
|
|
/// * `Err(_)` - Error (too many markers, unexpected marker in strict mode, etc.)
|
|
fn advance_to_next_sos(
|
|
&mut self,
|
|
first_marker: Marker,
|
|
stream: &mut BitStream
|
|
) -> Result<bool, DecodeErrors> {
|
|
// Limit iterations to prevent DoS from malicious files.
|
|
const MAX_INTER_SCAN_MARKERS: usize = 64;
|
|
|
|
// Parse the first marker that triggered this call
|
|
self.parse_marker_inner(first_marker)?;
|
|
stream.reset();
|
|
|
|
for _ in 0..MAX_INTER_SCAN_MARKERS {
|
|
let marker = get_marker(&mut self.stream, stream)?;
|
|
|
|
match marker {
|
|
Marker::SOS => {
|
|
self.parse_marker_inner(Marker::SOS)?;
|
|
stream.reset();
|
|
trace!("Found SOS marker, continuing decode");
|
|
return Ok(true);
|
|
}
|
|
Marker::EOI => {
|
|
stream.seen_eoi = true;
|
|
trace!("Found EOI marker");
|
|
return Ok(false);
|
|
}
|
|
Marker::DHT | Marker::DQT | Marker::DRI | Marker::COM => {
|
|
trace!("Parsing inter-scan marker {:?}", marker);
|
|
self.parse_marker_inner(marker)?;
|
|
}
|
|
Marker::APP(_) => {
|
|
trace!("Parsing inter-scan APP marker {:?}", marker);
|
|
self.parse_marker_inner(marker)?;
|
|
}
|
|
other => {
|
|
if self.options.strict_mode() {
|
|
return Err(DecodeErrors::Format(format!(
|
|
"Unexpected marker {:?} while scanning for SOS between scans",
|
|
other
|
|
)));
|
|
}
|
|
// Non-strict: skip unknown marker
|
|
warn!("Skipping unexpected marker {:?} between scans", other);
|
|
let length = self.stream.get_u16_be_err()?;
|
|
if length >= 2 {
|
|
self.stream.skip((length - 2) as usize)?;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
Err(DecodeErrors::FormatStatic(
|
|
"Too many markers between scans (exceeded limit of 64)"
|
|
))
|
|
}
|
|
|
|
// handle RST markers.
|
|
// No-op if not using restarts
|
|
// this routine is shared with mcu_prog
|
|
#[cold]
|
|
pub(crate) fn handle_rst(&mut self, stream: &mut BitStream) -> Result<(), DecodeErrors> {
|
|
self.todo = self.restart_interval;
|
|
|
|
if let Some(marker) = stream.marker {
|
|
// Found a marker
|
|
// Read stream and see what marker is stored there
|
|
match marker {
|
|
Marker::RST(_) => {
|
|
// reset stream
|
|
stream.reset();
|
|
// Initialize dc predictions to zero for all components
|
|
self.components.iter_mut().for_each(|x| x.dc_pred = 0);
|
|
// Start iterating again. from position.
|
|
}
|
|
Marker::EOI => {
|
|
// silent pass
|
|
}
|
|
// Valid markers that can appear between scans at a restart boundary
|
|
// (restart interval aligns with end of scan). Leave for caller.
|
|
Marker::SOS | Marker::DHT | Marker::DQT | Marker::DRI | Marker::COM
|
|
| Marker::APP(_) => {}
|
|
_ => {
|
|
if self.options.strict_mode() {
|
|
return Err(DecodeErrors::MCUError(format!(
|
|
"Unexpected marker {marker:?} at restart boundary"
|
|
)));
|
|
}
|
|
warn!("Unexpected marker {:?} at restart boundary", marker);
|
|
}
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
#[allow(clippy::too_many_lines, clippy::too_many_arguments)]
|
|
pub(crate) fn post_process(
|
|
&mut self, pixels: &mut [u8], i: usize, mcu_height: usize, width: usize,
|
|
padded_width: usize, pixels_written: &mut usize, upsampler_scratch_space: &mut [i16]
|
|
) -> Result<(), DecodeErrors> {
|
|
let out_colorspace_components = self.options.jpeg_get_out_colorspace().num_components();
|
|
|
|
let mut px = *pixels_written;
|
|
// indicates whether image is vertically up-sampled
|
|
let is_vertically_sampled = self
|
|
.components
|
|
.iter()
|
|
.any(|c| c.sample_ratio == SampleRatios::HV || c.sample_ratio == SampleRatios::V);
|
|
|
|
let mut comp_len = self.components.len();
|
|
|
|
// If we are moving from YCbCr -> Luma, we do not allocate storage for other components, so we
|
|
// will panic when we are trying to read samples, so for that case,
|
|
// hardcode it so that we don't panic when doing
|
|
// *samp = &samples[j][pos * padded_width..(pos + 1) * padded_width]
|
|
if out_colorspace_components < comp_len && self.options.jpeg_get_out_colorspace() == Luma {
|
|
comp_len = out_colorspace_components;
|
|
}
|
|
let mut color_conv_function =
|
|
|num_iters: usize, samples: [&[i16]; 4]| -> Result<(), DecodeErrors> {
|
|
for (pos, output) in pixels[px..]
|
|
.chunks_exact_mut(width * out_colorspace_components)
|
|
.take(num_iters)
|
|
.enumerate()
|
|
{
|
|
let mut raw_samples: [&[i16]; 4] = [&[], &[], &[], &[]];
|
|
|
|
// iterate over each line, since color-convert needs only
|
|
// one line
|
|
for (j, samp) in raw_samples.iter_mut().enumerate().take(comp_len) {
|
|
let temp = &samples[j].get(pos * padded_width..(pos + 1) * padded_width);
|
|
if temp.is_none() {
|
|
return Err(DecodeErrors::FormatStatic("Missing samples"));
|
|
}
|
|
*samp = temp.unwrap();
|
|
}
|
|
color_convert(
|
|
&raw_samples,
|
|
self.color_convert_16,
|
|
self.input_colorspace,
|
|
self.options.jpeg_get_out_colorspace(),
|
|
output,
|
|
width,
|
|
padded_width
|
|
)?;
|
|
px += width * out_colorspace_components;
|
|
}
|
|
Ok(())
|
|
};
|
|
|
|
let comps = &mut self.components[..];
|
|
|
|
if self.is_interleaved && self.options.jpeg_get_out_colorspace() != ColorSpace::Luma {
|
|
for comp in comps.iter_mut() {
|
|
upsample(
|
|
comp,
|
|
mcu_height,
|
|
i,
|
|
upsampler_scratch_space,
|
|
is_vertically_sampled
|
|
)?;
|
|
}
|
|
|
|
if is_vertically_sampled {
|
|
if i > 0 {
|
|
// write the last line, it wasn't up-sampled as we didn't have row_down
|
|
// yet
|
|
let mut samples: [&[i16]; 4] = [&[], &[], &[], &[]];
|
|
|
|
for (samp, component) in samples.iter_mut().zip(comps.iter()) {
|
|
*samp = &component.first_row_upsample_dest;
|
|
}
|
|
|
|
// ensure length matches for all samples
|
|
let _first_len = samples[0].len();
|
|
|
|
// This was a good check, but can be caused to panic, esp on invalid/corrupt images.
|
|
// See one in issue https://github.com/etemesi254/zune-image/issues/262, so for now
|
|
// we just ignore and generate invalid images at the end.
|
|
|
|
//
|
|
//
|
|
// for samp in samples.iter().take(comp_len) {
|
|
// assert_eq!(first_len, samp.len());
|
|
// }
|
|
let num_iters = self.coeff * self.v_max;
|
|
|
|
color_conv_function(num_iters, samples)?;
|
|
}
|
|
|
|
// After up-sampling the last row, save any row that can be used for
|
|
// a later up-sampling,
|
|
//
|
|
// E.g the Y sample is not sampled but we haven't finished upsampling the last row of
|
|
// the previous mcu, since we don't have the down row, so save it
|
|
for component in comps.iter_mut() {
|
|
if component.sample_ratio != SampleRatios::H {
|
|
// We don't care about H sampling factors, since it's copied in the workers function
|
|
|
|
// copy last row to be used for the next color conversion
|
|
let size = component.vertical_sample
|
|
* component.width_stride
|
|
* component.sample_ratio.sample();
|
|
|
|
let last_bytes =
|
|
component.raw_coeff.rchunks_exact_mut(size).next().unwrap();
|
|
|
|
component
|
|
.first_row_upsample_dest
|
|
.copy_from_slice(last_bytes);
|
|
}
|
|
}
|
|
}
|
|
|
|
let mut samples: [&[i16]; 4] = [&[], &[], &[], &[]];
|
|
|
|
for (samp, component) in samples.iter_mut().zip(comps.iter()) {
|
|
*samp = if component.sample_ratio == SampleRatios::None {
|
|
&component.raw_coeff
|
|
} else {
|
|
&component.upsample_dest
|
|
};
|
|
}
|
|
|
|
// we either do 7 or 8 MCU's depending on the state, this only applies to
|
|
// vertically sampled images
|
|
//
|
|
// for rows up until the last MCU, we do not upsample the last stride of the MCU
|
|
// which means that the number of iterations should take that into account is one less the
|
|
// up-sampled size
|
|
//
|
|
// For the last MCU, we upsample the last stride, meaning that if we hit the last MCU, we
|
|
// should sample full raw coeffs
|
|
let is_last_considered = is_vertically_sampled && (i != mcu_height.saturating_sub(1));
|
|
|
|
let num_iters = (8 - usize::from(is_last_considered)) * self.coeff * self.v_max;
|
|
|
|
color_conv_function(num_iters, samples)?;
|
|
} else {
|
|
let mut channels_ref: [&[i16]; MAX_COMPONENTS] = [&[]; MAX_COMPONENTS];
|
|
|
|
self.components
|
|
.iter()
|
|
.enumerate()
|
|
.for_each(|(pos, x)| channels_ref[pos] = &x.raw_coeff);
|
|
|
|
if let SampleRatios::Generic(_, v) = self.info.sample_ratio {
|
|
color_conv_function(8 * v * self.coeff, channels_ref)?;
|
|
} else {
|
|
color_conv_function(8 * self.coeff, channels_ref)?;
|
|
}
|
|
}
|
|
|
|
*pixels_written = px;
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
enum McuContinuation {
|
|
Ok,
|
|
AnotherSos,
|
|
/// Found an inter-scan marker (DHT/DQT/DRI/COM/APP) that needs handling.
|
|
/// The caller should parse it and scan for the next SOS.
|
|
InterScanMarker(Marker),
|
|
Terminate
|
|
}
|