From f34f5afb165e9caaba4dd96c38c574d3c260f05a Mon Sep 17 00:00:00 2001 From: GF Date: Sat, 26 Sep 2026 03:59:10 -0400 Subject: [PATCH 1/5] perf(jpeg): keep entropy state in registers and widen NEON colour conversion - Decode MCU rows against local copies of the bit reader and DC predictors so the accumulator stays in registers; cold refill paths take and return the reader by value. - Refill the bit accumulator eight bytes at a time when the window has no 0xFF, falling back to the byte-wise path for stuffing, markers, and the input tail. - Replace the per-coefficient touched-index list and activity state machine with one AC flag and an OR of natural indices. - Run NEON YCbCr->RGB in exact 16-bit lanes (constants split as 65536*k + remainder), sixteen pixels per iteration, and test it against the scalar converter for every chroma pair. --- crates/j2k-jpeg/src/backend/neon.rs | 174 ++++++++------- crates/j2k-jpeg/src/backend/tests.rs | 20 ++ crates/j2k-jpeg/src/entropy/block.rs | 154 ++++++-------- .../src/entropy/sequential/deposit.rs | 53 +++++ .../src/entropy/sequential/fast420/rows.rs | 30 ++- .../src/entropy/sequential/generic/row.rs | 40 +++- .../src/entropy/sequential/restart.rs | 11 +- .../j2k-jpeg/src/entropy/sequential/rgb444.rs | 47 ++++- crates/j2k-jpeg/src/internal/bit_reader.rs | 199 ++++++++++++++++-- .../src/internal/bit_reader/terminal.rs | 2 +- crates/j2k-jpeg/src/simd/neon_memory.rs | 32 ++- 11 files changed, 563 insertions(+), 199 deletions(-) diff --git a/crates/j2k-jpeg/src/backend/neon.rs b/crates/j2k-jpeg/src/backend/neon.rs index f1d52c88d..89c30b018 100644 --- a/crates/j2k-jpeg/src/backend/neon.rs +++ b/crates/j2k-jpeg/src/backend/neon.rs @@ -1,18 +1,17 @@ // SPDX-License-Identifier: MIT OR Apache-2.0 use core::arch::aarch64::{ - int32x4_t, uint16x8_t, uint8x16_t, uint8x8_t, vaddq_s32, vaddq_u16, vcombine_u16, vcombine_u8, - vdupq_n_s32, vdupq_n_u16, vget_high_u16, vget_low_u16, vmovl_u16, vmovl_u8, vmulq_n_s32, - vqmovn_u16, vqmovun_s32, vreinterpretq_s32_u32, vshrq_n_s32, vshrq_n_u16, vsubq_s32, vzip_u8, - vzipq_u16, + int16x8_t, int32x4_t, uint16x8_t, uint8x16_t, uint8x8_t, vaddq_s16, vaddq_u16, vcombine_s16, + vcombine_u8, vdup_n_u8, vdupq_n_s16, vdupq_n_u16, vdupq_n_u8, vget_low_s16, vget_low_u8, + vmlal_high_n_s16, vmlal_n_s16, vmovl_high_u8, vmovl_u8, vmull_high_n_s16, vmull_n_s16, + vqmovn_u16, vqmovun_high_s16, vqmovun_s16, vreinterpretq_s16_u16, vrshrn_n_s32, vshrq_n_u16, + vsubl_high_u8, vsubl_u8, vsubq_s16, vzip_u8, vzipq_u16, }; use super::row_pair::{normalize_simd_row_pair, normalize_ycbcr_row}; use super::{scalar, Rgb420ChromaRows, Rgb420Crop, Rgb420CroppedRowPair, Rgb420RowPair}; use crate::color::upsample::h2v2_fancy_sample_for_width; -use crate::color::ycbcr::{ - ycbcr_to_rgb, FIX_0_34414, FIX_0_71414, FIX_1_40200, FIX_1_77200, ROUND, -}; +use crate::color::ycbcr::{ycbcr_to_rgb, FIX_0_34414, FIX_0_71414, FIX_1_40200, FIX_1_77200}; use crate::simd::neon_memory; pub(crate) fn fill_rgb_row_from_gray(neon: fearless_simd::Neon, gray_row: &[u8], dst: &mut [u8]) { @@ -1384,30 +1383,85 @@ fn fill_chunk_from_vectors(y: uint8x8_t, cb: uint8x8_t, cr: uint8x8_t, dst_chunk fill_chunk_from_vectors_u16(y, cb16, cr16, dst_chunk); } +// The libjpeg chroma multipliers exceed `i16`, so each is split into a +// multiple of 65536 plus a remainder that fits a 16-bit lane. For integer `k`, +// `(65536·k·x + p + ROUND) >> 16 == k·x + ((p + ROUND) >> 16)` exactly, so the +// split keeps `ycbcr_to_rgb`'s rounding bit for bit while every multiply stays +// a 16→32-bit widening one and the rest of the math runs eight lanes wide. +const R_CR_REMAINDER: i16 = narrow_multiplier(FIX_1_40200 - 65_536); +const G_CB: i16 = narrow_multiplier(FIX_0_34414); +const G_CR_REMAINDER: i16 = narrow_multiplier(FIX_0_71414 - 65_536); +const B_CB_REMAINDER: i16 = narrow_multiplier(FIX_1_77200 - 2 * 65_536); + +const fn narrow_multiplier(value: i32) -> i16 { + assert!(value >= i16::MIN as i32 && value <= i16::MAX as i32); + #[expect( + clippy::cast_possible_truncation, + reason = "the range is asserted at compile time" + )] + let narrowed = value as i16; + narrowed +} + +#[target_feature(enable = "neon")] +fn rounded_high_half(low: int32x4_t, high: int32x4_t) -> int16x8_t { + // `vrshrn` computes `(x + ROUND) >> 16` without intermediate overflow. + vcombine_s16(vrshrn_n_s32::<16>(low), vrshrn_n_s32::<16>(high)) +} + +/// Convert eight pixels whose chroma lanes hold 8-bit sample values. #[target_feature(enable = "neon")] fn fill_chunk_from_vectors_u16(y: uint8x8_t, cb: uint16x8_t, cr: uint16x8_t, dst_chunk: &mut [u8]) { let Some(dst) = dst_chunk.first_chunk_mut::<{ LANES * 3 }>() else { return; }; - let y16 = vmovl_u8(y); - - let y_lo = widen_low(y16); - let y_hi = widen_high(y16); - let cb_lo = subtract_bias(widen_low(cb)); - let cb_hi = subtract_bias(widen_high(cb)); - let cr_lo = subtract_bias(widen_low(cr)); - let cr_hi = subtract_bias(widen_high(cr)); - - let (r_lo, g_lo, b_lo) = convert_half(y_lo, cb_lo, cr_lo); - let (r_hi, g_hi, b_hi) = convert_half(y_hi, cb_hi, cr_hi); + let bias = vdupq_n_s16(128); + let (r, g, b) = ycbcr_lanes_to_rgb( + vreinterpretq_s16_u16(vmovl_u8(y)), + vsubq_s16(vreinterpretq_s16_u16(cb), bias), + vsubq_s16(vreinterpretq_s16_u16(cr), bias), + ); + neon_memory::store_rgb8(dst, vqmovun_s16(r), vqmovun_s16(g), vqmovun_s16(b)); +} - let r_bytes = pack_eight_u8(r_lo, r_hi); - let g_bytes = pack_eight_u8(g_lo, g_hi); - let b_bytes = pack_eight_u8(b_lo, b_hi); +/// `ycbcr_to_rgb` on eight lanes of luma and centered chroma, before the +/// final unsigned saturation. +#[target_feature(enable = "neon")] +fn ycbcr_lanes_to_rgb( + y: int16x8_t, + cb: int16x8_t, + cr: int16x8_t, +) -> (int16x8_t, int16x8_t, int16x8_t) { + // R = Y + Cr + ((26345·Cr + ROUND) >> 16) + let r_fraction = rounded_high_half( + vmull_n_s16(vget_low_s16(cr), R_CR_REMAINDER), + vmull_high_n_s16(cr, R_CR_REMAINDER), + ); + let r = vaddq_s16(vaddq_s16(y, cr), r_fraction); + + // G = Y - Cr - ((22554·Cb - 18734·Cr + ROUND) >> 16) + let g_fraction = rounded_high_half( + vmlal_n_s16( + vmull_n_s16(vget_low_s16(cb), G_CB), + vget_low_s16(cr), + G_CR_REMAINDER, + ), + vmlal_high_n_s16(vmull_high_n_s16(cb, G_CB), cr, G_CR_REMAINDER), + ); + let g = vsubq_s16(vsubq_s16(y, cr), g_fraction); - neon_memory::store_rgb8(dst, r_bytes, g_bytes, b_bytes); + // B = Y + 2·Cb + ((-14942·Cb + ROUND) >> 16) + let b_fraction = rounded_high_half( + vmull_n_s16(vget_low_s16(cb), B_CB_REMAINDER), + vmull_high_n_s16(cb, B_CB_REMAINDER), + ); + let b = vaddq_s16(vaddq_s16(y, vaddq_s16(cb, cb)), b_fraction); + (r, g, b) } +/// Sixteen pixels per iteration so loads and the interleaving `st3` store use +/// full 128-bit registers; two 8-pixel halves would issue twice the +/// multi-structure store µops. #[target_feature(enable = "neon")] fn fill_rgb_row_from_ycbcr_chunk16_neon( y: &[u8; UPSAMPLED_LANES], @@ -1415,14 +1469,28 @@ fn fill_rgb_row_from_ycbcr_chunk16_neon( cr: &[u8; UPSAMPLED_LANES], dst: &mut [u8; UPSAMPLED_LANES * 3], ) { - let (y, y_remainder) = y.as_chunks::(); - let (cb, cb_remainder) = cb.as_chunks::(); - let (cr, cr_remainder) = cr.as_chunks::(); - let (dst, dst_remainder) = dst.as_chunks_mut::<{ LANES * 3 }>(); - let _ = (y_remainder, cb_remainder, cr_remainder, dst_remainder); - - fill_chunk(&y[0], &cb[0], &cr[0], &mut dst[0]); - fill_chunk(&y[1], &cb[1], &cr[1], &mut dst[1]); + let y = neon_memory::load_u8x16(y); + let cb = neon_memory::load_u8x16(cb); + let cr = neon_memory::load_u8x16(cr); + let bias = vdup_n_u8(128); + // Widening subtraction wraps in `u16`; reinterpreting as `i16` yields the + // exact centered chroma for 8-bit samples. + let (r_lo, g_lo, b_lo) = ycbcr_lanes_to_rgb( + vreinterpretq_s16_u16(vmovl_u8(vget_low_u8(y))), + vreinterpretq_s16_u16(vsubl_u8(vget_low_u8(cb), bias)), + vreinterpretq_s16_u16(vsubl_u8(vget_low_u8(cr), bias)), + ); + let (r_hi, g_hi, b_hi) = ycbcr_lanes_to_rgb( + vreinterpretq_s16_u16(vmovl_high_u8(y)), + vreinterpretq_s16_u16(vsubl_high_u8(cb, vdupq_n_u8(128))), + vreinterpretq_s16_u16(vsubl_high_u8(cr, vdupq_n_u8(128))), + ); + neon_memory::store_rgb8x16( + dst, + vqmovun_high_s16(vqmovun_s16(r_lo), r_hi), + vqmovun_high_s16(vqmovun_s16(g_lo), g_hi), + vqmovun_high_s16(vqmovun_s16(b_lo), b_hi), + ); } #[target_feature(enable = "neon")] @@ -1434,52 +1502,6 @@ fn load_eight(src: &[u8], offset: usize) -> uint8x8_t { neon_memory::load_u8x8(chunk) } -#[target_feature(enable = "neon")] -fn widen_low(values: core::arch::aarch64::uint16x8_t) -> int32x4_t { - vreinterpretq_s32_u32(vmovl_u16(vget_low_u16(values))) -} - -#[target_feature(enable = "neon")] -fn widen_high(values: core::arch::aarch64::uint16x8_t) -> int32x4_t { - vreinterpretq_s32_u32(vmovl_u16(vget_high_u16(values))) -} - -#[target_feature(enable = "neon")] -fn subtract_bias(values: int32x4_t) -> int32x4_t { - vsubq_s32(values, vdupq_n_s32(128)) -} - -#[target_feature(enable = "neon")] -fn fixed_mul_shift(values: int32x4_t, coefficient: i32) -> int32x4_t { - vshrq_n_s32( - vaddq_s32(vmulq_n_s32(values, coefficient), vdupq_n_s32(ROUND)), - 16, - ) -} - -#[target_feature(enable = "neon")] -fn convert_half(y: int32x4_t, cb: int32x4_t, cr: int32x4_t) -> (int32x4_t, int32x4_t, int32x4_t) { - let r = vaddq_s32(y, fixed_mul_shift(cr, FIX_1_40200)); - let g = vsubq_s32( - y, - vshrq_n_s32( - vaddq_s32( - vaddq_s32(vmulq_n_s32(cb, FIX_0_34414), vmulq_n_s32(cr, FIX_0_71414)), - vdupq_n_s32(ROUND), - ), - 16, - ), - ); - let b = vaddq_s32(y, fixed_mul_shift(cb, FIX_1_77200)); - (r, g, b) -} - -#[target_feature(enable = "neon")] -fn pack_eight_u8(low: int32x4_t, high: int32x4_t) -> uint8x8_t { - let words = vcombine_u16(vqmovun_s32(low), vqmovun_s32(high)); - vqmovn_u16(words) -} - #[target_feature(enable = "neon")] fn fill_upsampled_420_chunk( near: &[u8], diff --git a/crates/j2k-jpeg/src/backend/tests.rs b/crates/j2k-jpeg/src/backend/tests.rs index 65be6f8aa..e111fdf19 100644 --- a/crates/j2k-jpeg/src/backend/tests.rs +++ b/crates/j2k-jpeg/src/backend/tests.rs @@ -362,6 +362,26 @@ fn gray_rows_expand_to_equal_rgb_channels() { assert_eq!(dst, vec![10, 10, 10, 40, 40, 40, 90, 90, 90, 200, 200, 200]); } +/// The SIMD converters restructure the fixed-point arithmetic, so check them +/// against the scalar reference for every (Cb, Cr) pair, with luma swept +/// across the range so both saturation edges are exercised. +#[test] +fn detected_backend_ycbcr_matches_scalar_for_every_chroma_pair() { + let backend = super::Backend::detect(); + let cb: Vec = (0..=255u8).collect(); + let mut expected = vec![0u8; 256 * 3]; + let mut actual = vec![0u8; 256 * 3]; + for cr_value in 0..=255u8 { + let cr = vec![cr_value; 256]; + for y_value in (0..=255u8).step_by(15).chain([1, 128, 254, 255]) { + let y = vec![y_value; 256]; + scalar::fill_rgb_row_from_ycbcr(&y, &cb, &cr, &mut expected); + backend.fill_rgb_row_from_ycbcr(&y, &cb, &cr, &mut actual); + assert_eq!(actual, expected, "y={y_value} cr={cr_value}"); + } + } +} + #[test] fn ycbcr_rows_match_per_pixel_reference() { let y = [16u8, 40, 90, 200]; diff --git a/crates/j2k-jpeg/src/entropy/block.rs b/crates/j2k-jpeg/src/entropy/block.rs index 0b32d46c7..ee9de6cbd 100644 --- a/crates/j2k-jpeg/src/entropy/block.rs +++ b/crates/j2k-jpeg/src/entropy/block.rs @@ -25,8 +25,6 @@ use crate::entropy::ZIGZAG; use crate::error::{HuffmanFailure, JpegError}; use crate::internal::bit_reader::BitReader; -const DENSE_CLEAR_THRESHOLD: usize = 4; - #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(crate) enum BlockActivity { DcOnly, @@ -56,27 +54,21 @@ impl ReducedIdctCoefficients { } } -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -enum ClearMode { - Sparse, - Dense, -} - #[derive(Debug, Clone)] pub(crate) struct CoefficientBlock { coeffs: [i16; 64], - touched: [u8; 64], - touched_len: usize, - clear_mode: ClearMode, + /// Set by any AC store. Clearing a block that holds AC coefficients + /// zeroes all 128 bytes (a handful of vector stores); a DC-only block + /// clears just the DC. This keeps per-coefficient bookkeeping to a single + /// flag store instead of an index list maintained in memory. + ac_written: bool, } impl Default for CoefficientBlock { fn default() -> Self { Self { coeffs: [0; 64], - touched: [0; 64], - touched_len: 0, - clear_mode: ClearMode::Sparse, + ac_written: false, } } } @@ -88,16 +80,12 @@ impl CoefficientBlock { )] #[inline(always)] fn clear_touched(&mut self) { - match self.clear_mode { - ClearMode::Sparse => { - for &idx in &self.touched[..self.touched_len] { - self.coeffs[idx as usize] = 0; - } - } - ClearMode::Dense => self.coeffs.fill(0), + if self.ac_written { + self.coeffs = [0; 64]; + self.ac_written = false; + } else { + self.coeffs[0] = 0; } - self.touched_len = 0; - self.clear_mode = ClearMode::Sparse; } #[expect( @@ -105,20 +93,18 @@ impl CoefficientBlock { reason = "measured entropy-block hot path requires cross-helper inlining" )] #[inline(always)] + fn store_dc(&mut self, value: i16) { + self.coeffs[0] = value; + } + #[expect( - clippy::cast_possible_truncation, - reason = "coefficient indices are bounded to the 64-entry JPEG block" + clippy::inline_always, + reason = "measured entropy-block hot path requires cross-helper inlining" )] + #[inline(always)] fn store(&mut self, idx: usize, value: i16) { self.coeffs[idx] = value; - if self.clear_mode == ClearMode::Sparse { - if self.touched_len < DENSE_CLEAR_THRESHOLD { - self.touched[self.touched_len] = idx as u8; - self.touched_len += 1; - } else { - self.clear_mode = ClearMode::Dense; - } - } + self.ac_written = true; } #[expect( @@ -140,17 +126,20 @@ impl CoefficientBlock { } } +/// Classify a block from the OR of every stored AC coefficient's natural +/// index. AC indices are nonzero, so any AC store makes the OR nonzero, and +/// bit 5 is set exactly when some index is at least 32 (rows 4-7). Keeping +/// one OR per coefficient replaces a per-coefficient state-machine update. #[expect( clippy::inline_always, reason = "measured entropy-block hot path requires cross-helper inlining" )] #[inline(always)] -fn extend_activity(activity: BlockActivity, natural_idx: usize) -> BlockActivity { - if natural_idx < 32 { - match activity { - BlockActivity::DcOnly | BlockActivity::BottomHalfZero => BlockActivity::BottomHalfZero, - BlockActivity::General => BlockActivity::General, - } +fn activity_from_ac_indices(ac_index_or: usize) -> BlockActivity { + if ac_index_or == 0 { + BlockActivity::DcOnly + } else if ac_index_or & 32 == 0 { + BlockActivity::BottomHalfZero } else { BlockActivity::General } @@ -196,9 +185,9 @@ pub(crate) fn decode_block_with_activity( *prev_dc = prev_dc.wrapping_add(diff); // Dequant the DC in natural-order position 0 (zigzag index 0 → natural 0). let dc_dequant = (*prev_dc).wrapping_mul(i32::from(quant[0])); - block.store(0, clamp_i16(dc_dequant)); + block.store_dc(clamp_i16(dc_dequant)); - let mut activity = BlockActivity::DcOnly; + let mut ac_index_or = 0usize; drive_ac_fast::(br, ac_table, |k, ac| { let natural_idx = ZIGZAG[k] as usize; // Quant table entries are stored in zigzag order per T.81 §B.2.4.1, @@ -206,10 +195,10 @@ pub(crate) fn decode_block_with_activity( let value = ac_decoded_value(ac); let dequant = value.wrapping_mul(i32::from(quant[k])); block.store(natural_idx, clamp_i16(dequant)); - activity = extend_activity(activity, natural_idx); + ac_index_or |= natural_idx; Ok(()) })?; - Ok(activity) + Ok(activity_from_ac_indices(ac_index_or)) } /// Decode one dequantized block directly into a freshly zeroed output block. @@ -265,21 +254,21 @@ pub(crate) fn decode_block_quantized_and_dequantized_with_activity( // DC. let diff = dc_table.decode_fast_dc(br)?; *prev_dc = prev_dc.wrapping_add(diff); - quantized_block.store(0, clamp_i16(*prev_dc)); + quantized_block.store_dc(clamp_i16(*prev_dc)); let dc_dequant = (*prev_dc).wrapping_mul(i32::from(quant[0])); - dequantized_block.store(0, clamp_i16(dc_dequant)); + dequantized_block.store_dc(clamp_i16(dc_dequant)); - let mut activity = BlockActivity::DcOnly; + let mut ac_index_or = 0usize; drive_ac_fast::(br, ac_table, |k, ac| { let natural_idx = ZIGZAG[k] as usize; let value = ac_decoded_value(ac); quantized_block.store(natural_idx, clamp_i16(value)); let dequant = value.wrapping_mul(i32::from(quant[k])); dequantized_block.store(natural_idx, clamp_i16(dequant)); - activity = extend_activity(activity, natural_idx); + ac_index_or |= natural_idx; Ok(()) })?; - Ok(activity) + Ok(activity_from_ac_indices(ac_index_or)) } #[expect( @@ -301,7 +290,7 @@ pub(crate) fn decode_block_with_dc_status( let diff = dc_table.decode_fast_dc(br)?; *prev_dc = prev_dc.wrapping_add(diff); let dc_dequant = (*prev_dc).wrapping_mul(i32::from(quant[0])); - block.store(0, clamp_i16(dc_dequant)); + block.store_dc(clamp_i16(dc_dequant)); let mut dc_only = true; drive_ac_fast::(br, ac_table, |k, ac| { @@ -334,7 +323,7 @@ pub(crate) fn decode_block_for_reduced_idct( let diff = dc_table.decode_fast_dc(br)?; *prev_dc = prev_dc.wrapping_add(diff); let dc_dequant = (*prev_dc).wrapping_mul(i32::from(quant[0])); - block.store(0, clamp_i16(dc_dequant)); + block.store_dc(clamp_i16(dc_dequant)); let mut dc_only_for_reduced_idct = true; drive_ac_fast::(br, ac_table, |k, ac| { @@ -368,7 +357,7 @@ pub(crate) fn decode_block_for_1x1_idct( let diff = dc_table.decode_fast_dc(br)?; *prev_dc = prev_dc.wrapping_add(diff); let dc_dequant = (*prev_dc).wrapping_mul(i32::from(quant[0])); - block.store(0, clamp_i16(dc_dequant)); + block.store_dc(clamp_i16(dc_dequant)); drive_ac_fast::(br, ac_table, |_, _| Ok(()))?; Ok(()) @@ -793,58 +782,47 @@ mod tests { } #[test] - fn extend_activity_promotes_top_half_ac_without_marking_general() { - assert_eq!( - extend_activity(BlockActivity::DcOnly, 31), - BlockActivity::BottomHalfZero - ); - assert_eq!( - extend_activity(BlockActivity::BottomHalfZero, 7), - BlockActivity::BottomHalfZero - ); + fn top_half_ac_indices_classify_as_bottom_half_zero() { + for indices in [&[31usize][..], &[7], &[1, 8, 31], &[3, 17, 26]] { + let or = indices.iter().fold(0, |acc, &idx| acc | idx); + assert_eq!( + activity_from_ac_indices(or), + BlockActivity::BottomHalfZero, + "{indices:?}" + ); + } + assert_eq!(activity_from_ac_indices(0), BlockActivity::DcOnly); } #[test] - fn extend_activity_marks_bottom_half_ac_as_general() { - assert_eq!( - extend_activity(BlockActivity::DcOnly, 32), - BlockActivity::General - ); - assert_eq!( - extend_activity(BlockActivity::BottomHalfZero, 40), - BlockActivity::General - ); + fn any_bottom_half_ac_index_classifies_as_general() { + for indices in [&[32usize][..], &[40], &[1, 63], &[7, 31, 32]] { + let or = indices.iter().fold(0, |acc, &idx| acc | idx); + assert_eq!( + activity_from_ac_indices(or), + BlockActivity::General, + "{indices:?}" + ); + } } #[test] - fn switches_to_dense_clear_after_threshold_and_zeroes_full_block() { + fn clear_zeroes_every_coefficient_after_ac_stores() { let mut block = CoefficientBlock::default(); - for (i, idx) in [0usize, 1, 8, 16, 24].into_iter().enumerate() { + block.store_dc(9); + for (i, idx) in [1usize, 8, 16, 24, 63].into_iter().enumerate() { block.store( idx, i16::try_from(i + 1).expect("fixture coefficient fits in i16"), ); } - assert_eq!(block.clear_mode, ClearMode::Dense); block.clear_touched(); - assert!(block.coefficients().iter().all(|&c| c == 0)); - assert_eq!(block.touched_len, 0); - assert_eq!(block.clear_mode, ClearMode::Sparse); - } - #[test] - fn stays_sparse_below_dense_clear_threshold() { - let mut block = CoefficientBlock::default(); - for (i, idx) in [0usize, 2, 4, 6].into_iter().enumerate() { - block.store( - idx, - i16::try_from(i + 1).expect("fixture coefficient fits in i16"), - ); - } - - assert_eq!(block.clear_mode, ClearMode::Sparse); - assert_eq!(block.touched_len, DENSE_CLEAR_THRESHOLD); + // A DC-only block that follows must also clear back to all zeroes. + block.store_dc(-4); + block.clear_touched(); + assert!(block.coefficients().iter().all(|&c| c == 0)); } } diff --git a/crates/j2k-jpeg/src/entropy/sequential/deposit.rs b/crates/j2k-jpeg/src/entropy/sequential/deposit.rs index e307b3caa..05e6b919b 100644 --- a/crates/j2k-jpeg/src/entropy/sequential/deposit.rs +++ b/crates/j2k-jpeg/src/entropy/sequential/deposit.rs @@ -169,6 +169,59 @@ pub(super) struct FastTile420EntropyState<'a, 'b> { pub(super) coeff: &'a mut CoefficientBlock, } +/// Register-friendly copy of a [`FastTile420EntropyState`]'s reader and DC +/// predictors. Behind `&mut` fields the accumulator round-trips through +/// memory on every Huffman symbol; decoding a row against this local (with +/// the row kernel inlined, so it never escapes) lets LLVM keep it in +/// registers. Callers restore it even when the row fails. +pub(super) struct LocalEntropy<'b> { + br: BitReader<'b>, + dc: [i32; 3], +} + +impl<'b> LocalEntropy<'b> { + #[expect( + clippy::inline_always, + reason = "must inline so the borrowed locals never escape the row frame" + )] + #[inline(always)] + pub(super) fn state<'a>( + &'a mut self, + coeff: &'a mut CoefficientBlock, + ) -> FastTile420EntropyState<'a, 'b> { + let [y, cb, cr] = &mut self.dc; + FastTile420EntropyState { + br: &mut self.br, + dc: FastTile420DcState { y, cb, cr }, + coeff, + } + } +} + +impl<'b> FastTile420EntropyState<'_, 'b> { + #[expect( + clippy::inline_always, + reason = "must inline so the copied reader stays register-resident" + )] + #[inline(always)] + pub(super) fn take_local(&self) -> LocalEntropy<'b> { + LocalEntropy { + br: self.br.clone(), + dc: [*self.dc.y, *self.dc.cb, *self.dc.cr], + } + } + + #[expect( + clippy::inline_always, + reason = "must inline so the copied reader stays register-resident" + )] + #[inline(always)] + pub(super) fn restore_local(&mut self, local: LocalEntropy<'b>) { + *self.br = local.br; + [*self.dc.y, *self.dc.cb, *self.dc.cr] = local.dc; + } +} + #[derive(Clone, Copy)] pub(super) struct FastTile420Window { pub(super) mcus_per_row: u32, diff --git a/crates/j2k-jpeg/src/entropy/sequential/fast420/rows.rs b/crates/j2k-jpeg/src/entropy/sequential/fast420/rows.rs index 45d35ec57..5a7e5d909 100644 --- a/crates/j2k-jpeg/src/entropy/sequential/fast420/rows.rs +++ b/crates/j2k-jpeg/src/entropy/sequential/fast420/rows.rs @@ -209,11 +209,39 @@ pub(super) fn decode_mcu_row_fast_tile_420_scaled( Ok(()) } +pub(in crate::entropy::sequential) fn decode_mcu_row_fast_tile_420( + components: FastTile420Components<'_>, + backend: Backend, + state: &mut FastTile420EntropyState<'_, '_>, + pixels: &mut [u8; 64], + window: FastTile420Window, + stripe: &mut StripeBuffer, + profiler: &mut impl Fast420Profiler, +) -> Result<(), JpegError> { + let mut local = state.take_local(); + let result = decode_mcu_row_fast_tile_420_local( + components, + backend, + &mut local.state(&mut *state.coeff), + pixels, + window, + stripe, + profiler, + ); + state.restore_local(local); + result +} + #[expect( clippy::too_many_lines, reason = "the MCU-row kernel keeps six block decodes and plane deposits in JPEG sampling order" )] -pub(in crate::entropy::sequential) fn decode_mcu_row_fast_tile_420( +#[expect( + clippy::inline_always, + reason = "inlining into the row wrapper is what keeps the local reader out of memory" +)] +#[inline(always)] +fn decode_mcu_row_fast_tile_420_local( components: FastTile420Components<'_>, backend: Backend, state: &mut FastTile420EntropyState<'_, '_>, diff --git a/crates/j2k-jpeg/src/entropy/sequential/generic/row.rs b/crates/j2k-jpeg/src/entropy/sequential/generic/row.rs index 576980465..7a18ec38c 100644 --- a/crates/j2k-jpeg/src/entropy/sequential/generic/row.rs +++ b/crates/j2k-jpeg/src/entropy/sequential/generic/row.rs @@ -40,11 +40,49 @@ pub(super) struct McuRowState<'a, 'b> { pub(super) expected_rst: &'a mut u8, } +pub(super) fn decode_mcu_row( + context: &McuRowContext<'_>, + state: &mut McuRowState<'_, '_>, + mcu_y: u32, + stripe: &mut StripeBuffer, +) -> Result<(), JpegError> { + // Decode against locals so the bit accumulator and DC predictors live in + // registers for the whole row instead of round-tripping through `state` + // on every symbol. They are written back even when the row fails. + let mut prev_dc = [0i32; 4]; + let Some(local_dc) = prev_dc.get_mut(..state.prev_dc.len()) else { + return decode_mcu_row_local(context, state, mcu_y, stripe); + }; + local_dc.copy_from_slice(state.prev_dc); + let mut br = state.br.clone(); + let result = decode_mcu_row_local( + context, + &mut McuRowState { + br: &mut br, + prev_dc: &mut *local_dc, + coeff: &mut *state.coeff, + pixels: &mut *state.pixels, + mcus_since_restart: &mut *state.mcus_since_restart, + expected_rst: &mut *state.expected_rst, + }, + mcu_y, + stripe, + ); + *state.br = br; + state.prev_dc.copy_from_slice(local_dc); + result +} + #[expect( clippy::too_many_lines, reason = "the MCU kernel traverses component sampling factors while preserving entropy, predictor, IDCT, and deposit order" )] -pub(super) fn decode_mcu_row( +#[expect( + clippy::inline_always, + reason = "inlining into the row wrapper is what keeps the local reader out of memory" +)] +#[inline(always)] +fn decode_mcu_row_local( context: &McuRowContext<'_>, state: &mut McuRowState<'_, '_>, mcu_y: u32, diff --git a/crates/j2k-jpeg/src/entropy/sequential/restart.rs b/crates/j2k-jpeg/src/entropy/sequential/restart.rs index ae2f3b28b..55df2bd07 100644 --- a/crates/j2k-jpeg/src/entropy/sequential/restart.rs +++ b/crates/j2k-jpeg/src/entropy/sequential/restart.rs @@ -148,6 +148,11 @@ pub(super) struct McuPosition { pub(super) total: u32, } +#[expect( + clippy::inline_always, + reason = "the due check runs once per MCU and must inline so callers' local readers stay in registers" +)] +#[inline(always)] pub(super) fn consume_restart_marker_if_due( br: &mut BitReader<'_>, restart: u16, @@ -159,7 +164,11 @@ pub(super) fn consume_restart_marker_if_due( return Ok(false); } - *expected_rst = br.consume_restart_marker(*expected_rst, position.current, position.total)?; + let (restarted, next_rst) = + br.clone() + .restarted(*expected_rst, position.current, position.total)?; + *br = restarted; + *expected_rst = next_rst; Ok(true) } diff --git a/crates/j2k-jpeg/src/entropy/sequential/rgb444.rs b/crates/j2k-jpeg/src/entropy/sequential/rgb444.rs index 9bfca5456..e002f1c1c 100644 --- a/crates/j2k-jpeg/src/entropy/sequential/rgb444.rs +++ b/crates/j2k-jpeg/src/entropy/sequential/rgb444.rs @@ -115,22 +115,47 @@ struct FastRgb444McuRowState<'a, 'b> { expected_rst: &'a mut u8, } +fn decode_mcu_row_fast_rgb_444( + context: &FastRgb444McuRowContext<'_>, + state: &mut FastRgb444McuRowState<'_, '_>, + mcu_y: u32, + stripe: &mut StripeBuffer, +) -> Result<(), JpegError> { + // Decode against locals so the bit accumulator and DC predictors live in + // registers for the whole row instead of round-tripping through `state` + // on every symbol. They are written back even when the row fails. + let mut br = state.br.clone(); + let mut dc = [*state.y_dc, *state.cb_dc, *state.cr_dc]; + let result = decode_mcu_row_fast_rgb_444_local(context, state, &mut br, &mut dc, mcu_y, stripe); + *state.br = br; + [*state.y_dc, *state.cb_dc, *state.cr_dc] = dc; + result +} + #[expect( clippy::too_many_lines, reason = "the 4:4:4 MCU kernel keeps three block decodes and direct RGB row deposition in sampling order" )] -fn decode_mcu_row_fast_rgb_444( +#[expect( + clippy::inline_always, + reason = "inlining into the row wrapper is what keeps the local reader out of memory" +)] +#[inline(always)] +fn decode_mcu_row_fast_rgb_444_local( context: &FastRgb444McuRowContext<'_>, state: &mut FastRgb444McuRowState<'_, '_>, + br: &mut BitReader<'_>, + dc: &mut [i32; 3], mcu_y: u32, stripe: &mut StripeBuffer, ) -> Result<(), JpegError> { + let [y_dc, cb_dc, cr_dc] = dc; for plane_idx in 0..3 { assert_stripe_deposit_capacity(stripe, plane_idx, 1, 1, context.mcus_per_row, 8); } for mx in 0..context.mcus_per_row { if consume_restart_marker_if_due( - state.br, + br, context.restart, *state.mcus_since_restart, state.expected_rst, @@ -139,19 +164,19 @@ fn decode_mcu_row_fast_rgb_444( total: context.mcu_rows * context.mcus_per_row, }, )? { - *state.y_dc = 0; - *state.cb_dc = 0; - *state.cr_dc = 0; + *y_dc = 0; + *cb_dc = 0; + *cr_dc = 0; *state.mcus_since_restart = 0; } let block_x = mx * 8; let y_activity = decode_block_with_activity( - state.br, + br, context.y_comp.dc_table, context.y_comp.ac_table, - state.y_dc, + y_dc, context.y_comp.quant, state.coeff, )?; @@ -177,10 +202,10 @@ fn decode_mcu_row_fast_rgb_444( ); let cb_activity = decode_block_with_activity( - state.br, + br, context.cb_comp.dc_table, context.cb_comp.ac_table, - state.cb_dc, + cb_dc, context.cb_comp.quant, state.coeff, )?; @@ -206,10 +231,10 @@ fn decode_mcu_row_fast_rgb_444( ); let cr_activity = decode_block_with_activity( - state.br, + br, context.cr_comp.dc_table, context.cr_comp.ac_table, - state.cr_dc, + cr_dc, context.cr_comp.quant, state.coeff, )?; diff --git a/crates/j2k-jpeg/src/internal/bit_reader.rs b/crates/j2k-jpeg/src/internal/bit_reader.rs index 4069d080c..2e83f1109 100644 --- a/crates/j2k-jpeg/src/internal/bit_reader.rs +++ b/crates/j2k-jpeg/src/internal/bit_reader.rs @@ -34,18 +34,27 @@ pub(crate) struct BitReaderSnapshot { pub(crate) bits: u8, } +/// Hot loops copy the reader into a local, decode, and write it back so the +/// accumulator stays in registers; `Clone` (not `Copy`) keeps those copies +/// explicit. +#[derive(Clone)] pub(crate) struct BitReader<'a> { bytes: &'a [u8], /// Cursor into `bytes`. Always either (a) past the last consumed byte, or /// (b) pointing at the leading `0xFF` of a marker the refill paused at. pos: usize, /// MSB-first bit accumulator. The `bits` most-significant bits contain - /// the next coded bits; lower (`64 - bits`) bits are zero. + /// the next coded bits. The lower (`64 - bits`) bits are either zero or a + /// prefix of the not-yet-counted byte at `pos`, preloaded by the bulk + /// refill; every later refill ORs identical values into those positions. acc: u64, /// Number of valid bits in `acc`, 0..=64. bits: u8, - /// Synthetic trailing one bits appended only for terminal Huffman lookahead. - synthetic_bits: u8, + /// Synthetic one bits appended for terminal Huffman lookahead since the + /// last reset. Real bits are always consumed before synthetic ones, so + /// the unread synthetic count is `min(synthetic_pushed, bits)`. Deriving + /// it keeps this bookkeeping out of `consume_bits`. + synthetic_pushed: u8, /// Set when refill stopped at a marker. Cleared by [`Self::take_marker`]. marker: Option, /// Cursor immediately after the observed marker code, including FF fill. @@ -68,7 +77,7 @@ impl<'a> BitReader<'a> { pos: 0, acc: 0, bits: 0, - synthetic_bits: 0, + synthetic_pushed: 0, marker: None, marker_end: 0, allow_eof_padding, @@ -83,18 +92,56 @@ impl<'a> BitReader<'a> { )] #[inline(always)] pub(crate) fn ensure_bits(&mut self, n: u8) -> Result<(), JpegError> { + if self.bits >= n || self.try_refill_bulk() { + return Ok(()); + } + *self = self.clone().ensure_bits_slow(n)?; + Ok(()) + } + + /// Byte-at-a-time refill for windows holding `0xFF`, markers, or the + /// input tail. Taking and returning the reader by value keeps callers' + /// local readers from escaping, so their accumulators stay in registers. + #[cold] + #[inline(never)] + fn ensure_bits_slow(mut self, n: u8) -> Result { while self.bits < n { if !self.refill_one_byte() { - if self.bits >= n { - return Ok(()); - } return Err(JpegError::HuffmanDecode { mcu: 0, reason: HuffmanFailure::TableExhausted, }); } } - Ok(()) + Ok(self) + } + + /// Branch-light refill: when the next eight input bytes contain no `0xFF` + /// (no stuffing and no marker), OR them in with one big-endian load and + /// count every whole byte that fits, leaving `bits` in `56..=63`. Returns + /// `false`, changing nothing, when the window is short or holds `0xFF`. + #[inline(always)] + fn try_refill_bulk(&mut self) -> bool { + const ONES: u64 = 0x0101_0101_0101_0101; + const HIGHS: u64 = 0x8080_8080_8080_8080; + let Some(window) = self + .bytes + .get(self.pos..) + .and_then(|tail| tail.first_chunk::<8>()) + else { + return false; + }; + let word = u64::from_be_bytes(*window); + // A byte of `word` is 0xFF exactly when that byte of `!word` is zero. + let inverted = !word; + if inverted.wrapping_sub(ONES) & word & HIGHS != 0 { + return false; + } + self.acc |= word >> self.bits; + let whole_bytes = (ACC_BITS - 1 - self.bits) >> 3; + self.pos += usize::from(whole_bytes); + self.bits += whole_bytes << 3; + true } /// Fill the accumulator to at least `n` bits, padding with `1` bits when @@ -114,6 +161,16 @@ impl<'a> BitReader<'a> { )] #[inline(always)] pub(crate) fn ensure_bits_padded(&mut self, n: u8) -> Result<(), JpegError> { + if self.bits >= n || self.try_refill_bulk() { + return Ok(()); + } + *self = self.clone().ensure_bits_padded_slow(n)?; + Ok(()) + } + + #[cold] + #[inline(never)] + fn ensure_bits_padded_slow(mut self, n: u8) -> Result { let mut refilled = false; while self.bits < n { if !self.refill_one_byte() { @@ -128,16 +185,16 @@ impl<'a> BitReader<'a> { while self.bits < n { self.acc |= 1u64 << (ACC_BITS - 1 - self.bits); self.bits += 1; - self.synthetic_bits += 1; + self.synthetic_pushed = self.synthetic_pushed.saturating_add(1); } - return Ok(()); + return Ok(self); } refilled = true; } if refilled { self.refill_to_threshold(); } - Ok(()) + Ok(self) } /// Refill one byte of data into the accumulator. Returns `true` if a @@ -228,10 +285,6 @@ impl<'a> BitReader<'a> { "consume_bits({n}) with only {} buffered", self.bits ); - let real_bits = self.bits - self.synthetic_bits; - if n > real_bits { - self.synthetic_bits -= n - real_bits; - } self.acc <<= n; self.bits -= n; } @@ -248,7 +301,9 @@ impl<'a> BitReader<'a> { /// After consuming bits, top up the accumulator so the next Huffman peek /// can always examine 16 bits without further refill. fn refill_to_threshold(&mut self) { - while self.bits < REFILL_THRESHOLD && self.refill_one_byte() {} + if self.bits < REFILL_THRESHOLD && !self.try_refill_bulk() { + while self.bits < REFILL_THRESHOLD && self.refill_one_byte() {} + } } /// Signed-value extension per T.81 §F.2.2.1 ("EXTEND" procedure). `ssss` @@ -288,6 +343,21 @@ impl<'a> BitReader<'a> { Some(m) } + /// By-value [`Self::consume_restart_marker`] for hot loops that keep the + /// reader in a local: passing `&mut` to an out-of-line call would force + /// the whole reader into memory for the entire loop. + #[cold] + #[inline(never)] + pub(crate) fn restarted( + mut self, + expected_rst: u8, + mcu_at: u32, + mcu_total: u32, + ) -> Result<(Self, u8), JpegError> { + let next_rst = self.consume_restart_marker(expected_rst, mcu_at, mcu_total)?; + Ok((self, next_rst)) + } + /// Consume the next restart marker and return the following RST index. /// /// Entropy decoders normally prefetch far enough to observe the marker. @@ -327,7 +397,7 @@ impl<'a> BitReader<'a> { } pub(crate) fn snapshot(&self) -> BitReaderSnapshot { - let real_bits = self.bits.saturating_sub(self.synthetic_bits); + let real_bits = self.unread_real_bits(); let acc = if real_bits == 0 { 0 } else { @@ -345,7 +415,7 @@ impl<'a> BitReader<'a> { pub(crate) fn reset_at_restart(&mut self) { self.acc = 0; self.bits = 0; - self.synthetic_bits = 0; + self.synthetic_pushed = 0; } pub(crate) fn from_snapshot(bytes: &'a [u8], snapshot: BitReaderSnapshot) -> Self { @@ -354,7 +424,7 @@ impl<'a> BitReader<'a> { pos: snapshot.pos, acc: snapshot.acc, bits: snapshot.bits, - synthetic_bits: 0, + synthetic_pushed: 0, marker: None, marker_end: 0, allow_eof_padding: false, @@ -594,4 +664,95 @@ mod tests { assert_eq!(br.consume_restart_marker(0, 1, 2).unwrap(), 1); assert_eq!(br.snapshot().bits, 0); } + + /// Deterministic xorshift so failures reproduce from the printed seed. + fn next_random(state: &mut u64) -> u64 { + *state ^= *state << 13; + *state ^= *state >> 7; + *state ^= *state << 17; + *state + } + + /// Byte-stuffs `payload` (every data `0xFF` becomes `0xFF 0x00`) and + /// terminates it with `marker`, like an entropy segment. + fn stuffed_segment(payload: &[u8], marker: u8) -> Vec { + let mut segment = Vec::with_capacity(payload.len() * 2 + 2); + for &byte in payload { + segment.push(byte); + if byte == 0xff { + segment.push(0x00); + } + } + segment.extend_from_slice(&[0xff, marker]); + segment + } + + fn reference_bits(payload: &[u8], bit_pos: usize, n: u8) -> u32 { + (0..usize::from(n)).fold(0u32, |value, i| { + let bit = payload + .get((bit_pos + i) / 8) + .map_or(1, |byte| (byte >> (7 - (bit_pos + i) % 8)) & 1); + (value << 1) | u32::from(bit) + }) + } + + #[test] + fn bulk_refill_matches_bitwise_reference_across_stuffing_and_markers() { + for seed in 1..=400u64 { + let mut rng = seed.wrapping_mul(0x9e37_79b9_7f4a_7c15) | 1; + let len = (next_random(&mut rng) % 48) as usize; + // Dense 0xFF so stuffing lands at every offset of the 8-byte window. + let payload: Vec = (0..len) + .map(|_| match next_random(&mut rng) % 4 { + 0 => 0xff, + _ => next_random(&mut rng).to_le_bytes()[0], + }) + .collect(); + let marker = if seed % 2 == 0 { 0xd9 } else { 0xd0 }; + let segment = stuffed_segment(&payload, marker); + let mut br = BitReader::new(&segment); + + let total_bits = payload.len() * 8; + let mut bit_pos = 0usize; + while bit_pos < total_bits + 24 { + let n = 1 + (next_random(&mut rng) % 16) as u8; + br.ensure_bits_padded(n) + .unwrap_or_else(|err| panic!("seed {seed}: padded ensure failed: {err:?}")); + assert_eq!( + br.peek_bits(n), + reference_bits(&payload, bit_pos, n), + "seed {seed}: {n} bits at bit {bit_pos}" + ); + br.consume_bits(n); + bit_pos += usize::from(n); + assert_eq!( + usize::from(br.unread_real_bits()), + usize::from(br.bits).min(total_bits.saturating_sub(bit_pos)), + "seed {seed}: real-bit accounting at bit {bit_pos}" + ); + } + assert_eq!(br.take_marker(), Some(marker), "seed {seed}"); + assert_eq!(br.position(), segment.len(), "seed {seed}"); + } + } + + #[test] + fn bulk_refill_snapshot_restores_only_counted_bits() { + let payload: Vec = (0u8..=40).map(|i| i.wrapping_mul(37)).collect(); + let segment = stuffed_segment(&payload, 0xd9); + let mut br = BitReader::new(&segment); + for n in [3u8, 11, 7, 16, 1, 9, 13] { + br.ensure_bits_padded(n).unwrap(); + br.consume_bits(n); + let snapshot = br.snapshot(); + let mut original = br.clone(); + let mut restored = BitReader::from_snapshot(&segment, snapshot); + for width in [5u8, 12, 16, 2] { + assert_eq!( + restored.read_bits(width).unwrap(), + original.read_bits(width).unwrap() + ); + } + } + } } diff --git a/crates/j2k-jpeg/src/internal/bit_reader/terminal.rs b/crates/j2k-jpeg/src/internal/bit_reader/terminal.rs index 4f6a7b278..910d31791 100644 --- a/crates/j2k-jpeg/src/internal/bit_reader/terminal.rs +++ b/crates/j2k-jpeg/src/internal/bit_reader/terminal.rs @@ -16,7 +16,7 @@ impl BitReader<'_> { } pub(crate) fn unread_real_bits(&self) -> u8 { - self.bits - self.synthetic_bits + self.bits.saturating_sub(self.synthetic_pushed) } pub(crate) fn unread_real_bits_are_ones(&self) -> bool { diff --git a/crates/j2k-jpeg/src/simd/neon_memory.rs b/crates/j2k-jpeg/src/simd/neon_memory.rs index 10aafad9f..7826e1a3f 100644 --- a/crates/j2k-jpeg/src/simd/neon_memory.rs +++ b/crates/j2k-jpeg/src/simd/neon_memory.rs @@ -3,7 +3,8 @@ //! Fixed-size `AArch64` NEON memory operations. use core::arch::aarch64::{ - int16x8_t, uint8x16_t, uint8x8_t, uint8x8x3_t, vld1_u8, vld1q_s16, vst1_u8, vst1q_u8, vst3_u8, + int16x8_t, uint8x16_t, uint8x16x3_t, uint8x8_t, uint8x8x3_t, vld1_u8, vld1q_s16, vld1q_u8, + vst1_u8, vst1q_u8, vst3_u8, vst3q_u8, }; #[inline] @@ -81,3 +82,32 @@ pub(crate) fn store_rgb8(dst: &mut [u8; 24], red: uint8x8_t, green: uint8x8_t, b // - Initialization: the store initializes all twenty-four output bytes. unsafe { vst3_u8(dst.as_mut_ptr(), uint8x8x3_t(red, green, blue)) }; } + +#[inline] +#[target_feature(enable = "neon")] +pub(crate) fn load_u8x16(src: &[u8; 16]) -> uint8x16_t { + // SAFETY: + // - Feature availability: callers run inside a `Neon` token kernel. + // - Bounds: the array reference proves sixteen readable bytes. + // - Alignment: AArch64 `vld1q_u8` supports unaligned byte addresses. + // - Aliasing: the shared reference permits reads and no writes occur. + // - Initialization: all bytes behind a Rust reference are initialized. + unsafe { vld1q_u8(src.as_ptr()) } +} + +#[inline] +#[target_feature(enable = "neon")] +pub(crate) fn store_rgb8x16( + dst: &mut [u8; 48], + red: uint8x16_t, + green: uint8x16_t, + blue: uint8x16_t, +) { + // SAFETY: + // - Feature availability: callers run inside a `Neon` token kernel. + // - Bounds: the array reference proves space for sixteen three-byte pixels. + // - Alignment: AArch64 `vst3q_u8` supports unaligned byte addresses. + // - Aliasing: the exclusive reference prevents overlapping live access. + // - Initialization: the store initializes all forty-eight output bytes. + unsafe { vst3q_u8(dst.as_mut_ptr(), uint8x16x3_t(red, green, blue)) }; +} From d6aa6d47d26804ae6222e98d28a92acd935a3866 Mon Sep 17 00:00:00 2001 From: GF Date: Sat, 26 Sep 2026 03:59:10 -0400 Subject: [PATCH 2/5] docs: note fearless_simd 1.0 upgrade in changelog --- CHANGELOG.md | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3ee558d9f..90d2a2b7c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,6 +3,14 @@ This changelog tracks the current release line. Historical phase notes and stale roadmap entries have been removed from the public documentation set. +## [Unreleased] + +- Upgrades `fearless_simd` to 1.0.0 for the `j2k-jpeg` AArch64 NEON kernels + and the optional `j2k-native` `simd` feature. The dependency is private, so + public APIs are unchanged. `j2k-jpeg` keeps its private exact-AVX2 token on + x86-64 because `fearless_simd::Avx2` still represents the broader x86-64-v3 + feature set. + ## [0.11.1] - 2026-09-12 - Adds full-image and batched Metal reconstruction for sampled RGB components without MCT, including odd dimensions and 4:2:2/4:2:0 sampling. From 6ee8cf91ac27b40cc0d90c12bf0fd2d4a958e8ce Mon Sep 17 00:00:00 2001 From: GF Date: Sat, 26 Sep 2026 20:19:38 -0400 Subject: [PATCH 3/5] perf(jpeg): optimize extended, lossless, and progressive decoding --- crates/j2k-jpeg/src/backend/mod.rs | 51 ++ crates/j2k-jpeg/src/backend/neon_upsample.rs | 72 ++ crates/j2k-jpeg/src/backend/tests.rs | 86 ++ crates/j2k-jpeg/src/color/upsample.rs | 49 +- crates/j2k-jpeg/src/context.rs | 71 +- crates/j2k-jpeg/src/decoder.rs | 13 +- crates/j2k-jpeg/src/decoder/allocation.rs | 17 +- .../j2k-jpeg/src/decoder/extended12/planes.rs | 12 +- .../src/decoder/extended12/progressive.rs | 2 +- .../extended12/progressive/color444.rs | 6 +- .../extended12/progressive/four_component.rs | 2 +- .../extended12/progressive/subsampled.rs | 2 +- .../src/decoder/extended12/sequential.rs | 1 + .../decoder/extended12/sequential/color444.rs | 1 + .../extended12/sequential/four_component.rs | 9 +- .../extended12/sequential/subsampled.rs | 8 +- .../j2k-jpeg/src/decoder/extended12/state.rs | 18 +- .../src/decoder/extended12/writers.rs | 512 +++++++---- .../src/decoder/extended12/writers/tests.rs | 261 ++++++ .../j2k-jpeg/src/decoder/lossless_helpers.rs | 359 +------- .../src/decoder/lossless_helpers/tests.rs | 53 +- .../j2k-jpeg/src/decoder/lossless_render.rs | 809 +++++++----------- crates/j2k-jpeg/src/decoder/plan.rs | 6 +- .../j2k-jpeg/src/decoder/plan/construction.rs | 5 +- crates/j2k-jpeg/src/decoder/rows.rs | 11 +- crates/j2k-jpeg/src/decoder/scratch.rs | 19 +- crates/j2k-jpeg/src/entropy/huffman.rs | 43 + crates/j2k-jpeg/src/entropy/progressive.rs | 2 +- .../src/entropy/progressive/allocation.rs | 30 + .../src/entropy/progressive/render.rs | 44 +- .../j2k-jpeg/src/entropy/progressive/scan.rs | 513 ++++++++--- .../j2k-jpeg/src/entropy/progressive/tests.rs | 76 +- .../src/entropy/sequential/plan/resolved.rs | 9 +- crates/j2k-jpeg/src/idct/mod.rs | 1 - crates/j2k-jpeg/src/idct/neon.rs | 336 +++++--- crates/j2k-jpeg/src/internal/bit_reader.rs | 59 +- crates/j2k-jpeg/src/lossless/mod.rs | 121 ++- crates/j2k-jpeg/src/lossless/scan.rs | 510 +++++++++++ crates/j2k-jpeg/src/lossless/scan/tests.rs | 393 +++++++++ crates/j2k-jpeg/src/simd/neon_memory.rs | 16 +- .../j2k-jpeg/tests/libjpeg_turbo_compare.rs | 269 ++++++ .../src/jpeg_fixtures/builders.rs | 146 +++- 42 files changed, 3572 insertions(+), 1451 deletions(-) create mode 100644 crates/j2k-jpeg/src/backend/neon_upsample.rs create mode 100644 crates/j2k-jpeg/src/decoder/extended12/writers/tests.rs create mode 100644 crates/j2k-jpeg/src/lossless/scan.rs create mode 100644 crates/j2k-jpeg/src/lossless/scan/tests.rs diff --git a/crates/j2k-jpeg/src/backend/mod.rs b/crates/j2k-jpeg/src/backend/mod.rs index 350a1ca2a..eb0ee8f1a 100644 --- a/crates/j2k-jpeg/src/backend/mod.rs +++ b/crates/j2k-jpeg/src/backend/mod.rs @@ -15,6 +15,9 @@ mod x86; #[cfg(target_arch = "aarch64")] mod neon; +#[cfg(target_arch = "aarch64")] +mod neon_upsample; + #[cfg(any(target_arch = "x86_64", target_arch = "aarch64"))] mod row_pair; @@ -307,6 +310,41 @@ impl Backend { } } + /// One visible row of fancy 4:2:0 chroma upsampling; bit-exact with + /// [`crate::color::upsample::upsample_h2v2_fancy_row`]. + pub(crate) fn upsample_h2v2_fancy_row( + self, + rows: [&[u8]; 3], + output_width: usize, + output_is_bottom: bool, + out: &mut [u8], + ) { + let [prev, curr, next] = rows; + match self { + #[cfg(target_arch = "aarch64")] + Self::Neon(neon) => { + let near = if output_is_bottom { next } else { prev }; + crate::color::upsample::emit_h2v2_row_with( + near, + curr, + output_width, + out, + |curr, near, interior| { + neon_upsample::h2v2_fancy_pairs(neon, curr, near, interior) + }, + ); + } + _ => crate::color::upsample::upsample_h2v2_fancy_row( + prev, + curr, + next, + output_width, + output_is_bottom, + out, + ), + } + } + /// 8×8 inverse DCT of a dequantized coefficient block. Output is /// level-shifted by +128 and clamped to `[0, 255]` — bit-exact with /// [`idct::scalar::idct_islow`] on every legal JPEG input. @@ -320,6 +358,19 @@ impl Backend { } } + /// 8×8 inverse DCT of a dequantized 12-bit coefficient block, level-shifted + /// by +2048 and clamped to `[0, 4095]`; bit-exact with + /// [`idct::scalar::idct_islow_12bit`] on every input. + pub(crate) fn idct_12bit(self, input: &[i16; 64], output: &mut [u16; 64]) { + match self { + // The NEON transform declines blocks that could overflow its + // 32-bit lanes; those fall through to the 64-bit scalar one. + #[cfg(target_arch = "aarch64")] + Self::Neon(neon) if idct::neon::idct_islow_12bit(neon, input, output) => {} + _ => idct::scalar::idct_islow_12bit(input, output), + } + } + pub(crate) fn idct_bottom_half_zero(self, input: &[i16; 64], output: &mut [u8; 64]) { match self { Self::Scalar => idct::scalar::idct_islow_bottom_half_zero(input, output), diff --git a/crates/j2k-jpeg/src/backend/neon_upsample.rs b/crates/j2k-jpeg/src/backend/neon_upsample.rs new file mode 100644 index 000000000..b04e38c87 --- /dev/null +++ b/crates/j2k-jpeg/src/backend/neon_upsample.rs @@ -0,0 +1,72 @@ +// SPDX-License-Identifier: MIT OR Apache-2.0 + +//! NEON fancy 4:2:0 chroma row upsampling. + +use core::arch::aarch64::{ + vaddq_u16, vcombine_u8, vdup_n_u8, vdupq_n_u16, vdupq_n_u8, vget_low_u8, vmlal_high_u8, + vmlal_u8, vmlaq_n_u16, vmovl_high_u8, vmovl_u8, vshrn_n_u16, vzip1q_u8, vzip2q_u8, +}; + +use crate::simd::neon_memory; + +fearless_simd::kernel! { + /// Fill the leading interior pairs of a fancy 4:2:0 chroma row, sixteen + /// chroma samples at a time. With column sums `cs(i) = 3 * curr[i] + + /// near[i]`, pair `i` is `interior[2i] = (3 * cs(i) + cs(i + 1) + 7) >> 4` + /// and `interior[2i + 1] = (3 * cs(i + 1) + cs(i) + 8) >> 4`, exactly as + /// libjpeg-turbo's `h2v2_fancy_upsample`. Returns how many pairs were + /// filled; the caller finishes the rest. + pub(crate) fn h2v2_fancy_pairs( + neon: Neon, + curr: &[u8], + near: &[u8], + interior: &mut [u8], + ) -> usize { + h2v2_fancy_pairs_neon(curr, near, interior) + } +} + +#[target_feature(enable = "neon")] +fn h2v2_fancy_pairs_neon(curr: &[u8], near: &[u8], interior: &mut [u8]) -> usize { + let pairs = interior.len() / 2; + let mut done = 0; + // Each step reads samples `done..done + 17` and writes 32 outputs. + while done + 16 <= pairs { + let (Some(c0), Some(c1), Some(n0), Some(n1)) = ( + curr.get(done..).and_then(<[u8]>::first_chunk::<16>), + curr.get(done + 1..).and_then(<[u8]>::first_chunk::<16>), + near.get(done..).and_then(<[u8]>::first_chunk::<16>), + near.get(done + 1..).and_then(<[u8]>::first_chunk::<16>), + ) else { + break; + }; + let Some((lo, rest)) = interior[2 * done..].split_first_chunk_mut::<16>() else { + break; + }; + let Some(hi) = rest.first_chunk_mut::<16>() else { + break; + }; + let (c0, c1) = (neon_memory::load_u8x16(c0), neon_memory::load_u8x16(c1)); + let (n0, n1) = (neon_memory::load_u8x16(n0), neon_memory::load_u8x16(n1)); + let (three, three_q) = (vdup_n_u8(3), vdupq_n_u8(3)); + // Column sums of samples `i` (left) and `i + 1` (right), in u16. + let left_lo = vmlal_u8(vmovl_u8(vget_low_u8(n0)), vget_low_u8(c0), three); + let left_hi = vmlal_high_u8(vmovl_high_u8(n0), c0, three_q); + let right_lo = vmlal_u8(vmovl_u8(vget_low_u8(n1)), vget_low_u8(c1), three); + let right_hi = vmlal_high_u8(vmovl_high_u8(n1), c1, three_q); + // At most 3 * 1020 + 1020 + 8 = 4088, so the u16 lanes cannot wrap. + let (seven, eight) = (vdupq_n_u16(7), vdupq_n_u16(8)); + let odd = vcombine_u8( + vshrn_n_u16::<4>(vaddq_u16(vmlaq_n_u16(right_lo, left_lo, 3), seven)), + vshrn_n_u16::<4>(vaddq_u16(vmlaq_n_u16(right_hi, left_hi, 3), seven)), + ); + let even = vcombine_u8( + vshrn_n_u16::<4>(vaddq_u16(vmlaq_n_u16(left_lo, right_lo, 3), eight)), + vshrn_n_u16::<4>(vaddq_u16(vmlaq_n_u16(left_hi, right_hi, 3), eight)), + ); + neon_memory::store_u8x16(lo, vzip1q_u8(odd, even)); + neon_memory::store_u8x16(hi, vzip2q_u8(odd, even)); + done += 16; + } + done +} diff --git a/crates/j2k-jpeg/src/backend/tests.rs b/crates/j2k-jpeg/src/backend/tests.rs index e111fdf19..701e07cb7 100644 --- a/crates/j2k-jpeg/src/backend/tests.rs +++ b/crates/j2k-jpeg/src/backend/tests.rs @@ -1195,3 +1195,89 @@ fn neon_420_cropped_row_pair_matches_scalar_reference_across_chunks() { crop_width, ); } + +#[test] +fn every_backend_matches_the_64_bit_scalar_12bit_idct() { + // Includes blocks beyond the NEON 32-bit guard, which must fall back. + let blocks = [ + [0i16; 64], + core::array::from_fn(|i| if i == 0 { i16::MAX } else { 0 }), + core::array::from_fn(|i| if i % 3 == 0 { i16::MIN } else { 1_000 }), + core::array::from_fn(|i| i16::try_from(i * 37 % 211).expect("small") - 100), + core::array::from_fn(|i| if i % 2 == 0 { i16::MAX } else { i16::MIN }), + ]; + let kinds = [ + super::BackendKind::Scalar, + #[cfg(target_arch = "aarch64")] + super::BackendKind::Neon, + #[cfg(target_arch = "x86_64")] + super::BackendKind::Avx2, + ]; + for kind in kinds { + let Some(backend) = super::Backend::for_test_kind(kind) else { + continue; + }; + for (index, block) in blocks.iter().enumerate() { + let mut expected = [0u16; 64]; + crate::idct::scalar::idct_islow_12bit(block, &mut expected); + let mut actual = [0u16; 64]; + backend.idct_12bit(block, &mut actual); + assert_eq!(actual, expected, "{kind:?} block {index}"); + } + } +} + +#[test] +fn fancy_h2v2_rows_match_the_per_sample_reference_on_every_backend() { + let kinds = [ + super::BackendKind::Scalar, + #[cfg(target_arch = "aarch64")] + super::BackendKind::Neon, + #[cfg(target_arch = "x86_64")] + super::BackendKind::Avx2, + ]; + let mut state = 0x1357_9bdf_u32; + let mut row = |len: usize| -> Vec { + (0..len) + .map(|_| { + state ^= state << 13; + state ^= state >> 17; + state ^= state << 5; + u8::try_from(state >> 24).expect("byte") + }) + .collect() + }; + for chroma_width in 1..=70usize { + let (prev, curr, next) = (row(chroma_width), row(chroma_width), row(chroma_width)); + for output_width in [2 * chroma_width, 2 * chroma_width - 1, chroma_width] { + for bottom in [false, true] { + let near = if bottom { &next } else { &prev }; + let expected: Vec = (0..output_width) + .map(|x| crate::color::upsample::h2v2_fancy_sample(near, &curr, x)) + .collect(); + for kind in kinds { + let Some(backend) = super::Backend::for_test_kind(kind) else { + continue; + }; + let mut actual = vec![0xAA; output_width + 3]; + backend.upsample_h2v2_fancy_row( + [&prev, &curr, &next], + output_width, + bottom, + &mut actual, + ); + assert_eq!( + &actual[..output_width], + &expected[..], + "{kind:?} chroma {chroma_width} width {output_width} bottom {bottom}" + ); + assert_eq!( + &actual[output_width..], + &[0xAA; 3], + "{kind:?} wrote past width" + ); + } + } + } + } +} diff --git a/crates/j2k-jpeg/src/color/upsample.rs b/crates/j2k-jpeg/src/color/upsample.rs index 3641a7396..0d54acfba 100644 --- a/crates/j2k-jpeg/src/color/upsample.rs +++ b/crates/j2k-jpeg/src/color/upsample.rs @@ -179,10 +179,55 @@ pub(crate) fn upsample_h2v2_fancy_row( /// vertical weighting, then run the 3:1 horizontal blend and emit one /// upsampled luma row. `near` = chroma row above (for the top output) or /// below (for the bottom output) the current chroma row. +/// +/// Equal to [`h2v2_fancy_sample_with_len`] at every visible column, but one +/// pass over column-sum pairs: each pair of chroma samples `(a, b)` yields the +/// odd output right of `a` and the even output left of `b`. fn emit_h2v2_row(near: &[u8], curr: &[u8], output_width: usize, out: &mut [u8]) { + emit_h2v2_row_with(near, curr, output_width, out, |_, _, _| 0); +} + +/// [`emit_h2v2_row`] whose interior pairs may start with a SIMD prefix: +/// `vector_pairs(curr, near, interior)` fills a leading run of pairs with +/// the same arithmetic and returns its length. +#[expect( + clippy::cast_possible_truncation, + reason = "libjpeg-compatible weighted sums are shifted back into the u8 chroma range" +)] +pub(crate) fn emit_h2v2_row_with( + near: &[u8], + curr: &[u8], + output_width: usize, + out: &mut [u8], + vector_pairs: impl FnOnce(&[u8], &[u8], &mut [u8]) -> usize, +) { let n = curr.len(); - for (x, slot) in out.iter_mut().enumerate().take(output_width) { - *slot = h2v2_fancy_sample_with_len(near, curr, n, n * 2, x); + let visible_width = output_width.min(out.len()); + let visible = &mut out[..visible_width]; + if n <= 1 || visible.is_empty() { + for (x, slot) in visible.iter_mut().enumerate() { + *slot = h2v2_fancy_sample_with_len(near, curr, n, n * 2, x); + } + return; + } + let colsum = |i: usize| 3 * u16::from(curr[i]) + u16::from(near[i]); + visible[0] = ((4 * colsum(0) + 8) >> 4) as u8; + // Interior pairs, as an indexed loop over pre-sized slices so the + // compiler drops the bounds checks and vectorizes the interleaved stores. + let pairs = (visible_width - 1) / 2; + let (curr_band, near_band) = (&curr[..=pairs], &near[..=pairs]); + let interior = &mut visible[1..=2 * pairs]; + let done = vector_pairs(curr_band, near_band, interior).min(pairs); + for i in done..pairs { + let left = 3 * u16::from(curr_band[i]) + u16::from(near_band[i]); + let right = 3 * u16::from(curr_band[i + 1]) + u16::from(near_band[i + 1]); + interior[2 * i] = ((3 * left + right + 7) >> 4) as u8; + interior[2 * i + 1] = ((3 * right + left + 8) >> 4) as u8; + } + if visible_width.is_multiple_of(2) { + // The last odd column of a full row, or an odd visible width. + visible[visible_width - 1] = + h2v2_fancy_sample_with_len(near, curr, n, n * 2, visible_width - 1); } } diff --git a/crates/j2k-jpeg/src/context.rs b/crates/j2k-jpeg/src/context.rs index 8b19a9801..574ce2c3c 100644 --- a/crates/j2k-jpeg/src/context.rs +++ b/crates/j2k-jpeg/src/context.rs @@ -213,13 +213,16 @@ impl DecoderContext { table } + /// The compiled table for `raw`, from the cache or compiled into it. + /// Returned by reference so a caller copies the ~24 KiB table at most + /// once, into its own arena. pub(crate) fn resolve_huffman_table_with_live_budget( &mut self, raw: &RawHuffmanTable, role: HuffmanTableRole, live_bytes: &mut usize, cap: usize, - ) -> Result { + ) -> Result<&HuffmanTable, JpegError> { let digest = digest_huffman_table(raw, role); self.resolve_huffman_table_with_digest_and_live_budget(raw, role, digest, live_bytes, cap) } @@ -235,44 +238,46 @@ impl DecoderContext { digest: u64, live_bytes: &mut usize, cap: usize, - ) -> Result { + ) -> Result<&HuffmanTable, JpegError> { self.ensure_huffman_cache_slots(live_bytes, cap)?; - let start = (digest as usize) % self.huffman_tables.len(); - for probe in 0..self.huffman_tables.len() { - let slot = (start + probe) % self.huffman_tables.len(); - match &self.huffman_tables[slot] { + let slots = self.huffman_tables.len(); + let start = (digest as usize) % slots; + // The first matching or free slot along the probe sequence; when + // every slot holds another table, `start` is evicted. + let (slot, hit) = (0..slots) + .map(|probe| (start + probe) % slots) + .find_map(|slot| match &self.huffman_tables[slot] { Some(cached) if cached.digest == digest && cached.role == role && &cached.raw == raw => { - self.cache_hits = self.cache_hits.saturating_add(1); - return Ok(cached.table.clone()); - } - None => { - let table = HuffmanTable::from_raw(raw, role)?; - self.huffman_tables[slot] = Some(CachedHuffmanTable { - digest, - role, - raw: raw.clone(), - table: table.clone(), - }); - self.cache_misses = self.cache_misses.saturating_add(1); - return Ok(table); + Some((slot, true)) } - Some(_) => {} + None => Some((slot, false)), + Some(_) => None, + }) + .unwrap_or((start, false)); + if hit { + self.cache_hits = self.cache_hits.saturating_add(1); + } else { + let evicting = self.huffman_tables[slot].is_some(); + let table = HuffmanTable::from_raw(raw, role)?; + self.huffman_tables[slot] = Some(CachedHuffmanTable { + digest, + role, + raw: raw.clone(), + table, + }); + self.cache_misses = self.cache_misses.saturating_add(1); + if evicting { + self.cache_evictions = self.cache_evictions.saturating_add(1); } } - - let slot = start; - let table = HuffmanTable::from_raw(raw, role)?; - self.huffman_tables[slot] = Some(CachedHuffmanTable { - digest, - role, - raw: raw.clone(), - table: table.clone(), - }); - self.cache_misses = self.cache_misses.saturating_add(1); - self.cache_evictions = self.cache_evictions.saturating_add(1); - Ok(table) + self.huffman_tables[slot] + .as_ref() + .map(|cached| &cached.table) + .ok_or(JpegError::InternalInvariant { + reason: "Huffman cache slot emptied while resolving a table", + }) } fn ensure_huffman_cache_slots( @@ -627,6 +632,7 @@ mod tests { &mut live_bytes, MAX_DECODER_CONTEXT_ALLOCATION_BYTES, ) + .cloned() } fn resolve_huffman_table_with_digest( @@ -642,6 +648,7 @@ mod tests { &mut live_bytes, MAX_DECODER_CONTEXT_ALLOCATION_BYTES, ) + .cloned() } #[test] diff --git a/crates/j2k-jpeg/src/decoder.rs b/crates/j2k-jpeg/src/decoder.rs index db6ccb655..70a4b450c 100644 --- a/crates/j2k-jpeg/src/decoder.rs +++ b/crates/j2k-jpeg/src/decoder.rs @@ -27,7 +27,7 @@ use crate::info::{ use crate::internal::bit_reader::BitReader; use crate::internal::checkpoint::{checkpoint_before_mcu, CpuCheckpointCache, DeviceCheckpoint}; use crate::internal::scratch::ScratchPool; -use crate::lossless::{lossless_predict, LosslessSample}; +use crate::lossless::LosslessSample; use crate::output::{ validate_buffer, Gray8Writer, InterleavedRgbWriter, OutputWriter, Rgb8Writer, Rgba8Writer, }; @@ -69,12 +69,9 @@ use self::extended12::{lossless_color_sampling, upsample_h2v1_sample_at, upsampl mod lossless_helpers; pub(crate) use self::lossless_helpers::restart_index_allocation_bytes; use self::lossless_helpers::{ - decode_lossless_color_sample, decode_lossless_sampled_color_mcu, emit_decode_scan_profile, - lossless_predictor_gray_rows, lossless_predictor_value, lossless_predictor_value_u16, - resolve_lossless_color_components, restart_index_for_stream, validate_lossless_color_plan, - write_lossless_color16_sampled_output, write_lossless_color8_sampled_output, - LosslessColorIntoSample, LosslessColorPlanes, LosslessColorRowSample, LosslessRestartTracker, - LosslessSampledColorPlanesMut, LosslessSampledMcu, + emit_decode_scan_profile, lossless_output_order, restart_index_for_stream, + validate_lossless_color_plan, write_lossless_color16_sampled_output, + write_lossless_color8_sampled_output, LosslessColorPlanes, }; mod color_convert; use self::color_convert::{ @@ -317,7 +314,7 @@ struct PreparedDecoderMetadata { struct PreparedLosslessPlan { predictor: u8, bit_depth: u8, - dc_table: PreparedHuffmanTableId, + point_transform: u8, dimensions: (u32, u32), scan_offset: usize, } diff --git a/crates/j2k-jpeg/src/decoder/allocation.rs b/crates/j2k-jpeg/src/decoder/allocation.rs index 33e9076bf..956c8dbbd 100644 --- a/crates/j2k-jpeg/src/decoder/allocation.rs +++ b/crates/j2k-jpeg/src/decoder/allocation.rs @@ -164,6 +164,11 @@ pub(super) fn compute_progressive_scratch_bytes( checked_usize_product(&[components.len(), COMPONENT_IMAGE_METADATA_BYTES], cap)?; let dc_predictors = checked_usize_product(&[components.len(), core::mem::size_of::()], cap)?; + // Scan-phase nonzero masks: one outer vector per component and one u64 + // per block, live next to the DC predictors. + let mask_outer = + checked_usize_product(&[components.len(), core::mem::size_of::>()], cap)?; + let mut mask_payload = 0usize; let mut coefficient_payload = 0usize; let mut plane_samples = 0usize; for component in components { @@ -173,6 +178,8 @@ pub(super) fn compute_progressive_scratch_bytes( )?; let coeffs = checked_usize_product(&[blocks, core::mem::size_of::<[i32; 64]>()], cap)?; coefficient_payload = checked_workspace_add(coefficient_payload, coeffs, cap)?; + let masks = checked_usize_product(&[blocks, core::mem::size_of::()], cap)?; + mask_payload = checked_workspace_add(mask_payload, masks, cap)?; let samples = checked_usize_product( &[ component.block_cols as usize, @@ -199,7 +206,12 @@ pub(super) fn compute_progressive_scratch_bytes( } else { eight_bit_render }; - let phase_peak = dc_predictors.max(render_phase); + let scan_phase = checked_workspace_add( + checked_workspace_add(mask_outer, mask_payload, cap)?, + dc_predictors, + cap, + )?; + let phase_peak = scan_phase.max(render_phase); checked_workspace_add( checked_workspace_add(coefficient_outer, coefficient_payload, cap)?, phase_peak, @@ -342,7 +354,8 @@ mod tests { fn progressive_scratch_counts_every_simultaneous_outer_and_payload() { let expected = size_of::>() + size_of::<[i32; 64]>() - + (COMPONENT_IMAGE_METADATA_BYTES + 64 + 8 * 7).max(size_of::()); + + (COMPONENT_IMAGE_METADATA_BYTES + 64 + 8 * 7) + .max(size_of::>() + size_of::() + size_of::()); assert_eq!( compute_progressive_scratch_bytes( &[progressive_component()], diff --git a/crates/j2k-jpeg/src/decoder/extended12/planes.rs b/crates/j2k-jpeg/src/decoder/extended12/planes.rs index 8c14687b6..8360f42b8 100644 --- a/crates/j2k-jpeg/src/decoder/extended12/planes.rs +++ b/crates/j2k-jpeg/src/decoder/extended12/planes.rs @@ -5,7 +5,7 @@ use alloc::vec::Vec; use super::super::{ - finish_scan, BitReader, CoefficientBlock, JpegError, PreparedDecodePlan, + finish_scan, Backend, BitReader, CoefficientBlock, JpegError, PreparedDecodePlan, PreparedProgressivePlan, SofKind, Warning, ZIGZAG, }; use super::state::{decode_extended12_block_pixels, Extended12RestartTracker}; @@ -21,6 +21,7 @@ use self::allocation::{ }; pub(super) fn decode_extended12_color_planes( + backend: Backend, plan: &PreparedDecodePlan, scan_bytes: &[u8], sof: SofKind, @@ -59,6 +60,7 @@ pub(super) fn decode_extended12_color_planes( for by in 0..u32::from(component.v) { for bx in 0..u32::from(component.h) { decode_extended12_block_pixels( + backend, &mut br, resolved, &mut prev_dc[output_index], @@ -83,6 +85,7 @@ pub(super) fn decode_extended12_color_planes( } pub(super) fn decode_extended12_four_component_planes( + backend: Backend, plan: &PreparedDecodePlan, scan_bytes: &[u8], sof: SofKind, @@ -121,6 +124,7 @@ pub(super) fn decode_extended12_four_component_planes( for by in 0..u32::from(component.v) { for bx in 0..u32::from(component.h) { decode_extended12_block_pixels( + backend, &mut br, resolved, &mut prev_dc[output_index], @@ -145,6 +149,7 @@ pub(super) fn decode_extended12_four_component_planes( } pub(super) fn render_progressive12_color_planes( + backend: Backend, plan: &PreparedProgressivePlan, dct_blocks: &ProgressiveDctBlocks, ) -> Result<[Extended12Plane; 3], JpegError> { @@ -171,7 +176,7 @@ pub(super) fn render_progressive12_color_planes( if dequant[1..].iter().all(|&coeff| coeff == 0) { pixels.fill(crate::idct::idct_islow_12bit_dc_only_sample(dequant[0])); } else { - crate::idct::idct_islow_12bit(&dequant, &mut pixels); + backend.idct_12bit(&dequant, &mut pixels); } deposit_extended12_block(&mut planes[output_index], bx * 8, by * 8, &pixels); } @@ -182,6 +187,7 @@ pub(super) fn render_progressive12_color_planes( } pub(super) fn render_progressive12_four_component_planes( + backend: Backend, plan: &PreparedProgressivePlan, dct_blocks: &ProgressiveDctBlocks, ) -> Result<[Extended12Plane; 4], JpegError> { @@ -208,7 +214,7 @@ pub(super) fn render_progressive12_four_component_planes( if dequant[1..].iter().all(|&coeff| coeff == 0) { pixels.fill(crate::idct::idct_islow_12bit_dc_only_sample(dequant[0])); } else { - crate::idct::idct_islow_12bit(&dequant, &mut pixels); + backend.idct_12bit(&dequant, &mut pixels); } deposit_extended12_block(&mut planes[output_index], bx * 8, by * 8, &pixels); } diff --git a/crates/j2k-jpeg/src/decoder/extended12/progressive.rs b/crates/j2k-jpeg/src/decoder/extended12/progressive.rs index 3f3d050f3..60a1a9dd0 100644 --- a/crates/j2k-jpeg/src/decoder/extended12/progressive.rs +++ b/crates/j2k-jpeg/src/decoder/extended12/progressive.rs @@ -169,7 +169,7 @@ impl Decoder<'_> { if dequant[1..].iter().all(|&coeff| coeff == 0) { pixels.fill(crate::idct::idct_islow_12bit_dc_only_sample(dequant[0])); } else { - crate::idct::idct_islow_12bit(&dequant, &mut pixels); + self.backend.idct_12bit(&dequant, &mut pixels); } write_extended12_block_region( out, diff --git a/crates/j2k-jpeg/src/decoder/extended12/progressive/color444.rs b/crates/j2k-jpeg/src/decoder/extended12/progressive/color444.rs index 3d55a9a1b..f5d8b21c4 100644 --- a/crates/j2k-jpeg/src/decoder/extended12/progressive/color444.rs +++ b/crates/j2k-jpeg/src/decoder/extended12/progressive/color444.rs @@ -80,10 +80,8 @@ impl Decoder<'_> { dequant[output_index][0], )); } else { - crate::idct::idct_islow_12bit( - &dequant[output_index], - &mut pixels[output_index], - ); + self.backend + .idct_12bit(&dequant[output_index], &mut pixels[output_index]); } } write_extended12_rgb_block_region( diff --git a/crates/j2k-jpeg/src/decoder/extended12/progressive/four_component.rs b/crates/j2k-jpeg/src/decoder/extended12/progressive/four_component.rs index 6d6b7f517..3f93a7c91 100644 --- a/crates/j2k-jpeg/src/decoder/extended12/progressive/four_component.rs +++ b/crates/j2k-jpeg/src/decoder/extended12/progressive/four_component.rs @@ -35,7 +35,7 @@ impl Decoder<'_> { let output_rect = scaled_rect_covering(roi, downscale)?; let dct_blocks = decode_progressive_dct_blocks(plan, self.bytes, 0)?; - let planes = render_progressive12_four_component_planes(plan, &dct_blocks)?; + let planes = render_progressive12_four_component_planes(self.backend, plan, &dct_blocks)?; write_extended12_four_component_planes_region( out, stride, diff --git a/crates/j2k-jpeg/src/decoder/extended12/progressive/subsampled.rs b/crates/j2k-jpeg/src/decoder/extended12/progressive/subsampled.rs index be5b13e75..10e856a02 100644 --- a/crates/j2k-jpeg/src/decoder/extended12/progressive/subsampled.rs +++ b/crates/j2k-jpeg/src/decoder/extended12/progressive/subsampled.rs @@ -41,7 +41,7 @@ impl Decoder<'_> { let output_rect = scaled_rect_covering(roi, downscale)?; let dct_blocks = decode_progressive_dct_blocks(plan, self.bytes, 0)?; - let planes = render_progressive12_color_planes(plan, &dct_blocks)?; + let planes = render_progressive12_color_planes(self.backend, plan, &dct_blocks)?; let write_region = Extended12WriteRegion { output_rect, dimensions: self.info.dimensions, diff --git a/crates/j2k-jpeg/src/decoder/extended12/sequential.rs b/crates/j2k-jpeg/src/decoder/extended12/sequential.rs index 7b952ec98..ddc2e886a 100644 --- a/crates/j2k-jpeg/src/decoder/extended12/sequential.rs +++ b/crates/j2k-jpeg/src/decoder/extended12/sequential.rs @@ -217,6 +217,7 @@ impl Decoder<'_> { prev_dc = 0; } decode_extended12_block_pixels( + self.backend, &mut br, component, &mut prev_dc, diff --git a/crates/j2k-jpeg/src/decoder/extended12/sequential/color444.rs b/crates/j2k-jpeg/src/decoder/extended12/sequential/color444.rs index 1426dee57..fd9e17d3f 100644 --- a/crates/j2k-jpeg/src/decoder/extended12/sequential/color444.rs +++ b/crates/j2k-jpeg/src/decoder/extended12/sequential/color444.rs @@ -54,6 +54,7 @@ impl Decoder<'_> { let output_index = component.output_index; let component = self.plan.resolve_component(component)?; decode_extended12_block_pixels( + self.backend, &mut br, component, &mut prev_dc[output_index], diff --git a/crates/j2k-jpeg/src/decoder/extended12/sequential/four_component.rs b/crates/j2k-jpeg/src/decoder/extended12/sequential/four_component.rs index 5364795c2..dce32968c 100644 --- a/crates/j2k-jpeg/src/decoder/extended12/sequential/four_component.rs +++ b/crates/j2k-jpeg/src/decoder/extended12/sequential/four_component.rs @@ -57,6 +57,7 @@ impl Decoder<'_> { let output_index = component.output_index; let component = self.plan.resolve_component(component)?; decode_extended12_block_pixels( + self.backend, &mut br, component, &mut prev_dc[output_index], @@ -103,8 +104,12 @@ impl Decoder<'_> { let output_rect = scaled_rect_covering(roi, downscale)?; let scan_bytes = &self.bytes[self.plan.scan_offset..]; - let (planes, scan_warnings) = - decode_extended12_four_component_planes(&self.plan, scan_bytes, self.info.sof_kind)?; + let (planes, scan_warnings) = decode_extended12_four_component_planes( + self.backend, + &self.plan, + scan_bytes, + self.info.sof_kind, + )?; write_extended12_four_component_planes_region( out, stride, diff --git a/crates/j2k-jpeg/src/decoder/extended12/sequential/subsampled.rs b/crates/j2k-jpeg/src/decoder/extended12/sequential/subsampled.rs index dc86faee9..a94a26f23 100644 --- a/crates/j2k-jpeg/src/decoder/extended12/sequential/subsampled.rs +++ b/crates/j2k-jpeg/src/decoder/extended12/sequential/subsampled.rs @@ -34,8 +34,12 @@ impl Decoder<'_> { let output_rect = scaled_rect_covering(roi, downscale)?; let scan_bytes = &self.bytes[self.plan.scan_offset..]; - let (planes, scan_warnings) = - decode_extended12_color_planes(&self.plan, scan_bytes, self.info.sof_kind)?; + let (planes, scan_warnings) = decode_extended12_color_planes( + self.backend, + &self.plan, + scan_bytes, + self.info.sof_kind, + )?; let write_region = Extended12WriteRegion { output_rect, dimensions: self.info.dimensions, diff --git a/crates/j2k-jpeg/src/decoder/extended12/state.rs b/crates/j2k-jpeg/src/decoder/extended12/state.rs index 3d5c46ab2..3d1f3b8a3 100644 --- a/crates/j2k-jpeg/src/decoder/extended12/state.rs +++ b/crates/j2k-jpeg/src/decoder/extended12/state.rs @@ -4,25 +4,33 @@ pub(super) use super::super::lossless_helpers::Extended12RestartTracker; use super::super::{ - decode_block_with_activity, BitReader, BlockActivity, CoefficientBlock, JpegError, + decode_block_with_activity, Backend, BitReader, BlockActivity, CoefficientBlock, JpegError, ResolvedPreparedComponentPlan, }; +/// Decode one block and inverse-transform it to 12-bit samples. pub(super) fn decode_extended12_block_pixels( + backend: Backend, br: &mut BitReader<'_>, component: ResolvedPreparedComponentPlan<'_>, prev_dc: &mut i32, coeff: &mut CoefficientBlock, pixels: &mut [u16; 64], ) -> Result<(), JpegError> { - let activity = decode_block_with_activity( - br, + // Decode through a local copy so the accumulator stays in registers for + // the whole block instead of being stored back after every symbol; it is + // written back even on error. + let mut local = br.clone(); + let decoded = decode_block_with_activity( + &mut local, component.dc_table, component.ac_table, prev_dc, component.quant, coeff, - )?; + ); + *br = local; + let activity = decoded?; match activity { BlockActivity::DcOnly => { pixels.fill(crate::idct::idct_islow_12bit_dc_only_sample( @@ -30,7 +38,7 @@ pub(super) fn decode_extended12_block_pixels( )); } BlockActivity::BottomHalfZero | BlockActivity::General => { - crate::idct::idct_islow_12bit(coeff.coefficients(), pixels); + backend.idct_12bit(coeff.coefficients(), pixels); } } Ok(()) diff --git a/crates/j2k-jpeg/src/decoder/extended12/writers.rs b/crates/j2k-jpeg/src/decoder/extended12/writers.rs index b9dd0160c..0c6e02819 100644 --- a/crates/j2k-jpeg/src/decoder/extended12/writers.rs +++ b/crates/j2k-jpeg/src/decoder/extended12/writers.rs @@ -2,15 +2,19 @@ //! Scaling- and ROI-aware extended-precision output writers. +use core::ops::Range; + use super::super::{ColorSpace, DownscaleFactor, Rect}; use super::planes::Extended12Plane; use super::sampling::Extended12ColorSampling; use super::upsample::{ extended12_plane_row, sample_extended12_plane_at, upsample_extended12_h2v1_at, - upsample_extended12_plane_h2v1_at, upsample_extended12_plane_h2v2_at, - upsample_h2v2_u16_rows_at, + upsample_extended12_plane_h2v1_at, upsample_extended12_plane_h2v2_at, upsample_h2v2_rows_at, }; +#[cfg(test)] +mod tests; + #[derive(Debug, Clone, Copy)] pub(super) enum Extended12Output { Gray16, @@ -31,6 +35,50 @@ pub(super) struct Extended12WriteRegion { pub(super) output: Extended12Output, } +/// Output coordinates in `start..start + len` whose clamped source +/// coordinate `min(o * denom, extent - 1)` lies in the block span `lo..hi`. +/// The mapping is monotonic, so the matches form one contiguous range. +fn block_output_span( + start: u32, + len: u32, + denom: u32, + extent: u32, + lo: u32, + hi: u32, +) -> Range { + let end = start + len; + let first = lo.div_ceil(denom).clamp(start, end); + let last = if hi >= extent { + end + } else { + hi.div_ceil(denom).clamp(first, end) + }; + first..last +} + +/// Visit each output pixel that samples the 8x8 block at `block_origin`, as +/// `(dst_row, dst_col, src_index)`. +fn for_each_block_output( + region: Extended12WriteRegion, + block_origin: (u32, u32), + mut visit: impl FnMut(usize, usize, usize), +) { + let (width, height) = region.dimensions; + let (x0, y0) = block_origin; + let denom = region.downscale.denominator(); + let rect = region.output_rect; + let rows = block_output_span(rect.y, rect.h, denom, height, y0, (y0 + 8).min(height)); + let cols = block_output_span(rect.x, rect.w, denom, width, x0, (x0 + 8).min(width)); + for output_y in rows { + let src_row = (output_y.saturating_mul(denom).min(height - 1) - y0) as usize; + let dst_row = (output_y - rect.y) as usize; + for output_x in cols.clone() { + let src_col = (output_x.saturating_mul(denom).min(width - 1) - x0) as usize; + visit(dst_row, (output_x - rect.x) as usize, src_row * 8 + src_col); + } + } +} + pub(super) fn write_extended12_rgb_block_region( out: &mut [u8], stride: usize, @@ -39,46 +87,25 @@ pub(super) fn write_extended12_rgb_block_region( block_origin: (u32, u32), pixels: &[[u16; 64]; 3], ) { - let (width, height) = region.dimensions; - let (x0, y0) = block_origin; - let block_x1 = (x0 + 8).min(width); - let block_y1 = (y0 + 8).min(height); - let denom = region.downscale.denominator(); - let output_rect = region.output_rect; - for output_y in output_rect.y..output_rect.y + output_rect.h { - let source_y = output_y.saturating_mul(denom).min(height - 1); - if source_y < y0 || source_y >= block_y1 { - continue; - } - let src_row = (source_y - y0) as usize; - let dst_row = (output_y - output_rect.y) as usize; - for output_x in output_rect.x..output_rect.x + output_rect.w { - let source_x = output_x.saturating_mul(denom).min(width - 1); - if source_x < x0 || source_x >= block_x1 { - continue; - } - let src_col = (source_x - x0) as usize; - let src_index = src_row * 8 + src_col; - let dst_col = (output_x - output_rect.x) as usize; - let dst_start = dst_row * stride + dst_col * 6; - let dst = &mut out[dst_start..dst_start + 6]; - let (r, g, b) = match projection { - Extended12RgbProjection::Identity => ( - pixels[0][src_index], - pixels[1][src_index], - pixels[2][src_index], - ), - Extended12RgbProjection::YCbCr => crate::color::ycbcr::ycbcr12_to_rgb16( - pixels[0][src_index], - pixels[1][src_index], - pixels[2][src_index], - ), - }; - dst[0..2].copy_from_slice(&r.to_le_bytes()); - dst[2..4].copy_from_slice(&g.to_le_bytes()); - dst[4..6].copy_from_slice(&b.to_le_bytes()); - } + if region.downscale.denominator() == 1 { + write_rgb16_block_full(out, stride, region, projection, block_origin, pixels); + return; } + for_each_block_output(region, block_origin, |dst_row, dst_col, src_index| { + let (r, g, b) = match projection { + Extended12RgbProjection::Identity => ( + pixels[0][src_index], + pixels[1][src_index], + pixels[2][src_index], + ), + Extended12RgbProjection::YCbCr => crate::color::ycbcr::ycbcr12_to_rgb16( + pixels[0][src_index], + pixels[1][src_index], + pixels[2][src_index], + ), + }; + write_rgb16(out, dst_row * stride + dst_col * 6, (r, g, b)); + }); } pub(super) fn write_extended12_four_component_block_region( @@ -89,49 +116,31 @@ pub(super) fn write_extended12_four_component_block_region( block_origin: (u32, u32), pixels: &[[u16; 64]; 4], ) { - let (width, height) = region.dimensions; - let (x0, y0) = block_origin; - let block_x1 = (x0 + 8).min(width); - let block_y1 = (y0 + 8).min(height); - let denom = region.downscale.denominator(); - let output_rect = region.output_rect; - for output_y in output_rect.y..output_rect.y + output_rect.h { - let source_y = output_y.saturating_mul(denom).min(height - 1); - if source_y < y0 || source_y >= block_y1 { - continue; - } - let src_row = (source_y - y0) as usize; - let dst_row = (output_y - output_rect.y) as usize; - for output_x in output_rect.x..output_rect.x + output_rect.w { - let source_x = output_x.saturating_mul(denom).min(width - 1); - if source_x < x0 || source_x >= block_x1 { - continue; - } - let src_col = (source_x - x0) as usize; - let src_index = src_row * 8 + src_col; - let (r, g, b) = match color_space { - ColorSpace::Cmyk => crate::color::cmyk::inverted_cmyk12_to_rgb16( - pixels[0][src_index], - pixels[1][src_index], - pixels[2][src_index], - pixels[3][src_index], - ), - ColorSpace::Ycck => crate::color::cmyk::ycck12_to_rgb16( - pixels[0][src_index], - pixels[1][src_index], - pixels[2][src_index], - pixels[3][src_index], - ), - _ => unreachable!("12-bit four-component path only accepts CMYK/YCCK"), - }; - let dst_col = (output_x - output_rect.x) as usize; - let dst_start = dst_row * stride + dst_col * 6; - let dst = &mut out[dst_start..dst_start + 6]; - dst[0..2].copy_from_slice(&r.to_le_bytes()); - dst[2..4].copy_from_slice(&g.to_le_bytes()); - dst[4..6].copy_from_slice(&b.to_le_bytes()); - } - } + for_each_block_output(region, block_origin, |dst_row, dst_col, src_index| { + let (r, g, b) = match color_space { + ColorSpace::Cmyk => crate::color::cmyk::inverted_cmyk12_to_rgb16( + pixels[0][src_index], + pixels[1][src_index], + pixels[2][src_index], + pixels[3][src_index], + ), + ColorSpace::Ycck => crate::color::cmyk::ycck12_to_rgb16( + pixels[0][src_index], + pixels[1][src_index], + pixels[2][src_index], + pixels[3][src_index], + ), + _ => unreachable!("12-bit four-component path only accepts CMYK/YCCK"), + }; + write_rgb16(out, dst_row * stride + dst_col * 6, (r, g, b)); + }); +} + +fn write_rgb16(out: &mut [u8], offset: usize, (r, g, b): (u16, u16, u16)) { + let dst = &mut out[offset..offset + 6]; + dst[0..2].copy_from_slice(&r.to_le_bytes()); + dst[2..4].copy_from_slice(&g.to_le_bytes()); + dst[4..6].copy_from_slice(&b.to_le_bytes()); } pub(super) fn write_extended12_color422_planes_region( @@ -144,29 +153,35 @@ pub(super) fn write_extended12_color422_planes_region( let (width, height) = region.dimensions; let denom = region.downscale.denominator(); let output_rect = region.output_rect; + let chroma_rows = planes[1].pixels.len() / planes[1].stride; for output_y in output_rect.y..output_rect.y + output_rect.h { let source_y = output_y.saturating_mul(denom).min(height - 1) as usize; - let dst_row = (output_y - output_rect.y) as usize; - for output_x in output_rect.x..output_rect.x + output_rect.w { + let luma = extended12_plane_row(&planes[0], source_y); + let chroma_y = source_y.min(chroma_rows - 1); + let cb_row = extended12_plane_row(&planes[1], chroma_y); + let cr_row = extended12_plane_row(&planes[2], chroma_y); + let dst_row = &mut out[(output_y - output_rect.y) as usize * stride..]; + if denom == 1 { + let span = output_rect.x as usize..(output_rect.x + output_rect.w) as usize; + let rows = [cb_row, cr_row]; + write_fancy_full_row( + dst_row, + projection, + luma, + span, + (FANCY_H2V1, cb_row.len() - 1), + |channel, sample| u32::from(rows[channel][sample]), + ); + continue; + } + for (dst, output_x) in dst_row + .chunks_exact_mut(6) + .zip(output_rect.x..output_rect.x + output_rect.w) + { let source_x = output_x.saturating_mul(denom).min(width - 1) as usize; - let y = planes[0].pixels[source_y * planes[0].stride + source_x]; - let chroma_y = source_y.min(planes[1].pixels.len() / planes[1].stride - 1); - let cb_row = &planes[1].pixels - [chroma_y * planes[1].stride..chroma_y * planes[1].stride + planes[1].width]; - let cr_row = &planes[2].pixels - [chroma_y * planes[2].stride..chroma_y * planes[2].stride + planes[2].width]; let c1 = upsample_extended12_h2v1_at(cb_row, source_x); let c2 = upsample_extended12_h2v1_at(cr_row, source_x); - let (r, g, b) = match projection { - Extended12RgbProjection::Identity => (y, c1, c2), - Extended12RgbProjection::YCbCr => crate::color::ycbcr::ycbcr12_to_rgb16(y, c1, c2), - }; - let dst_col = (output_x - output_rect.x) as usize; - let dst_start = dst_row * stride + dst_col * 6; - let dst = &mut out[dst_start..dst_start + 6]; - dst[0..2].copy_from_slice(&r.to_le_bytes()); - dst[2..4].copy_from_slice(&g.to_le_bytes()); - dst[4..6].copy_from_slice(&b.to_le_bytes()); + write_projected_rgb16(dst, projection, luma[source_x], c1, c2); } } } @@ -185,39 +200,162 @@ pub(super) fn write_extended12_color420_planes_region( let chroma_height = (height as usize).div_ceil(2); for output_y in output_rect.y..output_rect.y + output_rect.h { let source_y = output_y.saturating_mul(denom).min(height - 1) as usize; - let dst_row = (output_y - output_rect.y) as usize; - for output_x in output_rect.x..output_rect.x + output_rect.w { - let source_x = output_x.saturating_mul(denom).min(width - 1) as usize; - let y = planes[0].pixels[source_y * planes[0].stride + source_x]; - let chroma_y = (source_y / 2).min(chroma_height - 1); - let prev_y = chroma_y.saturating_sub(1); - let next_y = (chroma_y + 1).min(chroma_height - 1); - let c1 = upsample_h2v2_u16_rows_at( - extended12_plane_row(&planes[1], prev_y), - extended12_plane_row(&planes[1], chroma_y), - extended12_plane_row(&planes[1], next_y), - source_x, - !source_y.is_multiple_of(2), + let luma = extended12_plane_row(&planes[0], source_y); + let chroma_y = (source_y / 2).min(chroma_height - 1); + // Odd rows blend toward the chroma row below, even rows toward the + // one above. + let near_y = if source_y.is_multiple_of(2) { + chroma_y.saturating_sub(1) + } else { + (chroma_y + 1).min(chroma_height - 1) + }; + let (cb, cb_near) = ( + extended12_plane_row(&planes[1], chroma_y), + extended12_plane_row(&planes[1], near_y), + ); + let (cr, cr_near) = ( + extended12_plane_row(&planes[2], chroma_y), + extended12_plane_row(&planes[2], near_y), + ); + let dst_row = &mut out[(output_y - output_rect.y) as usize * stride..]; + if denom == 1 && cb.len() > 1 { + let span = output_rect.x as usize..(output_rect.x + output_rect.w) as usize; + let rows = [(cb, cb_near), (cr, cr_near)]; + write_fancy_full_row( + dst_row, + projection, + luma, + span, + (FANCY_H2V2, cb.len() - 1), + |channel, sample| { + let (curr, near) = rows[channel]; + 3 * u32::from(curr[sample]) + u32::from(near[sample]) + }, ); - let c2 = upsample_h2v2_u16_rows_at( - extended12_plane_row(&planes[2], prev_y), - extended12_plane_row(&planes[2], chroma_y), - extended12_plane_row(&planes[2], next_y), - source_x, - !source_y.is_multiple_of(2), + continue; + } + for (dst, output_x) in dst_row + .chunks_exact_mut(6) + .zip(output_rect.x..output_rect.x + output_rect.w) + { + let source_x = output_x.saturating_mul(denom).min(width - 1) as usize; + let c1 = upsample_h2v2_rows_at(cb, cb_near, cb.len() * 2, source_x); + let c2 = upsample_h2v2_rows_at(cr, cr_near, cr.len() * 2, source_x); + write_projected_rgb16(dst, projection, luma[source_x], c1, c2); + } + } +} + +/// Output pixels a [`write_fancy_full_row`] chunk covers. +const FANCY_CHUNK: usize = 64; + +/// Fancy horizontal upsampling rule: each chroma sample `c` yields an even +/// output `(3 * sum(c) + sum(left) + even_round) >> shift` and an odd output +/// `(3 * sum(c) + sum(right) + odd_round) >> shift`, with the column sums +/// replicated past the row ends. This reproduces libjpeg-turbo's fancy +/// upsamplers, including their end-of-row cases. +#[derive(Clone, Copy)] +struct FancyRule { + even_round: u32, + odd_round: u32, + shift: u32, +} + +/// `h2v1_fancy_upsample` for DCT-coded samples ([`upsample_extended12_h2v1_at`]). +const FANCY_H2V1: FancyRule = FancyRule { + even_round: 1, + odd_round: 2, + shift: 2, +}; + +/// `h2v2_fancy_upsample` over `3 * curr + near` column sums +/// ([`upsample_h2v2_rows_at`]). +const FANCY_H2V2: FancyRule = FancyRule { + even_round: 8, + odd_round: 7, + shift: 4, +}; + +/// Full-resolution color row over source columns `span`, upsampling chroma a +/// chunk at a time. `column_sum(channel, sample)` gives the chroma column sum +/// of channel 0 (Cb) or 1 (Cr). +fn write_fancy_full_row( + dst_row: &mut [u8], + projection: Extended12RgbProjection, + luma: &[u16], + span: Range, + (rule, last_sample): (FancyRule, usize), + column_sum: impl Fn(usize, usize) -> u32, +) { + let mut up = [[0u16; FANCY_CHUNK + 2]; 2]; + let mut start = span.start; + let mut dst = dst_row.chunks_exact_mut(6); + while start < span.end { + let end = (start + FANCY_CHUNK).min(span.end); + let first_sample = start / 2; + let pairs = end.div_ceil(2) - first_sample; + for (channel, up) in up.iter_mut().enumerate() { + fancy_pairs( + rule, + first_sample, + last_sample, + &mut up[..2 * pairs], + |sample| column_sum(channel, sample), ); - let (r, g, b) = match projection { - Extended12RgbProjection::Identity => (y, c1, c2), - Extended12RgbProjection::YCbCr => crate::color::ycbcr::ycbcr12_to_rgb16(y, c1, c2), - }; - let dst_col = (output_x - output_rect.x) as usize; - let dst_start = dst_row * stride + dst_col * 6; - let dst = &mut out[dst_start..dst_start + 6]; - dst[0..2].copy_from_slice(&r.to_le_bytes()); - dst[2..4].copy_from_slice(&g.to_le_bytes()); - dst[4..6].copy_from_slice(&b.to_le_bytes()); } + let offset = start - 2 * first_sample; + let chunk = luma[start..end] + .iter() + .zip(&up[0][offset..]) + .zip(&up[1][offset..]); + for (((&y, &c1), &c2), dst) in chunk.zip(&mut dst) { + write_projected_rgb16(dst, projection, y, c1, c2); + } + start = end; + } +} + +/// Upsample chroma samples `first..first + out.len() / 2` into interleaved +/// even/odd output pairs under `rule`. +#[expect( + clippy::cast_possible_truncation, + reason = "a weighted mean of 12-bit samples fits u16" +)] +fn fancy_pairs( + rule: FancyRule, + first: usize, + last: usize, + out: &mut [u16], + column_sum: impl Fn(usize) -> u32, +) { + // Column sums for samples `first - 1 ..= first + pairs`, replicated at + // the row ends. + let mut sums = [0u32; FANCY_CHUNK / 2 + 3]; + let pairs = out.len() / 2; + for (offset, sum) in sums[..pairs + 2].iter_mut().enumerate() { + *sum = column_sum((first + offset).saturating_sub(1).min(last)); } + for (pair, window) in out.chunks_exact_mut(2).zip(sums.windows(3)) { + let this = 3 * window[1]; + pair[0] = ((this + window[0] + rule.even_round) >> rule.shift) as u16; + pair[1] = ((this + window[2] + rule.odd_round) >> rule.shift) as u16; + } +} + +fn write_projected_rgb16( + dst: &mut [u8], + projection: Extended12RgbProjection, + c0: u16, + c1: u16, + c2: u16, +) { + let (r, g, b) = match projection { + Extended12RgbProjection::Identity => (c0, c1, c2), + Extended12RgbProjection::YCbCr => crate::color::ycbcr::ycbcr12_to_rgb16(c0, c1, c2), + }; + dst[0..2].copy_from_slice(&r.to_le_bytes()); + dst[2..4].copy_from_slice(&g.to_le_bytes()); + dst[4..6].copy_from_slice(&b.to_le_bytes()); } pub(super) fn write_extended12_four_component_planes_region( @@ -293,43 +431,85 @@ pub(super) fn write_extended12_block_region( block_origin: (u32, u32), pixels: &[u16; 64], ) { - let (width, height) = region.dimensions; - let (x0, y0) = block_origin; - let block_x1 = (x0 + 8).min(width); - let block_y1 = (y0 + 8).min(height); - let denom = region.downscale.denominator(); - let bytes_per_pixel = match region.output { - Extended12Output::Gray16 => 2, - Extended12Output::Rgb16 => 6, - }; - let output_rect = region.output_rect; - for output_y in output_rect.y..output_rect.y + output_rect.h { - let source_y = output_y.saturating_mul(denom).min(height - 1); - if source_y < y0 || source_y >= block_y1 { - continue; - } - let src_row = (source_y - y0) as usize; - let dst_row = (output_y - output_rect.y) as usize; - for output_x in output_rect.x..output_rect.x + output_rect.w { - let source_x = output_x.saturating_mul(denom).min(width - 1); - if source_x < x0 || source_x >= block_x1 { - continue; + if matches!(region.output, Extended12Output::Gray16) && region.downscale.denominator() == 1 { + write_gray16_block_full(out, stride, region, block_origin, pixels); + return; + } + for_each_block_output(region, block_origin, |dst_row, dst_col, src_index| { + let sample = pixels[src_index]; + match region.output { + Extended12Output::Gray16 => { + let offset = dst_row * stride + dst_col * 2; + out[offset..offset + 2].copy_from_slice(&sample.to_le_bytes()); } - let src_col = (source_x - x0) as usize; - let sample = pixels[src_row * 8 + src_col].to_le_bytes(); - let dst_col = (output_x - output_rect.x) as usize; - let dst_start = dst_row * stride + dst_col * bytes_per_pixel; - let dst = &mut out[dst_start..dst_start + bytes_per_pixel]; - match region.output { - Extended12Output::Gray16 => { - dst.copy_from_slice(&sample); - } - Extended12Output::Rgb16 => { - dst[0..2].copy_from_slice(&sample); - dst[2..4].copy_from_slice(&sample); - dst[4..6].copy_from_slice(&sample); - } + Extended12Output::Rgb16 => { + write_rgb16( + out, + dst_row * stride + dst_col * 6, + (sample, sample, sample), + ); } } + }); +} + +/// Full-resolution block rows, clipped to the image and the output +/// rectangle: `(block row, source column, output row, output column, length)` +/// for each visible row of the block at `(x0, y0)`. +fn full_block_rows( + region: Extended12WriteRegion, + (x0, y0): (u32, u32), +) -> impl Iterator { + let (width, height) = region.dimensions; + let rect = region.output_rect; + let rows = block_output_span(rect.y, rect.h, 1, height, y0, (y0 + 8).min(height)); + let cols = block_output_span(rect.x, rect.w, 1, width, x0, (x0 + 8).min(width)); + let (src_col, dst_col, len) = ( + (cols.start - x0) as usize, + (cols.start - rect.x) as usize, + cols.len(), + ); + rows.map(move |output_y| { + ( + (output_y - y0) as usize, + src_col, + (output_y - rect.y) as usize, + dst_col, + len, + ) + }) +} + +fn write_gray16_block_full( + out: &mut [u8], + stride: usize, + region: Extended12WriteRegion, + block_origin: (u32, u32), + pixels: &[u16; 64], +) { + for (row, src_col, dst_row, dst_col, len) in full_block_rows(region, block_origin) { + let src = &pixels[row * 8 + src_col..][..len]; + let dst = &mut out[dst_row * stride + dst_col * 2..][..len * 2]; + for (dst, &sample) in dst.chunks_exact_mut(2).zip(src) { + dst.copy_from_slice(&sample.to_le_bytes()); + } + } +} + +fn write_rgb16_block_full( + out: &mut [u8], + stride: usize, + region: Extended12WriteRegion, + projection: Extended12RgbProjection, + block_origin: (u32, u32), + pixels: &[[u16; 64]; 3], +) { + for (row, src_col, dst_row, dst_col, len) in full_block_rows(region, block_origin) { + let start = row * 8 + src_col; + let [c0, c1, c2] = pixels.each_ref().map(|plane| &plane[start..start + len]); + let dst = &mut out[dst_row * stride + dst_col * 6..][..len * 6]; + for (dst, ((&c0, &c1), &c2)) in dst.chunks_exact_mut(6).zip(c0.iter().zip(c1).zip(c2)) { + write_projected_rgb16(dst, projection, c0, c1, c2); + } } } diff --git a/crates/j2k-jpeg/src/decoder/extended12/writers/tests.rs b/crates/j2k-jpeg/src/decoder/extended12/writers/tests.rs new file mode 100644 index 000000000..79e5e7027 --- /dev/null +++ b/crates/j2k-jpeg/src/decoder/extended12/writers/tests.rs @@ -0,0 +1,261 @@ +// SPDX-License-Identifier: MIT OR Apache-2.0 + +use alloc::vec; +use alloc::vec::Vec; + +use super::super::planes::Extended12Plane; +use super::*; + +fn xorshift(state: &mut u32) -> u32 { + *state ^= *state << 13; + *state ^= *state >> 17; + *state ^= *state << 5; + *state +} + +fn samples(seed: u32, len: usize) -> Vec { + let mut state = seed | 1; + (0..len) + .map(|_| u16::try_from(xorshift(&mut state) % 4096).expect("12-bit sample")) + .collect() +} + +#[test] +fn block_output_span_matches_brute_force() { + for extent in 1..=20u32 { + for denom in [1u32, 2, 4, 8] { + let outputs = extent.div_ceil(denom); + for start in 0..outputs { + for len in 0..=outputs - start { + for lo in (0..extent).step_by(8) { + let hi = (lo + 8).min(extent); + let expected: Vec = (start..start + len) + .filter(|&o| { + let source = o.saturating_mul(denom).min(extent - 1); + (lo..hi).contains(&source) + }) + .collect(); + let actual: Vec = + block_output_span(start, len, denom, extent, lo, hi).collect(); + assert_eq!( + actual, expected, + "extent {extent} denom {denom} start {start} len {len} block {lo}" + ); + } + } + } + } + } +} + +#[test] +fn fancy_pairs_match_the_per_pixel_upsamplers() { + for len in 1..=40usize { + let curr = samples(u32::try_from(len).expect("len"), len); + let near = samples(u32::try_from(len * 7).expect("len"), len); + for first in 0..len { + let pairs = (len - first).min(FANCY_CHUNK / 2); + let mut h2v1 = vec![0u16; 2 * pairs]; + fancy_pairs(FANCY_H2V1, first, len - 1, &mut h2v1, |i| { + u32::from(curr[i]) + }); + let mut h2v2 = vec![0u16; 2 * pairs]; + fancy_pairs(FANCY_H2V2, first, len - 1, &mut h2v2, |i| { + 3 * u32::from(curr[i]) + u32::from(near[i]) + }); + for (offset, (&up1, &up2)) in h2v1.iter().zip(&h2v2).enumerate() { + let x = 2 * first + offset; + assert_eq!( + up1, + upsample_extended12_h2v1_at(&curr, x), + "h2v1 len {len} x {x}" + ); + // Single-sample rows take the dedicated path in the writers. + if len > 1 { + let expected = upsample_h2v2_rows_at(&curr, &near, 2 * len, x); + assert_eq!(up2, expected, "h2v2 len {len} x {x}"); + } + } + } + } +} + +fn plane(seed: u32, width: usize, rows: usize) -> Extended12Plane { + let stride = width.next_multiple_of(8); + Extended12Plane { + pixels: samples(seed, stride * rows), + stride, + width, + } +} + +fn region(dimensions: (u32, u32), output_rect: Rect) -> Extended12WriteRegion { + Extended12WriteRegion { + output_rect, + dimensions, + downscale: DownscaleFactor::Full, + output: Extended12Output::Rgb16, + } +} + +fn rects(width: u32, height: u32) -> [Rect; 3] { + [ + Rect::full((width, height)), + Rect { + x: 1, + y: 1, + w: width - 2, + h: height - 2, + }, + Rect { + x: 3, + y: 2, + w: width - 4, + h: 1, + }, + ] +} + +/// Per-pixel reference for the full-resolution subsampled writers. +fn expected_subsampled( + planes: &[Extended12Plane; 3], + dimensions: (u32, u32), + rect: Rect, + h2v2: bool, +) -> Vec { + let stride = rect.w as usize * 6; + let mut out = vec![0u8; stride * rect.h as usize]; + let chroma_height = (dimensions.1 as usize).div_ceil(2); + for y in 0..rect.h as usize { + let source_y = rect.y as usize + y; + for x in 0..rect.w as usize { + let source_x = rect.x as usize + x; + let luma = planes[0].pixels[source_y * planes[0].stride + source_x]; + let [c1, c2] = [&planes[1], &planes[2]].map(|plane| { + if h2v2 { + super::super::upsample::upsample_extended12_plane_h2v2_at( + plane, + chroma_height, + source_x, + source_y, + ) + } else { + super::super::upsample::upsample_extended12_plane_h2v1_at( + plane, source_x, source_y, + ) + } + }); + let (r, g, b) = crate::color::ycbcr::ycbcr12_to_rgb16(luma, c1, c2); + let dst = &mut out[y * stride + x * 6..][..6]; + dst[0..2].copy_from_slice(&r.to_le_bytes()); + dst[2..4].copy_from_slice(&g.to_le_bytes()); + dst[4..6].copy_from_slice(&b.to_le_bytes()); + } + } + out +} + +#[test] +fn full_resolution_subsampled_writers_match_per_pixel_upsampling() { + for (width, height) in [(5u32, 6u32), (67, 9), (130, 4), (4, 3)] { + for h2v2 in [true, false] { + let chroma_width = (width as usize).div_ceil(2); + let chroma_rows = if h2v2 { + (height as usize).div_ceil(2) + } else { + height as usize + }; + let planes = [ + plane(width * 3, width as usize, height as usize), + plane(width * 5, chroma_width, chroma_rows), + plane(width * 11, chroma_width, chroma_rows), + ]; + for rect in rects(width, height) { + let stride = rect.w as usize * 6; + let mut actual = vec![0u8; stride * rect.h as usize]; + let write = if h2v2 { + write_extended12_color420_planes_region + } else { + write_extended12_color422_planes_region + }; + write( + &mut actual, + stride, + region((width, height), rect), + Extended12RgbProjection::YCbCr, + &planes, + ); + assert_eq!( + actual, + expected_subsampled(&planes, (width, height), rect, h2v2), + "{width}x{height} h2v2={h2v2} {rect:?}" + ); + } + } + } +} + +#[test] +fn full_resolution_block_writers_match_the_per_pixel_path() { + let (width, height) = (21u32, 13u32); + let pixels: [[u16; 64]; 3] = [1, 2, 3].map(|seed| { + let values = samples(seed, 64); + core::array::from_fn(|i| values[i]) + }); + for rect in rects(width, height) { + for output in [Extended12Output::Gray16, Extended12Output::Rgb16] { + let bytes = if matches!(output, Extended12Output::Gray16) { + 2 + } else { + 6 + }; + let stride = rect.w as usize * bytes; + let write_region = Extended12WriteRegion { + output, + ..region((width, height), rect) + }; + let mut gray = vec![0u8; stride * rect.h as usize]; + let mut expected_gray = gray.clone(); + let mut color = vec![0u8; rect.w as usize * 6 * rect.h as usize]; + let mut expected_color = color.clone(); + for origin in [(0, 0), (8, 0), (16, 8), (8, 8)] { + write_extended12_block_region(&mut gray, stride, write_region, origin, &pixels[0]); + for_each_block_output(write_region, origin, |row, col, src| { + let sample = pixels[0][src]; + match output { + Extended12Output::Gray16 => { + let at = row * stride + col * 2; + expected_gray[at..at + 2].copy_from_slice(&sample.to_le_bytes()); + } + Extended12Output::Rgb16 => { + write_rgb16( + &mut expected_gray, + row * stride + col * 6, + (sample, sample, sample), + ); + } + } + }); + let color_stride = rect.w as usize * 6; + write_extended12_rgb_block_region( + &mut color, + color_stride, + write_region, + Extended12RgbProjection::YCbCr, + origin, + &pixels, + ); + for_each_block_output(write_region, origin, |row, col, src| { + let rgb = crate::color::ycbcr::ycbcr12_to_rgb16( + pixels[0][src], + pixels[1][src], + pixels[2][src], + ); + write_rgb16(&mut expected_color, row * color_stride + col * 6, rgb); + }); + } + assert_eq!(gray, expected_gray, "{output:?} {rect:?}"); + assert_eq!(color, expected_color, "{rect:?}"); + } + } +} diff --git a/crates/j2k-jpeg/src/decoder/lossless_helpers.rs b/crates/j2k-jpeg/src/decoder/lossless_helpers.rs index 10917b6b1..00a28f7b4 100644 --- a/crates/j2k-jpeg/src/decoder/lossless_helpers.rs +++ b/crates/j2k-jpeg/src/decoder/lossless_helpers.rs @@ -1,13 +1,12 @@ // SPDX-License-Identifier: MIT OR Apache-2.0 use super::{ - downscale_profile_name, emit_jpeg_profile_fields, lossless_predict, upsample_h2v1_sample_at, + downscale_profile_name, emit_jpeg_profile_fields, upsample_h2v1_sample_at, upsample_h2v2_rows_at, BitReader, ColorSpace, DownscaleFactor, Duration, Info, JpegError, LosslessColorSampling, LosslessSample, MarkerKind, PreparedLosslessPlan, ProfileField, Rect, RestartIndex, RestartSegment, SofKind, }; use crate::allocation::{checked_allocation_bytes, try_vec_with_capacity}; -use crate::entropy::huffman::DcHuffmanTable; use crate::entropy::sequential::PreparedDecodePlan; pub(crate) fn restart_index_allocation_bytes( @@ -171,51 +170,6 @@ pub(super) fn emit_decode_scan_profile( }); } -pub(super) fn consume_lossless_restart( - br: &mut BitReader<'_>, - sample_index: u32, - total_samples: u32, - expected_rst: &mut u8, -) -> Result<(), JpegError> { - br.reset_at_restart(); - *expected_rst = br.consume_restart_marker(*expected_rst, sample_index, total_samples)?; - Ok(()) -} - -pub(super) struct LosslessRestartTracker { - restart_interval: u32, - total_units: u32, - units_since_restart: u32, - expected_rst: u8, -} - -impl LosslessRestartTracker { - pub(super) fn new(restart_interval: Option, total_units: u32) -> Self { - Self { - restart_interval: u32::from(restart_interval.unwrap_or(0)), - total_units, - units_since_restart: 0, - expected_rst: 0, - } - } - - pub(super) fn begin_unit( - &mut self, - br: &mut BitReader<'_>, - unit_index: u32, - ) -> Result { - if self.restart_interval > 0 && self.units_since_restart == self.restart_interval { - consume_lossless_restart(br, unit_index, self.total_units, &mut self.expected_rst)?; - self.units_since_restart = 0; - } - Ok(self.restart_interval > 0 && self.units_since_restart == 0) - } - - pub(super) fn finish_unit(&mut self) { - self.units_since_restart += 1; - } -} - pub(super) fn consume_extended12_restart( br: &mut BitReader<'_>, mcu_index: u32, @@ -262,172 +216,27 @@ impl Extended12RestartTracker { } } -pub(super) fn lossless_predictor_value( - predictor: u8, - out: &[u8], - stride: usize, - x: usize, - y: usize, -) -> i32 { - lossless_predict(predictor, 128, x, y, |sx, sy| { - i32::from(out[sy * stride + sx]) - }) -} - -pub(super) fn lossless_predictor_color_into( - predictor: u8, - out: &[u8], - stride: usize, - x: usize, - y: usize, - component: usize, -) -> i32 { - lossless_predict(predictor, P::RESTART_PREDICTOR, x, y, |sx, sy| { - P::read_le(&out[sy * stride + (sx * 3 + component) * P::BYTES..]) - }) -} - -pub(super) fn lossless_predictor_gray_rows( - predictor: u8, - curr_row: &[u8], - prev_row: &[u8], - x: usize, - y: usize, -) -> i32 { - lossless_predict(predictor, P::RESTART_PREDICTOR, x, y, |sx, sy| { - let row = if sy == y { curr_row } else { prev_row }; - P::read_le(&row[sx * P::BYTES..]) - }) -} - -pub(super) fn lossless_predictor_color_rows( - predictor: u8, - curr_row: &[u8], - prev_row: &[u8], - x: usize, - y: usize, - component: usize, -) -> i32 { - lossless_predict(predictor, P::RESTART_PREDICTOR, x, y, |sx, sy| { - let row = if sy == y { curr_row } else { prev_row }; - P::read_le(&row[(sx * 3 + component) * P::BYTES..]) - }) -} - -pub(super) trait LosslessColorSampleTarget { - fn predict(&self, predictor: u8, component: usize) -> i32; - fn write(&mut self, component: usize, sample: P); -} - -pub(super) struct LosslessColorIntoSample<'a> { - pub(super) out: &'a mut [u8], - pub(super) stride: usize, - pub(super) x: usize, - pub(super) y: usize, -} - -impl LosslessColorSampleTarget

for LosslessColorIntoSample<'_> { - fn predict(&self, predictor: u8, component: usize) -> i32 { - lossless_predictor_color_into::

( - predictor, - self.out, - self.stride, - self.x, - self.y, - component, - ) - } - - fn write(&mut self, component: usize, sample: P) { - let offset = self.y * self.stride + (self.x * 3 + component) * P::BYTES; - sample.write_le(&mut self.out[offset..]); - } -} - -pub(super) struct LosslessColorRowSample<'a> { - pub(super) curr_row: &'a mut [u8], - pub(super) prev_row: &'a [u8], - pub(super) x: usize, - pub(super) y: usize, -} - -impl LosslessColorSampleTarget

for LosslessColorRowSample<'_> { - fn predict(&self, predictor: u8, component: usize) -> i32 { - lossless_predictor_color_rows::

( - predictor, - self.curr_row, - self.prev_row, - self.x, - self.y, - component, - ) - } - - fn write(&mut self, component: usize, sample: P) { - let offset = (self.x * 3 + component) * P::BYTES; - sample.write_le(&mut self.curr_row[offset..]); - } -} - -#[expect( - clippy::cast_possible_truncation, - reason = "component indices are validated against the JPEG maximum component count" -)] -pub(super) fn resolve_lossless_color_components( +/// Scan-component index feeding each of the three output channels. +pub(super) fn lossless_output_order( decode_plan: &PreparedDecodePlan, -) -> Result<[ResolvedLosslessColorComponent<'_>; 3], JpegError> { +) -> Result<[usize; 3], JpegError> { + let unsupported = || JpegError::UnsupportedComponentCount { + count: u8::try_from(decode_plan.components.len()).unwrap_or(u8::MAX), + }; if decode_plan.components.len() != 3 { - return Err(JpegError::UnsupportedComponentCount { - count: decode_plan.components.len() as u8, - }); + return Err(unsupported()); } - let resolve = |index: usize| { - let component = &decode_plan.components[index]; - if component.output_index >= 3 { - return Err(JpegError::UnsupportedComponentCount { - count: decode_plan.components.len() as u8, - }); + let mut order = [usize::MAX; 3]; + for (index, component) in decode_plan.components.iter().enumerate() { + let slot = order + .get_mut(component.output_index) + .ok_or_else(unsupported)?; + if *slot != usize::MAX { + return Err(unsupported()); } - Ok(ResolvedLosslessColorComponent { - h: usize::from(component.h), - v: usize::from(component.v), - output_index: component.output_index, - dc_table: decode_plan.dc_table(component)?, - }) - }; - Ok([resolve(0)?, resolve(1)?, resolve(2)?]) -} - -#[derive(Clone, Copy)] -pub(super) struct ResolvedLosslessColorComponent<'a> { - h: usize, - v: usize, - output_index: usize, - dc_table: DcHuffmanTable<'a>, -} - -pub(super) fn decode_lossless_color_sample( - br: &mut BitReader<'_>, - components: &[ResolvedLosslessColorComponent<'_>; 3], - predictor: u8, - restart_first_sample: bool, - target: &mut T, -) -> Result<(), JpegError> -where - P: LosslessSample, - T: LosslessColorSampleTarget

, -{ - for component in components { - let predicted = if restart_first_sample { - P::RESTART_PREDICTOR - } else { - target.predict(predictor, component.output_index) - }; - let diff = component.dc_table.decode_fast_dc(br)?; - let sample = P::from_i32(predicted + diff)?; - target.write(component.output_index, sample); + *slot = index; } - Ok(()) + Ok(order) } #[expect( @@ -461,124 +270,6 @@ pub(super) fn validate_lossless_color_plan( Ok(()) } -#[derive(Clone, Copy)] -pub(super) struct LosslessPlaneSample { - pub(super) x: usize, - pub(super) y: usize, - pub(super) restart_first_sample: bool, -} - -pub(super) fn decode_lossless_plane_sample( - br: &mut BitReader<'_>, - table: DcHuffmanTable<'_>, - predictor: u8, - plane: &mut [P], - width: usize, - sample: LosslessPlaneSample, -) -> Result<(), JpegError> { - let predicted = if sample.restart_first_sample { - P::RESTART_PREDICTOR - } else { - lossless_predictor_plane(predictor, plane, width, sample.x, sample.y) - }; - let diff = table.decode_fast_dc(br)?; - plane[sample.y * width + sample.x] = P::from_i32(predicted + diff)?; - Ok(()) -} - -pub(super) struct LosslessSampledColorPlanesMut<'a, P> { - pub(super) c0: &'a mut [P], - pub(super) c1: &'a mut [P], - pub(super) c2: &'a mut [P], - pub(super) dimensions: (usize, usize), - pub(super) chroma_dimensions: (usize, usize), -} - -impl

LosslessSampledColorPlanesMut<'_, P> { - fn component_plane(&mut self, output_index: usize) -> Option<(&mut [P], usize, usize)> { - match output_index { - 0 => Some((&mut *self.c0, self.dimensions.0, self.dimensions.1)), - 1 => Some(( - &mut *self.c1, - self.chroma_dimensions.0, - self.chroma_dimensions.1, - )), - 2 => Some(( - &mut *self.c2, - self.chroma_dimensions.0, - self.chroma_dimensions.1, - )), - _ => None, - } - } -} - -#[derive(Clone, Copy)] -pub(super) struct LosslessSampledMcu { - pub(super) x: usize, - pub(super) y: usize, - pub(super) restart_first_mcu: bool, -} - -#[expect( - clippy::cast_possible_truncation, - reason = "component indices are validated against the JPEG maximum component count" -)] -pub(super) fn decode_lossless_sampled_color_mcu

( - br: &mut BitReader<'_>, - components: &[ResolvedLosslessColorComponent<'_>; 3], - predictor: u8, - mcu: LosslessSampledMcu, - planes: &mut LosslessSampledColorPlanesMut<'_, P>, -) -> Result<(), JpegError> -where - P: LosslessSample, -{ - for component in components { - let Some((plane, plane_width, plane_height)) = - planes.component_plane(component.output_index) - else { - return Err(JpegError::UnsupportedComponentCount { - count: components.len() as u8, - }); - }; - for local_y in 0..component.v { - for local_x in 0..component.h { - let x = mcu.x * component.h + local_x; - let y = mcu.y * component.v + local_y; - if x >= plane_width || y >= plane_height { - continue; - } - decode_lossless_plane_sample( - br, - component.dc_table, - predictor, - plane, - plane_width, - LosslessPlaneSample { - x, - y, - restart_first_sample: mcu.restart_first_mcu && local_x == 0 && local_y == 0, - }, - )?; - } - } - } - Ok(()) -} - -pub(super) fn lossless_predictor_plane( - predictor: u8, - plane: &[P], - width: usize, - x: usize, - y: usize, -) -> i32 { - lossless_predict(predictor, P::RESTART_PREDICTOR, x, y, |sx, sy| { - plane[sy * width + sx].into() - }) -} - pub(super) struct LosslessColorPlanes<'a, P> { pub(super) c0: &'a [P], pub(super) c1: &'a [P], @@ -757,21 +448,5 @@ pub(super) fn upsample_h2v1_u16_at(row: &[u16], output_x: usize) -> u16 { upsample_h2v1_sample_at(row, output_x) } -pub(super) fn lossless_predictor_value_u16( - predictor: u8, - out: &[u8], - stride: usize, - x: usize, - y: usize, -) -> i32 { - lossless_predict(predictor, 32768, x, y, |sx, sy| { - i32::from(read_gray16_sample(out, sy * stride + sx * 2)) - }) -} - -pub(super) fn read_gray16_sample(out: &[u8], offset: usize) -> u16 { - u16::from_le_bytes([out[offset], out[offset + 1]]) -} - #[cfg(test)] mod tests; diff --git a/crates/j2k-jpeg/src/decoder/lossless_helpers/tests.rs b/crates/j2k-jpeg/src/decoder/lossless_helpers/tests.rs index 914f0bb77..5534f101f 100644 --- a/crates/j2k-jpeg/src/decoder/lossless_helpers/tests.rs +++ b/crates/j2k-jpeg/src/decoder/lossless_helpers/tests.rs @@ -5,9 +5,7 @@ use alloc::vec; use crate::{ColorSpace, Decoder, JpegError, MarkerKind}; use super::{ - lossless_predictor_color_into, lossless_predictor_color_rows, lossless_predictor_gray_rows, - lossless_predictor_plane, lossless_predictor_value, lossless_predictor_value_u16, - read_gray16_sample, restart_index_allocation_bytes, restart_index_for_stream, + lossless_output_order, restart_index_allocation_bytes, restart_index_for_stream, restart_segment_capacity, upsample_h2v1_u16_at, upsample_h2v1_u8_at, upsample_h2v2_u16_at, upsample_h2v2_u8_at, write_lossless_color16_sampled_output, write_lossless_color8_sampled_output, LosslessColorPlanes, LosslessColorSampling, @@ -128,41 +126,20 @@ fn restart_index_rejects_more_markers_than_geometry_allows() { } #[test] -fn predictors_read_interleaved_rows_and_planes_consistently() { - let gray = [10, 20, 30, 40]; - assert_eq!(lossless_predictor_value(1, &gray, 2, 1, 1), 30); - assert_eq!( - lossless_predictor_gray_rows::(2, &[30, 40], &[10, 20], 1, 1), - 20 - ); - - let color = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]; - assert_eq!( - lossless_predictor_color_into::(1, &color, 6, 1, 1, 2), - 9 - ); - assert_eq!( - lossless_predictor_color_rows::( - 2, - &[7, 8, 9, 10, 11, 12], - &[1, 2, 3, 4, 5, 6], - 1, - 1, - 1 - ), - 5 - ); - assert_eq!( - lossless_predictor_plane(7, &[10u8, 20, 30, 40], 2, 1, 1), - 25 - ); -} - -#[test] -fn gray16_helpers_use_little_endian_samples() { - let samples = [0x34, 0x12, 0x78, 0x56, 0xbc, 0x9a, 0xf0, 0xde]; - assert_eq!(read_gray16_sample(&samples, 2), 0x5678); - assert_eq!(lossless_predictor_value_u16(1, &samples, 4, 1, 1), 0x9abc); +fn output_order_follows_frame_component_order() { + for bytes in [ + j2k_test_support::lossless_predictor_rgb_3x3_jpeg(1), + j2k_test_support::lossless_predictor_ycbcr_3x3_jpeg(1), + ] { + let decoder = Decoder::new(&bytes).expect("lossless color fixture"); + assert_eq!(lossless_output_order(&decoder.plan), Ok([0, 1, 2])); + } + let gray_bytes = j2k_test_support::lossless_predictor_grayscale_3x3_jpeg(1); + let gray = Decoder::new(&gray_bytes).expect("lossless gray fixture"); + assert!(matches!( + lossless_output_order(&gray.plan), + Err(JpegError::UnsupportedComponentCount { count: 1 }) + )); } #[test] diff --git a/crates/j2k-jpeg/src/decoder/lossless_render.rs b/crates/j2k-jpeg/src/decoder/lossless_render.rs index 40d58b27f..71a113805 100644 --- a/crates/j2k-jpeg/src/decoder/lossless_render.rs +++ b/crates/j2k-jpeg/src/decoder/lossless_render.rs @@ -1,22 +1,22 @@ // SPDX-License-Identifier: MIT OR Apache-2.0 -//! Lossless grayscale and color rendering paths. +//! Lossless grayscale and color rendering paths. Every path pulls whole MCU +//! rows from one [`LosslessScan`] and stores them into its output. use super::{ allocate_output_buffer_with_live_budget, checked_scratch_len, convert_ycbcr16_to_rgb16_in_place, convert_ycbcr8_to_rgb8_in_place, copy_gray16_scaled_rect, copy_gray8_scaled_rect, copy_ycbcr16_row_to_rgb16, copy_ycbcr8_row_to_rgb8, - decode_lossless_color_sample, decode_lossless_sampled_color_mcu, finish_scan, - lossless_color_sampling, lossless_predictor_gray_rows, lossless_predictor_value, - lossless_predictor_value_u16, lossless_sampled_plane_layout, merged_warnings, - resolve_lossless_color_components, scaled_rect_covering, validate_lossless_color_plan, - write_lossless_color16_sampled_output, write_lossless_color8_sampled_output, BitReader, - ColorSpace, DecodeOutcome, Decoder, DownscaleFactor, JpegError, LosslessColorIntoSample, - LosslessColorPlanes, LosslessColorRowSample, LosslessColorSampling, LosslessRestartTracker, - LosslessSample, LosslessSampledColorPlanesMut, LosslessSampledMcu, LosslessSampledPlaneLayout, - Rect, RowSink, Vec, + lossless_color_sampling, lossless_output_order, lossless_sampled_plane_layout, merged_warnings, + scaled_rect_covering, validate_lossless_color_plan, write_lossless_color16_sampled_output, + write_lossless_color8_sampled_output, ColorSpace, DecodeOutcome, Decoder, DownscaleFactor, + JpegError, LosslessColorPlanes, LosslessColorSampling, LosslessSample, + LosslessSampledPlaneLayout, PreparedLosslessPlan, Rect, RowSink, Vec, }; -use crate::allocation::{try_reserve_for_len_with_live_budget, try_resize_filled}; +use crate::allocation::{ + try_reserve_for_len_with_live_budget, try_resize_filled, try_vec_with_capacity, +}; +use crate::lossless::scan::{LosslessComponentSpec, LosslessScan, LosslessScanSpec}; #[cfg(test)] mod tests; @@ -44,23 +44,78 @@ fn allocate_lossless_sampled_planes( Ok(OwnedLosslessSampledPlanes { c0, c1, c2 }) } -impl Decoder<'_> { - #[expect( - clippy::cast_possible_truncation, - reason = "row and column indices are bounded by validated u32 image dimensions" - )] - pub(super) fn decode_lossless_gray8_into( - &self, +/// Color output steps that depend on the sample container. +trait LosslessColorOutput: LosslessSample { + fn convert_ycbcr_in_place(out: &mut [u8], stride: usize, dimensions: (u32, u32)); + + fn write_sampled( out: &mut [u8], stride: usize, - ) -> Result { + color_space: ColorSpace, + sampling: LosslessColorSampling, + dimensions: (usize, usize), + planes: LosslessColorPlanes<'_, Self>, + ); +} + +impl LosslessColorOutput for u8 { + fn convert_ycbcr_in_place(out: &mut [u8], stride: usize, dimensions: (u32, u32)) { + convert_ycbcr8_to_rgb8_in_place(out, stride, dimensions); + } + + fn write_sampled( + out: &mut [u8], + stride: usize, + color_space: ColorSpace, + sampling: LosslessColorSampling, + dimensions: (usize, usize), + planes: LosslessColorPlanes<'_, Self>, + ) { + write_lossless_color8_sampled_output( + out, + stride, + color_space, + sampling, + dimensions, + planes, + ); + } +} + +impl LosslessColorOutput for u16 { + fn convert_ycbcr_in_place(out: &mut [u8], stride: usize, dimensions: (u32, u32)) { + convert_ycbcr16_to_rgb16_in_place(out, stride, dimensions); + } + + fn write_sampled( + out: &mut [u8], + stride: usize, + color_space: ColorSpace, + sampling: LosslessColorSampling, + dimensions: (usize, usize), + planes: LosslessColorPlanes<'_, Self>, + ) { + write_lossless_color16_sampled_output( + out, + stride, + color_space, + sampling, + dimensions, + planes, + ); + } +} + +impl Decoder<'_> { + /// The lossless plan, checked against the `P` output container. + fn lossless_plan_for(&self) -> Result<&PreparedLosslessPlan, JpegError> { let plan = self .lossless_plan .as_ref() .ok_or(JpegError::NotImplemented { sof: self.info.sof_kind, })?; - if plan.bit_depth != 8 { + if plan.bit_depth != P::BIT_DEPTH { return Err(JpegError::UnsupportedBitDepth { depth: plan.bit_depth, }); @@ -70,37 +125,76 @@ impl Decoder<'_> { predictor: plan.predictor, }); } - let dc_table = self.plan.dc_table_by_id(plan.dc_table)?; - - let (width, height) = plan.dimensions; - let scan_bytes = &self.bytes[plan.scan_offset..]; - let mut br = BitReader::new(scan_bytes); - let total_samples = width.saturating_mul(height); - let mut restart_tracker = - LosslessRestartTracker::new(self.plan.restart_interval, total_samples); - for y in 0..height as usize { - for x in 0..width as usize { - let sample_index = y as u32 * width + x as u32; - let restart_first_sample = restart_tracker.begin_unit(&mut br, sample_index)?; - let predictor = if restart_first_sample { - 128 - } else { - lossless_predictor_value(plan.predictor, out, stride, x, y) - }; - let diff = dc_table.decode_fast_dc(&mut br)?; - let sample = ::from_i32(predictor + diff)?; - out[y * stride + x] = sample; - restart_tracker.finish_unit(); - } + Ok(plan) + } + + /// Start the scan over every plan component, in scan order. + fn lossless_scan( + &self, + plan: &PreparedLosslessPlan, + ) -> Result, JpegError> { + let mut components = try_vec_with_capacity(self.plan.components.len())?; + for component in &self.plan.components { + components.push(LosslessComponentSpec { + h: component.h, + v: component.v, + table: self.plan.dc_table(component)?, + }); } + LosslessScan::new( + &self.bytes[plan.scan_offset..], + LosslessScanSpec { + dimensions: plan.dimensions, + predictor: plan.predictor, + precision: plan.bit_depth, + point_transform: plan.point_transform, + restart_interval: self.plan.restart_interval, + }, + &components, + ) + } - let scan_warnings = finish_scan(&mut br, true)?; + fn lossless_gray_scan(&self) -> Result, JpegError> { + let plan = self.lossless_plan_for::

()?; + if self.plan.components.len() != 1 { + return Err(JpegError::NotImplemented { + sof: self.info.sof_kind, + }); + } + self.lossless_scan(plan) + } + + fn finish_lossless_scan(&self, scan: LosslessScan<'_, '_>) -> Result { + let scan_warnings = scan.finish()?; Ok(DecodeOutcome { decoded: Rect::full(self.info.dimensions), warnings: merged_warnings(&self.warnings, scan_warnings)?, }) } + fn decode_lossless_gray_into( + &self, + out: &mut [u8], + stride: usize, + ) -> Result { + let mut scan = self.lossless_gray_scan::

()?; + let shift = scan.point_transform(); + let row_bytes = checked_scratch_len(&[self.info.dimensions.0 as usize, P::BYTES])?; + while scan.decode_mcu_row()? { + let y = scan.row_base(0); + P::store_row(scan.row(0, 0), shift, &mut out[y * stride..][..row_bytes]); + } + self.finish_lossless_scan(scan) + } + + pub(super) fn decode_lossless_gray8_into( + &self, + out: &mut [u8], + stride: usize, + ) -> Result { + self.decode_lossless_gray_into::(out, stride) + } + pub(super) fn decode_lossless_gray8_region_scaled_into( &self, out: &mut [u8], @@ -133,91 +227,47 @@ impl Decoder<'_> { Ok(outcome) } + /// Stream grayscale rows through `emit_row`, assembling each in `row`. #[expect( clippy::cast_possible_truncation, - reason = "row and column indices are bounded by validated u32 image dimensions" + reason = "row indices are bounded by validated u32 image dimensions" )] - pub(super) fn decode_lossless_gray_rows( + fn decode_lossless_gray_rows( &self, sink: &mut S, - prev_row: &mut Vec, - curr_row: &mut Vec, + row: &mut Vec, mut emit_row: impl FnMut(&mut S, u32, &[u8]) -> Result<(), JpegError>, ) -> Result where P: LosslessSample, S: RowSink, { - let plan = self - .lossless_plan - .as_ref() - .ok_or(JpegError::NotImplemented { - sof: self.info.sof_kind, - })?; - if plan.bit_depth != P::BIT_DEPTH { - return Err(JpegError::UnsupportedBitDepth { - depth: plan.bit_depth, - }); - } if self.info.color_space != ColorSpace::Grayscale { return Err(JpegError::NotImplemented { sof: self.info.sof_kind, }); } - if !(1..=7).contains(&plan.predictor) { - return Err(JpegError::UnsupportedPredictor { - predictor: plan.predictor, - }); + let mut scan = self.lossless_gray_scan::

()?; + let shift = scan.point_transform(); + let row_len = checked_scratch_len(&[self.info.dimensions.0 as usize, P::BYTES])?; + try_resize_filled(row, row_len, 0)?; + while scan.decode_mcu_row()? { + P::store_row(scan.row(0, 0), shift, &mut row[..row_len]); + emit_row(sink, scan.row_base(0) as u32, &row[..row_len])?; } - let dc_table = self.plan.dc_table_by_id(plan.dc_table)?; - - let (width, height) = plan.dimensions; - let width = width as usize; - let row_len = checked_scratch_len(&[width, P::BYTES])?; - try_resize_filled(prev_row, row_len, 0)?; - try_resize_filled(curr_row, row_len, 0)?; - - let scan_bytes = &self.bytes[plan.scan_offset..]; - let mut br = BitReader::new(scan_bytes); - let total_samples = plan.dimensions.0.saturating_mul(height); - let mut restart_tracker = - LosslessRestartTracker::new(self.plan.restart_interval, total_samples); - for y in 0..height as usize { - for x in 0..width { - let sample_index = y as u32 * plan.dimensions.0 + x as u32; - let restart_first_sample = restart_tracker.begin_unit(&mut br, sample_index)?; - let predictor = if restart_first_sample { - P::RESTART_PREDICTOR - } else { - lossless_predictor_gray_rows::

(plan.predictor, curr_row, prev_row, x, y) - }; - let diff = dc_table.decode_fast_dc(&mut br)?; - let sample = P::from_i32(predictor + diff)?; - sample.write_le(&mut curr_row[x * P::BYTES..]); - restart_tracker.finish_unit(); - } - emit_row(sink, y as u32, &curr_row[..row_len])?; - core::mem::swap(prev_row, curr_row); - } - - let scan_warnings = finish_scan(&mut br, true)?; - Ok(DecodeOutcome { - decoded: Rect::full(self.info.dimensions), - warnings: merged_warnings(&self.warnings, scan_warnings)?, - }) + self.finish_lossless_scan(scan) } pub(super) fn decode_lossless_gray8_rows( &self, sink: &mut S, - prev_row: &mut Vec, - curr_row: &mut Vec, + row: &mut Vec, rgb_row: &mut [u8], ) -> Result where S: RowSink, { - self.decode_lossless_gray_rows::(sink, prev_row, curr_row, |sink, y, gray_row| { + self.decode_lossless_gray_rows::(sink, row, |sink, y, gray_row| { let rgb_len = gray_row.len().saturating_mul(3); if rgb_row.len() < rgb_len { return Err(JpegError::OutputBufferTooSmall { @@ -232,15 +282,58 @@ impl Decoder<'_> { }) } - pub(super) fn decode_lossless_rgb8_into( + pub(super) fn decode_lossless_gray16_rows( + &self, + sink: &mut S, + row: &mut Vec, + ) -> Result + where + S: RowSink, + { + self.decode_lossless_gray_rows::(sink, row, S::write_row) + } + + pub(super) fn decode_lossless_gray16_into( &self, out: &mut [u8], stride: usize, ) -> Result { - self.decode_lossless_color8_output_into(out, stride, ColorSpace::Rgb) + self.decode_lossless_gray_into::(out, stride) } - pub(super) fn decode_lossless_color8_output_into( + pub(super) fn decode_lossless_gray16_region_scaled_into( + &self, + out: &mut [u8], + stride: usize, + roi: Rect, + downscale: DownscaleFactor, + external_live_bytes: usize, + ) -> Result { + if roi == Rect::full(self.info.dimensions) && downscale == DownscaleFactor::Full { + return self.decode_lossless_gray16_into(out, stride); + } + + let (width, height) = self.info.dimensions; + let full_stride = width as usize * 2; + let full_len = checked_scratch_len(&[full_stride, height as usize])?; + let (mut live_bytes, workspace_cap) = self.decode_phase_live_bytes(external_live_bytes)?; + let mut full = + allocate_output_buffer_with_live_budget(full_len, &mut live_bytes, workspace_cap)?; + let mut outcome = self.decode_lossless_gray16_into(&mut full, full_stride)?; + let output_rect = scaled_rect_covering(roi, downscale)?; + copy_gray16_scaled_rect( + &full, + (width, height), + output_rect, + downscale.denominator(), + out, + stride, + ); + outcome.decoded = roi; + Ok(outcome) + } + + fn decode_lossless_color_output_into( &self, out: &mut [u8], stride: usize, @@ -249,14 +342,14 @@ impl Decoder<'_> { match lossless_color_sampling(&self.info) { Some(LosslessColorSampling::S444) => { let outcome = - self.decode_lossless_color8_components_into(out, stride, color_space)?; + self.decode_lossless_color_components_into::

(out, stride, color_space)?; if color_space == ColorSpace::YCbCr { - convert_ycbcr8_to_rgb8_in_place(out, stride, self.info.dimensions); + P::convert_ycbcr_in_place(out, stride, self.info.dimensions); } Ok(outcome) } - Some(LosslessColorSampling::S422 | LosslessColorSampling::S420) => { - self.decode_lossless_color8_sampled_into(out, stride, color_space) + Some(sampling @ (LosslessColorSampling::S422 | LosslessColorSampling::S420)) => { + self.decode_lossless_color_sampled_into::

(out, stride, color_space, sampling) } None => Err(JpegError::NotImplemented { sof: self.info.sof_kind, @@ -264,114 +357,43 @@ impl Decoder<'_> { } } - #[expect( - clippy::cast_possible_truncation, - reason = "component row indices are bounded by validated u32 image dimensions" - )] - pub(super) fn decode_lossless_color_components_into

( + /// 4:4:4 color: each MCU row is one pixel row of all three components. + fn decode_lossless_color_components_into( &self, out: &mut [u8], stride: usize, color_space: ColorSpace, - ) -> Result - where - P: LosslessSample, - { - let plan = self - .lossless_plan - .as_ref() - .ok_or(JpegError::NotImplemented { - sof: self.info.sof_kind, - })?; + ) -> Result { + let plan = self.lossless_plan_for::

()?; validate_lossless_color_plan::

(plan, &self.plan, &self.info, color_space)?; - let components = resolve_lossless_color_components(&self.plan)?; - - let (width, height) = plan.dimensions; - let scan_bytes = &self.bytes[plan.scan_offset..]; - let mut br = BitReader::new(scan_bytes); - let total_pixels = width.saturating_mul(height); - let mut restart_tracker = - LosslessRestartTracker::new(self.plan.restart_interval, total_pixels); - for y in 0..height as usize { - for x in 0..width as usize { - let pixel_index = y as u32 * width + x as u32; - let restart_first_pixel = restart_tracker.begin_unit(&mut br, pixel_index)?; - decode_lossless_color_sample::( - &mut br, - &components, - plan.predictor, - restart_first_pixel, - &mut LosslessColorIntoSample { - out: &mut *out, - stride, - x, - y, - }, - )?; - restart_tracker.finish_unit(); - } + let [r, g, b] = lossless_output_order(&self.plan)?; + let mut scan = self.lossless_scan(plan)?; + let shift = scan.point_transform(); + let row_bytes = checked_scratch_len(&[self.info.dimensions.0 as usize, 3, P::BYTES])?; + while scan.decode_mcu_row()? { + let y = scan.row_base(0); + P::store_interleaved( + [scan.row(r, 0), scan.row(g, 0), scan.row(b, 0)], + shift, + &mut out[y * stride..][..row_bytes], + ); } - - let scan_warnings = finish_scan(&mut br, true)?; - Ok(DecodeOutcome { - decoded: Rect::full(self.info.dimensions), - warnings: merged_warnings(&self.warnings, scan_warnings)?, - }) + self.finish_lossless_scan(scan) } - pub(super) fn decode_lossless_color8_components_into( + /// 4:2:2 and 4:2:0 color: reconstruct full component planes, then + /// upsample and convert them into `out`. + fn decode_lossless_color_sampled_into( &self, out: &mut [u8], stride: usize, color_space: ColorSpace, + sampling: LosslessColorSampling, ) -> Result { - self.decode_lossless_color_components_into::(out, stride, color_space) - } - - #[expect( - clippy::cast_possible_truncation, - reason = "MCU and sample indices are bounded by validated u32 image dimensions" - )] - pub(super) fn decode_lossless_color_sampled_into

( - &self, - out: &mut [u8], - stride: usize, - color_space: ColorSpace, - write_output: impl FnOnce( - &mut [u8], - usize, - ColorSpace, - LosslessColorSampling, - (usize, usize), - LosslessColorPlanes<'_, P>, - ), - ) -> Result - where - P: LosslessSample, - { - let plan = self - .lossless_plan - .as_ref() - .ok_or(JpegError::NotImplemented { - sof: self.info.sof_kind, - })?; - let sampling = lossless_color_sampling(&self.info).ok_or(JpegError::NotImplemented { - sof: self.info.sof_kind, - })?; - if !matches!( - sampling, - LosslessColorSampling::S422 | LosslessColorSampling::S420 - ) { - return Err(JpegError::NotImplemented { - sof: self.info.sof_kind, - }); - } + let plan = self.lossless_plan_for::

()?; validate_lossless_color_plan::

(plan, &self.plan, &self.info, color_space)?; - let components = resolve_lossless_color_components(&self.plan)?; - - let (width, height) = plan.dimensions; - let width = width as usize; - let height = height as usize; + let order = lossless_output_order(&self.plan)?; + let (width, height) = (plan.dimensions.0 as usize, plan.dimensions.1 as usize); let layout = lossless_sampled_plane_layout(&self.info, super::DEFAULT_MAX_DECODE_BYTES)? .ok_or(JpegError::NotImplemented { sof: self.info.sof_kind, @@ -395,72 +417,53 @@ impl Decoder<'_> { cap: self.plan.scratch_bytes, }); } - let (chroma_width, chroma_height) = layout.chroma_dimensions; + let chroma_width = layout.chroma_dimensions.0; let OwnedLosslessSampledPlanes { mut c0, mut c1, mut c2, } = allocate_lossless_sampled_planes::

(layout, self.plan.scratch_bytes)?; - let mut planes = LosslessSampledColorPlanesMut { - c0: &mut c0, - c1: &mut c1, - c2: &mut c2, - dimensions: (width, height), - chroma_dimensions: (chroma_width, chroma_height), - }; - - let scan_bytes = &self.bytes[plan.scan_offset..]; - let mut br = BitReader::new(scan_bytes); - let total_mcus = (chroma_width * chroma_height) as u32; - let mut restart_tracker = - LosslessRestartTracker::new(self.plan.restart_interval, total_mcus); - for mcu_y in 0..chroma_height { - for mcu_x in 0..chroma_width { - let mcu_index = (mcu_y * chroma_width + mcu_x) as u32; - let restart_first_mcu = restart_tracker.begin_unit(&mut br, mcu_index)?; - decode_lossless_sampled_color_mcu::

( - &mut br, - &components, - plan.predictor, - LosslessSampledMcu { - x: mcu_x, - y: mcu_y, - restart_first_mcu, - }, - &mut planes, - )?; - restart_tracker.finish_unit(); + + let mut scan = self.lossless_scan(plan)?; + let shift = scan.point_transform(); + while scan.decode_mcu_row()? { + for (output, &component) in order.iter().enumerate() { + let (plane, plane_width) = match output { + 0 => (&mut c0, width), + 1 => (&mut c1, chroma_width), + _ => (&mut c2, chroma_width), + }; + let base = scan.row_base(component); + for index in 0..scan.rows_ready(component) { + let dst = &mut plane[(base + index) * plane_width..][..plane_width]; + for (dst, &sample) in dst.iter_mut().zip(scan.row(component, index)) { + *dst = P::from_sample(sample << shift); + } + } } } - - let scan_warnings = finish_scan(&mut br, true)?; - let LosslessSampledColorPlanesMut { c0, c1, c2, .. } = planes; - write_output( + let outcome = self.finish_lossless_scan(scan)?; + P::write_sampled( out, stride, color_space, sampling, (width, height), - LosslessColorPlanes { c0, c1, c2 }, + LosslessColorPlanes { + c0: &c0, + c1: &c1, + c2: &c2, + }, ); - Ok(DecodeOutcome { - decoded: Rect::full(self.info.dimensions), - warnings: merged_warnings(&self.warnings, scan_warnings)?, - }) + Ok(outcome) } - pub(super) fn decode_lossless_color8_sampled_into( + pub(super) fn decode_lossless_rgb8_into( &self, out: &mut [u8], stride: usize, - color_space: ColorSpace, ) -> Result { - self.decode_lossless_color_sampled_into::( - out, - stride, - color_space, - write_lossless_color8_sampled_output, - ) + self.decode_lossless_color_output_into::(out, stride, ColorSpace::Rgb) } pub(super) fn decode_lossless_ycbcr8_into( @@ -468,19 +471,36 @@ impl Decoder<'_> { out: &mut [u8], stride: usize, ) -> Result { - self.decode_lossless_color8_output_into(out, stride, ColorSpace::YCbCr) + self.decode_lossless_color_output_into::(out, stride, ColorSpace::YCbCr) } + pub(super) fn decode_lossless_rgb16_into( + &self, + out: &mut [u8], + stride: usize, + ) -> Result { + self.decode_lossless_color_output_into::(out, stride, ColorSpace::Rgb) + } + + pub(super) fn decode_lossless_ycbcr16_into( + &self, + out: &mut [u8], + stride: usize, + ) -> Result { + self.decode_lossless_color_output_into::(out, stride, ColorSpace::YCbCr) + } + + /// Stream 4:4:4 color rows, assembling each in `row` and converting + /// YCbCr through `conversion_row`. #[expect( clippy::cast_possible_truncation, - reason = "row and column indices are bounded by validated u32 image dimensions" + reason = "row indices are bounded by validated u32 image dimensions" )] - pub(super) fn decode_lossless_color_rows( + fn decode_lossless_color_rows( &self, sink: &mut S, - prev_row: &mut Vec, - curr_row: &mut Vec, - conversion_row: Option<&mut [u8]>, + row: &mut Vec, + mut conversion_row: Option<&mut [u8]>, color_space: ColorSpace, convert_row: impl Fn(&[u8], &mut [u8]), ) -> Result @@ -488,79 +508,55 @@ impl Decoder<'_> { P: LosslessSample, S: RowSink, { - let plan = self - .lossless_plan - .as_ref() - .ok_or(JpegError::NotImplemented { - sof: self.info.sof_kind, - })?; + let plan = self.lossless_plan_for::

()?; validate_lossless_color_plan::

(plan, &self.plan, &self.info, color_space)?; - let components = resolve_lossless_color_components(&self.plan)?; - - let (width, height) = plan.dimensions; - let width = width as usize; - let row_len = checked_scratch_len(&[width, 3, P::BYTES])?; - try_resize_filled(prev_row, row_len, 0)?; - try_resize_filled(curr_row, row_len, 0)?; - - let scan_bytes = &self.bytes[plan.scan_offset..]; - let mut br = BitReader::new(scan_bytes); - let total_pixels = plan.dimensions.0.saturating_mul(height); - let mut restart_tracker = - LosslessRestartTracker::new(self.plan.restart_interval, total_pixels); - let mut conversion_row = conversion_row; - for y in 0..height as usize { - for x in 0..width { - let pixel_index = y as u32 * plan.dimensions.0 + x as u32; - let restart_first_pixel = restart_tracker.begin_unit(&mut br, pixel_index)?; - decode_lossless_color_sample::( - &mut br, - &components, - plan.predictor, - restart_first_pixel, - &mut LosslessColorRowSample { - curr_row: &mut *curr_row, - prev_row: &*prev_row, - x, - y, - }, - )?; - restart_tracker.finish_unit(); - } - let row = if color_space == ColorSpace::YCbCr { - let row = conversion_row - .as_deref_mut() - .ok_or(JpegError::OutputBufferTooSmall { - required: row_len, - provided: 0, - })?; - if row.len() < row_len { - return Err(JpegError::OutputBufferTooSmall { - required: row_len, - provided: row.len(), - }); - } - convert_row(&curr_row[..row_len], &mut row[..row_len]); - &row[..row_len] + if lossless_color_sampling(&self.info) != Some(LosslessColorSampling::S444) { + return Err(JpegError::NotImplemented { + sof: self.info.sof_kind, + }); + } + let [r, g, b] = lossless_output_order(&self.plan)?; + let row_len = checked_scratch_len(&[plan.dimensions.0 as usize, 3, P::BYTES])?; + try_resize_filled(row, row_len, 0)?; + + let mut scan = self.lossless_scan(plan)?; + let shift = scan.point_transform(); + while scan.decode_mcu_row()? { + P::store_interleaved( + [scan.row(r, 0), scan.row(g, 0), scan.row(b, 0)], + shift, + &mut row[..row_len], + ); + let emitted = if color_space == ColorSpace::YCbCr { + let converted = + conversion_row + .as_deref_mut() + .ok_or(JpegError::OutputBufferTooSmall { + required: row_len, + provided: 0, + })?; + let provided = converted.len(); + let converted = + converted + .get_mut(..row_len) + .ok_or(JpegError::OutputBufferTooSmall { + required: row_len, + provided, + })?; + convert_row(&row[..row_len], converted); + &*converted } else { - &curr_row[..row_len] + &row[..row_len] }; - sink.write_row(y as u32, row)?; - core::mem::swap(prev_row, curr_row); + sink.write_row(scan.row_base(0) as u32, emitted)?; } - - let scan_warnings = finish_scan(&mut br, true)?; - Ok(DecodeOutcome { - decoded: Rect::full(self.info.dimensions), - warnings: merged_warnings(&self.warnings, scan_warnings)?, - }) + self.finish_lossless_scan(scan) } pub(super) fn decode_lossless_color8_rows( &self, sink: &mut S, - prev_row: &mut Vec, - curr_row: &mut Vec, + row: &mut Vec, conversion_row: Option<&mut [u8]>, color_space: ColorSpace, ) -> Result @@ -569,33 +565,17 @@ impl Decoder<'_> { { self.decode_lossless_color_rows::( sink, - prev_row, - curr_row, + row, conversion_row, color_space, copy_ycbcr8_row_to_rgb8, ) } - pub(super) fn decode_lossless_gray16_rows( - &self, - sink: &mut S, - prev_row: &mut Vec, - curr_row: &mut Vec, - ) -> Result - where - S: RowSink, - { - self.decode_lossless_gray_rows::(sink, prev_row, curr_row, |sink, y, row| { - sink.write_row(y, row) - }) - } - pub(super) fn decode_lossless_color16_rows( &self, sink: &mut S, - prev_row: &mut Vec, - curr_row: &mut Vec, + row: &mut Vec, conversion_row: Option<&mut [u8]>, color_space: ColorSpace, ) -> Result @@ -604,163 +584,10 @@ impl Decoder<'_> { { self.decode_lossless_color_rows::( sink, - prev_row, - curr_row, + row, conversion_row, color_space, copy_ycbcr16_row_to_rgb16, ) } - - pub(super) fn decode_lossless_rgb16_into( - &self, - out: &mut [u8], - stride: usize, - ) -> Result { - self.decode_lossless_color16_output_into(out, stride, ColorSpace::Rgb) - } - - pub(super) fn decode_lossless_color16_output_into( - &self, - out: &mut [u8], - stride: usize, - color_space: ColorSpace, - ) -> Result { - match lossless_color_sampling(&self.info) { - Some(LosslessColorSampling::S444) => { - let outcome = - self.decode_lossless_color16_components_into(out, stride, color_space)?; - if color_space == ColorSpace::YCbCr { - convert_ycbcr16_to_rgb16_in_place(out, stride, self.info.dimensions); - } - Ok(outcome) - } - Some(LosslessColorSampling::S422 | LosslessColorSampling::S420) => { - self.decode_lossless_color16_sampled_into(out, stride, color_space) - } - None => Err(JpegError::NotImplemented { - sof: self.info.sof_kind, - }), - } - } - - pub(super) fn decode_lossless_color16_components_into( - &self, - out: &mut [u8], - stride: usize, - color_space: ColorSpace, - ) -> Result { - self.decode_lossless_color_components_into::(out, stride, color_space) - } - - pub(super) fn decode_lossless_color16_sampled_into( - &self, - out: &mut [u8], - stride: usize, - color_space: ColorSpace, - ) -> Result { - self.decode_lossless_color_sampled_into::( - out, - stride, - color_space, - write_lossless_color16_sampled_output, - ) - } - - pub(super) fn decode_lossless_ycbcr16_into( - &self, - out: &mut [u8], - stride: usize, - ) -> Result { - self.decode_lossless_color16_output_into(out, stride, ColorSpace::YCbCr) - } - - #[expect( - clippy::cast_possible_truncation, - reason = "row and column indices are bounded by validated u32 image dimensions" - )] - pub(super) fn decode_lossless_gray16_into( - &self, - out: &mut [u8], - stride: usize, - ) -> Result { - let plan = self - .lossless_plan - .as_ref() - .ok_or(JpegError::NotImplemented { - sof: self.info.sof_kind, - })?; - if plan.bit_depth != 16 { - return Err(JpegError::UnsupportedBitDepth { - depth: plan.bit_depth, - }); - } - if !(1..=7).contains(&plan.predictor) { - return Err(JpegError::UnsupportedPredictor { - predictor: plan.predictor, - }); - } - let dc_table = self.plan.dc_table_by_id(plan.dc_table)?; - - let (width, height) = plan.dimensions; - let scan_bytes = &self.bytes[plan.scan_offset..]; - let mut br = BitReader::new(scan_bytes); - let total_samples = width.saturating_mul(height); - let mut restart_tracker = - LosslessRestartTracker::new(self.plan.restart_interval, total_samples); - for y in 0..height as usize { - for x in 0..width as usize { - let sample_index = y as u32 * width + x as u32; - let restart_first_sample = restart_tracker.begin_unit(&mut br, sample_index)?; - let predictor = if restart_first_sample { - 32768 - } else { - lossless_predictor_value_u16(plan.predictor, out, stride, x, y) - }; - let diff = dc_table.decode_fast_dc(&mut br)?; - let sample = ::from_i32(predictor + diff)?; - let offset = y * stride + x * 2; - sample.write_le(&mut out[offset..offset + 2]); - restart_tracker.finish_unit(); - } - } - - let scan_warnings = finish_scan(&mut br, true)?; - Ok(DecodeOutcome { - decoded: Rect::full(self.info.dimensions), - warnings: merged_warnings(&self.warnings, scan_warnings)?, - }) - } - - pub(super) fn decode_lossless_gray16_region_scaled_into( - &self, - out: &mut [u8], - stride: usize, - roi: Rect, - downscale: DownscaleFactor, - external_live_bytes: usize, - ) -> Result { - if roi == Rect::full(self.info.dimensions) && downscale == DownscaleFactor::Full { - return self.decode_lossless_gray16_into(out, stride); - } - - let (width, height) = self.info.dimensions; - let full_stride = width as usize * 2; - let full_len = checked_scratch_len(&[full_stride, height as usize])?; - let (mut live_bytes, workspace_cap) = self.decode_phase_live_bytes(external_live_bytes)?; - let mut full = - allocate_output_buffer_with_live_budget(full_len, &mut live_bytes, workspace_cap)?; - let mut outcome = self.decode_lossless_gray16_into(&mut full, full_stride)?; - let output_rect = scaled_rect_covering(roi, downscale)?; - copy_gray16_scaled_rect( - &full, - (width, height), - output_rect, - downscale.denominator(), - out, - stride, - ); - outcome.decoded = roi; - Ok(outcome) - } } diff --git a/crates/j2k-jpeg/src/decoder/plan.rs b/crates/j2k-jpeg/src/decoder/plan.rs index f6bb2c7a9..c0063a355 100644 --- a/crates/j2k-jpeg/src/decoder/plan.rs +++ b/crates/j2k-jpeg/src/decoder/plan.rs @@ -249,7 +249,6 @@ impl Decoder<'_> { ensure_prepared_construction_fits(header, prepared_bytes)?; let workspace_cap = decode_workspace_cap(header, prepared_bytes)?; let mut components = construction.try_vec(scan.components.len())?; - let mut first_dc_table = None; for scan_component in scan.components.iter().copied() { let component_index = find_component_index(&header.component_ids, scan_component.id) .ok_or(JpegError::UnknownScanComponent { @@ -274,7 +273,6 @@ impl Decoder<'_> { class: 0, id: scan_component.dc_table, })?; - first_dc_table.get_or_insert(dc_table); components.push(PreparedComponentPlan { h, v, @@ -304,9 +302,7 @@ impl Decoder<'_> { let lossless = PreparedLosslessPlan { predictor: scan.ss, bit_depth: info.bit_depth, - dc_table: first_dc_table.ok_or(JpegError::MissingMarker { - marker: MarkerKind::Sos, - })?, + point_transform: scan.al, dimensions: info.dimensions, scan_offset: plan.scan_offset, }; diff --git a/crates/j2k-jpeg/src/decoder/plan/construction.rs b/crates/j2k-jpeg/src/decoder/plan/construction.rs index 2d65ea9d6..e585aed93 100644 --- a/crates/j2k-jpeg/src/decoder/plan/construction.rs +++ b/crates/j2k-jpeg/src/decoder/plan/construction.rs @@ -92,8 +92,9 @@ impl PreparedConstructionBudget { reason: "prepared construction context bytes exceed its live-byte ledger", }, )?; - let table = - ctx.resolve_huffman_table_with_live_budget(raw, role, &mut self.live_bytes, self.cap)?; + let table = ctx + .resolve_huffman_table_with_live_budget(raw, role, &mut self.live_bytes, self.cap)? + .clone(); let context_bytes = ctx.retained_allocation_bytes(); let expected = checked_add_under_cap(non_context_bytes, context_bytes, self.cap)?; if self.live_bytes != expected { diff --git a/crates/j2k-jpeg/src/decoder/rows.rs b/crates/j2k-jpeg/src/decoder/rows.rs index 85cf5ddcf..e5b487b24 100644 --- a/crates/j2k-jpeg/src/decoder/rows.rs +++ b/crates/j2k-jpeg/src/decoder/rows.rs @@ -86,7 +86,6 @@ impl Decoder<'_> { pool.prepare_lossless_rows(predictor_len, sink_len, scratch_bytes)?; let result = self.decode_lossless_gray8_rows( sink, - &mut pool.lossless_prev_row, &mut pool.lossless_curr_row, &mut rows.top_row, ); @@ -96,18 +95,13 @@ impl Decoder<'_> { (ColorSpace::Grayscale, 16) => { let predictor_len = checked_scratch_len(&[width, 2])?; let _rows = pool.prepare_lossless_rows(predictor_len, 0, scratch_bytes)?; - self.decode_lossless_gray16_rows( - sink, - &mut pool.lossless_prev_row, - &mut pool.lossless_curr_row, - ) + self.decode_lossless_gray16_rows(sink, &mut pool.lossless_curr_row) } (ColorSpace::Rgb, 8) => { let predictor_len = checked_scratch_len(&[width, 3])?; let _rows = pool.prepare_lossless_rows(predictor_len, 0, scratch_bytes)?; self.decode_lossless_color8_rows( sink, - &mut pool.lossless_prev_row, &mut pool.lossless_curr_row, None, ColorSpace::Rgb, @@ -118,7 +112,6 @@ impl Decoder<'_> { let mut rows = pool.prepare_lossless_rows(row_len, row_len, scratch_bytes)?; let result = self.decode_lossless_color8_rows( sink, - &mut pool.lossless_prev_row, &mut pool.lossless_curr_row, Some(&mut rows.top_row), ColorSpace::YCbCr, @@ -131,7 +124,6 @@ impl Decoder<'_> { let _rows = pool.prepare_lossless_rows(predictor_len, 0, scratch_bytes)?; self.decode_lossless_color16_rows( sink, - &mut pool.lossless_prev_row, &mut pool.lossless_curr_row, None, ColorSpace::Rgb, @@ -142,7 +134,6 @@ impl Decoder<'_> { let mut rows = pool.prepare_lossless_rows(row_len, row_len, scratch_bytes)?; let result = self.decode_lossless_color16_rows( sink, - &mut pool.lossless_prev_row, &mut pool.lossless_curr_row, Some(&mut rows.top_row), ColorSpace::YCbCr, diff --git a/crates/j2k-jpeg/src/decoder/scratch.rs b/crates/j2k-jpeg/src/decoder/scratch.rs index 163675e9e..5c0c85734 100644 --- a/crates/j2k-jpeg/src/decoder/scratch.rs +++ b/crates/j2k-jpeg/src/decoder/scratch.rs @@ -7,6 +7,7 @@ use super::{ }; use crate::entropy::sequential::StripeLayout; use crate::info::{ColorSpace, Info, SamplingFactors}; +use crate::lossless::scan::lossless_scan_allocation_bytes; #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub(super) struct LosslessSampledPlaneLayout { @@ -137,11 +138,27 @@ pub(super) fn additional_decode_scratch_bytes( Ok(additional) } +/// Largest lossless workspace: the scan's own row buffers plus either the +/// sampled color planes or the streaming row buffers. pub(super) fn compute_lossless_scratch_bytes(info: &Info, cap: usize) -> Result { let sampled_planes = lossless_sampled_plane_layout(info, cap)?.map_or(0, |layout| layout.total_bytes); let row_scratch = compute_lossless_row_scratch_bytes(info, cap)?; - Ok(sampled_planes.max(row_scratch)) + let scan_rows = lossless_scan_allocation_bytes(info.dimensions, info.sampling.components())?; + let total = sampled_planes + .max(row_scratch) + .checked_add(scan_rows) + .ok_or(JpegError::MemoryCapExceeded { + requested: usize::MAX, + cap, + })?; + if total > cap { + return Err(JpegError::MemoryCapExceeded { + requested: total, + cap, + }); + } + Ok(total) } pub(super) fn lossless_sampled_plane_layout( diff --git a/crates/j2k-jpeg/src/entropy/huffman.rs b/crates/j2k-jpeg/src/entropy/huffman.rs index 7479a7048..ce287a96f 100644 --- a/crates/j2k-jpeg/src/entropy/huffman.rs +++ b/crates/j2k-jpeg/src/entropy/huffman.rs @@ -361,9 +361,47 @@ impl DcHuffmanTable<'_> { self.0.decode_fast_dc(br) } + #[expect( + clippy::inline_always, + reason = "progressive DC scans decode one symbol per block; an out-of-line call spills the reader" + )] + #[inline(always)] pub(crate) fn decode(self, br: &mut BitReader<'_>) -> Result { self.0.decode(br) } + + /// Decode one lossless (SOF3) difference, T.81 H.1.2.2, modulo 2^16. + /// + /// Categories 0..=15 extend like DC differences. Category 16 carries no + /// extra bits and always means 32768. Reconstruction is modulo 2^16 + /// (T.81 H.2.1), so the wrapped `u16` is the exact difference. + #[expect( + clippy::inline_always, + reason = "measured Huffman lookup hot path requires cross-helper inlining" + )] + #[inline(always)] + #[expect( + clippy::cast_possible_truncation, + clippy::cast_sign_loss, + reason = "lossless differences are defined modulo 2^16, so wrapping into u16 is exact" + )] + pub(crate) fn decode_lossless_diff(self, br: &mut BitReader<'_>) -> Result { + br.ensure_bits_padded(FAST_BITS)?; + let peek = br.peek_bits(FAST_BITS) as usize; + let packed = self.0.packed[peek]; + if packed != 0 { + br.consume_bits((packed & DC_FAST_LEN_MASK) as u8); + return Ok((packed >> DC_FAST_VALUE_SHIFT) as u16); + } + match self.0.decode(br)? { + 16 => Ok(0x8000), + ssss @ 0..=15 => Ok(br.receive_extend(ssss)? as u16), + _ => Err(JpegError::HuffmanDecode { + mcu: 0, + reason: HuffmanFailure::InvalidSymbol, + }), + } + } } impl AcHuffmanTable<'_> { @@ -385,6 +423,11 @@ impl AcHuffmanTable<'_> { self.0.skip_fast_ac(br) } + #[expect( + clippy::inline_always, + reason = "progressive AC scans decode every symbol through this; an out-of-line call spills the reader" + )] + #[inline(always)] pub(crate) fn decode(self, br: &mut BitReader<'_>) -> Result { self.0.decode(br) } diff --git a/crates/j2k-jpeg/src/entropy/progressive.rs b/crates/j2k-jpeg/src/entropy/progressive.rs index f62242226..fa5fd8154 100644 --- a/crates/j2k-jpeg/src/entropy/progressive.rs +++ b/crates/j2k-jpeg/src/entropy/progressive.rs @@ -46,7 +46,7 @@ pub(crate) fn decode_progressive( component_image_capacity_bytes(images.capacity(), &images)?, plan.scratch_bytes, )?; - emit_component_images(plan, &images, image_live_bytes, writer)?; + emit_component_images(plan, backend, &images, image_live_bytes, writer)?; Ok(Vec::new()) } diff --git a/crates/j2k-jpeg/src/entropy/progressive/allocation.rs b/crates/j2k-jpeg/src/entropy/progressive/allocation.rs index b8b66c33e..234abc037 100644 --- a/crates/j2k-jpeg/src/entropy/progressive/allocation.rs +++ b/crates/j2k-jpeg/src/entropy/progressive/allocation.rs @@ -82,6 +82,36 @@ pub(super) fn coefficient_capacity_bytes( Ok(total) } +/// Per-block nonzero masks for the scan phase, parallel to `coeffs` and +/// charged on top of `initial_live_bytes`. Returns the masks and the live +/// bytes after them. +pub(super) fn allocate_nonzero_masks( + plan: &PreparedProgressivePlan, + coeffs: &[Vec<[i32; 64]>], + initial_live_bytes: usize, +) -> Result<(Vec>, usize), JpegError> { + let mut live_bytes = initial_live_bytes; + let mut masks = Vec::new(); + try_reserve_for_len_with_live_budget( + &mut masks, + coeffs.len(), + &mut live_bytes, + plan.scratch_bytes, + )?; + for component_coeffs in coeffs { + let mut component_masks = Vec::new(); + try_reserve_for_len_with_live_budget( + &mut component_masks, + component_coeffs.len(), + &mut live_bytes, + plan.scratch_bytes, + )?; + component_masks.resize(component_coeffs.len(), 0u64); + masks.push(component_masks); + } + Ok((masks, live_bytes)) +} + pub(super) fn checked_phase_capacity( initial: usize, additional: usize, diff --git a/crates/j2k-jpeg/src/entropy/progressive/render.rs b/crates/j2k-jpeg/src/entropy/progressive/render.rs index cce31a1e0..1b4d0eac7 100644 --- a/crates/j2k-jpeg/src/entropy/progressive/render.rs +++ b/crates/j2k-jpeg/src/entropy/progressive/render.rs @@ -6,9 +6,7 @@ use alloc::vec::Vec; use crate::allocation::{checked_allocation_len, try_reserve_for_len_with_live_budget}; use crate::backend::Backend; -use crate::color::upsample::{ - upsample_h1v2_fancy_row, upsample_h2v1_fancy_row, upsample_h2v2_fancy_row, -}; +use crate::color::upsample::{upsample_h1v2_fancy_row, upsample_h2v1_fancy_row}; use crate::entropy::block::clamp_i16; use crate::entropy::ZIGZAG; use crate::error::JpegError; @@ -38,14 +36,11 @@ pub(super) fn render_component_images( { let mut dequant = [0i16; 64]; let mut pixels = [0u8; 64]; + let natural_quant = natural_order_quant(&component.quant); for by in 0..component.block_rows as usize { for bx in 0..component.block_cols as usize { let block_index = by * component.block_cols as usize + bx; - dequantize_block( - &component_coeffs[block_index], - &component.quant, - &mut dequant, - ); + dequantize_block(&component_coeffs[block_index], &natural_quant, &mut dequant); backend.idct(&dequant, &mut pixels); deposit_block(&mut image.plane, image.stride, bx * 8, by * 8, &pixels); } @@ -54,12 +49,21 @@ pub(super) fn render_component_images( Ok(images) } -fn dequantize_block(coeffs: &[i32; 64], quant: &[u16; 64], out: &mut [i16; 64]) { - out.fill(0); - for k in 0..64 { - let natural_idx = usize::from(ZIGZAG[k]); - let value = coeffs[natural_idx].wrapping_mul(i32::from(quant[k])); - out[natural_idx] = clamp_i16(value); +/// Quantization table (stored in zigzag order, T.81 B.2.4.1) permuted to the +/// natural order the coefficients are kept in. +fn natural_order_quant(quant: &[u16; 64]) -> [u16; 64] { + let mut natural = [0u16; 64]; + for (&q, &natural_idx) in quant.iter().zip(&ZIGZAG) { + natural[usize::from(natural_idx)] = q; + } + natural +} + +/// Dequantize a natural-order block with a natural-order quant table; a +/// straight lane-wise loop the compiler vectorizes. +fn dequantize_block(coeffs: &[i32; 64], natural_quant: &[u16; 64], out: &mut [i16; 64]) { + for ((out, &coeff), &q) in out.iter_mut().zip(coeffs).zip(natural_quant) { + *out = clamp_i16(coeff.wrapping_mul(i32::from(q))); } } @@ -73,6 +77,7 @@ fn deposit_block(plane: &mut [u8], stride: usize, x: usize, y: usize, block: &[u pub(super) fn emit_component_images( plan: &PreparedProgressivePlan, + backend: Backend, images: &[ComponentImage], image_live_bytes: usize, writer: &mut W, @@ -94,7 +99,7 @@ pub(super) fn emit_component_images( )?; gray.resize(width_usize, 0u8); for y in 0..height { - upsample_component_row(plan, 0, image, y, &mut gray); + upsample_component_row(plan, backend, 0, image, y, &mut gray); writer.write_gray_row(y, &gray)?; } return Ok(()); @@ -116,9 +121,9 @@ pub(super) fn emit_component_images( try_reserve_for_len_with_live_budget(&mut c, width_usize, &mut live_bytes, plan.scratch_bytes)?; c.resize(width_usize, 0u8); for y in 0..height { - upsample_component_row(plan, first, &images[first], y, &mut a); - upsample_component_row(plan, second, &images[second], y, &mut b); - upsample_component_row(plan, third, &images[third], y, &mut c); + upsample_component_row(plan, backend, first, &images[first], y, &mut a); + upsample_component_row(plan, backend, second, &images[second], y, &mut b); + upsample_component_row(plan, backend, third, &images[third], y, &mut c); match plan.color_space { ColorSpace::YCbCr => writer.write_ycbcr_row(y, &a, &b, &c)?, ColorSpace::Rgb => writer.write_rgb_row(y, &a, &b, &c)?, @@ -145,6 +150,7 @@ fn component_slot(plan: &PreparedProgressivePlan, output_index: usize) -> Result fn upsample_component_row( plan: &PreparedProgressivePlan, + backend: Backend, component_index: usize, image: &ComponentImage, y: u32, @@ -168,7 +174,7 @@ fn upsample_component_row( let prev = component_row(component, image, prev_y); let curr = component_row(component, image, sample_y); let next = component_row(component, image, next_y); - upsample_h2v2_fancy_row(prev, curr, next, out.len(), y % 2 == 1, out); + backend.upsample_h2v2_fancy_row([prev, curr, next], out.len(), y % 2 == 1, out); } else if h_ratio == 1 && v_ratio == 2 { let sample_y = ((y / 2) as usize).min(component.sample_height.saturating_sub(1) as usize); let prev_y = sample_y.saturating_sub(1); diff --git a/crates/j2k-jpeg/src/entropy/progressive/scan.rs b/crates/j2k-jpeg/src/entropy/progressive/scan.rs index 15c818604..1d6bd6220 100644 --- a/crates/j2k-jpeg/src/entropy/progressive/scan.rs +++ b/crates/j2k-jpeg/src/entropy/progressive/scan.rs @@ -5,12 +5,14 @@ use alloc::vec::Vec; use crate::allocation::try_reserve_for_len_with_live_budget; +use crate::entropy::huffman::{AcHuffmanTable, DcHuffmanTable}; use crate::entropy::ZIGZAG; use crate::error::{HuffmanFailure, JpegError}; use crate::internal::bit_reader::BitReader; use super::allocation::{ - allocate_coefficients, checked_phase_capacity, coefficient_capacity_bytes, + allocate_coefficients, allocate_nonzero_masks, checked_phase_capacity, + coefficient_capacity_bytes, }; use super::model::{ PreparedProgressiveComponentPlan, PreparedProgressivePlan, PreparedProgressiveScan, @@ -18,13 +20,52 @@ use super::model::{ }; use super::terminal::finish_progressive_scan; -struct ProgressiveBlockTarget<'a> { +struct ProgressiveBlockTarget<'a, 't> { component: &'a PreparedProgressiveComponentPlan, scan_component: &'a PreparedProgressiveScanComponent, + tables: ScanTables<'t>, block_x: u32, block_y: u32, } +/// Huffman tables one scan component uses, resolved once per scan: DC first +/// scans need the DC table, AC scans the AC table, DC refinement neither. +#[derive(Clone, Copy, Default)] +struct ScanTables<'t> { + dc: Option>, + ac: Option>, +} + +/// Largest component count of one JPEG scan (T.81 B.2.3). +const MAX_SCAN_COMPONENTS: usize = 4; + +fn resolve_scan_tables<'t>( + plan: &'t PreparedProgressivePlan, + scan: &PreparedProgressiveScan, + scan_components: &[PreparedProgressiveScanComponent], +) -> Result<[ScanTables<'t>; MAX_SCAN_COMPONENTS], JpegError> { + let mut tables = [ScanTables::default(); MAX_SCAN_COMPONENTS]; + if scan_components.len() > MAX_SCAN_COMPONENTS { + return Err(JpegError::InternalInvariant { + reason: "prepared progressive scan has more than four components", + }); + } + for (slot, scan_component) in tables.iter_mut().zip(scan_components) { + if scan.ss != 0 { + slot.ac = Some(plan.ac_table(scan_component.ac_table)?); + } else if scan.ah == 0 { + slot.dc = Some(plan.dc_table(scan_component.dc_table)?); + } + } + Ok(tables) +} + +fn missing_table() -> JpegError { + JpegError::InternalInvariant { + reason: "progressive scan references a missing prepared Huffman table", + } +} + pub(crate) fn decode_progressive_dct_blocks( plan: &PreparedProgressivePlan, bytes: &[u8], @@ -36,18 +77,32 @@ pub(crate) fn decode_progressive_dct_blocks( coefficient_capacity_bytes(coeffs.capacity(), &coeffs)?, plan.scratch_bytes, )?; + let (mut nonzero, scan_live_bytes) = + allocate_nonzero_masks(plan, &coeffs, coefficient_live_bytes)?; for scan in &plan.scans { - decode_progressive_scan(plan, scan, bytes, &mut coeffs, coefficient_live_bytes)?; + let mut state = CoefficientState { + blocks: &mut coeffs, + nonzero: &mut nonzero, + }; + decode_progressive_scan(plan, scan, bytes, &mut state, scan_live_bytes)?; } Ok(ProgressiveDctBlocks { quantized: coeffs }) } +/// Coefficient blocks, plus a mask per block whose bit `k` is set when the +/// coefficient at zigzag position `k` (1..=63) is nonzero. The masks let AC +/// refinement visit only nonzero coefficients. +struct CoefficientState<'a> { + blocks: &'a mut [Vec<[i32; 64]>], + nonzero: &'a mut [Vec], +} + fn decode_progressive_scan( plan: &PreparedProgressivePlan, scan: &PreparedProgressiveScan, bytes: &[u8], - coeffs: &mut [Vec<[i32; 64]>], - coefficient_live_bytes: usize, + coeffs: &mut CoefficientState<'_>, + scan_live_bytes: usize, ) -> Result<(), JpegError> { let scan_bytes = bytes .get(scan.entropy_offset..) @@ -56,7 +111,7 @@ fn decode_progressive_scan( expected: 1, })?; let mut br = BitReader::new_with_eof_padding(scan_bytes, scan.terminal_code == 0); - let mut live_bytes = coefficient_live_bytes; + let mut live_bytes = scan_live_bytes; let mut dc_predictors = Vec::new(); try_reserve_for_len_with_live_budget( &mut dc_predictors, @@ -67,36 +122,109 @@ fn decode_progressive_scan( dc_predictors.resize(plan.components.len(), 0i32); let mut eob_run = 0u32; let restart = u32::from(scan.restart_interval.unwrap_or(0)); - let mut mcus_since_restart = 0u32; - let mut expected_rst = 0u8; let total_mcus = scan_mcu_count(plan, scan)?; + let scan_components = plan.scan_components(scan)?; + let tables = if total_mcus == 0 { + [ScanTables::default(); MAX_SCAN_COMPONENTS] + } else { + resolve_scan_tables(plan, scan, scan_components)? + }; + let context = ScanContext { + plan, + scan, + scan_components, + tables, + }; + + let state = (&mut dc_predictors[..], &mut eob_run); + let counts = (restart, total_mcus); + match (scan.ss == 0, scan.ah == 0) { + (true, true) => decode_scan_mcus::(&context, &mut br, coeffs, state, counts), + (true, false) => decode_scan_mcus::(&context, &mut br, coeffs, state, counts), + (false, true) => decode_scan_mcus::(&context, &mut br, coeffs, state, counts), + (false, false) => decode_scan_mcus::(&context, &mut br, coeffs, state, counts), + }?; + + finish_progressive_scan(&mut br, scan_bytes, scan, eob_run) +} + +/// Scan kinds (T.81 G.1.2): `Ss == 0` selects DC, `Ah == 0` a first pass. +/// Each kind gets its own monomorphized MCU loop, so a loop only carries +/// the one block decoder it needs. +const DC_FIRST: u8 = 0; +const DC_REFINE: u8 = 1; +const AC_FIRST: u8 = 2; +const AC_REFINE: u8 = 3; + +/// Decode every MCU of one scan of kind `KIND`. Kept out of line per kind: +/// its callee chain is inlined so the reader stays in registers. +fn decode_scan_mcus( + context: &ScanContext<'_>, + br: &mut BitReader<'_>, + coeffs: &mut CoefficientState<'_>, + (dc_predictors, eob_run): (&mut [i32], &mut u32), + (restart, total_mcus): (u32, u32), +) -> Result<(), JpegError> { + // Decode through a local copy so the accumulator stays in registers; it + // is written back even on error. + let mut local = br.clone(); + let decoded = decode_scan_mcus_local::( + context, + &mut local, + coeffs, + (dc_predictors, eob_run), + (restart, total_mcus), + ); + *br = local; + decoded +} +#[expect( + clippy::inline_always, + reason = "inlining keeps the caller's local bit reader in registers for the whole scan" +)] +#[inline(always)] +fn decode_scan_mcus_local( + context: &ScanContext<'_>, + br: &mut BitReader<'_>, + coeffs: &mut CoefficientState<'_>, + (dc_predictors, eob_run): (&mut [i32], &mut u32), + (restart, total_mcus): (u32, u32), +) -> Result<(), JpegError> { + if let [scan_component] = context.scan_components { + return decode_single_component_scan::( + context, + scan_component, + br, + coeffs, + (dc_predictors, eob_run), + restart, + ); + } + let mut mcus_since_restart = 0u32; + let mut expected_rst = 0u8; for mcu_index in 0..total_mcus { if restart > 0 && mcus_since_restart == restart { - consume_restart( - &mut br, - mcu_index, - total_mcus, - &mut expected_rst, - &mut dc_predictors, - &mut eob_run, - )?; + let (restarted, next_rst) = + br.clone().restarted(expected_rst, mcu_index, total_mcus)?; + *br = restarted; + expected_rst = next_rst; + dc_predictors.fill(0); + *eob_run = 0; mcus_since_restart = 0; } - - decode_progressive_mcu( - plan, - scan, - &mut br, - coeffs, - &mut dc_predictors, - &mut eob_run, - mcu_index, - )?; + decode_progressive_mcu::(context, br, coeffs, dc_predictors, eob_run, mcu_index)?; mcus_since_restart += 1; } + Ok(()) +} - finish_progressive_scan(&mut br, scan_bytes, scan, eob_run) +/// Per-scan state shared by every block of the scan. +struct ScanContext<'p> { + plan: &'p PreparedProgressivePlan, + scan: &'p PreparedProgressiveScan, + scan_components: &'p [PreparedProgressiveScanComponent], + tables: [ScanTables<'p>; MAX_SCAN_COMPONENTS], } fn scan_mcu_count( @@ -122,34 +250,106 @@ fn scan_mcu_count( } } -fn consume_restart( +/// A non-interleaved scan, whose MCU is one block: walk the coded blocks row +/// by row (the order of `mcu_index`) without per-block index arithmetic. +#[expect( + clippy::inline_always, + reason = "inlining keeps the caller's local bit reader in registers for the whole scan" +)] +#[inline(always)] +fn decode_single_component_scan( + context: &ScanContext<'_>, + scan_component: &PreparedProgressiveScanComponent, br: &mut BitReader<'_>, - mcu_index: u32, - total_mcus: u32, - expected_rst: &mut u8, - dc_predictors: &mut [i32], - eob_run: &mut u32, + coeffs: &mut CoefficientState<'_>, + (dc_predictors, eob_run): (&mut [i32], &mut u32), + restart: u32, ) -> Result<(), JpegError> { - *expected_rst = br.consume_restart_marker(*expected_rst, mcu_index, total_mcus)?; - dc_predictors.fill(0); - *eob_run = 0; + let component = context + .plan + .components + .get(scan_component.component_index) + .ok_or(JpegError::InternalInvariant { + reason: "prepared progressive scan references an unknown component", + })?; + let (coded_cols, rows) = ( + progressive_coded_block_cols(component), + progressive_coded_block_rows(component), + ); + let total_mcus = coded_cols.saturating_mul(rows); + let cols = coded_cols as usize; + let block_cols = component.block_cols as usize; + let (Some(component_coeffs), Some(component_masks)) = ( + coeffs.blocks.get_mut(scan_component.component_index), + coeffs.nonzero.get_mut(scan_component.component_index), + ) else { + return Err(invalid_symbol()); + }; + let tables = context.tables[0]; + let mut mcus_since_restart = 0u32; + let mut expected_rst = 0u8; + let mut mcu_index = 0u32; + for by in 0..rows as usize { + let row_start = by * block_cols; + let (Some(row), Some(masks)) = ( + component_coeffs.get_mut(row_start..row_start + cols), + component_masks.get_mut(row_start..row_start + cols), + ) else { + return Err(invalid_symbol()); + }; + for (block, nonzero) in row.iter_mut().zip(masks) { + if restart > 0 && mcus_since_restart == restart { + let (restarted, next_rst) = + br.clone().restarted(expected_rst, mcu_index, total_mcus)?; + *br = restarted; + expected_rst = next_rst; + dc_predictors.fill(0); + *eob_run = 0; + mcus_since_restart = 0; + } + if KIND == DC_FIRST || KIND == AC_FIRST { + decode_progressive_block_first::( + context.scan, + tables, + br, + (block, nonzero), + &mut dc_predictors[scan_component.component_index], + eob_run, + )?; + } else { + decode_progressive_block_refine::( + context.scan, + tables, + br, + (block, nonzero), + eob_run, + )?; + } + mcus_since_restart += 1; + mcu_index += 1; + } + } Ok(()) } -fn decode_progressive_mcu( - plan: &PreparedProgressivePlan, - scan: &PreparedProgressiveScan, +#[expect( + clippy::inline_always, + reason = "inlining keeps the caller's local bit reader in registers for the whole scan" +)] +#[inline(always)] +fn decode_progressive_mcu( + context: &ScanContext<'_>, br: &mut BitReader<'_>, - coeffs: &mut [Vec<[i32; 64]>], + coeffs: &mut CoefficientState<'_>, dc_predictors: &mut [i32], eob_run: &mut u32, mcu_index: u32, ) -> Result<(), JpegError> { - let scan_components = plan.scan_components(scan)?; - if scan_components.len() > 1 { + let plan = context.plan; + if context.scan_components.len() > 1 { let mcu_x = mcu_index % plan.mcu_cols; let mcu_y = mcu_index / plan.mcu_cols; - for scan_component in scan_components { + for (scan_component, &tables) in context.scan_components.iter().zip(&context.tables) { let component = plan.components.get(scan_component.component_index).ok_or( JpegError::InternalInvariant { reason: "prepared progressive scan references an unknown component", @@ -160,12 +360,12 @@ fn decode_progressive_mcu( let target = ProgressiveBlockTarget { component, scan_component, + tables, block_x: mcu_x * u32::from(component.h) + bx, block_y: mcu_y * u32::from(component.v) + by, }; - decode_progressive_block_at( - plan, - scan, + decode_progressive_block_at::( + context.scan, &target, br, coeffs, @@ -176,11 +376,13 @@ fn decode_progressive_mcu( } } } else { - let scan_component = scan_components - .first() - .ok_or(JpegError::InternalInvariant { - reason: "prepared progressive scan has no components", - })?; + let scan_component = + context + .scan_components + .first() + .ok_or(JpegError::InternalInvariant { + reason: "prepared progressive scan has no components", + })?; let component = plan.components.get(scan_component.component_index).ok_or( JpegError::InternalInvariant { reason: "prepared progressive scan references an unknown component", @@ -190,21 +392,33 @@ fn decode_progressive_mcu( let target = ProgressiveBlockTarget { component, scan_component, + tables: context.tables[0], block_x: mcu_index % coded_cols, block_y: mcu_index / coded_cols, }; - decode_progressive_block_at(plan, scan, &target, br, coeffs, dc_predictors, eob_run)?; + decode_progressive_block_at::( + context.scan, + &target, + br, + coeffs, + dc_predictors, + eob_run, + )?; } Ok(()) } -fn decode_progressive_block_at( - plan: &PreparedProgressivePlan, +#[expect( + clippy::inline_always, + reason = "inlining keeps the caller's local bit reader in registers for the whole scan" +)] +#[inline(always)] +fn decode_progressive_block_at( scan: &PreparedProgressiveScan, - target: &ProgressiveBlockTarget<'_>, + target: &ProgressiveBlockTarget<'_, '_>, br: &mut BitReader<'_>, - coeffs: &mut [Vec<[i32; 64]>], + coeffs: &mut CoefficientState<'_>, dc_predictors: &mut [i32], eob_run: &mut u32, ) -> Result<(), JpegError> { @@ -215,40 +429,49 @@ fn decode_progressive_block_at( mcu: 0, reason: HuffmanFailure::InvalidSymbol, })?; - let block = coeffs - .get_mut(target.scan_component.component_index) - .and_then(|component_coeffs| component_coeffs.get_mut(block_index)) - .ok_or(JpegError::HuffmanDecode { - mcu: 0, - reason: HuffmanFailure::InvalidSymbol, - })?; + let component = target.scan_component.component_index; + let (Some(block), Some(nonzero)) = ( + coeffs + .blocks + .get_mut(component) + .and_then(|blocks| blocks.get_mut(block_index)), + coeffs + .nonzero + .get_mut(component) + .and_then(|masks| masks.get_mut(block_index)), + ) else { + return Err(invalid_symbol()); + }; - if scan.ah == 0 { - decode_progressive_block_first( - plan, + if KIND == DC_FIRST || KIND == AC_FIRST { + decode_progressive_block_first::( scan, - target.scan_component, + target.tables, br, - block, - &mut dc_predictors[target.scan_component.component_index], + (block, nonzero), + &mut dc_predictors[component], eob_run, ) } else { - decode_progressive_block_refine(plan, scan, target.scan_component, br, block, eob_run) + decode_progressive_block_refine::(scan, target.tables, br, (block, nonzero), eob_run) } } -fn decode_progressive_block_first( - plan: &PreparedProgressivePlan, +#[expect( + clippy::inline_always, + reason = "inlining keeps the caller's local bit reader in registers for the whole scan" +)] +#[inline(always)] +fn decode_progressive_block_first( scan: &PreparedProgressiveScan, - scan_component: &PreparedProgressiveScanComponent, + tables: ScanTables<'_>, br: &mut BitReader<'_>, - block: &mut [i32; 64], + (block, nonzero): (&mut [i32; 64], &mut u64), dc_predictor: &mut i32, eob_run: &mut u32, ) -> Result<(), JpegError> { - if scan.ss == 0 { - let dc_table = plan.dc_table(scan_component.dc_table)?; + if KIND == DC_FIRST { + let dc_table = tables.dc.ok_or_else(missing_table)?; let ssss = dc_table.decode(br)?; if ssss > 15 { return Err(invalid_symbol()); @@ -259,31 +482,35 @@ fn decode_progressive_block_first( return Ok(()); } - let ac_table = plan.ac_table(scan_component.ac_table)?; + let ac_table = tables.ac.ok_or_else(missing_table)?; if *eob_run > 0 { *eob_run -= 1; return Ok(()); } - let mut k = scan.ss; - while k <= scan.se { + // `usize` band positions: `k <= se <= 63` wherever a coefficient is + // stored, so the mask only drops a bounds check. + let (mut k, se) = (usize::from(scan.ss), usize::from(scan.se)); + while k <= se { let symbol = ac_table.decode(br)?; let run = symbol >> 4; let ssss = symbol & 0x0F; if ssss == 0 { if run == 15 { - k = k.saturating_add(16); + k += 16; } else { *eob_run = decode_eob_run(br, run)?; break; } } else { - k = k.saturating_add(run); - if k > scan.se { + k += usize::from(run); + if k > se { return Err(invalid_symbol()); } let value = br.receive_extend(ssss)?.wrapping_shl(u32::from(scan.al)); - block[usize::from(ZIGZAG[k as usize])] = value; + block[usize::from(ZIGZAG[k & 63]) & 63] = value; + // `receive_extend` values are nonzero for SSSS >= 1. + *nonzero |= 1 << (k & 63); k += 1; } } @@ -307,26 +534,34 @@ fn progressive_coded_block_rows(component: &PreparedProgressiveComponentPlan) -> .min(component.block_rows) } -fn decode_progressive_block_refine( - plan: &PreparedProgressivePlan, +#[expect( + clippy::inline_always, + reason = "inlining keeps the caller's local bit reader in registers for the whole scan" +)] +#[inline(always)] +fn decode_progressive_block_refine( scan: &PreparedProgressiveScan, - scan_component: &PreparedProgressiveScanComponent, + tables: ScanTables<'_>, br: &mut BitReader<'_>, - block: &mut [i32; 64], + (block, nonzero): (&mut [i32; 64], &mut u64), eob_run: &mut u32, ) -> Result<(), JpegError> { let bit = 1i32 << scan.al; - if scan.ss == 0 { + if KIND == DC_REFINE { if br.read_bits(1)? != 0 { block[0] |= bit; } return Ok(()); } - let ac_table = plan.ac_table(scan_component.ac_table)?; + let ac_table = tables.ac.ok_or_else(missing_table)?; + // Correction bits are read with deferred top-ups; every path finishes + // them before the next Huffman decode or before returning. if *eob_run > 0 { *eob_run -= 1; - refine_non_zeroes(br, block, scan.ss, scan.se, 64, bit)?; + let refined = refine_band(br, block, *nonzero, scan.ss, scan.se, 64, bit); + br.finish_deferred_bits(); + refined?; return Ok(()); } @@ -348,17 +583,24 @@ fn decode_progressive_block_refine( } } 1 => { - value = if br.read_bits(1)? != 0 { bit } else { -bit }; + let positive = br.read_bit_deferred(); + if positive.is_err() { + br.finish_deferred_bits(); + } + value = if positive? { bit } else { -bit }; } _ => return Err(invalid_symbol()), } - k = refine_non_zeroes(br, block, k, scan.se, zero_run_length, bit)?; + let refined = refine_band(br, block, *nonzero, k, scan.se, zero_run_length, bit); + br.finish_deferred_bits(); + k = refined?; if value != 0 { if k > scan.se { return Err(invalid_symbol()); } - block[usize::from(ZIGZAG[k as usize])] = value; + block[usize::from(ZIGZAG[usize::from(k)])] = value; + *nonzero |= 1 << (k & 63); } k += 1; } @@ -366,6 +608,11 @@ fn decode_progressive_block_refine( Ok(()) } +#[expect( + clippy::inline_always, + reason = "inlining keeps the caller's local bit reader in registers for the whole scan" +)] +#[inline(always)] pub(super) fn decode_eob_run(br: &mut BitReader<'_>, run_bits: u8) -> Result { let mut eob_run = (1u32 << run_bits) - 1; if run_bits > 0 { @@ -374,6 +621,9 @@ pub(super) fn decode_eob_run(br: &mut BitReader<'_>, run_bits: u8) -> Result, block: &mut [i32; 64], @@ -382,15 +632,19 @@ pub(super) fn refine_non_zeroes( mut zero_run_length: usize, bit: i32, ) -> Result { - for k in start..=end { - let idx = usize::from(ZIGZAG[k as usize]); - let coeff = &mut block[idx]; + // Reads correction bits deferred; the caller finishes them. `ZIGZAG` + // entries are below 64, so the mask only drops a bounds check. + let band = ZIGZAG + .get(usize::from(start)..=usize::from(end)) + .unwrap_or(&[]); + for (offset, &natural) in band.iter().enumerate() { + let coeff = &mut block[usize::from(natural) & 63]; if *coeff == 0 { if zero_run_length == 0 { - return Ok(k); + return Ok(start + band_offset(offset)); } zero_run_length -= 1; - } else if br.read_bits(1)? != 0 && (*coeff & bit) == 0 { + } else if br.read_bit_deferred()? && (*coeff & bit) == 0 { if *coeff > 0 { *coeff = coeff.wrapping_add(bit); } else { @@ -401,6 +655,69 @@ pub(super) fn refine_non_zeroes( Ok(end) } +/// Refine a band (T.81 G.1.2.3) driven by the block's nonzero mask, like the +/// coefficient-by-coefficient `refine_non_zeroes` test oracle: it visits only +/// nonzero coefficients, and finds the zero that ends the run with bit +/// arithmetic. Returns the same position and reads the same correction bits, +/// in the same order. +#[expect( + clippy::inline_always, + reason = "inlining keeps the caller's local bit reader in registers for the whole scan" +)] +#[inline(always)] +pub(super) fn refine_band( + br: &mut BitReader<'_>, + block: &mut [i32; 64], + nonzero: u64, + start: u8, + end: u8, + zero_run_length: usize, + bit: i32, +) -> Result { + let (first, last) = (usize::from(start), usize::from(end)); + if first > last || last > 63 { + return Ok(end); + } + let band = (u64::MAX >> (63 - last)) & (u64::MAX << first); + // The run ends at the zero that follows `zero_run_length` skipped zeros. + let mut zeros = !nonzero & band; + let stop = if zero_run_length < zeros.count_ones() as usize { + for _ in 0..zero_run_length { + zeros &= zeros - 1; + } + Some(zeros.trailing_zeros()) + } else { + None + }; + let mut pending = nonzero + & match stop { + Some(position) => band & !(u64::MAX << position), + None => band, + }; + while pending != 0 { + let k = pending.trailing_zeros() as usize; + pending &= pending - 1; + let coeff = &mut block[usize::from(ZIGZAG[k]) & 63]; + if br.read_bit_deferred()? && (*coeff & bit) == 0 { + if *coeff > 0 { + *coeff = coeff.wrapping_add(bit); + } else { + *coeff = coeff.wrapping_sub(bit); + } + } + } + Ok(stop.map_or(end, |position| band_offset(position as usize))) +} + +/// Offset within a spectral band, which has at most 64 coefficients. +#[expect( + clippy::cast_possible_truncation, + reason = "band offsets index a 64-entry zigzag table" +)] +fn band_offset(offset: usize) -> u8 { + offset as u8 +} + fn invalid_symbol() -> JpegError { JpegError::HuffmanDecode { mcu: 0, diff --git a/crates/j2k-jpeg/src/entropy/progressive/tests.rs b/crates/j2k-jpeg/src/entropy/progressive/tests.rs index 06eee2f6d..f61a32df8 100644 --- a/crates/j2k-jpeg/src/entropy/progressive/tests.rs +++ b/crates/j2k-jpeg/src/entropy/progressive/tests.rs @@ -2,7 +2,7 @@ use super::allocation::{checked_phase_capacity, validate_coefficient_workspace}; use super::model::PreparedProgressiveComponentPlan; -use super::scan::{decode_eob_run, refine_non_zeroes}; +use super::scan::{decode_eob_run, refine_band, refine_non_zeroes}; use crate::allocation::checked_allocation_bytes; use crate::entropy::ZIGZAG; use crate::error::JpegError; @@ -85,3 +85,77 @@ fn refine_non_zeroes_stops_at_requested_zero_run() { assert_eq!(index, 2); } + +fn xorshift(state: &mut u32) -> u32 { + *state ^= *state << 13; + *state ^= *state >> 17; + *state ^= *state << 5; + *state +} + +/// Zigzag-order nonzero mask of a natural-order block's AC coefficients. +fn nonzero_mask(block: &[i32; 64]) -> u64 { + (1..64) + .filter(|&k| block[usize::from(ZIGZAG[k])] != 0) + .fold(0, |mask, k| mask | 1 << k) +} + +#[test] +fn mask_driven_refinement_matches_the_coefficient_walk() { + let mut state = 0x2468_ace1; + for case in 0..20_000 { + let mut block = [0i32; 64]; + let density = xorshift(&mut state) % 5; + for coefficient in block.iter_mut().skip(1) { + if xorshift(&mut state) % 5 < density { + let magnitude = i32::try_from(1 + xorshift(&mut state) % 40).expect("small"); + *coefficient = if xorshift(&mut state) & 1 == 0 { + magnitude + } else { + -magnitude + }; + } + } + let start = u8::try_from(1 + xorshift(&mut state) % 63).expect("band start"); + let end = start + u8::try_from(xorshift(&mut state) % u32::from(64 - start)).expect("end"); + let zero_run_length = if case % 7 == 0 { + 64 + } else { + usize::try_from(xorshift(&mut state) % 16).expect("run") + }; + let bit = 1i32 << (xorshift(&mut state) % 3); + let bytes: Vec = (0..16) + .map(|_| u8::try_from(xorshift(&mut state) % 255).expect("byte")) + .collect(); + + let mut expected_block = block; + let mut expected_reader = BitReader::new(&bytes); + let expected = refine_non_zeroes( + &mut expected_reader, + &mut expected_block, + start, + end, + zero_run_length, + bit, + ); + let mut actual_block = block; + let mut actual_reader = BitReader::new(&bytes); + let actual = refine_band( + &mut actual_reader, + &mut actual_block, + nonzero_mask(&block), + start, + end, + zero_run_length, + bit, + ); + let context = format!("case {case}: band {start}..={end}, run {zero_run_length}"); + assert_eq!(actual, expected, "{context}"); + assert_eq!(actual_block, expected_block, "{context}"); + assert_eq!( + actual_reader.read_bits(8).ok(), + expected_reader.read_bits(8).ok(), + "{context}: bits consumed" + ); + } +} diff --git a/crates/j2k-jpeg/src/entropy/sequential/plan/resolved.rs b/crates/j2k-jpeg/src/entropy/sequential/plan/resolved.rs index abec9c6df..c234d21a0 100644 --- a/crates/j2k-jpeg/src/entropy/sequential/plan/resolved.rs +++ b/crates/j2k-jpeg/src/entropy/sequential/plan/resolved.rs @@ -3,7 +3,7 @@ //! Checked prepared-table resolution outside entropy hot loops. use super::{PreparedComponentPlan, PreparedDecodePlan}; -use crate::entropy::huffman::{AcHuffmanTable, DcHuffmanTable, PreparedHuffmanTableId}; +use crate::entropy::huffman::{AcHuffmanTable, DcHuffmanTable}; use crate::error::JpegError; /// Borrowed component metadata with table IDs resolved once before entering an @@ -26,13 +26,6 @@ impl PreparedDecodePlan { })?) } - pub(crate) fn dc_table_by_id( - &self, - id: PreparedHuffmanTableId, - ) -> Result, JpegError> { - self.huffman_tables.get_dc(id) - } - pub(crate) fn ac_table( &self, component: &PreparedComponentPlan, diff --git a/crates/j2k-jpeg/src/idct/mod.rs b/crates/j2k-jpeg/src/idct/mod.rs index 0a23d73e5..7372ca0ee 100644 --- a/crates/j2k-jpeg/src/idct/mod.rs +++ b/crates/j2k-jpeg/src/idct/mod.rs @@ -14,7 +14,6 @@ pub(crate) mod neon; pub(crate) mod avx2; pub(crate) use scalar::idct_islow; -pub(crate) use scalar::idct_islow_12bit; pub(crate) use scalar::idct_islow_12bit_dc_only_sample; pub(crate) use scalar::idct_islow_dc_only; pub(crate) use scalar::idct_islow_dc_only_pixel; diff --git a/crates/j2k-jpeg/src/idct/neon.rs b/crates/j2k-jpeg/src/idct/neon.rs index 6a7e44c30..a4da446ea 100644 --- a/crates/j2k-jpeg/src/idct/neon.rs +++ b/crates/j2k-jpeg/src/idct/neon.rs @@ -12,9 +12,11 @@ //! between passes even on adversarial coefficients. use core::arch::aarch64::{ - int16x8_t, int32x4_t, vaddq_s32, vcombine_s16, vcombine_s32, vdupq_n_s32, vget_high_s16, - vget_high_s32, vget_low_s16, vget_low_s32, vgetq_lane_u64, vmovl_s16, vmulq_n_s32, vorrq_u64, - vqmovn_s32, vqmovun_s16, vreinterpretq_u64_s16, vshlq_n_s32, vshrq_n_s32, vsubq_s32, vtrnq_s32, + int16x8_t, int32x4_t, vaddq_s32, vaddvq_s32, vcombine_s16, vcombine_s32, vcombine_u16, + vdupq_n_s32, vdupq_n_u16, vget_high_s16, vget_high_s32, vget_low_s16, vget_low_s32, + vgetq_lane_u64, vminq_u16, vmovl_s16, vmulq_n_s32, vorrq_u64, vpadalq_s16, vqabsq_s16, + vqmovn_s32, vqmovun_s16, vqmovun_s32, vreinterpretq_u64_s16, vshlq_n_s32, vshrq_n_s32, + vsubq_s32, vtrnq_s32, }; use j2k_codec_math::jpeg::idct; @@ -92,68 +94,8 @@ fn idct_islow_kernel(input: &[i16; 64], output: &mut [u8; 64]) { let cw_lo = idct_1d_x4::([r0l, r1l, r2l, r3l, r4l, r5l, r6l, r7l], round1); let cw_hi = idct_1d_x4::([r0h, r1h, r2h, r3h, r4h, r5h, r6h, r7h], round1); - // Transpose to pass-2 input. We need `q_c[l] = (row l, col c)`, meaning - // 8 int32x4_t pairs where lane = row. Split into four independent 4×4 - // i32 transposes because `cw_lo` covers cols 0..3 and `cw_hi` covers - // cols 4..7, each across two row halves. - // - // q_c_lo: 4 low rows (0..3) of col c. - // q_c_hi: 4 high rows (4..7) of col c. - let [q0l, q1l, q2l, q3l] = transpose_4x4_i32(cw_lo[0], cw_lo[1], cw_lo[2], cw_lo[3]); - let [q4l, q5l, q6l, q7l] = transpose_4x4_i32(cw_hi[0], cw_hi[1], cw_hi[2], cw_hi[3]); - let [q0h, q1h, q2h, q3h] = transpose_4x4_i32(cw_lo[4], cw_lo[5], cw_lo[6], cw_lo[7]); - let [q4h, q5h, q6h, q7h] = transpose_4x4_i32(cw_hi[4], cw_hi[5], cw_hi[6], cw_hi[7]); - - // Pass 2: row IDCT. - let round2 = vdupq_n_s32(1 << (PASS2_SHIFT - 1)); - - let rw_lo = idct_1d_x4::([q0l, q1l, q2l, q3l, q4l, q5l, q6l, q7l], round2); - let rw_hi = idct_1d_x4::([q0h, q1h, q2h, q3h, q4h, q5h, q6h, q7h], round2); - - // `rw_lo[k]` lane `l` = (row l, col k) for rows 0..3; `rw_hi[k]` for - // rows 4..7. Transpose back to row-major, applying the +128 level - // shift, then pack to u8. - let bias = vdupq_n_s32(128); - - // Transpose rw_lo (cols 0..7 × rows 0..3) into row-major low (cols 0..3). - let [fll0, fll1, fll2, fll3] = transpose_4x4_i32( - vaddq_s32(rw_lo[0], bias), - vaddq_s32(rw_lo[1], bias), - vaddq_s32(rw_lo[2], bias), - vaddq_s32(rw_lo[3], bias), - ); - let [flh0, flh1, flh2, flh3] = transpose_4x4_i32( - vaddq_s32(rw_lo[4], bias), - vaddq_s32(rw_lo[5], bias), - vaddq_s32(rw_lo[6], bias), - vaddq_s32(rw_lo[7], bias), - ); - let [fhl0, fhl1, fhl2, fhl3] = transpose_4x4_i32( - vaddq_s32(rw_hi[0], bias), - vaddq_s32(rw_hi[1], bias), - vaddq_s32(rw_hi[2], bias), - vaddq_s32(rw_hi[3], bias), - ); - let [fhh0, fhh1, fhh2, fhh3] = transpose_4x4_i32( - vaddq_s32(rw_hi[4], bias), - vaddq_s32(rw_hi[5], bias), - vaddq_s32(rw_hi[6], bias), - vaddq_s32(rw_hi[7], bias), - ); - - // `fll_r` = row r (0..3), cols 0..3 as i32x4. - // `flh_r` = row r (0..3), cols 4..7. - // `fhl_r` = row r (4..7), cols 0..3. - // `fhh_r` = row r (4..7), cols 4..7. - let (rows, _) = output.as_chunks_mut::<8>(); - store_row(&mut rows[0], fll0, flh0); - store_row(&mut rows[1], fll1, flh1); - store_row(&mut rows[2], fll2, flh2); - store_row(&mut rows[3], fll3, flh3); - store_row(&mut rows[4], fhl0, fhh0); - store_row(&mut rows[5], fhl1, fhh1); - store_row(&mut rows[6], fhl2, fhh2); - store_row(&mut rows[7], fhl3, fhh3); + let (q_lo, q_hi) = transpose_pass1(cw_lo, cw_hi); + store_rows_u8(output, pass2_row_major::(q_lo, q_hi, 128)); } fearless_simd::kernel! { @@ -198,16 +140,74 @@ fn idct_islow_bottom_half_zero_rows( let cw_lo = idct_1d_x4_bottom_half_zero::(r0l, r1l, r2l, r3l, round1); let cw_hi = idct_1d_x4_bottom_half_zero::(r0h, r1h, r2h, r3h, round1); + let (q_lo, q_hi) = transpose_pass1(cw_lo, cw_hi); + store_rows_u8(output, pass2_row_major::(q_lo, q_hi, 128)); +} + +fearless_simd::kernel! { + #[inline] + #[cfg(test)] + fn bottom_half_is_zero(neon: Neon, input: &[i16; 64]) -> bool { + let (rows, _) = input.as_chunks::<8>(); + bottom_half_rows_are_zero( + neon_memory::load_i16x8(&rows[4]), + neon_memory::load_i16x8(&rows[5]), + neon_memory::load_i16x8(&rows[6]), + neon_memory::load_i16x8(&rows[7]), + ) + } +} + +#[inline] +#[target_feature(enable = "neon")] +fn bottom_half_rows_are_zero( + row4: int16x8_t, + row5: int16x8_t, + row6: int16x8_t, + row7: int16x8_t, +) -> bool { + let bottom = vorrq_u64( + vorrq_u64(vreinterpretq_u64_s16(row4), vreinterpretq_u64_s16(row5)), + vorrq_u64(vreinterpretq_u64_s16(row6), vreinterpretq_u64_s16(row7)), + ); + vgetq_lane_u64::<0>(bottom) == 0 && vgetq_lane_u64::<1>(bottom) == 0 +} + +/// Transpose pass-1 column output into pass-2 row input. `cw_lo`/`cw_hi` +/// hold columns 0..3/4..7 with one row per vector; the result holds rows +/// 0..3/4..7 with one column per vector, so lane `l` of `lo[c]` is +/// (row l, col c). +#[inline] +#[target_feature(enable = "neon")] +fn transpose_pass1( + cw_lo: [int32x4_t; 8], + cw_hi: [int32x4_t; 8], +) -> ([int32x4_t; 8], [int32x4_t; 8]) { let [q0l, q1l, q2l, q3l] = transpose_4x4_i32(cw_lo[0], cw_lo[1], cw_lo[2], cw_lo[3]); let [q4l, q5l, q6l, q7l] = transpose_4x4_i32(cw_hi[0], cw_hi[1], cw_hi[2], cw_hi[3]); let [q0h, q1h, q2h, q3h] = transpose_4x4_i32(cw_lo[4], cw_lo[5], cw_lo[6], cw_lo[7]); let [q4h, q5h, q6h, q7h] = transpose_4x4_i32(cw_hi[4], cw_hi[5], cw_hi[6], cw_hi[7]); + ( + [q0l, q1l, q2l, q3l, q4l, q5l, q6l, q7l], + [q0h, q1h, q2h, q3h, q4h, q5h, q6h, q7h], + ) +} - let round2 = vdupq_n_s32(1 << (PASS2_SHIFT - 1)); - let rw_lo = idct_1d_x4::([q0l, q1l, q2l, q3l, q4l, q5l, q6l, q7l], round2); - let rw_hi = idct_1d_x4::([q0h, q1h, q2h, q3h, q4h, q5h, q6h, q7h], round2); - - let bias = vdupq_n_s32(128); +/// Pass 2 (row IDCT) on transposed input, then transpose back to row-major +/// and add the level shift `bias`. Returns `rows[r] = (cols 0..3, cols 4..7)`. +#[inline] +#[target_feature(enable = "neon")] +fn pass2_row_major( + q_lo: [int32x4_t; 8], + q_hi: [int32x4_t; 8], + bias: i32, +) -> [(int32x4_t, int32x4_t); 8] { + let round2 = vdupq_n_s32(1 << (SHIFT - 1)); + let rw_lo = idct_1d_x4::(q_lo, round2); + let rw_hi = idct_1d_x4::(q_hi, round2); + // `rw_lo[k]` lane `l` = (row l, col k) for rows 0..3; `rw_hi[k]` for + // rows 4..7. + let bias = vdupq_n_s32(bias); let [fll0, fll1, fll2, fll3] = transpose_4x4_i32( vaddq_s32(rw_lo[0], bias), vaddq_s32(rw_lo[1], bias), @@ -232,45 +232,93 @@ fn idct_islow_bottom_half_zero_rows( vaddq_s32(rw_hi[6], bias), vaddq_s32(rw_hi[7], bias), ); + [ + (fll0, flh0), + (fll1, flh1), + (fll2, flh2), + (fll3, flh3), + (fhl0, fhh0), + (fhl1, fhh1), + (fhl2, fhh2), + (fhl3, fhh3), + ] +} - let (rows, _) = output.as_chunks_mut::<8>(); - store_row(&mut rows[0], fll0, flh0); - store_row(&mut rows[1], fll1, flh1); - store_row(&mut rows[2], fll2, flh2); - store_row(&mut rows[3], fll3, flh3); - store_row(&mut rows[4], fhl0, fhh0); - store_row(&mut rows[5], fhl1, fhh1); - store_row(&mut rows[6], fhl2, fhh2); - store_row(&mut rows[7], fhl3, fhh3); +#[inline] +#[target_feature(enable = "neon")] +fn store_rows_u8(output: &mut [u8; 64], rows: [(int32x4_t, int32x4_t); 8]) { + let (dst, _) = output.as_chunks_mut::<8>(); + for (dst, (lo, hi)) in dst.iter_mut().zip(rows) { + store_row(dst, lo, hi); + } } +/// Pass-1 precision bits of the 12-bit ISLOW transform (`jidctint.c`). +const PASS1_BITS_12: i32 = 1; + +/// Largest `sum(|coefficient|)` for which every 12-bit ISLOW intermediate fits +/// in `i32`. Pass 1 cannot overflow for any `i16` input (its largest L1 gain +/// is 61,214, and 61,214 * 32,768 < 2^31). A pass-2 intermediate is a linear +/// form of the 64 coefficients whose largest single-coefficient gain is +/// 58,244, plus at most 30,607 of pass-1 rounding and the 2^16 descale +/// rounding, so `58,244 * sum + 96,143 < 2^31` holds up to a sum of 36,869. +/// The margin below that absorbs `vqabsq_s16` reporting 32,767 for -32,768. +const IDCT12_I32_SAFE_ABS_SUM: i32 = 36_800; + fearless_simd::kernel! { - #[inline] - #[cfg(test)] - fn bottom_half_is_zero(neon: Neon, input: &[i16; 64]) -> bool { - let (rows, _) = input.as_chunks::<8>(); - bottom_half_rows_are_zero( - neon_memory::load_i16x8(&rows[4]), - neon_memory::load_i16x8(&rows[5]), - neon_memory::load_i16x8(&rows[6]), - neon_memory::load_i16x8(&rows[7]), - ) + /// 12-bit ISLOW IDCT of one block into samples level-shifted by +2048 and + /// clamped to `[0, 4095]`, bit-exact with the 64-bit + /// [`super::scalar::idct_islow_12bit`]. Returns `false`, leaving `output` + /// untouched, when the coefficients are large enough that 32-bit lanes + /// could overflow; the caller then runs the 64-bit scalar transform. + pub(crate) fn idct_islow_12bit(neon: Neon, input: &[i16; 64], output: &mut [u16; 64]) -> bool { + idct_islow_12bit_kernel(input, output) } } -#[inline] #[target_feature(enable = "neon")] -fn bottom_half_rows_are_zero( - row4: int16x8_t, - row5: int16x8_t, - row6: int16x8_t, - row7: int16x8_t, -) -> bool { - let bottom = vorrq_u64( - vorrq_u64(vreinterpretq_u64_s16(row4), vreinterpretq_u64_s16(row5)), - vorrq_u64(vreinterpretq_u64_s16(row6), vreinterpretq_u64_s16(row7)), - ); - vgetq_lane_u64::<0>(bottom) == 0 && vgetq_lane_u64::<1>(bottom) == 0 +fn idct_islow_12bit_kernel(input: &[i16; 64], output: &mut [u16; 64]) -> bool { + const PASS1_SHIFT: i32 = CONST_BITS - PASS1_BITS_12; + const PASS2_SHIFT: i32 = CONST_BITS + PASS1_BITS_12 + 3; + let (rows, _) = input.as_chunks::<8>(); + let row0 = neon_memory::load_i16x8(&rows[0]); + let row1 = neon_memory::load_i16x8(&rows[1]); + let row2 = neon_memory::load_i16x8(&rows[2]); + let row3 = neon_memory::load_i16x8(&rows[3]); + let row4 = neon_memory::load_i16x8(&rows[4]); + let row5 = neon_memory::load_i16x8(&rows[5]); + let row6 = neon_memory::load_i16x8(&rows[6]); + let row7 = neon_memory::load_i16x8(&rows[7]); + let mut magnitude = vdupq_n_s32(0); + for row in [row0, row1, row2, row3, row4, row5, row6, row7] { + magnitude = vpadalq_s16(magnitude, vqabsq_s16(row)); + } + if vaddvq_s32(magnitude) > IDCT12_I32_SAFE_ABS_SUM { + return false; + } + + let (r0l, r0h) = widen(row0); + let (r1l, r1h) = widen(row1); + let (r2l, r2h) = widen(row2); + let (r3l, r3h) = widen(row3); + let (r4l, r4h) = widen(row4); + let (r5l, r5h) = widen(row5); + let (r6l, r6h) = widen(row6); + let (r7l, r7h) = widen(row7); + let round1 = vdupq_n_s32(1 << (PASS1_SHIFT - 1)); + let cw_lo = idct_1d_x4::([r0l, r1l, r2l, r3l, r4l, r5l, r6l, r7l], round1); + let cw_hi = idct_1d_x4::([r0h, r1h, r2h, r3h, r4h, r5h, r6h, r7h], round1); + let (q_lo, q_hi) = transpose_pass1(cw_lo, cw_hi); + let rows = pass2_row_major::(q_lo, q_hi, 2048); + + let max_sample = vdupq_n_u16(4095); + let (dst, _) = output.as_chunks_mut::<8>(); + for (dst, (lo, hi)) in dst.iter_mut().zip(rows) { + // Saturating narrows clamp below at 0; the min clamps above at 4095. + let packed = vcombine_u16(vqmovun_s32(lo), vqmovun_s32(hi)); + neon_memory::store_u16x8(dst, vminq_u16(packed, max_sample)); + } + true } /// Saturating-narrow an (i32x4, i32x4) pair to u8x8 and store at `dst`. @@ -484,6 +532,92 @@ mod tests { assert_eq!(scalar_out, neon_out); } + fn neon_token() -> fearless_simd::Neon { + fearless_simd::Level::new() + .as_neon() + .expect("AArch64 test host must provide NEON") + } + + fn assert_12bit_matches_scalar(input: &[i16; 64], label: &str) { + let mut expected = [0u16; 64]; + crate::idct::scalar::idct_islow_12bit(input, &mut expected); + let mut actual = [0xffffu16; 64]; + assert!( + idct_islow_12bit(neon_token(), input, &mut actual), + "{label}: guard declined a block within the 32-bit bound" + ); + assert_eq!(actual, expected, "{label}"); + } + + fn xorshift(state: &mut u32) -> u32 { + *state ^= *state << 13; + *state ^= *state >> 17; + *state ^= *state << 5; + *state + } + + fn abs_sum(input: &[i16; 64]) -> i32 { + input.iter().map(|&c| i32::from(c).abs()).sum() + } + + #[test] + fn neon_12bit_matches_scalar_on_random_blocks() { + let mut state = 0x1234_5678; + for case in 0..4_000 { + let mut input = [0i16; 64]; + let magnitude = [64, 512, 4_096, 32_767][case % 4]; + let density = [2, 8, 64][case % 3]; + for coefficient in input.iter_mut().take(density) { + let raw = xorshift(&mut state) % (2 * magnitude + 1); + *coefficient = i16::try_from(i64::from(raw) - i64::from(magnitude)) + .expect("coefficient range"); + } + if abs_sum(&input) > IDCT12_I32_SAFE_ABS_SUM { + continue; + } + assert_12bit_matches_scalar(&input, &format!("case {case}")); + } + } + + #[test] + fn neon_12bit_matches_scalar_on_extreme_blocks_under_the_guard() { + // Whole guard budget on one coefficient: the largest single gains. + for index in 0..64 { + for value in [i16::MAX, i16::MIN] { + let mut input = [0i16; 64]; + input[index] = value; + assert_12bit_matches_scalar(&input, &format!("index {index} = {value}")); + } + } + // The budget spread across every coefficient with varied signs. + let mut state = 0x9e37_79b9; + for case in 0..2_000 { + let each = i16::try_from(IDCT12_I32_SAFE_ABS_SUM / 64).expect("share fits i16"); + let mut input = [0i16; 64]; + for coefficient in &mut input { + *coefficient = if xorshift(&mut state) & 1 == 0 { + each + } else { + -each + }; + } + assert_12bit_matches_scalar(&input, &format!("spread case {case}")); + } + } + + #[test] + fn neon_12bit_declines_blocks_that_could_overflow_32_bits() { + let mut input = [0i16; 64]; + input[0] = i16::MAX; + input[9] = 5_000; + let mut output = [7u16; 64]; + assert!(!idct_islow_12bit(neon_token(), &input, &mut output)); + assert_eq!( + output, [7u16; 64], + "a declined block leaves the output untouched" + ); + } + #[test] fn bottom_half_zero_detects_zero_and_nonzero_tails() { let neon = fearless_simd::Level::new() diff --git a/crates/j2k-jpeg/src/internal/bit_reader.rs b/crates/j2k-jpeg/src/internal/bit_reader.rs index 2e83f1109..59c262bd3 100644 --- a/crates/j2k-jpeg/src/internal/bit_reader.rs +++ b/crates/j2k-jpeg/src/internal/bit_reader.rs @@ -290,14 +290,50 @@ impl<'a> BitReader<'a> { } /// Combined peek + consume. Refills as needed. + #[expect( + clippy::inline_always, + reason = "progressive refinement reads single bits per coefficient; an out-of-line call spills the reader" + )] + #[inline(always)] pub(crate) fn read_bits(&mut self, n: u8) -> Result { self.ensure_bits(n)?; let v = self.peek_bits(n); self.consume_bits(n); - self.refill_to_threshold(); + self.top_up_to_threshold(); Ok(v) } + /// Read one bit like `read_bits(1)` but without the trailing top-up. + /// + /// A top-up loads the fewest bytes that leave at least 56 buffered bits + /// (or stops at a marker), and it never unloads bytes, so the loaded + /// frontier after a top-up depends only on the bits consumed so far. A + /// run of deferred reads followed by one [`Self::finish_deferred_bits`] + /// therefore leaves the same observable reader state as the same run of + /// `read_bits(1)` calls. Callers must finish before any Huffman decode, + /// restart marker, or end-of-scan check. + #[expect( + clippy::inline_always, + reason = "progressive refinement reads single bits per coefficient" + )] + #[inline(always)] + pub(crate) fn read_bit_deferred(&mut self) -> Result { + self.ensure_bits(1)?; + let bit = self.peek_bits(1) != 0; + self.consume_bits(1); + Ok(bit) + } + + /// Complete a run of [`Self::read_bit_deferred`] calls. + #[expect( + clippy::inline_always, + reason = "progressive refinement tops up once per run of correction bits" + )] + #[inline(always)] + pub(crate) fn finish_deferred_bits(&mut self) { + self.top_up_to_threshold(); + } + /// After consuming bits, top up the accumulator so the next Huffman peek /// can always examine 16 bits without further refill. fn refill_to_threshold(&mut self) { @@ -306,6 +342,27 @@ impl<'a> BitReader<'a> { } } + /// [`Self::refill_to_threshold`] for hot loops that keep the reader in a + /// local: the bulk path is inline and the byte-at-a-time path takes the + /// reader by value, so the local never escapes. + #[expect( + clippy::inline_always, + reason = "measured bit-buffer hot path requires cross-helper inlining" + )] + #[inline(always)] + fn top_up_to_threshold(&mut self) { + if self.bits < REFILL_THRESHOLD && !self.try_refill_bulk() { + *self = self.clone().top_up_to_threshold_slow(); + } + } + + #[cold] + #[inline(never)] + fn top_up_to_threshold_slow(mut self) -> Self { + while self.bits < REFILL_THRESHOLD && self.refill_one_byte() {} + self + } + /// Signed-value extension per T.81 §F.2.2.1 ("EXTEND" procedure). `ssss` /// is the category — a non-zero value in `1..=15` — and the return is the /// signed coefficient value. diff --git a/crates/j2k-jpeg/src/lossless/mod.rs b/crates/j2k-jpeg/src/lossless/mod.rs index 41f04464a..952864bd7 100644 --- a/crates/j2k-jpeg/src/lossless/mod.rs +++ b/crates/j2k-jpeg/src/lossless/mod.rs @@ -2,101 +2,84 @@ //! Shared lossless JPEG decode helpers. -use crate::error::{HuffmanFailure, JpegError}; +pub(crate) mod scan; -pub(crate) trait LosslessSample: Copy + Default + Into { - const RESTART_PREDICTOR: i32; +/// Output container for reconstructed lossless samples: `u8` for 8-bit and +/// `u16` (little-endian) for 16-bit images. +pub(crate) trait LosslessSample: Copy + Default { const BIT_DEPTH: u8; const BYTES: usize; - fn from_i32(value: i32) -> Result; - - fn read_le(src: &[u8]) -> i32; + /// Narrow a reconstructed sample. [`scan::LosslessScan`] has already + /// rejected values wider than the image precision. + fn from_sample(sample: u16) -> Self; fn write_le(self, dst: &mut [u8]); + + /// Store one component row, shifted left by the point transform. + fn store_row(samples: &[u16], shift: u8, dst: &mut [u8]) { + for (dst, &sample) in dst.chunks_exact_mut(Self::BYTES).zip(samples) { + Self::from_sample(sample << shift).write_le(dst); + } + } + + /// Interleave three component rows into one pixel row. + fn store_interleaved(rows: [&[u16]; 3], shift: u8, dst: &mut [u8]) { + let [c0, c1, c2] = rows; + for (pixel, ((&s0, &s1), &s2)) in dst + .chunks_exact_mut(3 * Self::BYTES) + .zip(c0.iter().zip(c1).zip(c2)) + { + let (first, rest) = pixel.split_at_mut(Self::BYTES); + let (second, third) = rest.split_at_mut(Self::BYTES); + Self::from_sample(s0 << shift).write_le(first); + Self::from_sample(s1 << shift).write_le(second); + Self::from_sample(s2 << shift).write_le(third); + } + } } impl LosslessSample for u8 { - const RESTART_PREDICTOR: i32 = 128; const BIT_DEPTH: u8 = 8; const BYTES: usize = 1; - fn from_i32(value: i32) -> Result { - u8::try_from(value).map_err(|_| invalid_lossless_symbol()) - } - - fn read_le(src: &[u8]) -> i32 { - i32::from(src[0]) + #[expect( + clippy::cast_possible_truncation, + reason = "8-bit lossless samples are range-checked to 0..=255 during reconstruction" + )] + fn from_sample(sample: u16) -> Self { + sample as u8 } fn write_le(self, dst: &mut [u8]) { dst[0] = self; } + + fn store_row(samples: &[u16], shift: u8, dst: &mut [u8]) { + for (dst, &sample) in dst.iter_mut().zip(samples) { + *dst = Self::from_sample(sample << shift); + } + } + + fn store_interleaved(rows: [&[u16]; 3], shift: u8, dst: &mut [u8]) { + let [c0, c1, c2] = rows; + for (pixel, ((&s0, &s1), &s2)) in dst.chunks_exact_mut(3).zip(c0.iter().zip(c1).zip(c2)) { + pixel[0] = Self::from_sample(s0 << shift); + pixel[1] = Self::from_sample(s1 << shift); + pixel[2] = Self::from_sample(s2 << shift); + } + } } impl LosslessSample for u16 { - const RESTART_PREDICTOR: i32 = 32_768; const BIT_DEPTH: u8 = 16; const BYTES: usize = 2; - fn from_i32(value: i32) -> Result { - u16::try_from(value).map_err(|_| invalid_lossless_symbol()) - } - - fn read_le(src: &[u8]) -> i32 { - i32::from(u16::from_le_bytes([src[0], src[1]])) + fn from_sample(sample: u16) -> Self { + sample } fn write_le(self, dst: &mut [u8]) { dst[..2].copy_from_slice(&self.to_le_bytes()); } } - -fn invalid_lossless_symbol() -> JpegError { - JpegError::HuffmanDecode { - mcu: 0, - reason: HuffmanFailure::InvalidSymbol, - } -} - -/// Spec predictor (ITU-T T.81 H.1.2.1) over a sample accessor. -/// -/// Edge rules: the first sample predicts `bias` (1 << (P - 1)); the first row -/// predicts Ra; the first column predicts Rb. `at(x, y)` must return the -/// reconstructed sample at the given coordinates; only `(x-1, y)`, `(x, y-1)` -/// and `(x-1, y-1)` are ever requested. -#[expect( - clippy::inline_always, - reason = "per-sample predictor hot path keeps the accessor closure monomorphized" -)] -#[inline(always)] -pub(crate) fn lossless_predict( - predictor: u8, - bias: i32, - x: usize, - y: usize, - at: impl Fn(usize, usize) -> i32, -) -> i32 { - if x == 0 && y == 0 { - return bias; - } - if y == 0 { - return at(x - 1, 0); - } - if x == 0 { - return at(0, y - 1); - } - let ra = at(x - 1, y); - let rb = at(x, y - 1); - let rc = at(x - 1, y - 1); - match predictor { - 1 => ra, - 2 => rb, - 3 => rc, - 4 => ra + rb - rc, - 5 => ra + ((rb - rc) >> 1), - 6 => rb + ((ra - rc) >> 1), - 7 => (ra + rb) >> 1, - _ => bias, - } -} diff --git a/crates/j2k-jpeg/src/lossless/scan.rs b/crates/j2k-jpeg/src/lossless/scan.rs new file mode 100644 index 000000000..b43c733b0 --- /dev/null +++ b/crates/j2k-jpeg/src/lossless/scan.rs @@ -0,0 +1,510 @@ +// SPDX-License-Identifier: MIT OR Apache-2.0 + +//! Row-at-a-time lossless (SOF3) scan decoding, ITU-T T.81 Annex H. +//! +//! Each MCU row decodes in two passes. The entropy pass reads every coded +//! difference in the row, including MCU padding samples that are coded but +//! never displayed. The reconstruction pass then undifferences each component +//! row with a loop specialized for the scan's predictor. As in libjpeg-turbo's +//! `jdlossls.c`: +//! - reconstruction is modulo 2^16, so differences are kept as wrapping `u16`; +//! - the first row of the scan and of every restart interval uses the 1-D +//! predictor (Ra), seeded with 2^(P-Pt-1); +//! - every later row starts from Rb, then uses the scan's predictor; +//! - restart intervals must cover whole MCU rows. + +use alloc::vec::Vec; + +use crate::allocation::{ + checked_add_allocation_bytes, checked_allocation_bytes, try_vec_filled, try_vec_with_capacity, +}; +use crate::entropy::huffman::DcHuffmanTable; +use crate::entropy::sequential::finish_scan; +use crate::error::{HuffmanFailure, JpegError, Warning}; +use crate::info::SofKind; +use crate::internal::bit_reader::BitReader; + +#[cfg(test)] +mod tests; + +/// One component of a lossless scan, in scan order. +#[derive(Clone, Copy)] +pub(crate) struct LosslessComponentSpec<'t> { + /// Horizontal sampling factor from the frame header. + pub(crate) h: u8, + /// Vertical sampling factor from the frame header. + pub(crate) v: u8, + pub(crate) table: DcHuffmanTable<'t>, +} + +/// Frame and scan parameters of a lossless scan. +#[derive(Clone, Copy)] +pub(crate) struct LosslessScanSpec { + pub(crate) dimensions: (u32, u32), + /// Predictor selection value `Ss`, 1..=7. + pub(crate) predictor: u8, + /// Sample precision `P`, 2..=16. + pub(crate) precision: u8, + /// Point transform `Pt` (`Al`), below `P`. + pub(crate) point_transform: u8, + pub(crate) restart_interval: Option, +} + +struct ComponentRows<'t> { + table: DcHuffmanTable<'t>, + mcu_h: usize, + mcu_v: usize, + width: usize, + height: usize, + /// Coded samples per row: MCUs per row times `mcu_h`. + coded_width: usize, + /// `mcu_v` rows of `coded_width` differences. + diffs: Vec, + /// `mcu_v + 1` rows of `width` samples. Row 0 holds the last row of the + /// previous MCU row; rows `1..=rows_ready` hold the current MCU row. + samples: Vec, + rows_ready: usize, + first_row: bool, +} + +/// Decoder state for one lossless scan; see the module docs. +pub(crate) struct LosslessScan<'a, 't> { + br: BitReader<'a>, + components: Vec>, + predictor: u8, + initial_predictor: u16, + /// Largest valid sample, or `u16::MAX` when every value is valid. + sample_limit: u16, + point_transform: u8, + mcus_per_row: usize, + mcu_rows: usize, + mcu_row: usize, + /// Restart interval in MCU rows; zero disables restarts. + restart_rows: usize, + rows_to_go: usize, + expected_rst: u8, +} + +/// Buffer layout of one scan component. +struct ComponentLayout { + /// Samples per MCU; 1x1 in a single-component (non-interleaved) scan. + mcu_h: usize, + mcu_v: usize, + /// Reconstructed extent, `ceil(X * h / Hmax)` by `ceil(Y * v / Vmax)`. + width: usize, + height: usize, + coded_width: usize, + diff_len: usize, + sample_len: usize, +} + +/// MCU grid of a scan: `(MCUs per row, MCU rows)`, plus each component's +/// layout. A single-component scan is non-interleaved, so its MCU is one sample. +fn scan_layout( + dimensions: (u32, u32), + sampling: &[(u8, u8)], +) -> Result<((usize, usize), Vec), JpegError> { + let max_h = usize::from(sampling.iter().map(|s| s.0).max().unwrap_or(0)); + let max_v = usize::from(sampling.iter().map(|s| s.1).max().unwrap_or(0)); + if max_h == 0 || max_v == 0 { + return Err(JpegError::NotImplemented { + sof: SofKind::Lossless, + }); + } + let (width, height) = (dimensions.0 as usize, dimensions.1 as usize); + let interleaved = sampling.len() > 1; + let grid = if interleaved { + (width.div_ceil(max_h), height.div_ceil(max_v)) + } else { + (width, height) + }; + let mut layouts = try_vec_with_capacity(sampling.len())?; + for &(h, v) in sampling { + let (h, v) = (usize::from(h), usize::from(v)); + let (mcu_h, mcu_v) = if interleaved { (h, v) } else { (1, 1) }; + let component_width = (width * h).div_ceil(max_h); + let coded_width = grid.0.checked_mul(mcu_h).ok_or_else(geometry_overflow)?; + let diff_len = coded_width + .checked_mul(mcu_v) + .ok_or_else(geometry_overflow)?; + let sample_len = component_width + .checked_mul(mcu_v + 1) + .ok_or_else(geometry_overflow)?; + layouts.push(ComponentLayout { + mcu_h, + mcu_v, + width: component_width, + height: (height * v).div_ceil(max_v), + coded_width, + diff_len, + sample_len, + }); + } + Ok((grid, layouts)) +} + +fn geometry_overflow() -> JpegError { + JpegError::MemoryCapExceeded { + requested: usize::MAX, + cap: j2k_core::DEFAULT_MAX_HOST_ALLOCATION_BYTES, + } +} + +/// Bytes [`LosslessScan::new`] allocates for components with this sampling. +pub(crate) fn lossless_scan_allocation_bytes( + dimensions: (u32, u32), + sampling: &[(u8, u8)], +) -> Result { + let (_, layouts) = scan_layout(dimensions, sampling)?; + let mut total = checked_allocation_bytes::>(layouts.len())?; + for layout in &layouts { + total = + checked_add_allocation_bytes(total, checked_allocation_bytes::(layout.diff_len)?)?; + total = checked_add_allocation_bytes( + total, + checked_allocation_bytes::(layout.sample_len)?, + )?; + } + Ok(total) +} + +impl<'a, 't> LosslessScan<'a, 't> { + /// Prepare to decode the scan whose entropy data starts at `scan_bytes`. + /// + /// # Errors + /// `UnsupportedPredictor` or `UnsupportedBitDepth` for parameters outside + /// T.81 Annex H, `NotImplemented` for a restart interval that does not + /// cover whole MCU rows (libjpeg-turbo rejects those too), or an + /// allocation error. + pub(crate) fn new( + scan_bytes: &'a [u8], + spec: LosslessScanSpec, + components: &[LosslessComponentSpec<'t>], + ) -> Result { + if !(1..=7).contains(&spec.predictor) { + return Err(JpegError::UnsupportedPredictor { + predictor: spec.predictor, + }); + } + if !(2..=16).contains(&spec.precision) || spec.point_transform >= spec.precision { + return Err(JpegError::UnsupportedBitDepth { + depth: spec.precision, + }); + } + let mut sampling = [(0u8, 0u8); 4]; + let sampling = + sampling + .get_mut(..components.len()) + .ok_or(JpegError::UnsupportedComponentCount { + count: u8::try_from(components.len()).unwrap_or(u8::MAX), + })?; + for (slot, component) in sampling.iter_mut().zip(components) { + *slot = (component.h, component.v); + } + let ((mcus_per_row, mcu_rows), layouts) = scan_layout(spec.dimensions, sampling)?; + let restart_mcus = usize::from(spec.restart_interval.unwrap_or(0)); + if restart_mcus != 0 && mcus_per_row != 0 && restart_mcus % mcus_per_row != 0 { + return Err(JpegError::NotImplemented { + sof: SofKind::Lossless, + }); + } + let restart_rows = restart_mcus.checked_div(mcus_per_row).unwrap_or(0); + + let mut rows = try_vec_with_capacity(components.len())?; + for (component, layout) in components.iter().zip(layouts) { + rows.push(ComponentRows { + table: component.table, + mcu_h: layout.mcu_h, + mcu_v: layout.mcu_v, + width: layout.width, + height: layout.height, + coded_width: layout.coded_width, + diffs: try_vec_filled(layout.diff_len, 0)?, + samples: try_vec_filled(layout.sample_len, 0)?, + rows_ready: 0, + first_row: true, + }); + } + + let significant_bits = spec.precision - spec.point_transform; + Ok(Self { + br: BitReader::new(scan_bytes), + components: rows, + predictor: spec.predictor, + initial_predictor: 1 << (significant_bits - 1), + sample_limit: if significant_bits >= 16 { + u16::MAX + } else { + (1 << significant_bits) - 1 + }, + point_transform: spec.point_transform, + mcus_per_row, + mcu_rows, + mcu_row: 0, + restart_rows, + rows_to_go: restart_rows, + expected_rst: 0, + }) + } + + /// Point transform the caller shifts samples left by on output. + pub(crate) fn point_transform(&self) -> u8 { + self.point_transform + } + + /// First component row held by the last decoded MCU row. + pub(crate) fn row_base(&self, component: usize) -> usize { + let c = &self.components[component]; + self.mcu_row.saturating_sub(1) * c.mcu_v + } + + /// Rows of `component` the last decoded MCU row reconstructed. + pub(crate) fn rows_ready(&self, component: usize) -> usize { + self.components[component].rows_ready + } + + /// Row `index` (below [`Self::rows_ready`]) of the last decoded MCU row. + pub(crate) fn row(&self, component: usize, index: usize) -> &[u16] { + let c = &self.components[component]; + let start = (index + 1) * c.width; + &c.samples[start..start + c.width] + } + + /// Decode and reconstruct the next MCU row. Returns `Ok(false)` once every + /// MCU row has been decoded. + /// + /// # Errors + /// Entropy, restart-marker, or out-of-range sample errors. + pub(crate) fn decode_mcu_row(&mut self) -> Result { + if self.mcu_row == self.mcu_rows { + return Ok(false); + } + if self.restart_rows != 0 { + if self.rows_to_go == 0 { + self.restart()?; + } + self.rows_to_go -= 1; + } + // The reader lives in a local for the whole row so its accumulator + // stays in registers; it is written back even on error. + let mut br = self.br.clone(); + let decoded = decode_differences(&mut br, &mut self.components, self.mcus_per_row); + self.br = br; + decoded?; + + for component in &mut self.components { + reconstruct_component( + component, + self.mcu_row, + self.predictor, + self.initial_predictor, + self.sample_limit, + )?; + } + self.mcu_row += 1; + Ok(true) + } + + #[expect( + clippy::cast_possible_truncation, + reason = "MCU indices are bounded by validated u32 image dimensions" + )] + fn restart(&mut self) -> Result<(), JpegError> { + let mcu_at = (self.mcu_row * self.mcus_per_row) as u32; + let mcu_total = (self.mcu_rows * self.mcus_per_row) as u32; + self.br.reset_at_restart(); + let (br, next_rst) = self + .br + .clone() + .restarted(self.expected_rst, mcu_at, mcu_total)?; + self.br = br; + self.expected_rst = next_rst; + self.rows_to_go = self.restart_rows; + for component in &mut self.components { + component.first_row = true; + } + Ok(()) + } + + /// Validate the end of the scan once every MCU row is decoded. + pub(crate) fn finish(mut self) -> Result, JpegError> { + finish_scan(&mut self.br, true) + } +} + +/// Entropy pass: every difference of one MCU row, in MCU order. +#[expect( + clippy::inline_always, + reason = "keeps the caller's local bit reader in registers across the row" +)] +#[inline(always)] +fn decode_differences( + br: &mut BitReader<'_>, + components: &mut [ComponentRows<'_>], + mcus_per_row: usize, +) -> Result<(), JpegError> { + match components { + [only] => { + let table = only.table; + for diff in &mut only.diffs[..mcus_per_row] { + *diff = table.decode_lossless_diff(br)?; + } + Ok(()) + } + [c0, c1, c2] if [c0.mcu_h, c0.mcu_v, c1.mcu_h, c1.mcu_v, c2.mcu_h, c2.mcu_v] == [1; 6] => { + let (t0, t1, t2) = (c0.table, c1.table, c2.table); + for ((d0, d1), d2) in c0.diffs[..mcus_per_row] + .iter_mut() + .zip(&mut c1.diffs[..mcus_per_row]) + .zip(&mut c2.diffs[..mcus_per_row]) + { + *d0 = t0.decode_lossless_diff(br)?; + *d1 = t1.decode_lossless_diff(br)?; + *d2 = t2.decode_lossless_diff(br)?; + } + Ok(()) + } + _ => decode_sampled_differences(br, components, mcus_per_row), + } +} + +#[expect( + clippy::inline_always, + reason = "keeps the caller's local bit reader in registers across the row" +)] +#[inline(always)] +fn decode_sampled_differences( + br: &mut BitReader<'_>, + components: &mut [ComponentRows<'_>], + mcus_per_row: usize, +) -> Result<(), JpegError> { + for mcu_x in 0..mcus_per_row { + for component in components.iter_mut() { + let table = component.table; + for row in 0..component.mcu_v { + let start = row * component.coded_width + mcu_x * component.mcu_h; + for diff in &mut component.diffs[start..start + component.mcu_h] { + *diff = table.decode_lossless_diff(br)?; + } + } + } + } + Ok(()) +} + +/// Reconstruction pass for one component's rows of an MCU row. +fn reconstruct_component( + component: &mut ComponentRows<'_>, + mcu_row: usize, + predictor: u8, + initial_predictor: u16, + sample_limit: u16, +) -> Result<(), JpegError> { + let width = component.width; + if mcu_row > 0 { + // Carry the last row of the previous MCU row into the predictor slot. + let last = component.rows_ready * width; + component.samples.copy_within(last..last + width, 0); + } + let first = mcu_row * component.mcu_v; + let rows = component.mcu_v.min(component.height.saturating_sub(first)); + component.rows_ready = rows; + for row in 0..rows { + let diff = &component.diffs[row * component.coded_width..][..width]; + let (above, current) = component.samples.split_at_mut((row + 1) * width); + let above = &above[row * width..]; + let current = &mut current[..width]; + if component.first_row { + undifference_first_row(diff, current, initial_predictor); + component.first_row = false; + } else { + undifference_row(predictor, diff, above, current); + } + if sample_limit != u16::MAX && current.iter().fold(0, |any, &s| any | s) > sample_limit { + return Err(JpegError::HuffmanDecode { + mcu: 0, + reason: HuffmanFailure::InvalidSymbol, + }); + } + } + Ok(()) +} + +/// First row of a scan or restart interval: predictor 1, seeded with +/// 2^(P-Pt-1). +fn undifference_first_row(diff: &[u16], out: &mut [u16], initial_predictor: u16) { + let mut ra = initial_predictor; + for (sample, &d) in out.iter_mut().zip(diff) { + ra = ra.wrapping_add(d); + *sample = ra; + } +} + +/// Any later row: Rb for the first sample, then the scan's predictor. +#[expect( + clippy::cast_possible_truncation, + clippy::cast_sign_loss, + reason = "predictions are reduced modulo 2^16 exactly like libjpeg-turbo's `& 0xFFFF`" +)] +fn undifference_row(predictor: u8, diff: &[u16], above: &[u16], out: &mut [u16]) { + let (Some((&d0, diff)), Some((&rb0, _)), Some((first, rest))) = ( + diff.split_first(), + above.split_first(), + out.split_first_mut(), + ) else { + return; + }; + *first = rb0.wrapping_add(d0); + let pairs = above.windows(2).zip(diff); + let mut ra = *first; + match predictor { + 1 => { + for (sample, &d) in rest.iter_mut().zip(diff) { + ra = ra.wrapping_add(d); + *sample = ra; + } + } + 2 => { + for (sample, (above, &d)) in rest.iter_mut().zip(pairs) { + *sample = above[1].wrapping_add(d); + } + } + 3 => { + for (sample, (above, &d)) in rest.iter_mut().zip(pairs) { + *sample = above[0].wrapping_add(d); + } + } + 4 => { + for (sample, (above, &d)) in rest.iter_mut().zip(pairs) { + ra = ra + .wrapping_add(above[1]) + .wrapping_sub(above[0]) + .wrapping_add(d); + *sample = ra; + } + } + 5 => { + for (sample, (above, &d)) in rest.iter_mut().zip(pairs) { + let (rb, rc) = (i32::from(above[1]), i32::from(above[0])); + let prediction = i32::from(ra) + ((rb - rc) >> 1); + ra = (prediction as u16).wrapping_add(d); + *sample = ra; + } + } + 6 => { + for (sample, (above, &d)) in rest.iter_mut().zip(pairs) { + let (rb, rc) = (i32::from(above[1]), i32::from(above[0])); + let prediction = rb + ((i32::from(ra) - rc) >> 1); + ra = (prediction as u16).wrapping_add(d); + *sample = ra; + } + } + _ => { + for (sample, (above, &d)) in rest.iter_mut().zip(pairs) { + let prediction = (u32::from(ra) + u32::from(above[1])) >> 1; + ra = (prediction as u16).wrapping_add(d); + *sample = ra; + } + } + } +} diff --git a/crates/j2k-jpeg/src/lossless/scan/tests.rs b/crates/j2k-jpeg/src/lossless/scan/tests.rs new file mode 100644 index 000000000..67134bdd2 --- /dev/null +++ b/crates/j2k-jpeg/src/lossless/scan/tests.rs @@ -0,0 +1,393 @@ +// SPDX-License-Identifier: MIT OR Apache-2.0 + +use alloc::vec; +use alloc::vec::Vec; + +use super::*; +use crate::entropy::huffman::HuffmanTable; +use crate::parse::tables::{HuffmanTableRole, HuffmanValues, RawHuffmanTable}; + +/// DC table giving every category 0..=16 a 5-bit code equal to the category. +fn category_table() -> HuffmanTable { + let mut bits = [0u8; 16]; + bits[4] = 17; + let values: Vec = (0..=16).collect(); + HuffmanTable::from_raw( + &RawHuffmanTable { + bits, + values: HuffmanValues::from_slice(&values), + }, + HuffmanTableRole::Dc, + ) + .expect("17 five-bit codes fit") +} + +#[derive(Default)] +struct BitWriter { + bytes: Vec, + acc: u32, + bits: u8, +} + +impl BitWriter { + fn put(&mut self, value: u32, count: u8) { + for bit in (0..count).rev() { + self.acc = (self.acc << 1) | ((value >> bit) & 1); + self.bits += 1; + if self.bits == 8 { + self.push_byte(); + } + } + } + + fn push_byte(&mut self) { + let byte = u8::try_from(self.acc & 0xff).expect("masked byte"); + self.bytes.push(byte); + if byte == 0xff { + self.bytes.push(0); + } + self.acc = 0; + self.bits = 0; + } + + /// Pad to a byte boundary with one bits, as T.81 F.1.2.3 requires. + fn align(&mut self) { + while self.bits != 0 { + self.put(1, 1); + } + } + + /// Encode one difference, interpreted modulo 2^16 like T.81 H.1.2.2. + fn diff(&mut self, diff: i32) { + let wrapped = diff.rem_euclid(1 << 16); + let signed = if wrapped > 0x8000 { + wrapped - 0x1_0000 + } else { + wrapped + }; + if signed == 0x8000 { + self.put(16, 5); + return; + } + let category = u8::try_from(32 - signed.unsigned_abs().leading_zeros()).expect("category"); + self.put(u32::from(category), 5); + if category != 0 { + let magnitude = if signed < 0 { signed - 1 } else { signed }; + let mask = (1u32 << category) - 1; + self.put(u32::from_ne_bytes(magnitude.to_ne_bytes()) & mask, category); + } + } + + fn finish(mut self) -> Vec { + self.align(); + self.bytes.extend_from_slice(&[0xff, 0xd9]); + self.bytes + } +} + +/// Reference T.81 H.1.2.1 prediction with libjpeg-turbo's modulo-2^16 +/// arithmetic, evaluated independently of the row kernels. +fn reference_predict(predictor: u8, ra: u16, rb: u16, rc: u16) -> u16 { + let (ra, rb, rc) = (i32::from(ra), i32::from(rb), i32::from(rc)); + let prediction = match predictor { + 1 => ra, + 2 => rb, + 3 => rc, + 4 => ra + rb - rc, + 5 => ra + ((rb - rc) >> 1), + 6 => rb + ((ra - rc) >> 1), + _ => (ra + rb) >> 1, + }; + u16::try_from(prediction.rem_euclid(1 << 16)).expect("reduced prediction") +} + +fn pseudo_random_row(seed: u32, len: usize) -> Vec { + let mut state = seed | 1; + (0..len) + .map(|_| { + state ^= state << 13; + state ^= state >> 17; + state ^= state << 5; + u16::try_from(state >> 16).expect("16-bit sample") + }) + .collect() +} + +#[test] +fn first_row_accumulates_from_the_initial_predictor_modulo_2_16() { + let mut out = [0u16; 4]; + undifference_first_row(&[5, 0xffff, 3, 0x8000], &mut out, 128); + assert_eq!(out, [133, 132, 135, 0x8087]); +} + +#[test] +fn later_rows_start_from_rb_then_apply_each_predictor() { + let above = pseudo_random_row(7, 33); + let diffs = pseudo_random_row(11, 33); + for predictor in 1..=7 { + let mut out = vec![0u16; 33]; + undifference_row(predictor, &diffs, &above, &mut out); + let mut expected = vec![0u16; 33]; + expected[0] = above[0].wrapping_add(diffs[0]); + for x in 1..33 { + let prediction = reference_predict(predictor, expected[x - 1], above[x], above[x - 1]); + expected[x] = prediction.wrapping_add(diffs[x]); + } + assert_eq!(out, expected, "predictor {predictor}"); + } +} + +#[test] +fn predictors_read_left_above_and_diagonal_neighbours() { + // Ra = 30, Rb = 20, Rc = 10 for the second sample of the lower row. + let above = [10u16, 20]; + let expected = [ + (1, 30), + (2, 20), + (3, 10), + (4, 40), + (5, 35), + (6, 30), + (7, 25), + ]; + for (predictor, prediction) in expected { + let mut out = [0u16; 2]; + undifference_row(predictor, &[20, 0], &above, &mut out); + assert_eq!(out, [30, prediction], "predictor {predictor}"); + } +} + +#[test] +fn predictions_wrap_modulo_2_16() { + // Rb = 0 plus 65535 gives Ra = 65535; predictor 4 then gives + // 65535 + 65535 - 0, which wraps to 65534, and +3 wraps to 1. + let mut out = [0u16; 2]; + undifference_row(4, &[0xffff, 3], &[0, 0xffff], &mut out); + assert_eq!(out, [0xffff, 1]); +} + +#[test] +fn category_16_decodes_as_32768_without_extra_bits() { + let table = category_table(); + let table = table.dc().expect("DC table"); + let mut writer = BitWriter::default(); + writer.diff(0x8000); + writer.diff(-1); + let bytes = writer.finish(); + let mut br = BitReader::new(&bytes); + assert_eq!(table.decode_lossless_diff(&mut br), Ok(0x8000)); + assert_eq!(table.decode_lossless_diff(&mut br), Ok(0xffff)); +} + +fn scan_samples( + table: &HuffmanTable, + bytes: &[u8], + spec: LosslessScanSpec, +) -> Result, JpegError> { + let dc = table.dc()?; + let mut scan = LosslessScan::new( + bytes, + spec, + &[LosslessComponentSpec { + h: 1, + v: 1, + table: dc, + }], + )?; + let mut samples = Vec::new(); + while scan.decode_mcu_row()? { + samples.extend_from_slice(scan.row(0, 0)); + } + scan.finish()?; + Ok(samples) +} + +/// Encode a grayscale image with T.81 prediction, restarting every +/// `restart_rows` rows (0 disables restarts). +fn encode_gray( + samples: &[u16], + width: usize, + predictor: u8, + bias: u16, + restart_rows: usize, +) -> Vec { + let mut writer = BitWriter::default(); + let rows: Vec<&[u16]> = samples.chunks(width).collect(); + let mut rst = 0u8; + for (y, row) in rows.iter().enumerate() { + let interval_row = if restart_rows == 0 { + y + } else { + y % restart_rows + }; + if restart_rows != 0 && y != 0 && interval_row == 0 { + writer.align(); + writer.bytes.extend_from_slice(&[0xff, 0xd0 | rst]); + rst = (rst + 1) & 7; + } + for (x, &sample) in row.iter().enumerate() { + let prediction = match (interval_row, x) { + (0, 0) => bias, + (0, _) => row[x - 1], + (_, 0) => rows[y - 1][0], + _ => reference_predict(predictor, row[x - 1], rows[y - 1][x], rows[y - 1][x - 1]), + }; + writer.diff(i32::from(sample) - i32::from(prediction)); + } + } + writer.finish() +} + +fn gray_spec(width: u32, height: u32, predictor: u8, precision: u8) -> LosslessScanSpec { + LosslessScanSpec { + dimensions: (width, height), + predictor, + precision, + point_transform: 0, + restart_interval: None, + } +} + +#[test] +fn sixteen_bit_scan_round_trips_wrapping_differences_for_every_predictor() { + let table = category_table(); + let (width, height) = (9usize, 5usize); + let mut samples = pseudo_random_row(3, width * height); + samples[0] = 0; + samples[1] = 0xffff; + samples[width] = 0x8000; + for predictor in 1..=7 { + let bytes = encode_gray(&samples, width, predictor, 0x8000, 0); + let decoded = scan_samples(&table, &bytes, gray_spec(9, 5, predictor, 16)); + assert_eq!(decoded, Ok(samples.clone()), "predictor {predictor}"); + } +} + +#[test] +fn restart_intervals_restart_the_one_dimensional_first_row() { + let table = category_table(); + let (width, height) = (6usize, 7usize); + let samples: Vec = pseudo_random_row(5, width * height) + .into_iter() + .map(|sample| sample >> 8) + .collect(); + for predictor in 1..=7 { + for restart_rows in [1usize, 2, 3] { + let bytes = encode_gray(&samples, width, predictor, 128, restart_rows); + let spec = LosslessScanSpec { + restart_interval: Some(u16::try_from(restart_rows * width).expect("interval")), + ..gray_spec(6, 7, predictor, 8) + }; + let decoded = scan_samples(&table, &bytes, spec); + assert_eq!( + decoded, + Ok(samples.clone()), + "predictor {predictor}, restart every {restart_rows} rows" + ); + } + } +} + +#[test] +fn restart_intervals_must_cover_whole_rows() { + let table = category_table(); + let spec = LosslessScanSpec { + restart_interval: Some(4), + ..gray_spec(6, 2, 1, 8) + }; + let dc = table.dc().expect("DC table"); + let result = LosslessScan::new( + &[0xff, 0xd9], + spec, + &[LosslessComponentSpec { + h: 1, + v: 1, + table: dc, + }], + ); + assert!(matches!( + result, + Err(JpegError::NotImplemented { + sof: SofKind::Lossless + }) + )); +} + +#[test] +fn samples_wider_than_the_precision_are_rejected() { + let table = category_table(); + // 128 + 200 = 328 does not fit an 8-bit sample. + let mut writer = BitWriter::default(); + writer.diff(200); + let bytes = writer.finish(); + let decoded = scan_samples(&table, &bytes, gray_spec(1, 1, 1, 8)); + assert_eq!( + decoded, + Err(JpegError::HuffmanDecode { + mcu: 0, + reason: HuffmanFailure::InvalidSymbol, + }) + ); +} + +#[test] +fn point_transform_scales_the_initial_predictor() { + let table = category_table(); + // P = 12, Pt = 4: the first sample predicts 2^(12-4-1) = 128. + let mut writer = BitWriter::default(); + writer.diff(1); + let bytes = writer.finish(); + let spec = LosslessScanSpec { + point_transform: 4, + ..gray_spec(1, 1, 1, 12) + }; + let dc = table.dc().expect("DC table"); + let mut scan = LosslessScan::new( + &bytes, + spec, + &[LosslessComponentSpec { + h: 1, + v: 1, + table: dc, + }], + ) + .expect("scan"); + assert_eq!(scan.decode_mcu_row(), Ok(true)); + assert_eq!(scan.row(0, 0), &[129]); + assert_eq!(scan.point_transform(), 4); +} + +#[test] +fn allocation_bytes_cover_padded_mcu_rows() { + let interleaved = + lossless_scan_allocation_bytes((5, 3), &[(2, 2), (1, 1), (1, 1)]).expect("layout"); + // Luma: 2 rows of 3 MCUs x 2 diffs, 3 rows of 5 samples; chroma: 1 row of + // 3 diffs, 2 rows of 3 samples each. + let samples = (12 + 15) + 2 * (3 + 6); + assert_eq!( + interleaved, + 3 * core::mem::size_of::>() + samples * 2 + ); + let gray = lossless_scan_allocation_bytes((5, 3), &[(2, 2)]).expect("layout"); + assert_eq!( + gray, + core::mem::size_of::>() + (5 + 10) * 2 + ); +} + +#[test] +fn stored_samples_are_little_endian_and_shifted_by_the_point_transform() { + use crate::lossless::LosslessSample; + + let mut gray16 = [0u8; 4]; + ::store_row(&[0x1234, 0x0abc], 4, &mut gray16); + assert_eq!(gray16, [0x40, 0x23, 0xc0, 0xab]); + + let mut rgb8 = [0u8; 6]; + ::store_interleaved([&[1, 4], &[2, 5], &[3, 6]], 1, &mut rgb8); + assert_eq!(rgb8, [2, 4, 6, 8, 10, 12]); + + let mut rgb16 = [0u8; 6]; + ::store_interleaved([&[0x0102], &[0x0304], &[0x0506]], 0, &mut rgb16); + assert_eq!(rgb16, [0x02, 0x01, 0x04, 0x03, 0x06, 0x05]); +} diff --git a/crates/j2k-jpeg/src/simd/neon_memory.rs b/crates/j2k-jpeg/src/simd/neon_memory.rs index 7826e1a3f..f6052d04f 100644 --- a/crates/j2k-jpeg/src/simd/neon_memory.rs +++ b/crates/j2k-jpeg/src/simd/neon_memory.rs @@ -3,8 +3,8 @@ //! Fixed-size `AArch64` NEON memory operations. use core::arch::aarch64::{ - int16x8_t, uint8x16_t, uint8x16x3_t, uint8x8_t, uint8x8x3_t, vld1_u8, vld1q_s16, vld1q_u8, - vst1_u8, vst1q_u8, vst3_u8, vst3q_u8, + int16x8_t, uint16x8_t, uint8x16_t, uint8x16x3_t, uint8x8_t, uint8x8x3_t, vld1_u8, vld1q_s16, + vld1q_u8, vst1_u8, vst1q_u16, vst1q_u8, vst3_u8, vst3q_u8, }; #[inline] @@ -59,6 +59,18 @@ pub(crate) fn store_u8x8(dst: &mut [u8; 8], values: uint8x8_t) { unsafe { vst1_u8(dst.as_mut_ptr(), values) }; } +#[inline] +#[target_feature(enable = "neon")] +pub(crate) fn store_u16x8(dst: &mut [u16; 8], values: uint16x8_t) { + // SAFETY: + // - Feature availability: callers run inside a `Neon` token kernel. + // - Bounds: the array reference proves eight writable u16 samples. + // - Alignment: AArch64 `vst1q_u16` supports unaligned u16 addresses. + // - Aliasing: the exclusive reference prevents overlapping live access. + // - Initialization: the store initializes every output sample. + unsafe { vst1q_u16(dst.as_mut_ptr(), values) }; +} + #[inline] #[target_feature(enable = "neon")] pub(crate) fn store_u8x16(dst: &mut [u8; 16], values: uint8x16_t) { diff --git a/crates/j2k-jpeg/tests/libjpeg_turbo_compare.rs b/crates/j2k-jpeg/tests/libjpeg_turbo_compare.rs index ac69ad6be..9b9e10dc6 100644 --- a/crates/j2k-jpeg/tests/libjpeg_turbo_compare.rs +++ b/crates/j2k-jpeg/tests/libjpeg_turbo_compare.rs @@ -370,6 +370,9 @@ mod turbo_codec { const TJPARAM_QUALITY: c_int = 3; const TJPARAM_SUBSAMP: c_int = 4; const TJPARAM_PROGRESSIVE: c_int = 12; + const TJPARAM_LOSSLESS: c_int = 15; + const TJPARAM_LOSSLESSPSV: c_int = 16; + const TJPARAM_RESTARTROWS: c_int = 19; pub(crate) const TJPF_RGB: c_int = 0; pub(crate) const TJPF_GRAY: c_int = 6; pub(crate) const TJSAMP_444: c_int = 0; @@ -412,6 +415,42 @@ mod turbo_codec { pitch: c_int, pixel_format: c_int, ) -> c_int; + fn tj3Compress16( + handle: *mut c_void, + src: *const u16, + width: c_int, + pitch: c_int, + height: c_int, + pixel_format: c_int, + jpeg: *mut *mut u8, + jpeg_size: *mut usize, + ) -> c_int; + fn tj3Decompress8( + handle: *mut c_void, + jpeg: *const u8, + jpeg_size: usize, + dst: *mut u8, + pitch: c_int, + pixel_format: c_int, + ) -> c_int; + fn tj3Decompress16( + handle: *mut c_void, + jpeg: *const u8, + jpeg_size: usize, + dst: *mut u16, + pitch: c_int, + pixel_format: c_int, + ) -> c_int; + } + + /// Lossless (SOF3) encoding parameters for a generated fixture. + #[derive(Clone, Copy, Debug)] + pub(crate) struct Lossless { + pub(crate) width: u32, + pub(crate) height: u32, + pub(crate) pixel_format: c_int, + pub(crate) predictor: c_int, + pub(crate) restart_rows: c_int, } /// Encoding parameters for a generated fixture. @@ -456,6 +495,19 @@ mod turbo_codec { } } + impl Handle { + fn configure_lossless(&self, spec: Lossless) { + for (param, value) in [ + (TJPARAM_LOSSLESS, 1), + (TJPARAM_LOSSLESSPSV, spec.predictor), + (TJPARAM_RESTARTROWS, spec.restart_rows), + ] { + // SAFETY: the handle is live and the parameter ids are TurboJPEG 3's. + self.check(unsafe { tj3Set(self.0, param, value) }, "tj3Set"); + } + } + } + impl Drop for Handle { fn drop(&mut self) { // SAFETY: the handle came from tj3Init and is destroyed once. @@ -531,6 +583,100 @@ mod turbo_codec { take_jpeg(jpeg, size) } + pub(crate) fn compress_lossless8(pixels: &[u8], spec: Lossless) -> Vec { + let handle = Handle::new(TJINIT_COMPRESS); + handle.configure_lossless(spec); + let pitch = spec.width as usize * channels(spec.pixel_format); + assert_eq!(pixels.len(), pitch * spec.height as usize); + let (mut jpeg, mut size) = (std::ptr::null_mut(), 0usize); + // SAFETY: as for compress8. + let rc = unsafe { + tj3Compress8( + handle.0, + pixels.as_ptr(), + c(spec.width as usize), + c(pitch), + c(spec.height as usize), + spec.pixel_format, + &raw mut jpeg, + &raw mut size, + ) + }; + handle.check(rc, "tj3Compress8 lossless"); + take_jpeg(jpeg, size) + } + + pub(crate) fn compress_lossless16(samples: &[u16], spec: Lossless) -> Vec { + let handle = Handle::new(TJINIT_COMPRESS); + handle.configure_lossless(spec); + let pitch = spec.width as usize * channels(spec.pixel_format); + assert_eq!(samples.len(), pitch * spec.height as usize); + let (mut jpeg, mut size) = (std::ptr::null_mut(), 0usize); + // SAFETY: as for compress8, with 16-bit samples. + let rc = unsafe { + tj3Compress16( + handle.0, + samples.as_ptr(), + c(spec.width as usize), + c(pitch), + c(spec.height as usize), + spec.pixel_format, + &raw mut jpeg, + &raw mut size, + ) + }; + handle.check(rc, "tj3Compress16"); + take_jpeg(jpeg, size) + } + + pub(crate) fn decompress8( + jpeg: &[u8], + width: u32, + height: u32, + pixel_format: c_int, + ) -> Vec { + let handle = Handle::new(TJINIT_DECOMPRESS); + let pitch = width as usize * channels(pixel_format); + let mut out = vec![0u8; pitch * height as usize]; + // SAFETY: `out` holds `height` rows of `pitch` samples. + let rc = unsafe { + tj3Decompress8( + handle.0, + jpeg.as_ptr(), + jpeg.len(), + out.as_mut_ptr(), + c(pitch), + pixel_format, + ) + }; + handle.check(rc, "tj3Decompress8"); + out + } + + pub(crate) fn decompress16( + jpeg: &[u8], + width: u32, + height: u32, + pixel_format: c_int, + ) -> Vec { + let handle = Handle::new(TJINIT_DECOMPRESS); + let pitch = width as usize * channels(pixel_format); + let mut out = vec![0u16; pitch * height as usize]; + // SAFETY: `out` holds `height` rows of `pitch` samples. + let rc = unsafe { + tj3Decompress16( + handle.0, + jpeg.as_ptr(), + jpeg.len(), + out.as_mut_ptr(), + c(pitch), + pixel_format, + ) + }; + handle.check(rc, "tj3Decompress16"); + out + } + /// Full-frame 12-bit decode with libjpeg-turbo's defaults (ISLOW IDCT, /// fancy upsampling). pub(crate) fn decompress12( @@ -676,3 +822,126 @@ fn extended12_decodes_match_turbo() { } } } + +/// Lossless streams from libjpeg-turbo's encoder, across predictors, restart +/// intervals, and sample sizes, decode exactly as libjpeg-turbo decodes them. +/// The 16-bit content jumps between extremes so differences wrap modulo 2^16 +/// and hit the 32768 (category 16) case. +#[cfg(all(has_libjpeg_turbo, has_libjpeg_turbo_v3))] +#[test] +fn lossless_decodes_match_turbo() { + use turbo_codec::{Lossless, TJPF_GRAY, TJPF_RGB}; + + if !turbo_available() { + return; + } + for (pixel_format, channels) in [(TJPF_GRAY, 1), (TJPF_RGB, 3)] { + for (width, height) in [(19, 7), (64, 16)] { + let smooth = textured_samples(width, height, channels, 65_535); + let wide = smooth + .iter() + .enumerate() + .map(|(index, &sample)| match index % 11 { + 3 => 0, + 4 => 65_535, + 5 => 32_768, + _ => u16::try_from(sample).expect("16-bit sample"), + }) + .collect::>(); + let narrow = smooth + .iter() + .map(|&sample| u8::try_from(sample >> 8).expect("8-bit sample")) + .collect::>(); + for predictor in 1..=7 { + for restart_rows in [0, 1, 3] { + let spec = Lossless { + width, + height, + pixel_format, + predictor, + restart_rows, + }; + let jpeg = turbo_codec::compress_lossless8(&narrow, spec); + let expected = turbo_codec::decompress8(&jpeg, width, height, pixel_format); + let format = if channels == 1 { + PixelFormat::Gray8 + } else { + PixelFormat::Rgb8 + }; + let (actual, _) = Decoder::new(&jpeg) + .expect("8-bit lossless decoder") + .decode_request(DecodeRequest::full(format)) + .unwrap_or_else(|err| panic!("{spec:?}: 8-bit lossless decode: {err}")); + assert!(actual == expected, "{spec:?}: 8-bit lossless mismatch"); + + let jpeg = turbo_codec::compress_lossless16(&wide, spec); + let expected = turbo_codec::decompress16(&jpeg, width, height, pixel_format); + let format = if channels == 1 { + PixelFormat::Gray16 + } else { + PixelFormat::Rgb16 + }; + let (actual, _) = Decoder::new(&jpeg) + .expect("16-bit lossless decoder") + .decode_request(DecodeRequest::full(format)) + .unwrap_or_else(|err| panic!("{spec:?}: 16-bit lossless decode: {err}")); + let actual = actual + .chunks_exact(2) + .map(|pair| u16::from_le_bytes([pair[0], pair[1]])) + .collect::>(); + assert!(actual == expected, "{spec:?}: 16-bit lossless mismatch"); + } + } + } + } +} + +/// Progressive streams from libjpeg-turbo's default script, which includes +/// successive-approximation refinement scans, decode exactly as turbo +/// decodes them across subsamplings, qualities, and non-MCU-aligned sizes. +#[cfg(all(has_libjpeg_turbo, has_libjpeg_turbo_v3))] +#[test] +fn progressive_decodes_match_turbo() { + use turbo_codec::{ + Encode, TJPF_GRAY, TJPF_RGB, TJSAMP_420, TJSAMP_422, TJSAMP_444, TJSAMP_GRAY, + }; + + if !turbo_available() { + return; + } + for (pixel_format, subsampling) in [ + (TJPF_GRAY, TJSAMP_GRAY), + (TJPF_RGB, TJSAMP_444), + (TJPF_RGB, TJSAMP_422), + (TJPF_RGB, TJSAMP_420), + ] { + let (channels, format) = if pixel_format == TJPF_GRAY { + (1, PixelFormat::Gray8) + } else { + (3, PixelFormat::Rgb8) + }; + for (width, height) in [(64, 48), (67, 45), (250, 3)] { + for quality in [50, 90, 100] { + let pixels = textured_samples(width, height, channels, 255) + .into_iter() + .map(|sample| u8::try_from(sample).expect("8-bit sample")) + .collect::>(); + let spec = Encode { + width, + height, + pixel_format, + subsampling, + quality, + progressive: true, + }; + let jpeg = turbo_codec::compress8(&pixels, spec); + let expected = turbo_codec::decompress8(&jpeg, width, height, pixel_format); + let (actual, _) = Decoder::new(&jpeg) + .expect("progressive decoder") + .decode_request(DecodeRequest::full(format)) + .unwrap_or_else(|err| panic!("{spec:?}: progressive decode: {err}")); + assert!(actual == expected, "{spec:?}: progressive mismatch"); + } + } + } +} diff --git a/crates/j2k-test-support/src/jpeg_fixtures/builders.rs b/crates/j2k-test-support/src/jpeg_fixtures/builders.rs index a8a4d98e9..a80ecd8c8 100644 --- a/crates/j2k-test-support/src/jpeg_fixtures/builders.rs +++ b/crates/j2k-test-support/src/jpeg_fixtures/builders.rs @@ -1926,6 +1926,8 @@ fn lossless_rgb_entropy_with_restarts( let sample = samples[pixel * 3 + component]; let predicted = if segment_offset == 0 { 128 + } else if y == segment_start / width { + i32::from(samples[(pixel - 1) * 3 + component]) } else { lossless_predicted_rgb_value(samples, width, x, y, component, predictor) }; @@ -1966,6 +1968,8 @@ fn lossless_entropy_with_restarts( let y = idx / width; let predicted = if segment_offset == 0 { 128 + } else if y == segment_start / width { + i32::from(samples[idx - 1]) } else { lossless_predicted_value(samples, width, x, y, predictor) }; @@ -2035,17 +2039,26 @@ fn lossless_422_entropy_8bit_with_restarts( let y = mcu / chroma_width; let mcu_x = mcu % chroma_width; let x0 = mcu_x * 2; + let restart = RestartPosition { + first_sample: segment_offset == 0, + first_row: Some(segment_start / chroma_width), + }; encode_lossless_component_sample_8bit_with_restart( - &mut bits, - c0, - width, - x0, - y, - predictor, - segment_offset == 0, + &mut bits, c0, width, x0, y, predictor, restart, ); if x0 + 1 < width { - encode_lossless_component_sample_8bit(&mut bits, c0, width, x0 + 1, y, predictor); + encode_lossless_component_sample_8bit_with_restart( + &mut bits, + c0, + width, + x0 + 1, + y, + predictor, + RestartPosition { + first_sample: false, + ..restart + }, + ); } encode_lossless_component_sample_8bit_with_restart( &mut bits, @@ -2054,7 +2067,7 @@ fn lossless_422_entropy_8bit_with_restarts( mcu_x, y, predictor, - segment_offset == 0, + restart, ); encode_lossless_component_sample_8bit_with_restart( &mut bits, @@ -2063,7 +2076,7 @@ fn lossless_422_entropy_8bit_with_restarts( mcu_x, y, predictor, - segment_offset == 0, + restart, ); } out.extend(pack_entropy_bits(bits)); @@ -2099,6 +2112,7 @@ fn lossless_420_entropy_8bit( mcu_x, mcu_y, restart_first_mcu: false, + restart_first_mcu_row: None, }, ); } @@ -2139,6 +2153,7 @@ fn lossless_420_entropy_8bit_with_restarts( mcu_x, mcu_y, restart_first_mcu: segment_offset == 0, + restart_first_mcu_row: Some(segment_start / chroma_width), }, ); } @@ -2232,17 +2247,26 @@ fn lossless_422_entropy_16bit_with_restarts( let y = mcu / chroma_width; let mcu_x = mcu % chroma_width; let x0 = mcu_x * 2; + let restart = RestartPosition { + first_sample: segment_offset == 0, + first_row: Some(segment_start / chroma_width), + }; encode_lossless_component_sample_16bit_with_restart( - &mut bits, - c0, - width, - x0, - y, - predictor, - segment_offset == 0, + &mut bits, c0, width, x0, y, predictor, restart, ); if x0 + 1 < width { - encode_lossless_component_sample_16bit(&mut bits, c0, width, x0 + 1, y, predictor); + encode_lossless_component_sample_16bit_with_restart( + &mut bits, + c0, + width, + x0 + 1, + y, + predictor, + RestartPosition { + first_sample: false, + ..restart + }, + ); } encode_lossless_component_sample_16bit_with_restart( &mut bits, @@ -2251,7 +2275,7 @@ fn lossless_422_entropy_16bit_with_restarts( mcu_x, y, predictor, - segment_offset == 0, + restart, ); encode_lossless_component_sample_16bit_with_restart( &mut bits, @@ -2260,7 +2284,7 @@ fn lossless_422_entropy_16bit_with_restarts( mcu_x, y, predictor, - segment_offset == 0, + restart, ); } out.extend(pack_entropy_bits(bits)); @@ -2296,6 +2320,7 @@ fn lossless_420_entropy_16bit( mcu_x, mcu_y, restart_first_mcu: false, + restart_first_mcu_row: None, }, ); } @@ -2336,6 +2361,7 @@ fn lossless_420_entropy_16bit_with_restarts( mcu_x, mcu_y, restart_first_mcu: segment_offset == 0, + restart_first_mcu_row: Some(segment_start / chroma_width), }, ); } @@ -2356,6 +2382,8 @@ struct Lossless420Mcu { mcu_x: usize, mcu_y: usize, restart_first_mcu: bool, + /// First MCU row of the current restart interval, if restarts are on. + restart_first_mcu_row: Option, } fn encode_lossless_420_mcu_8bit( @@ -2378,7 +2406,10 @@ fn encode_lossless_420_mcu_8bit( x, y, predictor, - mcu.restart_first_mcu && local_x == 0 && local_y == 0, + RestartPosition { + first_sample: mcu.restart_first_mcu && local_x == 0 && local_y == 0, + first_row: mcu.restart_first_mcu_row.map(|row| row * 2), + }, ); } } @@ -2390,7 +2421,10 @@ fn encode_lossless_420_mcu_8bit( mcu.mcu_x, mcu.mcu_y, predictor, - mcu.restart_first_mcu, + RestartPosition { + first_sample: mcu.restart_first_mcu, + first_row: mcu.restart_first_mcu_row, + }, ); encode_lossless_component_sample_8bit_with_restart( bits, @@ -2399,7 +2433,10 @@ fn encode_lossless_420_mcu_8bit( mcu.mcu_x, mcu.mcu_y, predictor, - mcu.restart_first_mcu, + RestartPosition { + first_sample: mcu.restart_first_mcu, + first_row: mcu.restart_first_mcu_row, + }, ); } @@ -2412,10 +2449,32 @@ fn encode_lossless_component_sample_8bit( predictor: u8, ) { encode_lossless_component_sample_8bit_with_restart( - bits, samples, width, x, y, predictor, false, + bits, + samples, + width, + x, + y, + predictor, + RestartPosition::NONE, ); } +/// Where a sample sits in its restart interval. T.81 H.1.2.1 predicts the +/// interval's first sample from 2^(P-1) and the rest of its first component +/// row from Ra alone, exactly like the first row of the scan. +#[derive(Clone, Copy)] +struct RestartPosition { + first_sample: bool, + first_row: Option, +} + +impl RestartPosition { + const NONE: Self = Self { + first_sample: false, + first_row: None, + }; +} + fn encode_lossless_component_sample_8bit_with_restart( bits: &mut Vec, samples: &[u8], @@ -2423,11 +2482,13 @@ fn encode_lossless_component_sample_8bit_with_restart( x: usize, y: usize, predictor: u8, - restart_first_sample: bool, + restart: RestartPosition, ) { let sample = samples[y * width + x]; - let predicted = if restart_first_sample { + let predicted = if restart.first_sample { 128 + } else if restart.first_row == Some(y) { + i32::from(samples[y * width + x - 1]) } else { lossless_predicted_value(samples, width, x, y, predictor) }; @@ -2459,7 +2520,10 @@ fn encode_lossless_420_mcu_16bit( x, y, predictor, - mcu.restart_first_mcu && local_x == 0 && local_y == 0, + RestartPosition { + first_sample: mcu.restart_first_mcu && local_x == 0 && local_y == 0, + first_row: mcu.restart_first_mcu_row.map(|row| row * 2), + }, ); } } @@ -2471,7 +2535,10 @@ fn encode_lossless_420_mcu_16bit( mcu.mcu_x, mcu.mcu_y, predictor, - mcu.restart_first_mcu, + RestartPosition { + first_sample: mcu.restart_first_mcu, + first_row: mcu.restart_first_mcu_row, + }, ); encode_lossless_component_sample_16bit_with_restart( bits, @@ -2480,7 +2547,10 @@ fn encode_lossless_420_mcu_16bit( mcu.mcu_x, mcu.mcu_y, predictor, - mcu.restart_first_mcu, + RestartPosition { + first_sample: mcu.restart_first_mcu, + first_row: mcu.restart_first_mcu_row, + }, ); } @@ -2493,7 +2563,13 @@ fn encode_lossless_component_sample_16bit( predictor: u8, ) { encode_lossless_component_sample_16bit_with_restart( - bits, samples, width, x, y, predictor, false, + bits, + samples, + width, + x, + y, + predictor, + RestartPosition::NONE, ); } @@ -2504,11 +2580,13 @@ fn encode_lossless_component_sample_16bit_with_restart( x: usize, y: usize, predictor: u8, - restart_first_sample: bool, + restart: RestartPosition, ) { let sample = samples[y * width + x]; - let predicted = if restart_first_sample { + let predicted = if restart.first_sample { 32768 + } else if restart.first_row == Some(y) { + i32::from(samples[y * width + x - 1]) } else { lossless_predicted_value_16bit(samples, width, x, y, predictor) }; @@ -2540,6 +2618,8 @@ fn lossless_entropy_16bit_with_restarts( let y = idx / width; let predicted = if segment_offset == 0 { 32768 + } else if y == segment_start / width { + i32::from(samples[idx - 1]) } else { lossless_predicted_value_16bit(samples, width, x, y, predictor) }; @@ -2602,6 +2682,8 @@ fn lossless_rgb_entropy_16bit_with_restarts( let sample = samples[pixel * 3 + component]; let predicted = if segment_offset == 0 { 32768 + } else if y == segment_start / width { + i32::from(samples[(pixel - 1) * 3 + component]) } else { lossless_predicted_rgb_value_16bit(samples, width, x, y, component, predictor) }; From 1839fb6336e9f92af291402334eedf88db9fa380 Mon Sep 17 00:00:00 2001 From: GF Date: Sat, 26 Sep 2026 20:19:38 -0400 Subject: [PATCH 4/5] perf(native): optimize bitplane scanning and sample placement --- .../j2k-native/src/j2c/arithmetic_decoder.rs | 80 ++-- crates/j2k-native/src/j2c/bitplane.rs | 3 + .../j2k-native/src/j2c/bitplane/arithmetic.rs | 178 ++------ crates/j2k-native/src/j2c/bitplane/bypass.rs | 10 +- crates/j2k-native/src/j2c/bitplane/context.rs | 42 -- crates/j2k-native/src/j2c/bitplane/facade.rs | 1 + .../src/j2c/bitplane/flag_passes.rs | 209 +++++++++ crates/j2k-native/src/j2c/bitplane/flags.rs | 139 ++++++ .../src/j2c/bitplane/golden_tests.rs | 5 +- .../src/j2c/bitplane/reconstruction.rs | 108 +++++ crates/j2k-native/src/j2c/bitplane/scan.rs | 53 +++ .../j2k-native/src/j2c/bitplane/schedule.rs | 49 +- crates/j2k-native/src/j2c/bitplane/state.rs | 135 ++---- .../src/j2c/bitplane/state/model.rs | 5 + .../src/j2c/bitplane/state/scan_masks.rs | 54 +++ .../src/j2c/bitplane/state/workspace.rs | 32 +- crates/j2k-native/src/j2c/bitplane/tests.rs | 427 ++++++++++++++---- crates/j2k-native/src/j2c/capabilities.rs | 16 +- crates/j2k-native/src/j2c/decode/store.rs | 19 + .../src/j2c/decode/store/upsample.rs | 166 +++++++ .../src/j2c/decode/subband/classic.rs | 24 +- .../src/j2c/ht_block_decode/magnitude.rs | 6 +- .../src/j2c/ht_block_decode/readers.rs | 136 +++--- .../src/j2c/ht_block_decode/readers/tests.rs | 119 +++++ .../j2k-native/src/scalar/classic_decode.rs | 18 +- crates/j2k-native/tests/ht_capabilities.rs | 55 ++- 26 files changed, 1558 insertions(+), 531 deletions(-) create mode 100644 crates/j2k-native/src/j2c/bitplane/flag_passes.rs create mode 100644 crates/j2k-native/src/j2c/bitplane/flags.rs create mode 100644 crates/j2k-native/src/j2c/bitplane/scan.rs create mode 100644 crates/j2k-native/src/j2c/bitplane/state/scan_masks.rs create mode 100644 crates/j2k-native/src/j2c/decode/store/upsample.rs create mode 100644 crates/j2k-native/src/j2c/ht_block_decode/readers/tests.rs diff --git a/crates/j2k-native/src/j2c/arithmetic_decoder.rs b/crates/j2k-native/src/j2c/arithmetic_decoder.rs index 64b3356a5..28e562468 100644 --- a/crates/j2k-native/src/j2c/arithmetic_decoder.rs +++ b/crates/j2k-native/src/j2c/arithmetic_decoder.rs @@ -6,6 +6,10 @@ use super::mq::PACKED_DECODER_STATES; +/// Tier-1 passes copy the decoder into a local, decode, and write it back so +/// the MQ registers stay in machine registers; `Clone` (not `Copy`) keeps +/// those copies explicit. +#[derive(Clone)] pub(crate) struct ArithmeticDecoder<'a> { /// The underlying encoded data. data: &'a [u8], @@ -139,54 +143,52 @@ impl<'a> ArithmeticDecoder<'a> { let qe = state as u32 & 0xffff; let mps = u32::from((state >> 16) & 1 != 0); - self.a -= qe; + let a = self.a - qe; - // This is a faster version that reduces branching, which has shown - // itself to be the main limiting factor for better performance. - // We short-circuit the case where just the most probably symbol is - // returned, and otherwise use a code path that works for both, - // MPS_EXCHANGE and LPS_EXCHANGE. - - if (self.c >> 16) < self.a && self.a & 0x8000 != 0 { - return mps; - } - - // Unified branchless MPS_EXCHANGE / LPS_EXCHANGE. In the Annex C.3.2 - // procedures, the only difference is that LPS flips the role of cond: + // Branchless: whether the symbol is the MPS and whether A needs + // renormalizing are data-dependent and poorly predicted, and the + // caller already branches on the decoded bit. The formulas below + // reduce to "return the MPS, change nothing" in the fast case. + // + // In the Annex C.3.2 procedures, the only difference between + // MPS_EXCHANGE and LPS_EXCHANGE is that LPS flips the role of cond: // exchange_mps: d = mps ^ cond, flip when cond, index = cond*nlps + inv*nmps // exchange_lps: d = mps ^ inv_cond, flip when inv_cond, index = cond*nmps + inv*nlps - // - // This is equivalent to XOR-ing cond with is_lps, so we can handle - // both paths with a single branchless computation. - // - // As can be seen above, renormalization is always performed. - let is_lps = u32::from((self.c >> 16) >= self.a); - - // LPS: C -= A << 16 (no-op when MPS). - let lps_mask = is_lps.wrapping_neg(); // 0xFFFF_FFFF if LPS, 0 if MPS - self.c -= (self.a << 16) & lps_mask; - - // Same condition as in exchange_mps / exchange_lps. - let cond = u32::from(self.a < qe); - - // LPS: a = qe (no-op when MPS, a stays as a - qe). - self.a = (self.a & !lps_mask) | (qe & lps_mask); - - // exchange_mps: d = mps ^ cond → cond ^ 0 - // exchange_lps: d = mps ^ inv_cond → cond ^ 1 - // unified: d = mps ^ (cond ^ is_lps) + // so both are handled by XOR-ing cond with is_lps. + let is_lps = u32::from((self.c >> 16) >= a); + // A probability-state transition happens exactly when A renormalizes: + // on every LPS, and on an MPS that leaves A below 0x8000. + let transition = is_lps | u32::from(a < 0x8000); + + // LPS: C -= A << 16 and A = Qe; MPS: A stays A - Qe. + let lps_mask = is_lps.wrapping_neg(); + self.c -= (a << 16) & lps_mask; + self.a = (a & !lps_mask) | (qe & lps_mask); + + // Same condition as in exchange_mps / exchange_lps. Without a + // transition, A >= 0x8000 > Qe, so this is zero and d is the MPS. + let cond = u32::from(a < qe); let pick_nlps = cond ^ is_lps; let d = mps ^ pick_nlps; - // exchange_mps: flip mps when cond & switch → (cond ^ 0) & switch - // exchange_lps: flip mps when inv_cond & switch → (cond ^ 1) & switch - // unified: flip mps when (cond ^ is_lps) & switch let pick_mask = u8::from(pick_nlps != 0).wrapping_neg(); let next_mps = (state >> 24) as u8; let next_lps = (state >> 32) as u8; - context.0 = next_mps ^ ((next_mps ^ next_lps) & pick_mask); - - self.renormalize(); + let next = next_mps ^ ((next_mps ^ next_lps) & pick_mask); + let keep_mask = u8::from(transition == 0).wrapping_neg(); + context.0 = next ^ ((next ^ context.0) & keep_mask); + + // Shift A back to at least 0x8000 in one step when the buffered bits + // suffice (zero shift in the fast case); otherwise take the byte-in + // loop. + let shifts_needed = self.a.leading_zeros() - 16; + if shifts_needed <= self.shift_count { + self.a <<= shifts_needed; + self.c <<= shifts_needed; + self.shift_count -= shifts_needed; + } else { + self.renormalize(); + } d } diff --git a/crates/j2k-native/src/j2c/bitplane.rs b/crates/j2k-native/src/j2c/bitplane.rs index 1c1c0714b..7d6af6ef4 100644 --- a/crates/j2k-native/src/j2c/bitplane.rs +++ b/crates/j2k-native/src/j2c/bitplane.rs @@ -4,8 +4,11 @@ mod arithmetic; mod bypass; mod context; mod facade; +mod flag_passes; +mod flags; mod observer; mod reconstruction; +mod scan; mod schedule; mod state; diff --git a/crates/j2k-native/src/j2c/bitplane/arithmetic.rs b/crates/j2k-native/src/j2c/bitplane/arithmetic.rs index 627a9a0f2..11600f7f4 100644 --- a/crates/j2k-native/src/j2c/bitplane/arithmetic.rs +++ b/crates/j2k-native/src/j2c/bitplane/arithmetic.rs @@ -1,13 +1,19 @@ // SPDX-License-Identifier: MIT OR Apache-2.0 +//! Per-coefficient arithmetic-coded passes, used when a code block cannot take +//! the packed-column passes in `flag_passes.rs`: vertically causal contexts, +//! or segments that mix in raw (bypass) coding. + use super::super::arithmetic_decoder::ArithmeticDecoder; use super::context::{ - context_label_magnitude_refinement_coding_from_state_lazy, - context_label_sign_coding_index_with_neighbors, context_label_zero_coding_from_neighbors, + context_label_magnitude_refinement_coding_from_state_lazy, context_label_sign_coding_index, + context_label_zero_coding_from_neighbors, +}; +use super::scan::{ + cleanup_candidate_scan_mask, cleanup_run_length_candidate, scan_unit_valid_mask, }; use super::state::{ - BitPlaneDecodeContext, COEFFICIENTS_PADDING, HAS_MAGNITUDE_REFINEMENT_MASK, - HAS_ZERO_CODING_MASK, SIGNIFICANCE_MASK, + BitPlaneDecodeContext, COEFFICIENTS_PADDING, HAS_MAGNITUDE_REFINEMENT_MASK, SIGNIFICANCE_MASK, }; #[expect( @@ -15,22 +21,25 @@ use super::state::{ reason = "Tier-1 coefficient helpers are measured inner-loop hot paths" )] #[inline(always)] -fn decode_sign_bit_arithmetic_with_neighbors( +fn decode_sign_bit_arithmetic( idx: usize, y: usize, ctx: &mut BitPlaneDecodeContext, decoder: &mut ArithmeticDecoder<'_>, ) { - let (ctx_label, xor_bit) = - context_label_sign_coding_index_with_neighbors::(idx, y, ctx); + let (ctx_label, xor_bit) = context_label_sign_coding_index(idx, y, ctx); let sign_bit = decoder.read_bit(ctx.arithmetic_decoder_context(ctx_label)) ^ u32::from(xor_bit); ctx.set_sign_index(idx, u8::from(sign_bit != 0)); } -pub(super) fn cleanup_pass_arithmetic_with_neighbors( +pub(super) fn cleanup_pass_arithmetic( ctx: &mut BitPlaneDecodeContext, - decoder: &mut ArithmeticDecoder<'_>, + shared_decoder: &mut ArithmeticDecoder<'_>, ) { + // Decode against a local copy: behind the `&mut` the MQ registers were + // stored and reloaded on every symbol. Written back after the pass. + let mut local_decoder = shared_decoder.clone(); + let decoder = &mut local_decoder; let width = ctx.width as usize; let height = ctx.height as usize; let padded_width = ctx.padded_width as usize; @@ -54,12 +63,7 @@ pub(super) fn cleanup_pass_arithmetic_with_neighbors( - ctx, - top_idx, - padded_width, - base_y, - ) + && cleanup_run_length_candidate(ctx, top_idx, padded_width, base_y) { // The four contiguous samples are all cleanup candidates // with zero context, so Annex D permits the RLC context. @@ -74,26 +78,12 @@ pub(super) fn cleanup_pass_arithmetic_with_neighbors( - significant_idx, - significant_y, - ctx, - decoder, - ); - ctx.set_significant_index_for_path::( - significant_idx, - padded_width, - ); + decode_sign_bit_arithmetic(significant_idx, significant_y, ctx, decoder); + ctx.set_significant_index(significant_idx, significant_y, padded_width); let mut idx = significant_idx + padded_width; for y in significant_y + 1..y_end { - cleanup_coefficient_arithmetic_with_neighbors::( - ctx, - decoder, - idx, - y, - padded_width, - ); + cleanup_coefficient_arithmetic(ctx, decoder, idx, y, padded_width); idx += padded_width; } continue; @@ -105,47 +95,20 @@ pub(super) fn cleanup_pass_arithmetic_with_neighbors( - ctx, - decoder, - idx, - y, - padded_width, - ); + cleanup_coefficient_arithmetic(ctx, decoder, idx, y, padded_width); } } } + *shared_decoder = local_decoder; } -#[expect( - clippy::inline_always, - reason = "Tier-1 coefficient helpers are measured inner-loop hot paths" -)] -#[inline(always)] -pub(super) fn cleanup_candidate_scan_mask( - ctx: &BitPlaneDecodeContext, - scan_unit: usize, - stripe_height: usize, -) -> u8 { - scan_unit_valid_mask(stripe_height) - & !(ctx.significant_scan_masks[scan_unit] | ctx.zero_coding_scan_masks[scan_unit]) -} - -#[expect( - clippy::inline_always, - reason = "Tier-1 coefficient helpers are measured inner-loop hot paths" -)] -#[inline(always)] -fn scan_unit_valid_mask(stripe_height: usize) -> u8 { - (1u8 << stripe_height) - 1 -} - -pub(super) fn significance_propagation_pass_arithmetic_with_neighbors< - const NORMAL_NEIGHBORS: bool, ->( +pub(super) fn significance_propagation_pass_arithmetic( ctx: &mut BitPlaneDecodeContext, - decoder: &mut ArithmeticDecoder<'_>, + shared_decoder: &mut ArithmeticDecoder<'_>, ) { + // Local MQ copy, as in the cleanup pass. + let mut local_decoder = shared_decoder.clone(); + let decoder = &mut local_decoder; let width = ctx.width as usize; let height = ctx.height as usize; let padded_width = ctx.padded_width as usize; @@ -159,8 +122,7 @@ pub(super) fn significance_propagation_pass_arithmetic_with_neighbors< for y in base_y..y_end { let state = ctx.coefficient_states[idx].0; - let neighbors = - neighborhood_significance_states_for_path::(ctx, idx, y); + let neighbors = ctx.neighborhood_significance_states_index(idx, y); // "The significance propagation pass only includes bits of coefficients // that were insignificant (the significance state has yet to be set) @@ -170,17 +132,15 @@ pub(super) fn significance_propagation_pass_arithmetic_with_neighbors< context_label_zero_coding_from_neighbors(neighbors, ctx.sub_band_type); let bit = decoder.read_bit(ctx.arithmetic_decoder_context(ctx_label)); ctx.push_magnitude_bit_index(idx, bit); - ctx.set_zero_coding_index(idx, padded_width); + ctx.set_zero_coding_index(idx, y, padded_width); // "If the value of this bit is 1 then the significance // state is set to 1 and the immediate next bit to be decoded is // the sign bit for the coefficient. Otherwise, the significance // state remains 0." if bit == 1 { - decode_sign_bit_arithmetic_with_neighbors::( - idx, y, ctx, decoder, - ); - ctx.set_significant_index_for_path::(idx, padded_width); + decode_sign_bit_arithmetic(idx, y, ctx, decoder); + ctx.set_significant_index(idx, y, padded_width); } } @@ -188,12 +148,16 @@ pub(super) fn significance_propagation_pass_arithmetic_with_neighbors< } } } + *shared_decoder = local_decoder; } -pub(super) fn magnitude_refinement_pass_arithmetic_with_neighbors( +pub(super) fn magnitude_refinement_pass_arithmetic( ctx: &mut BitPlaneDecodeContext, - decoder: &mut ArithmeticDecoder<'_>, + shared_decoder: &mut ArithmeticDecoder<'_>, ) { + // Local MQ copy, as in the cleanup pass. + let mut local_decoder = shared_decoder.clone(); + let decoder = &mut local_decoder; let width = ctx.width as usize; let height = ctx.height as usize; let padded_width = ctx.padded_width as usize; @@ -228,7 +192,7 @@ pub(super) fn magnitude_refinement_pass_arithmetic_with_neighbors(ctx, idx, y) + ctx.neighborhood_significance_states_index(idx, y) }); let bit = decoder.read_bit(ctx.arithmetic_decoder_context(ctx_label)); ctx.push_magnitude_bit_index(idx, bit); @@ -236,6 +200,7 @@ pub(super) fn magnitude_refinement_pass_arithmetic_with_neighbors( - ctx: &BitPlaneDecodeContext, - idx: usize, - y: usize, -) -> u8 { - if NORMAL_NEIGHBORS { - ctx.normal_neighborhood_significance_states_index(idx) - } else { - ctx.neighborhood_significance_states_index(idx, y) - } -} - -#[expect( - clippy::inline_always, - reason = "Tier-1 coefficient helpers are measured inner-loop hot paths" -)] -#[inline(always)] -fn cleanup_run_length_candidate_with_neighbors( - ctx: &BitPlaneDecodeContext, - top_idx: usize, - padded_width: usize, - base_y: usize, -) -> bool { - let mut idx = top_idx; - for y in base_y..base_y + 4 { - if ctx.coefficient_states[idx].0 & (SIGNIFICANCE_MASK | HAS_ZERO_CODING_MASK) != 0 - || neighborhood_significance_states_for_path::(ctx, idx, y) != 0 - { - return false; - } - idx += padded_width; - } - true -} -#[expect( - clippy::inline_always, - reason = "Tier-1 coefficient helpers are measured inner-loop hot paths" -)] -#[inline(always)] -fn cleanup_coefficient_arithmetic_with_neighbors( +fn cleanup_coefficient_arithmetic( ctx: &mut BitPlaneDecodeContext, decoder: &mut ArithmeticDecoder<'_>, idx: usize, y: usize, padded_width: usize, ) { - let neighbors = neighborhood_significance_states_for_path::(ctx, idx, y); + let neighbors = ctx.neighborhood_significance_states_index(idx, y); let ctx_label = context_label_zero_coding_from_neighbors(neighbors, ctx.sub_band_type); let bit = decoder.read_bit(ctx.arithmetic_decoder_context(ctx_label)); ctx.push_magnitude_bit_index(idx, bit); if bit == 1 { - decode_sign_bit_arithmetic_with_neighbors::(idx, y, ctx, decoder); - ctx.set_significant_index_for_path::(idx, padded_width); + decode_sign_bit_arithmetic(idx, y, ctx, decoder); + ctx.set_significant_index(idx, y, padded_width); } } - -#[expect( - clippy::inline_always, - reason = "Tier-1 coefficient helpers are measured inner-loop hot paths" -)] -#[inline(always)] -pub(super) fn cleanup_run_length_candidate( - ctx: &BitPlaneDecodeContext, - top_idx: usize, - padded_width: usize, - base_y: usize, -) -> bool { - cleanup_run_length_candidate_with_neighbors::(ctx, top_idx, padded_width, base_y) -} diff --git a/crates/j2k-native/src/j2c/bitplane/bypass.rs b/crates/j2k-native/src/j2c/bitplane/bypass.rs index a6527a278..e39a18fa0 100644 --- a/crates/j2k-native/src/j2c/bitplane/bypass.rs +++ b/crates/j2k-native/src/j2c/bitplane/bypass.rs @@ -1,11 +1,11 @@ // SPDX-License-Identifier: MIT OR Apache-2.0 use super::super::arithmetic_decoder::ArithmeticDecoderContext; -use super::arithmetic::cleanup_run_length_candidate; use super::context::{ context_label_magnitude_refinement_coding_from_state, context_label_sign_coding_index, context_label_zero_coding_from_neighbors, }; +use super::scan::cleanup_run_length_candidate; use super::state::{ BitPlaneDecodeContext, COEFFICIENTS_PADDING, HAS_MAGNITUDE_REFINEMENT_MASK, HAS_ZERO_CODING_MASK, SIGNIFICANCE_MASK, @@ -82,7 +82,7 @@ impl SafeScalarTier1 { let significant_idx = top_idx + first_significant * padded_width; ctx.push_magnitude_bit_index(significant_idx, 1); decode_sign_bit_bypass(significant_idx, significant_y, ctx, decoder)?; - ctx.set_significant_index(significant_idx, padded_width); + ctx.set_significant_index(significant_idx, significant_y, padded_width); let mut idx = significant_idx + padded_width; for y in significant_y + 1..y_end { @@ -127,11 +127,11 @@ impl SafeScalarTier1 { context_label_zero_coding_from_neighbors(neighbors, ctx.sub_band_type); let bit = decoder.read_bit(ctx.arithmetic_decoder_context(ctx_label))?; ctx.push_magnitude_bit_index(idx, bit); - ctx.set_zero_coding_index(idx, padded_width); + ctx.set_zero_coding_index(idx, y, padded_width); if bit == 1 { decode_sign_bit_bypass(idx, y, ctx, decoder)?; - ctx.set_significant_index(idx, padded_width); + ctx.set_significant_index(idx, y, padded_width); } } @@ -214,7 +214,7 @@ fn cleanup_coefficient_bypass( if bit == 1 { decode_sign_bit_bypass(idx, y, ctx, decoder)?; - ctx.set_significant_index(idx, padded_width); + ctx.set_significant_index(idx, y, padded_width); } } diff --git a/crates/j2k-native/src/j2c/bitplane/context.rs b/crates/j2k-native/src/j2c/bitplane/context.rs index 6c970e338..ddf77a6db 100644 --- a/crates/j2k-native/src/j2c/bitplane/context.rs +++ b/crates/j2k-native/src/j2c/bitplane/context.rs @@ -45,48 +45,6 @@ pub(super) fn context_label_sign_coding_index( SIGN_CONTEXT_LOOKUP[merged_significances as usize] } -#[expect( - clippy::inline_always, - reason = "Tier-1 context lookup is a measured coefficient-loop hot path" -)] -#[inline(always)] -pub(super) fn context_label_sign_coding_index_with_neighbors( - idx: usize, - y: usize, - ctx: &BitPlaneDecodeContext, -) -> (u8, u8) { - if NORMAL_NEIGHBORS { - context_label_sign_coding_index_normal(idx, ctx) - } else { - context_label_sign_coding_index(idx, y, ctx) - } -} - -#[expect( - clippy::inline_always, - reason = "Tier-1 context lookup is a measured coefficient-loop hot path" -)] -#[inline(always)] -pub(super) fn context_label_sign_coding_index_normal( - idx: usize, - ctx: &BitPlaneDecodeContext, -) -> (u8, u8) { - let significances = ctx.normal_neighborhood_significance_states_index(idx) & 0b0101_0101; - let padded_width = ctx.padded_width as usize; - - let top_sign = ctx.sign_index(idx - padded_width); - let left_sign = ctx.sign_index(idx - 1); - let right_sign = ctx.sign_index(idx + 1); - let bottom_sign = ctx.sign_index(idx + padded_width); - - let signs = (top_sign << 6) | (left_sign << 4) | (right_sign << 2) | bottom_sign; - let negative_significances = significances & signs; - let positive_significances = significances & !signs; - let merged_significances = (negative_significances << 1) | positive_significances; - - SIGN_CONTEXT_LOOKUP[merged_significances as usize] -} - /// Return the context label for zero coding (Section D.3.1). #[expect( clippy::inline_always, diff --git a/crates/j2k-native/src/j2c/bitplane/facade.rs b/crates/j2k-native/src/j2c/bitplane/facade.rs index 16595b850..63200b20c 100644 --- a/crates/j2k-native/src/j2c/bitplane/facade.rs +++ b/crates/j2k-native/src/j2c/bitplane/facade.rs @@ -149,6 +149,7 @@ pub(super) fn decode_code_block_segments_validated_with_observer { + decoder: ArithmeticDecoder<'d>, + contexts: &'c mut [ArithmeticDecoderContext; 19], + flags: &'c mut [u32], + coefficients: &'c mut [Coefficient], + zero_coding: &'static [u8; 512], + width: usize, + height: usize, + /// Row stride of both `flags` and `coefficients` (width plus padding). + stride: usize, + bit_position: u8, +} + +impl<'c, 'd> Scan<'c, 'd> { + fn new(ctx: &'c mut BitPlaneDecodeContext, decoder: &ArithmeticDecoder<'d>) -> Self { + Self { + decoder: decoder.clone(), + zero_coding: zero_coding_table(ctx.sub_band_type), + width: ctx.width as usize, + height: ctx.height as usize, + stride: ctx.padded_width as usize, + bit_position: ctx.current_bit_position, + contexts: &mut ctx.contexts, + flags: &mut ctx.flags, + coefficients: &mut ctx.coefficients, + } + } + + fn finish(self, decoder: &mut ArithmeticDecoder<'d>) { + *decoder = self.decoder; + } + + /// Flag and top-row coefficient indices of column 0 in `stripe`, and the + /// number of rows the stripe holds. + #[expect( + clippy::cast_possible_truncation, + reason = "a stripe holds at most four rows" + )] + fn stripe_start(&self, stripe: usize) -> (usize, usize, u32) { + ( + (stripe + 1) * self.stride + 1, + (4 * stripe + 1) * self.stride + 1, + (self.height - 4 * stripe).min(4) as u32, + ) + } + + #[expect(clippy::inline_always, reason = "per-symbol MQ decode")] + #[inline(always)] + fn read(&mut self, label: u8) -> u32 { + self.decoder.read_bit(&mut self.contexts[label as usize]) + } + + /// The coefficient at `row` just decoded a one while insignificant: record + /// the bit, decode its sign, and mark it significant. + #[expect(clippy::inline_always, reason = "Tier-1 coefficient-loop hot path")] + #[inline(always)] + fn become_significant(&mut self, own: &mut u32, flag_index: usize, top: usize, row: u32) { + let coefficient = &mut self.coefficients[top + row as usize * self.stride]; + coefficient.push_bit_at(1, self.bit_position); + let (label, xor) = sign_coding_label( + *own, + self.flags[flag_index - 1], + self.flags[flag_index + 1], + row, + ); + let sign = self.read(label) ^ u32::from(xor); + self.coefficients[top + row as usize * self.stride].set_sign(u8::from(sign != 0)); + set_significant(self.flags, own, flag_index, self.stride, row, sign); + } + + /// Zero-code the insignificant coefficient at `row`. + #[expect(clippy::inline_always, reason = "Tier-1 coefficient-loop hot path")] + #[inline(always)] + fn zero_code(&mut self, own: &mut u32, flag_index: usize, top: usize, row: u32) { + if self.read(zero_coding_label(self.zero_coding, *own, row)) != 0 { + self.become_significant(own, flag_index, top, row); + } + } +} + +pub(super) fn significance_propagation_pass_flags( + ctx: &mut BitPlaneDecodeContext, + decoder: &mut ArithmeticDecoder<'_>, +) { + let mut scan = Scan::new(ctx, decoder); + for stripe in 0..scan.height.div_ceil(4) { + let (flag_row, top_row, rows) = scan.stripe_start(stripe); + for x in 0..scan.width { + let flag_index = flag_row + x; + let mut own = scan.flags[flag_index]; + // Nothing significant in or around the column: no candidates. + if own == 0 { + continue; + } + for row in 0..rows { + let window = own >> (3 * row); + if window & (SIGMA_THIS | PI_THIS) == 0 && window & SIGMA_NEIGHBOURS != 0 { + own |= PI_THIS << (3 * row); + scan.zero_code(&mut own, flag_index, top_row + x, row); + } + } + scan.flags[flag_index] = own; + } + } + scan.finish(decoder); +} + +pub(super) fn magnitude_refinement_pass_flags( + ctx: &mut BitPlaneDecodeContext, + decoder: &mut ArithmeticDecoder<'_>, +) { + let mut scan = Scan::new(ctx, decoder); + for stripe in 0..scan.height.div_ceil(4) { + let (flag_row, top_row, rows) = scan.stripe_start(stripe); + for x in 0..scan.width { + let flag_index = flag_row + x; + let mut own = scan.flags[flag_index]; + if own & SIGMA_ROWS == 0 { + continue; + } + for row in 0..rows { + let window = own >> (3 * row); + // Significant before this bitplane's significance pass. + if window & (SIGMA_THIS | PI_THIS) != SIGMA_THIS { + continue; + } + let label = if window & MU_THIS != 0 { + 16 + } else if window & SIGMA_NEIGHBOURS != 0 { + 15 + } else { + 14 + }; + let bit = scan.read(label); + scan.coefficients[top_row + x + row as usize * scan.stride] + .push_bit_at(bit, scan.bit_position); + own |= MU_THIS << (3 * row); + } + scan.flags[flag_index] = own; + } + } + scan.finish(decoder); +} + +pub(super) fn cleanup_pass_flags( + ctx: &mut BitPlaneDecodeContext, + decoder: &mut ArithmeticDecoder<'_>, +) { + let mut scan = Scan::new(ctx, decoder); + // The run-length and uniform contexts are only used here; holding them in + // locals keeps consecutive run-length decodes off a store-to-load chain + // through the context array. + let mut run_length = scan.contexts[17]; + let mut uniform = scan.contexts[18]; + for stripe in 0..scan.height.div_ceil(4) { + let (flag_row, top_row, rows) = scan.stripe_start(stripe); + for x in 0..scan.width { + let flag_index = flag_row + x; + let top = top_row + x; + let mut own = scan.flags[flag_index]; + // A zero word is a full column of insignificant, unvisited + // coefficients with zero contexts: Annex D's run-length case. + if rows == 4 && own == 0 { + if scan.decoder.read_bit(&mut run_length) == 0 { + continue; + } + let first = (scan.decoder.read_bit(&mut uniform) << 1) + | scan.decoder.read_bit(&mut uniform); + scan.become_significant(&mut own, flag_index, top, first); + for row in first + 1..4 { + scan.zero_code(&mut own, flag_index, top, row); + } + } else { + for row in 0..rows { + if (own >> (3 * row)) & (SIGMA_THIS | PI_THIS) == 0 { + scan.zero_code(&mut own, flag_index, top, row); + } + } + } + // Visit marks last one bitplane. + scan.flags[flag_index] = own & !PI_ALL; + } + } + scan.contexts[17] = run_length; + scan.contexts[18] = uniform; + scan.finish(decoder); +} diff --git a/crates/j2k-native/src/j2c/bitplane/flags.rs b/crates/j2k-native/src/j2c/bitplane/flags.rs new file mode 100644 index 000000000..56c34b0fa --- /dev/null +++ b/crates/j2k-native/src/j2c/bitplane/flags.rs @@ -0,0 +1,139 @@ +// SPDX-License-Identifier: MIT OR Apache-2.0 + +//! Packed per-stripe-column state for the normal-neighbor arithmetic path. +//! +//! One `u32` per stripe column (four rows) carries everything the three +//! passes consult, so a coefficient's zero-coding context is a shift and a +//! table load, and a newly significant coefficient updates three words (six +//! at a stripe edge) instead of eight neighbor bytes plus scan masks. +//! +//! Bit layout, for window rows `-1..=4` (the stripe's four rows plus the row +//! above and below) and columns `-1..=1`: +//! - `3 * (row + 1) + (col + 1)`: significance of that neighbor (bits 0-17); +//! - 18: sign of row -1; `19 + 3 * row` for rows 0-4: sign of the column's +//! own coefficient in that row (bits 19, 22, 25, 28, 31); +//! - `20 + 3 * row`, rows 0-3: magnitude refinement has run (bits 20-29); +//! - `21 + 3 * row`, rows 0-3: coded in this bitplane's significance +//! propagation pass (bits 21-30). +//! +//! Shifting a word right by `3 * row` aligns row `row`'s 3x3 neighborhood +//! with bits 0-8 and its own refinement and visit bits with bits 20 and 21. +//! Sign bits are only ever set together with the matching significance bit. + +use super::super::build::SubBandType; +use j2k_codec_math::classic::{ + SIGN_CONTEXT_LOOKUP, ZERO_CTX_HH_LOOKUP, ZERO_CTX_HL_LOOKUP, ZERO_CTX_LL_LH_LOOKUP, +}; + +/// The coefficient's own significance within a row-aligned window. +pub(super) const SIGMA_THIS: u32 = 1 << 4; +/// The eight neighbors within a row-aligned window. +pub(super) const SIGMA_NEIGHBOURS: u32 = 0x1EF; +/// Magnitude refinement has run, row-aligned. +pub(super) const MU_THIS: u32 = 1 << 20; +/// Coded in this bitplane's significance propagation pass, row-aligned. +pub(super) const PI_THIS: u32 = 1 << 21; +/// All four rows' `PI_THIS` bits. +pub(super) const PI_ALL: u32 = PI_THIS | PI_THIS << 3 | PI_THIS << 6 | PI_THIS << 9; + +/// Zero-coding context label for a row-aligned 3x3 window (bits 0-8, the +/// centre ignored), derived from the byte-indexed Table D.1 lookups. +const fn zero_coding_window_table(byte_table: &[u8; 256]) -> [u8; 512] { + const fn bit(window: usize, index: usize) -> usize { + (window >> index) & 1 + } + let mut table = [0_u8; 512]; + let mut window = 0; + while window < 512 { + // Neighbor byte, MSB to LSB: TL, T, TR, L, BL, R, BR, B. + let byte = (bit(window, 0) << 7) + | (bit(window, 1) << 6) + | (bit(window, 2) << 5) + | (bit(window, 3) << 4) + | (bit(window, 6) << 3) + | (bit(window, 5) << 2) + | (bit(window, 8) << 1) + | bit(window, 7); + table[window] = byte_table[byte]; + window += 1; + } + table +} + +static ZERO_CODING_LL_LH: [u8; 512] = zero_coding_window_table(&ZERO_CTX_LL_LH_LOOKUP); +static ZERO_CODING_HL: [u8; 512] = zero_coding_window_table(&ZERO_CTX_HL_LOOKUP); +static ZERO_CODING_HH: [u8; 512] = zero_coding_window_table(&ZERO_CTX_HH_LOOKUP); + +pub(super) fn zero_coding_table(sub_band_type: SubBandType) -> &'static [u8; 512] { + match sub_band_type { + SubBandType::LowLow | SubBandType::LowHigh => &ZERO_CODING_LL_LH, + SubBandType::HighLow => &ZERO_CODING_HL, + SubBandType::HighHigh => &ZERO_CODING_HH, + } +} + +/// Zero-coding context label for row `row` of the column word `flags`. +#[expect( + clippy::inline_always, + reason = "Tier-1 context lookup is a measured coefficient-loop hot path" +)] +#[inline(always)] +pub(super) fn zero_coding_label(table: &[u8; 512], flags: u32, row: u32) -> u8 { + table[((flags >> (3 * row)) & SIGMA_NEIGHBOURS) as usize] +} + +/// Sign-coding context label and XOR bit (Table D.3) for row `row`, from the +/// column word and its left and right neighbors. +#[expect( + clippy::inline_always, + reason = "Tier-1 context lookup is a measured coefficient-loop hot path" +)] +#[inline(always)] +pub(super) fn sign_coding_label(flags: u32, left: u32, right: u32, row: u32) -> (u8, u8) { + let window = flags >> (3 * row); + let significances = (((window >> 1) & 1) << 6) + | (((window >> 3) & 1) << 4) + | (((window >> 5) & 1) << 2) + | ((window >> 7) & 1); + // The sign of row - 1 is bit 18 for row 0 and 19 + 3 * (row - 1) after. + let north_sign_bit = if row == 0 { 18 } else { 16 + 3 * row }; + let signs = (((flags >> north_sign_bit) & 1) << 6) + | (((left >> (19 + 3 * row)) & 1) << 4) + | (((right >> (19 + 3 * row)) & 1) << 2) + | ((flags >> (22 + 3 * row)) & 1); + let merged = ((significances & signs) << 1) | (significances & !signs); + SIGN_CONTEXT_LOOKUP[merged as usize] +} + +/// Record that row `row` of the column word at `index` became significant +/// with sign `sign`. The column's own word is the caller's register copy +/// `own`; neighbor words are updated in `flags`, whose row stride is +/// `stride` (one padding column and one padding stripe on every side). +#[expect(clippy::inline_always, reason = "Tier-1 coefficient-loop hot path")] +#[inline(always)] +pub(super) fn set_significant( + flags: &mut [u32], + own: &mut u32, + index: usize, + stride: usize, + row: u32, + sign: u32, +) { + let shift = 3 * row; + *own |= (SIGMA_THIS | sign << 19) << shift; + flags[index - 1] |= (1 << 5) << shift; + flags[index + 1] |= (1 << 3) << shift; + if row == 0 { + // Row 4 of the stripe above: significance bits 15-17, sign bit 31. + let north = index - stride; + flags[north - 1] |= 1 << 17; + flags[north] |= 1 << 16 | sign << 31; + flags[north + 1] |= 1 << 15; + } else if row == 3 { + // Row -1 of the stripe below: significance bits 0-2, sign bit 18. + let south = index + stride; + flags[south - 1] |= 1 << 2; + flags[south] |= 1 << 1 | sign << 18; + flags[south + 1] |= 1; + } +} diff --git a/crates/j2k-native/src/j2c/bitplane/golden_tests.rs b/crates/j2k-native/src/j2c/bitplane/golden_tests.rs index a79337256..9d253d137 100644 --- a/crates/j2k-native/src/j2c/bitplane/golden_tests.rs +++ b/crates/j2k-native/src/j2c/bitplane/golden_tests.rs @@ -427,12 +427,15 @@ fn strict_metadata_and_segment_failures_keep_exact_error_classes() { #[test] fn bitplane_modules_remain_focused_without_suppression_shortcuts() { const ROOT: &str = include_str!("../bitplane.rs"); - const MODULES: [(&str, &str, usize); 8] = [ + const MODULES: [(&str, &str, usize); 11] = [ ("arithmetic", include_str!("arithmetic.rs"), 320), ("bypass", include_str!("bypass.rs"), 280), ("context", include_str!("context.rs"), 210), ("facade", include_str!("facade.rs"), 170), + ("flag_passes", include_str!("flag_passes.rs"), 250), + ("flags", include_str!("flags.rs"), 170), ("observer", include_str!("observer.rs"), 100), + ("scan", include_str!("scan.rs"), 80), ("schedule", include_str!("schedule.rs"), 350), ("state", include_str!("state.rs"), 500), ("state/model", include_str!("state/model.rs"), 150), diff --git a/crates/j2k-native/src/j2c/bitplane/reconstruction.rs b/crates/j2k-native/src/j2c/bitplane/reconstruction.rs index 9d9aa111d..ab1391388 100644 --- a/crates/j2k-native/src/j2c/bitplane/reconstruction.rs +++ b/crates/j2k-native/src/j2c/bitplane/reconstruction.rs @@ -47,3 +47,111 @@ pub(super) fn reconstruct_irreversible_midpoint( reconstructed } } + +/// Per-code-block form of [`reconstruct_irreversible_midpoint`]. The pass +/// arithmetic is resolved once per block, and the common case (no ROI shift) +/// needs only a `u64` shift-or and a hardware integer-to-float conversion per +/// coefficient; `u128` has no hardware conversion. Every other case defers to +/// the reference function, so results are bit-identical. +#[derive(Clone, Copy)] +pub(crate) struct MidpointReconstructor { + /// Lowest decoded bit and whether the final pass was significance + /// propagation, when the fast path applies. + fast: Option<(u32, bool)>, + decoded_bitplanes: u8, + number_of_coding_passes: u8, + roi_shift: u8, +} + +impl MidpointReconstructor { + pub(crate) fn new(decoded_bitplanes: u8, number_of_coding_passes: u8, roi_shift: u8) -> Self { + let fast = if roi_shift != 0 || decoded_bitplanes == 0 || number_of_coding_passes == 0 { + None + } else { + let final_pass = u32::from(number_of_coding_passes) - 1; + u32::from(decoded_bitplanes) + .checked_sub(final_pass.div_ceil(3) + 1) + .filter(|&lowest| lowest < u64::BITS) + .map(|lowest| (lowest, final_pass % 3 == 1)) + }; + Self { + fast, + decoded_bitplanes, + number_of_coding_passes, + roi_shift, + } + } + + /// Reconstruct and dequantize a row: `output[i] = reconstruct(c[i]) * step` + /// for as many elements as both slices hold. + /// + /// With at most 31 decoded bitplanes, magnitudes and their doubled + /// midpoints fit in `u32`, and the per-coefficient branches become selects + /// that vectorize. Halving is exact, so `fixed * (step / 2)` rounds like + /// `(fixed / 2) * step`, and negation is a sign-bit flip. + #[expect( + clippy::cast_possible_truncation, + clippy::cast_precision_loss, + reason = "magnitudes are below 2^31 on the u32 path, which enters the f32 domain" + )] + pub(crate) fn dequantize_row( + self, + coefficients: &[Coefficient], + output: &mut [f32], + step: f32, + ) { + let Some((lowest, final_sigprop)) = self.fast.filter(|_| self.decoded_bitplanes <= 31) + else { + for (sample, &coefficient) in output.iter_mut().zip(coefficients) { + *sample = self.reconstruct(coefficient) * step; + } + return; + }; + let half_step = step * 0.5; + let refine = u32::from(final_sigprop); + for (sample, &coefficient) in output.iter_mut().zip(coefficients) { + let bits = coefficient.to_bits(); + let magnitude = bits as u32; + let negative = (bits >> 63) as u32; + let bit = lowest + (refine & !(magnitude >> lowest) & 1); + let value = ((magnitude << 1) | (1 << bit)) as f32 * half_step; + let value = f32::from_bits(value.to_bits() ^ (negative << 31)); + *sample = if magnitude == 0 { 0.0 } else { value }; + } + } + + #[expect( + clippy::cast_precision_loss, + clippy::inline_always, + reason = "irreversible coefficients enter the f32 domain here, once per coefficient" + )] + #[inline(always)] + pub(crate) fn reconstruct(self, coefficient: Coefficient) -> f32 { + let Some((lowest, final_sigprop)) = self.fast else { + return reconstruct_irreversible_midpoint( + coefficient, + self.decoded_bitplanes, + self.number_of_coding_passes, + self.roi_shift, + ); + }; + let signed = coefficient.get_i64(); + let magnitude = signed.unsigned_abs(); + if magnitude == 0 { + return 0.0; + } + // A final significance pass leaves already-significant coefficients + // unrefined; see `irreversible_midpoint_bit`. + let bit = lowest + u32::from(final_sigprop && magnitude & (1_u64 << lowest) == 0); + if bit >= u64::BITS { + return signed as f32; + } + // Magnitudes are below 2^63, so the doubled value fits in a u64. + let reconstructed = ((magnitude << 1) | (1_u64 << bit)) as f32 * 0.5; + if signed < 0 { + -reconstructed + } else { + reconstructed + } + } +} diff --git a/crates/j2k-native/src/j2c/bitplane/scan.rs b/crates/j2k-native/src/j2c/bitplane/scan.rs new file mode 100644 index 000000000..05940f6f1 --- /dev/null +++ b/crates/j2k-native/src/j2c/bitplane/scan.rs @@ -0,0 +1,53 @@ +// SPDX-License-Identifier: MIT OR Apache-2.0 + +//! Stripe-column scan helpers shared by the arithmetic and bypass passes. + +use super::state::{BitPlaneDecodeContext, HAS_ZERO_CODING_MASK, SIGNIFICANCE_MASK}; + +#[expect( + clippy::inline_always, + reason = "Tier-1 coefficient helpers are measured inner-loop hot paths" +)] +#[inline(always)] +pub(super) fn cleanup_candidate_scan_mask( + ctx: &BitPlaneDecodeContext, + scan_unit: usize, + stripe_height: usize, +) -> u8 { + scan_unit_valid_mask(stripe_height) + & !(ctx.significant_scan_masks[scan_unit] | ctx.zero_coding_scan_masks[scan_unit]) +} + +#[expect( + clippy::inline_always, + reason = "Tier-1 coefficient helpers are measured inner-loop hot paths" +)] +#[inline(always)] +pub(super) fn scan_unit_valid_mask(stripe_height: usize) -> u8 { + (1u8 << stripe_height) - 1 +} + +/// Whether a full stripe column qualifies for run-length cleanup coding: +/// no coefficient is significant, zero-coded, or has a significant neighbor. +#[expect( + clippy::inline_always, + reason = "Tier-1 coefficient helpers are measured inner-loop hot paths" +)] +#[inline(always)] +pub(super) fn cleanup_run_length_candidate( + ctx: &BitPlaneDecodeContext, + top_idx: usize, + padded_width: usize, + base_y: usize, +) -> bool { + let mut idx = top_idx; + for y in base_y..base_y + 4 { + if ctx.coefficient_states[idx].0 & (SIGNIFICANCE_MASK | HAS_ZERO_CODING_MASK) != 0 + || ctx.neighborhood_significance_states_index(idx, y) != 0 + { + return false; + } + idx += padded_width; + } + true +} diff --git a/crates/j2k-native/src/j2c/bitplane/schedule.rs b/crates/j2k-native/src/j2c/bitplane/schedule.rs index e8f2294ee..1e3fdb393 100644 --- a/crates/j2k-native/src/j2c/bitplane/schedule.rs +++ b/crates/j2k-native/src/j2c/bitplane/schedule.rs @@ -4,10 +4,13 @@ use super::super::arithmetic_decoder::ArithmeticDecoder; use super::super::build::CodeBlock; use super::super::decode::DecompositionStorage; use super::arithmetic::{ - cleanup_pass_arithmetic_with_neighbors, magnitude_refinement_pass_arithmetic_with_neighbors, - significance_propagation_pass_arithmetic_with_neighbors, + cleanup_pass_arithmetic, magnitude_refinement_pass_arithmetic, + significance_propagation_pass_arithmetic, }; use super::bypass::{BitDecoder, BypassDecoder, SafeScalarTier1}; +use super::flag_passes::{ + cleanup_pass_flags, magnitude_refinement_pass_flags, significance_propagation_pass_flags, +}; use super::observer::{J2kDecodeObserver, NoJ2kDecodeStats}; use super::state::{ extend_preallocated, push_preallocated, BitPlaneDecodeBuffers, BitPlaneDecodeContext, @@ -81,8 +84,8 @@ fn decode_inner_with_observer( let end = code_block .number_of_coding_passes .min(ctx.max_coding_passes); - if ctx.uses_normal_arithmetic_neighbor_path() { - handle_normal_arithmetic_coding_passes(0, end, ctx, &mut decoder, observer)?; + if ctx.uses_packed_columns() { + handle_packed_arithmetic_coding_passes(0, end, ctx, &mut decoder, observer)?; } else { handle_arithmetic_coding_passes(0, end, ctx, &mut decoder, observer)?; } @@ -136,27 +139,29 @@ fn decode_inner_with_observer( Some(()) } -fn handle_arithmetic_coding_passes( +pub(super) fn handle_arithmetic_coding_passes( start: u8, end: u8, ctx: &mut BitPlaneDecodeContext, decoder: &mut ArithmeticDecoder<'_>, observer: &mut impl J2kDecodeObserver, ) -> Option<()> { - handle_arithmetic_coding_passes_with_neighbors::(start, end, ctx, decoder, observer) + run_arithmetic_coding_passes::(start, end, ctx, decoder, observer) } -fn handle_normal_arithmetic_coding_passes( +pub(super) fn handle_packed_arithmetic_coding_passes( start: u8, end: u8, ctx: &mut BitPlaneDecodeContext, decoder: &mut ArithmeticDecoder<'_>, observer: &mut impl J2kDecodeObserver, ) -> Option<()> { - handle_arithmetic_coding_passes_with_neighbors::(start, end, ctx, decoder, observer) + run_arithmetic_coding_passes::(start, end, ctx, decoder, observer) } -fn handle_arithmetic_coding_passes_with_neighbors( +/// Arithmetic passes over the packed column words (`PACKED_COLUMNS`) or the +/// per-coefficient arrays; decisions and contexts are identical. +fn run_arithmetic_coding_passes( start: u8, end: u8, ctx: &mut BitPlaneDecodeContext, @@ -172,7 +177,11 @@ fn handle_arithmetic_coding_passes_with_neighbors( match coding_pass % 3 { 0 => { let phase_start = observer.phase_start(); - cleanup_pass_arithmetic_with_neighbors::(ctx, decoder); + if PACKED_COLUMNS { + cleanup_pass_flags(ctx, decoder); + } else { + cleanup_pass_arithmetic(ctx, decoder); + } if ctx.style.segmentation_symbols { let b0 = decoder.read_bit(ctx.arithmetic_decoder_context(18)); @@ -190,16 +199,20 @@ fn handle_arithmetic_coding_passes_with_neighbors( } 1 => { let phase_start = observer.phase_start(); - significance_propagation_pass_arithmetic_with_neighbors::( - ctx, decoder, - ); + if PACKED_COLUMNS { + significance_propagation_pass_flags(ctx, decoder); + } else { + significance_propagation_pass_arithmetic(ctx, decoder); + } observer.add_sigprop_us(phase_start); } 2 => { let phase_start = observer.phase_start(); - magnitude_refinement_pass_arithmetic_with_neighbors::( - ctx, decoder, - ); + if PACKED_COLUMNS { + magnitude_refinement_pass_flags(ctx, decoder); + } else { + magnitude_refinement_pass_arithmetic(ctx, decoder); + } observer.add_magref_us(phase_start); } _ => unreachable!(), @@ -286,8 +299,8 @@ pub(super) fn decode_code_block_segments_inner( if segment.use_arithmetic { let mut decoder = ArithmeticDecoder::new(segment_data); - if ctx.uses_normal_arithmetic_neighbor_path() { - handle_normal_arithmetic_coding_passes( + if ctx.uses_packed_columns() { + handle_packed_arithmetic_coding_passes( start_coding_pass, end_coding_pass, ctx, diff --git a/crates/j2k-native/src/j2c/bitplane/state.rs b/crates/j2k-native/src/j2c/bitplane/state.rs index 37879e19b..b3ae882c6 100644 --- a/crates/j2k-native/src/j2c/bitplane/state.rs +++ b/crates/j2k-native/src/j2c/bitplane/state.rs @@ -10,6 +10,7 @@ use crate::error::{bail, DecodingError, Result, ValidationError}; use crate::try_reserve_decode_elements; mod model; +mod scan_masks; mod workspace; pub(crate) use model::{Coefficient, CoefficientState, BITPLANE_BIT_SIZE}; @@ -98,6 +99,13 @@ pub(crate) struct BitPlaneDecodeContext { pub(super) significant_scan_masks: Vec, /// One 4-bit mask per scan stripe column for zero-coded coefficients in this bitplane. pub(super) zero_coding_scan_masks: Vec, + /// Packed per-stripe-column state (see `flags.rs`) used instead of the + /// per-coefficient state, neighbor, and scan-mask arrays when + /// `packed_columns` is set. One padding column and stripe on each side. + pub(super) flags: Vec, + /// Whether this block decodes through the packed-column passes: normal + /// neighbor contexts and arithmetic coding throughout. + pub(super) packed_columns: bool, /// The neighbor significances for each coefficient. pub(super) neighbor_significances: Vec, /// The magnitude and signs of each coefficient that is successively built @@ -137,6 +145,8 @@ impl BitPlaneDecodeContext { coefficient_states: Vec::new(), significant_scan_masks: Vec::new(), zero_coding_scan_masks: Vec::new(), + flags: Vec::new(), + packed_columns: false, coefficients: Vec::new(), neighbor_significances: Vec::new(), width: 0, @@ -160,8 +170,11 @@ impl BitPlaneDecodeContext { } } + /// Allocate a workspace able to decode `width` x `height` blocks through + /// either pass family. pub(crate) fn prepare(&mut self, width: u32, height: u32) -> Result<()> { - workspace::reset_decode_buffers(self, width, height).map(|_| ()) + workspace::reset_decode_buffers(self, width, height, workspace::StateLayout::Both) + .map(|_| ()) } pub(crate) fn allocated_bytes(&self) -> Result { @@ -174,6 +187,7 @@ impl BitPlaneDecodeContext { include_capacity::(&mut bytes, self.coefficient_states.capacity())?; include_capacity::(&mut bytes, self.significant_scan_masks.capacity())?; include_capacity::(&mut bytes, self.zero_coding_scan_masks.capacity())?; + include_capacity::(&mut bytes, self.flags.capacity())?; Ok(bytes) } @@ -192,8 +206,22 @@ impl BitPlaneDecodeContext { code_block_style: &CodeBlockStyle, total_bitplanes: u8, strict: bool, + all_segments_arithmetic: bool, ) -> Result<()> { - let padded_width = workspace::reset_decode_buffers(self, width, height)?; + // Packed and per-coefficient bookkeeping cannot be mixed in a block, + // so a raw (bypass) segment sends the whole block down the + // per-coefficient passes. + let packed_columns = all_segments_arithmetic + && !code_block_style.selective_arithmetic_coding_bypass + && !code_block_style.termination_on_each_pass + && !code_block_style.vertically_causal_context; + let layout = if packed_columns { + workspace::StateLayout::PackedColumns + } else { + workspace::StateLayout::PerCoefficient + }; + let padded_width = workspace::reset_decode_buffers(self, width, height, layout)?; + self.packed_columns = packed_columns; self.width = width; self.padded_width = padded_width; @@ -262,6 +290,9 @@ impl BitPlaneDecodeContext { code_block_style, total_bitplanes, strict, + // Codestream blocks take raw segments only under selective + // bypass, which already selects the per-coefficient passes. + true, ) } @@ -274,14 +305,13 @@ impl BitPlaneDecodeContext { .take(self.height as usize) } - pub(crate) fn reconstruct_irreversible_midpoint( + /// Midpoint reconstruction for this code block's decoded coefficients. + pub(crate) fn midpoint_reconstructor( &self, - coefficient: Coefficient, number_of_coding_passes: u8, roi_shift: u8, - ) -> f32 { - super::reconstruction::reconstruct_irreversible_midpoint( - coefficient, + ) -> super::reconstruction::MidpointReconstructor { + super::reconstruction::MidpointReconstructor::new( self.bitplanes, number_of_coding_passes, roi_shift, @@ -339,12 +369,12 @@ impl BitPlaneDecodeContext { #[expect(clippy::inline_always, reason = "Tier-1 coefficient-loop hot path")] #[inline(always)] - pub(super) fn set_significant_index(&mut self, idx: usize, padded_width: usize) { + pub(super) fn set_significant_index(&mut self, idx: usize, y: usize, padded_width: usize) { let is_significant = self.coefficient_states[idx].is_significant(); if !is_significant { self.coefficient_states[idx].set_significant(); - self.set_significant_scan_mask(idx, padded_width); + self.set_significant_scan_mask(idx, y, padded_width); // Update all neighbors so they know this coefficient is significant // now. @@ -359,87 +389,12 @@ impl BitPlaneDecodeContext { } } - #[expect(clippy::inline_always, reason = "Tier-1 coefficient-loop hot path")] - #[inline(always)] - pub(super) fn set_significant_index_for_path( - &mut self, - idx: usize, - padded_width: usize, - ) { - if NORMAL_NEIGHBORS { - self.set_significant_index_normal(idx, padded_width); - } else { - self.set_significant_index(idx, padded_width); - } - } - - #[expect(clippy::inline_always, reason = "Tier-1 coefficient-loop hot path")] - #[inline(always)] - pub(super) fn set_significant_index_normal(&mut self, idx: usize, padded_width: usize) { - if self.coefficient_states[idx].is_significant() { - return; - } - - self.coefficient_states[idx].set_significant(); - self.set_significant_scan_mask(idx, padded_width); - - let top_start = idx - padded_width - 1; - let top = &mut self.neighbor_significances[top_start..top_start + 3]; - top[0].set_bottom_right(); - top[1].set_bottom(); - top[2].set_bottom_left(); - - let middle_start = idx - 1; - let middle = &mut self.neighbor_significances[middle_start..middle_start + 3]; - middle[0].set_right(); - middle[2].set_left(); - - let bottom_start = idx + padded_width - 1; - let bottom = &mut self.neighbor_significances[bottom_start..bottom_start + 3]; - bottom[0].set_top_right(); - bottom[1].set_top(); - bottom[2].set_top_left(); - } - #[expect(clippy::inline_always, reason = "Tier-1 coefficient-loop hot path")] #[inline(always)] pub(super) fn push_magnitude_bit_index(&mut self, idx: usize, bit: u32) { self.coefficients[idx].push_bit_at(bit, self.current_bit_position); } - #[expect(clippy::inline_always, reason = "Tier-1 coefficient-loop hot path")] - #[inline(always)] - pub(super) fn set_zero_coding_index(&mut self, idx: usize, padded_width: usize) { - self.coefficient_states[idx].0 |= HAS_ZERO_CODING_MASK; - let (scan_unit, bit) = self.scan_unit_mask_index(idx, padded_width); - self.zero_coding_scan_masks[scan_unit] |= bit; - } - - #[expect(clippy::inline_always, reason = "Tier-1 coefficient-loop hot path")] - #[inline(always)] - pub(super) fn set_significant_scan_mask(&mut self, idx: usize, padded_width: usize) { - let (scan_unit, bit) = self.scan_unit_mask_index(idx, padded_width); - self.significant_scan_masks[scan_unit] |= bit; - } - - #[expect(clippy::inline_always, reason = "Tier-1 coefficient-loop hot path")] - #[inline(always)] - pub(super) fn scan_unit_mask_index(&self, idx: usize, padded_width: usize) -> (usize, u8) { - let row = idx / padded_width; - let col = idx - row * padded_width; - let pad = COEFFICIENTS_PADDING as usize; - debug_assert!(row >= pad); - debug_assert!(col >= pad); - - let y = row - pad; - let x = col - pad; - debug_assert!(y < self.height as usize); - debug_assert!(x < self.width as usize); - - let scan_unit = (y >> 2) * self.width as usize + x; - (scan_unit, 1u8 << (y & 3)) - } - #[expect(clippy::inline_always, reason = "Tier-1 coefficient-loop hot path")] #[inline(always)] pub(super) fn sign_index(&self, idx: usize) -> u8 { @@ -467,15 +422,7 @@ impl BitPlaneDecodeContext { #[expect(clippy::inline_always, reason = "Tier-1 coefficient-loop hot path")] #[inline(always)] - pub(super) fn normal_neighborhood_significance_states_index(&self, idx: usize) -> u8 { - self.neighbor_significances[idx].all() - } - - #[expect(clippy::inline_always, reason = "Tier-1 coefficient-loop hot path")] - #[inline(always)] - pub(super) fn uses_normal_arithmetic_neighbor_path(&self) -> bool { - !self.style.selective_arithmetic_coding_bypass - && !self.style.termination_on_each_pass - && !self.style.vertically_causal_context + pub(super) fn uses_packed_columns(&self) -> bool { + self.packed_columns } } diff --git a/crates/j2k-native/src/j2c/bitplane/state/model.rs b/crates/j2k-native/src/j2c/bitplane/state/model.rs index 6753d7c46..e5989f060 100644 --- a/crates/j2k-native/src/j2c/bitplane/state/model.rs +++ b/crates/j2k-native/src/j2c/bitplane/state/model.rs @@ -57,6 +57,11 @@ impl Coefficient { magnitude } + /// Sign in bit 63, magnitude in bits 0-62. + pub(crate) fn to_bits(self) -> u64 { + self.0 + } + pub(in crate::j2c::bitplane) fn set_sign(&mut self, sign: u8) { self.0 |= u64::from(sign) << 63; } diff --git a/crates/j2k-native/src/j2c/bitplane/state/scan_masks.rs b/crates/j2k-native/src/j2c/bitplane/state/scan_masks.rs new file mode 100644 index 000000000..3b4f368fc --- /dev/null +++ b/crates/j2k-native/src/j2c/bitplane/state/scan_masks.rs @@ -0,0 +1,54 @@ +// SPDX-License-Identifier: MIT OR Apache-2.0 + +//! Per-stripe-column bit masks that let the passes test four coefficients at +//! once instead of loading each coefficient's state. + +use super::{BitPlaneDecodeContext, COEFFICIENTS_PADDING, HAS_ZERO_CODING_MASK}; + +impl BitPlaneDecodeContext { + #[expect(clippy::inline_always, reason = "Tier-1 coefficient-loop hot path")] + #[inline(always)] + pub(in crate::j2c::bitplane) fn set_zero_coding_index( + &mut self, + idx: usize, + y: usize, + padded_width: usize, + ) { + self.coefficient_states[idx].0 |= HAS_ZERO_CODING_MASK; + let (scan_unit, bit) = self.scan_unit_mask_index(idx, y, padded_width); + self.zero_coding_scan_masks[scan_unit] |= bit; + } + + #[expect(clippy::inline_always, reason = "Tier-1 coefficient-loop hot path")] + #[inline(always)] + pub(in crate::j2c::bitplane) fn set_significant_scan_mask( + &mut self, + idx: usize, + y: usize, + padded_width: usize, + ) { + let (scan_unit, bit) = self.scan_unit_mask_index(idx, y, padded_width); + self.significant_scan_masks[scan_unit] |= bit; + } + + #[expect(clippy::inline_always, reason = "Tier-1 coefficient-loop hot path")] + #[inline(always)] + pub(in crate::j2c::bitplane) fn scan_unit_mask_index( + &self, + idx: usize, + y: usize, + padded_width: usize, + ) -> (usize, u8) { + // Callers pass the row they are scanning, so recovering the column + // costs a multiply; deriving the row from `idx` needed a division on + // every significance and zero-coding update. + let pad = COEFFICIENTS_PADDING as usize; + debug_assert!(idx / padded_width == y + pad); + let x = idx - (y + pad) * padded_width - pad; + debug_assert!(y < self.height as usize); + debug_assert!(x < self.width as usize); + + let scan_unit = (y >> 2) * self.width as usize + x; + (scan_unit, 1u8 << (y & 3)) + } +} diff --git a/crates/j2k-native/src/j2c/bitplane/state/workspace.rs b/crates/j2k-native/src/j2c/bitplane/state/workspace.rs index f1ef510e5..00e262b1a 100644 --- a/crates/j2k-native/src/j2c/bitplane/state/workspace.rs +++ b/crates/j2k-native/src/j2c/bitplane/state/workspace.rs @@ -15,6 +15,19 @@ struct ClassicWorkspaceLayout { padded_width: u32, coefficient_count: usize, scan_units: usize, + flag_words: usize, +} + +/// Which bookkeeping arrays a decode sizes and zeroes; the other set is +/// truncated, keeping its capacity. +#[derive(Clone, Copy, PartialEq, Eq)] +pub(super) enum StateLayout { + /// Per-coefficient state bytes, neighbor bytes, and scan masks. + PerCoefficient, + /// Packed per-stripe-column words. + PackedColumns, + /// Both, to size a workspace up front. + Both, } fn workspace_layout(width: u32, height: u32) -> Result { @@ -30,10 +43,14 @@ fn workspace_layout(width: u32, height: u32) -> Result { let coefficient_count = checked_decode_usize_product2(padded_width as usize, padded_height as usize)?; let scan_units = checked_decode_usize_product2(width as usize, height.div_ceil(4) as usize)?; + // One padding stripe above and below, one padding column on each side. + let flag_words = + checked_decode_usize_product2(padded_width as usize, height.div_ceil(4) as usize + 2)?; Ok(ClassicWorkspaceLayout { padded_width, coefficient_count, scan_units, + flag_words, }) } @@ -45,6 +62,7 @@ pub(crate) fn classic_decode_workspace_bytes(width: u32, height: u32) -> Result< include_elements::(&mut bytes, layout.coefficient_count)?; include_elements::(&mut bytes, layout.scan_units)?; include_elements::(&mut bytes, layout.scan_units)?; + include_elements::(&mut bytes, layout.flag_words)?; Ok(bytes) } @@ -65,6 +83,7 @@ pub(super) fn reset_decode_buffers( context: &mut BitPlaneDecodeContext, width: u32, height: u32, + state_layout: StateLayout, ) -> Result { let layout = workspace_layout(width, height)?; @@ -74,22 +93,29 @@ pub(super) fn reset_decode_buffers( layout.coefficient_count, Coefficient::default(), )?; + context.flags.clear(); + if state_layout != StateLayout::PerCoefficient { + try_resize_decode_elements(&mut context.flags, layout.flag_words, 0)?; + } context.neighbor_significances.clear(); + context.coefficient_states.clear(); + context.significant_scan_masks.clear(); + context.zero_coding_scan_masks.clear(); + if state_layout == StateLayout::PackedColumns { + return Ok(layout.padded_width); + } try_resize_decode_elements( &mut context.neighbor_significances, layout.coefficient_count, NeighborSignificances::default(), )?; - context.coefficient_states.clear(); try_resize_decode_elements( &mut context.coefficient_states, layout.coefficient_count, CoefficientState::default(), )?; - context.significant_scan_masks.clear(); try_resize_decode_elements(&mut context.significant_scan_masks, layout.scan_units, 0)?; - context.zero_coding_scan_masks.clear(); try_resize_decode_elements(&mut context.zero_coding_scan_masks, layout.scan_units, 0)?; Ok(layout.padded_width) } diff --git a/crates/j2k-native/src/j2c/bitplane/tests.rs b/crates/j2k-native/src/j2c/bitplane/tests.rs index 74416713e..e3644ef85 100644 --- a/crates/j2k-native/src/j2c/bitplane/tests.rs +++ b/crates/j2k-native/src/j2c/bitplane/tests.rs @@ -3,17 +3,23 @@ use alloc::{vec, vec::Vec}; use super::super::bitplane_encode; -use super::arithmetic::{cleanup_candidate_scan_mask, cleanup_run_length_candidate}; use super::bypass::{BitDecoder, BypassDecoder}; use super::context::{ context_label_magnitude_refinement_coding_from_state_lazy, context_label_sign_coding_index, - context_label_sign_coding_index_normal, + context_label_zero_coding_from_neighbors, }; use super::facade::decode_code_block_segments_validated; -use super::reconstruction::reconstruct_irreversible_midpoint; +use super::flags::{ + set_significant, sign_coding_label, zero_coding_label, zero_coding_table, SIGMA_NEIGHBOURS, + SIGMA_THIS, +}; +use super::observer::NoJ2kDecodeStats; +use super::reconstruction::{reconstruct_irreversible_midpoint, MidpointReconstructor}; +use super::scan::{cleanup_candidate_scan_mask, cleanup_run_length_candidate}; +use super::schedule::decode_code_block_segments_inner; use super::state::{ - BitPlaneDecodeContext, Coefficient, CoefficientState, NeighborSignificances, - COEFFICIENTS_PADDING, HAS_MAGNITUDE_REFINEMENT_MASK, HAS_ZERO_CODING_MASK, SIGNIFICANCE_MASK, + BitPlaneDecodeContext, Coefficient, NeighborSignificances, COEFFICIENTS_PADDING, + HAS_MAGNITUDE_REFINEMENT_MASK, HAS_ZERO_CODING_MASK, SIGNIFICANCE_MASK, }; use crate::j2c::build::SubBandType; use crate::j2c::codestream::CodeBlockStyle; @@ -61,6 +67,91 @@ fn classic_coefficient_state_preserves_38_bit_magnitude() { assert_eq!(coefficient.get(), i32::MIN); } +#[test] +fn per_block_midpoint_reconstruction_matches_the_reference_bit_for_bit() { + // Magnitudes cover zero, low bits, every single-bit position, dense + // patterns, and the 63-bit ceiling; the pass sweep hits every final-pass + // kind, and nonzero ROI shifts (decode caps bitplanes plus shift at 63) + // take the reference path. + let mut magnitudes = vec![0_u64, 1, 2, 3, 5, 0x55, 0xAA, 0xFFFF, (1 << 63) - 1]; + magnitudes.extend((0..63).map(|bit| 1_u64 << bit)); + magnitudes.extend((1..63).map(|bits| (1_u64 << bits) - 1)); + for decoded_bitplanes in [0_u8, 1, 2, 3, 8, 13, 31, 62, 63, 64, 70] { + for number_of_coding_passes in + 0..=u8::try_from((u32::from(decoded_bitplanes) * 3).min(200)).unwrap() + { + for roi_shift in [0_u8, 1, 7, 62, 63] { + let block = MidpointReconstructor::new( + decoded_bitplanes, + number_of_coding_passes, + roi_shift, + ); + for &magnitude in &magnitudes { + for sign in [0_u8, 1] { + let mut coefficient = Coefficient::default(); + for bit in 0..63 { + coefficient.push_bit_at(u32::from(magnitude >> bit & 1 == 1), bit); + } + coefficient.set_sign(sign); + assert_eq!( + block.reconstruct(coefficient).to_bits(), + reconstruct_irreversible_midpoint( + coefficient, + decoded_bitplanes, + number_of_coding_passes, + roi_shift, + ) + .to_bits(), + "bitplanes {decoded_bitplanes} passes {number_of_coding_passes} \ + roi {roi_shift} magnitude {magnitude:#x} sign {sign}" + ); + } + } + } + } + } +} + +#[test] +fn row_dequantization_matches_per_coefficient_reconstruction_bit_for_bit() { + let steps = [1.0_f32, 0.007_812_5, 3.3, 1.0e-3, 257.5, 0.123_456_7]; + for decoded_bitplanes in 1_u8..=40 { + let top = (1_u64 << decoded_bitplanes) - 1; + let mut magnitudes = vec![0, 1, 2, 3, top, top >> 1, top & 0x5555_5555_5555]; + magnitudes.extend((0..decoded_bitplanes).map(|bit| 1_u64 << bit)); + let coefficients = magnitudes + .iter() + .flat_map(|&magnitude| { + [0_u8, 1].map(|sign| { + let mut coefficient = Coefficient::default(); + for bit in 0..63 { + coefficient.push_bit_at(u32::from(magnitude >> bit & 1 == 1), bit); + } + coefficient.set_sign(sign); + coefficient + }) + }) + .collect::>(); + let max_passes = 1 + 3 * (decoded_bitplanes - 1); + for number_of_coding_passes in 1..=max_passes.min(96) { + let block = MidpointReconstructor::new(decoded_bitplanes, number_of_coding_passes, 0); + for step in steps { + let mut row = vec![f32::NAN; coefficients.len()]; + block.dequantize_row(&coefficients, &mut row, step); + for (value, &coefficient) in row.iter().zip(&coefficients) { + assert_eq!( + value.to_bits(), + (block.reconstruct(coefficient) * step).to_bits(), + "bitplanes {decoded_bitplanes} passes {number_of_coding_passes} \ + step {step} coefficient {:#x}", + coefficient.to_bits() + ); + } + } + } + } +} + #[test] fn irreversible_midpoint_reconstruction_tracks_the_last_decoded_pass() { let mut first_plane = Coefficient::default(); @@ -315,7 +406,7 @@ fn classic_bitplane_round_trips_seed_130_cb_block() { } #[test] -fn normal_neighborhood_significance_fast_path_returns_unmasked_neighbors() { +fn vertically_causal_context_masks_only_the_next_stripe_neighbor() { let mut ctx = BitPlaneDecodeContext { width: 1, height: 8, @@ -338,100 +429,232 @@ fn normal_neighborhood_significance_fast_path_returns_unmasked_neighbors() { ctx.neighbor_significances[idx].set_bottom(); assert_eq!(ctx.neighborhood_significance_states_index(idx, y), 1 << 6); + ctx.style.vertically_causal_context = false; assert_eq!( - ctx.normal_neighborhood_significance_states_index(idx), + ctx.neighborhood_significance_states_index(idx, y), (1 << 6) | 1 ); } #[test] -fn normal_sign_context_matches_generic_non_vertical_context() { - let mut ctx = BitPlaneDecodeContext { - width: 3, - height: 3, - padded_width: 5, - style: CodeBlockStyle::default(), - ..BitPlaneDecodeContext::default() +fn packed_column_contexts_match_per_coefficient_contexts() { + // Random significance and sign patterns, including stripe edges, a short + // final stripe, and block borders, recorded through both bookkeeping + // schemes; every coefficient must then see identical contexts. + let mut seed = 0x2545_f491_u32; + let mut next = move || { + seed ^= seed << 13; + seed ^= seed >> 17; + seed ^= seed << 5; + seed }; - let len = ctx.padded_width as usize * (ctx.height as usize + 2); - ctx.coefficients.resize(len, Coefficient::default()); - ctx.neighbor_significances - .resize(len, NeighborSignificances::default()); + for (width, height) in [(1, 1), (3, 5), (7, 10), (9, 12), (16, 3)] { + for density in [3_u32, 8, 14] { + let mut ctx = BitPlaneDecodeContext::default(); + let style = CodeBlockStyle::default(); + ctx.reset_for_job( + width, + height, + 0, + 1, + SubBandType::LowLow, + &style, + 8, + true, + false, + ) + .expect("per-coefficient context"); + let stride = ctx.padded_width as usize; + let mut flags = vec![0_u32; (height.div_ceil(4) as usize + 2) * stride]; + let flag_index = |x: usize, y: usize| (y / 4 + 1) * stride + x + 1; + let coefficient_index = |x: usize, y: usize| (y + 1) * stride + x + 1; + for y in 0..height as usize { + for x in 0..width as usize { + if next() % 16 >= density { + continue; + } + let sign = next() & 1; + let idx = coefficient_index(x, y); + ctx.set_sign_index(idx, u8::from(sign != 0)); + ctx.set_significant_index(idx, y, stride); + let index = flag_index(x, y); + let mut own = flags[index]; + let row = u32::try_from(y % 4).expect("row below four"); + set_significant(&mut flags, &mut own, index, stride, row, sign); + flags[index] = own; + } + } + for y in 0..height as usize { + for x in 0..width as usize { + let (idx, index, row) = ( + coefficient_index(x, y), + flag_index(x, y), + u32::try_from(y % 4).expect("row below four"), + ); + let word = flags[index]; + let neighbors = ctx.neighborhood_significance_states_index(idx, y); + for sub_band_type in [ + SubBandType::LowLow, + SubBandType::LowHigh, + SubBandType::HighLow, + SubBandType::HighHigh, + ] { + assert_eq!( + zero_coding_label(zero_coding_table(sub_band_type), word, row), + context_label_zero_coding_from_neighbors(neighbors, sub_band_type), + "{width}x{height} ({x}, {y})" + ); + } + assert_eq!( + sign_coding_label(word, flags[index - 1], flags[index + 1], row), + context_label_sign_coding_index(idx, y, &ctx), + "{width}x{height} ({x}, {y})" + ); + assert_eq!((word >> (3 * row)) & SIGMA_NEIGHBOURS != 0, neighbors != 0); + assert_eq!( + (word >> (3 * row)) & SIGMA_THIS != 0, + ctx.coefficient_states[idx].is_significant() + ); + } + } + } + } +} - let y = 1; - let idx = (y + COEFFICIENTS_PADDING as usize) * ctx.padded_width as usize - + COEFFICIENTS_PADDING as usize - + 1; - let padded_width = ctx.padded_width as usize; +/// Mostly small magnitudes with occasional large ones, like wavelet +/// subbands, so every pass and the run-length path occur. +fn random_subband_block(next: &mut impl FnMut() -> u32, len: u32) -> Vec { + (0..len) + .map(|_| { + let draw = next(); + let magnitude = match draw % 8 { + 0..=3 => 0, + 4 | 5 => (draw >> 8) % 4, + 6 => (draw >> 8) % 64, + _ => (draw >> 8) % 400, + }; + let magnitude = i32::try_from(magnitude).expect("bounded magnitude"); + if draw & 1 << 30 != 0 { + -magnitude + } else { + magnitude + } + }) + .collect() +} - ctx.neighbor_significances[idx].set_top(); - ctx.neighbor_significances[idx].set_left(); - ctx.neighbor_significances[idx].set_right(); - ctx.neighbor_significances[idx].set_bottom(); - ctx.set_sign_index(idx - padded_width, 1); - ctx.set_sign_index(idx - 1, 0); - ctx.set_sign_index(idx + 1, 1); - ctx.set_sign_index(idx + padded_width, 0); +/// Decode `passes` passes through the packed and the per-coefficient passes. +fn decode_through_both_pass_families( + encoded: &bitplane_encode::EncodedCodeBlockWithSegments, + (width, height): (u32, u32), + sub_band_type: SubBandType, + passes: u8, +) -> (Vec, Vec) { + let style = CodeBlockStyle::default(); + let segments = encoded + .segments + .iter() + .map(|segment| J2kCodeBlockSegment { + data_offset: segment.data_offset, + data_length: segment.data_length, + start_coding_pass: segment.start_coding_pass, + end_coding_pass: segment.end_coding_pass.min(passes), + use_arithmetic: segment.use_arithmetic, + }) + .collect::>(); + let mut packed = BitPlaneDecodeContext::default(); + decode_code_block_segments_validated( + &encoded.data, + &segments, + width, + height, + encoded.num_zero_bitplanes, + passes, + 10, + sub_band_type, + &style, + true, + &mut packed, + ) + .expect("packed decode"); + assert!(packed.uses_packed_columns()); + + let mut per_coefficient = BitPlaneDecodeContext::default(); + per_coefficient + .reset_for_job( + width, + height, + encoded.num_zero_bitplanes, + passes, + sub_band_type, + &style, + 10, + true, + false, + ) + .expect("per-coefficient context"); + decode_code_block_segments_inner( + &encoded.data, + &segments, + passes, + &mut per_coefficient, + &mut NoJ2kDecodeStats, + ) + .expect("per-coefficient decode"); - assert_eq!( - context_label_sign_coding_index_normal(idx, &ctx), - context_label_sign_coding_index(idx, y, &ctx) - ); + let rows = |ctx: &BitPlaneDecodeContext| { + ctx.coefficient_rows() + .flat_map(|row| row.iter().map(Coefficient::get_i64)) + .collect::>() + }; + (rows(&packed), rows(&per_coefficient)) } #[test] -fn normal_set_significant_index_matches_generic_neighbor_updates() { - let mut generic = BitPlaneDecodeContext { - width: 3, - height: 3, - padded_width: 5, - significant_scan_masks: vec![0; 3], - zero_coding_scan_masks: vec![0; 3], - style: CodeBlockStyle::default(), - ..BitPlaneDecodeContext::default() +fn packed_and_per_coefficient_passes_decode_identically() { + let mut seed = 0x9e37_79b9_u32; + let mut next = move || { + seed ^= seed << 13; + seed ^= seed >> 17; + seed ^= seed << 5; + seed }; - let len = generic.padded_width as usize * (generic.height as usize + 2); - generic - .coefficient_states - .resize(len, CoefficientState::default()); - generic - .neighbor_significances - .resize(len, NeighborSignificances::default()); - let mut normal = BitPlaneDecodeContext { - width: generic.width, - height: generic.height, - padded_width: generic.padded_width, - style: generic.style, - coefficient_states: generic.coefficient_states.clone(), - significant_scan_masks: vec![0; 3], - zero_coding_scan_masks: vec![0; 3], - neighbor_significances: generic.neighbor_significances.clone(), - ..BitPlaneDecodeContext::default() - }; - - let padded_width = generic.padded_width as usize; - let idx = - (1 + COEFFICIENTS_PADDING as usize) * padded_width + COEFFICIENTS_PADDING as usize + 1; - - generic.set_significant_index(idx, padded_width); - normal.set_significant_index_normal(idx, padded_width); - - assert_eq!( - normal.coefficient_states[idx].0, - generic.coefficient_states[idx].0 - ); - assert_eq!( - normal - .neighbor_significances - .iter() - .map(|neighbors| neighbors.0) - .collect::>(), - generic - .neighbor_significances - .iter() - .map(|neighbors| neighbors.0) - .collect::>() - ); + for (width, height) in [(1_u32, 1_u32), (4, 4), (5, 7), (13, 9), (32, 17), (64, 64)] { + for sub_band_type in [ + SubBandType::LowLow, + SubBandType::HighLow, + SubBandType::HighHigh, + ] { + let coefficients = random_subband_block(&mut next, width * height); + let encoded = bitplane_encode::encode_code_block_segments_with_style( + &coefficients, + width, + height, + sub_band_type, + 10, + CodeBlockStyle::default(), + ); + // Every truncation point, so each pass family's state after a + // partial bitplane is compared too. + for passes in 1..=encoded.num_coding_passes { + let (packed, per_coefficient) = decode_through_both_pass_families( + &encoded, + (width, height), + sub_band_type, + passes, + ); + assert_eq!( + packed, per_coefficient, + "{width}x{height} {passes}/{} passes", + encoded.num_coding_passes + ); + if passes == encoded.num_coding_passes { + let expected = coefficients.iter().map(|&value| i64::from(value)); + assert!(packed.iter().copied().eq(expected)); + } + } + } + } } #[test] @@ -462,7 +685,7 @@ fn first_magnitude_context_uses_neighbor_presence() { fn scan_unit_masks_track_significance_and_current_bitplane_zero_coding() { let mut ctx = BitPlaneDecodeContext::default(); let style = CodeBlockStyle::default(); - ctx.reset_for_job(5, 6, 0, 4, SubBandType::LowLow, &style, 8, true) + ctx.reset_for_job(5, 6, 0, 4, SubBandType::LowLow, &style, 8, true, false) .expect("reset context"); let padded_width = ctx.padded_width as usize; @@ -473,8 +696,8 @@ fn scan_unit_masks_track_significance_and_current_bitplane_zero_coding() { let scan_unit = (y >> 2) * ctx.width as usize + x; let bit = 1u8 << (y & 3); - ctx.set_significant_index(idx, padded_width); - ctx.set_zero_coding_index(idx, padded_width); + ctx.set_significant_index(idx, y, padded_width); + ctx.set_zero_coding_index(idx, y, padded_width); assert_eq!(ctx.significant_scan_masks[scan_unit], bit); assert_eq!(ctx.zero_coding_scan_masks[scan_unit], bit); @@ -493,9 +716,19 @@ fn scan_unit_masks_track_significance_and_current_bitplane_zero_coding() { selective_arithmetic_coding_bypass: true, ..CodeBlockStyle::default() }; - ctx.reset_for_job(5, 6, 0, 4, SubBandType::LowLow, &bypass_style, 8, true) - .expect("reset bypass context"); - ctx.set_zero_coding_index(idx, padded_width); + ctx.reset_for_job( + 5, + 6, + 0, + 4, + SubBandType::LowLow, + &bypass_style, + 8, + true, + false, + ) + .expect("reset bypass context"); + ctx.set_zero_coding_index(idx, y, padded_width); ctx.reset_for_next_bitplane(); assert_eq!(ctx.coefficient_states[idx].0 & HAS_ZERO_CODING_MASK, 0); } @@ -563,7 +796,7 @@ fn classic_selective_bypass_round_trips_padded_rgb8_cb_block() { fn cleanup_candidate_mask_excludes_significant_and_zero_coded_coefficients() { let mut ctx = BitPlaneDecodeContext::default(); let style = CodeBlockStyle::default(); - ctx.reset_for_job(5, 6, 0, 4, SubBandType::LowLow, &style, 8, true) + ctx.reset_for_job(5, 6, 0, 4, SubBandType::LowLow, &style, 8, true, false) .expect("reset context"); let padded_width = ctx.padded_width as usize; @@ -573,8 +806,8 @@ fn cleanup_candidate_mask_excludes_significant_and_zero_coded_coefficients() { (3 + COEFFICIENTS_PADDING as usize) * padded_width + COEFFICIENTS_PADDING as usize + 2; let scan_unit = 2; - ctx.set_significant_index(significant_idx, padded_width); - ctx.set_zero_coding_index(zero_coded_idx, padded_width); + ctx.set_significant_index(significant_idx, 1, padded_width); + ctx.set_zero_coding_index(zero_coded_idx, 3, padded_width); assert_eq!(cleanup_candidate_scan_mask(&ctx, scan_unit, 4), 0b0101); assert_eq!(cleanup_candidate_scan_mask(&ctx, scan_unit, 2), 0b0001); @@ -584,12 +817,12 @@ fn cleanup_candidate_mask_excludes_significant_and_zero_coded_coefficients() { fn cleanup_run_length_rejects_a_zero_coded_coefficient_without_significant_neighbors() { let mut ctx = BitPlaneDecodeContext::default(); let style = CodeBlockStyle::default(); - ctx.reset_for_job(1, 4, 0, 4, SubBandType::LowLow, &style, 8, true) + ctx.reset_for_job(1, 4, 0, 4, SubBandType::LowLow, &style, 8, true, false) .expect("reset context"); let padded_width = ctx.padded_width as usize; let top_idx = COEFFICIENTS_PADDING as usize * padded_width + COEFFICIENTS_PADDING as usize; - ctx.set_zero_coding_index(top_idx + padded_width, padded_width); + ctx.set_zero_coding_index(top_idx + padded_width, 1, padded_width); assert!(!cleanup_run_length_candidate( &ctx, diff --git a/crates/j2k-native/src/j2c/capabilities.rs b/crates/j2k-native/src/j2c/capabilities.rs index 66b572580..d1adc626e 100644 --- a/crates/j2k-native/src/j2c/capabilities.rs +++ b/crates/j2k-native/src/j2c/capabilities.rs @@ -197,15 +197,15 @@ impl<'a> CapabilityMarkerState<'a> { } pub(crate) fn validate_rsiz(&self, rsiz: u16) -> Result<(), CapabilityMarkerError> { - match (rsiz & HTJ2K_RSIZ_MASK != 0, self.htj2k.is_some()) { - (true, false) => Err(CapabilityMarkerError::Cap( + // Pcap15 without the Rsiz bit is accepted: CAP and the COD/COC block + // style fully describe HT decoding, and encoders before mid-2026 + // (this crate's included) omitted the bit. Rsiz without CAP leaves + // Ccap15 undefined, so that stays an error. + if rsiz & HTJ2K_RSIZ_MASK != 0 && self.htj2k.is_none() { + return Err(CapabilityMarkerError::Cap( "SIZ advertises Part 15 without Pcap15", - )), - (false, true) => Err(CapabilityMarkerError::Cap( - "Pcap15 is present without the Part 15 SIZ capability", - )), - _ => Ok(()), - }?; + )); + } if self.cpf_payload.is_some() && self.htj2k.is_none() { return Err(CapabilityMarkerError::Cpf( "CPF is present without Part 15 capabilities", diff --git a/crates/j2k-native/src/j2c/decode/store.rs b/crates/j2k-native/src/j2c/decode/store.rs index 61b29cfc7..fc6303bac 100644 --- a/crates/j2k-native/src/j2c/decode/store.rs +++ b/crates/j2k-native/src/j2c/decode/store.rs @@ -8,6 +8,8 @@ use super::{ use crate::j2c::roi::{output_grid_offset, output_region_rect}; use crate::math::round_ties_even_then_add; +mod upsample; + pub(super) fn apply_sign_shift_after_mct( tile_ctx: &mut TileDecodeContext, component_infos: &[ComponentInfo], @@ -192,6 +194,23 @@ pub(super) fn store<'a>( // Coefficients use the reconstructed resolution grid, which may be // reduced or clipped relative to the full component tile. let sample_rect = resolution_tile.rect.intersect(idwt_output.rect); + if x_shrink_factor == 1 && y_shrink_factor == 1 { + upsample::place_subsampled( + &upsample::SubsampledPlacement { + input: &idwt_output.coefficients, + input_stride: idwt_output.rect.width() as usize, + component_origin: (idwt_output.rect.x0, idwt_output.rect.y0), + columns: sample_rect.x0..sample_rect.x1, + rows: sample_rect.y0..sample_rect.y1, + scale: (u32::from(scale_x), u32::from(scale_y)), + image_size: (image_width, image_height), + image_offset: (x_offset, y_offset), + }, + &mut channel_data.container, + ); + return Ok(()); + } + for y in sample_rect.y0..sample_rect.y1 { let relative_y = (y - idwt_output.rect.y0) as usize; let reference_grid_y = (u32::from(scale_y) * y) / y_shrink_factor; diff --git a/crates/j2k-native/src/j2c/decode/store/upsample.rs b/crates/j2k-native/src/j2c/decode/store/upsample.rs new file mode 100644 index 000000000..23ad732f6 --- /dev/null +++ b/crates/j2k-native/src/j2c/decode/store/upsample.rs @@ -0,0 +1,166 @@ +// SPDX-License-Identifier: MIT OR Apache-2.0 + +//! Full-resolution placement of a subsampled component: every sample is +//! replicated over its `scale_x` x `scale_y` footprint on the reference grid, +//! clipped to the image. Footprints tile the grid without overlap, so each +//! output row of a footprint row is a copy of the first. + +use core::ops::Range; + +pub(super) struct SubsampledPlacement<'a> { + /// IDWT output, indexed relative to `component_origin` with row stride + /// `input_stride`. + pub(super) input: &'a [f32], + pub(super) input_stride: usize, + pub(super) component_origin: (u32, u32), + /// Component-grid columns and rows to place. + pub(super) columns: Range, + pub(super) rows: Range, + pub(super) scale: (u32, u32), + pub(super) image_size: (u32, u32), + /// Reference-grid position of the image origin. + pub(super) image_offset: (u32, u32), +} + +/// Place the samples into `output`, an image-sized row-major plane. +pub(super) fn place_subsampled(placement: &SubsampledPlacement<'_>, output: &mut [f32]) { + let (scale_x, scale_y) = placement.scale; + let (image_width, image_height) = placement.image_size; + let (x_offset, y_offset) = placement.image_offset; + let (x0, x1) = (placement.columns.start, placement.columns.end); + let first_column = (scale_x * x0).max(x_offset); + let end_column = (scale_x * x1).min(image_width + x_offset); + if first_column >= end_column { + return; + } + let width = image_width as usize; + let written = (first_column - x_offset) as usize..(end_column - x_offset) as usize; + let input_columns = + (x0 - placement.component_origin.0) as usize..(x1 - placement.component_origin.0) as usize; + for y in placement.rows.clone() { + let first_row = (scale_y * y).max(y_offset); + let end_row = (scale_y * y + scale_y).min(image_height + y_offset); + if first_row >= end_row { + continue; + } + let input_row = &placement.input + [(y - placement.component_origin.1) as usize * placement.input_stride..]; + let out_row = (first_row - y_offset) as usize; + fill_row( + &input_row[input_columns.clone()], + &mut output[out_row * width..][..width], + placement, + ); + for row in out_row + 1..(end_row - y_offset) as usize { + output.copy_within( + out_row * width + written.start..out_row * width + written.end, + row * width + written.start, + ); + } + } +} + +fn fill_row(samples: &[f32], out_row: &mut [f32], placement: &SubsampledPlacement<'_>) { + let scale_x = placement.scale.0; + let x_offset = placement.image_offset.0; + let (x0, x1) = (placement.columns.start, placement.columns.end); + if scale_x * x0 < x_offset || scale_x * x1 > placement.image_size.0 + x_offset { + // A footprint is clipped at the image edge. + for (x, &sample) in (x0..x1).zip(samples) { + let first = (scale_x * x).max(x_offset); + let end = (scale_x * x + scale_x).min(placement.image_size.0 + x_offset); + if first < end { + out_row[(first - x_offset) as usize..(end - x_offset) as usize].fill(sample); + } + } + return; + } + let out = &mut out_row[(scale_x * x0 - x_offset) as usize..(scale_x * x1 - x_offset) as usize]; + match scale_x { + 1 => out.copy_from_slice(samples), + 2 => { + for (pair, &sample) in out.chunks_exact_mut(2).zip(samples) { + pair[0] = sample; + pair[1] = sample; + } + } + _ => { + for (footprint, &sample) in out.chunks_exact_mut(scale_x as usize).zip(samples) { + footprint.fill(sample); + } + } + } +} + +#[cfg(test)] +mod tests { + use super::{place_subsampled, SubsampledPlacement}; + use alloc::vec; + use alloc::vec::Vec; + + /// The per-sample placement loop this module replaces at full resolution. + fn place_reference(placement: &SubsampledPlacement<'_>, output: &mut [f32]) { + let (scale_x, scale_y) = placement.scale; + let (image_width, image_height) = placement.image_size; + let (x_offset, y_offset) = placement.image_offset; + for y in placement.rows.clone() { + let relative_y = (y - placement.component_origin.1) as usize; + for x in placement.columns.clone() { + let relative_x = (x - placement.component_origin.0) as usize; + let sample = placement.input[relative_y * placement.input_stride + relative_x]; + for x_position in + (scale_x * x).max(x_offset)..(scale_x * x + scale_x).min(image_width + x_offset) + { + for y_position in (scale_y * y).max(y_offset) + ..(scale_y * y + scale_y).min(image_height + y_offset) + { + output[(y_position - y_offset) as usize * image_width as usize + + (x_position - x_offset) as usize] = sample; + } + } + } + } + } + + #[test] + fn placement_matches_the_per_sample_reference() { + for scale in [(1_u32, 1_u32), (2, 1), (1, 2), (2, 2), (3, 2)] { + // Offsets and component origins that clip the first and last + // footprints, alongside an aligned case. + for (image_offset, origin, columns, rows) in [ + ((0_u32, 0_u32), (0_u32, 0_u32), 0_u32..5, 0_u32..4), + ((1, 1), (0, 0), 0..5, 0..4), + ((0, 0), (2, 1), 2..6, 1..4), + ((3, 2), (1, 1), 1..6, 1..5), + ] { + let (scale_x, scale_y) = scale; + let image_size = ( + (scale_x * columns.end) + .saturating_sub(image_offset.0) + .max(1) + - 1, + (scale_y * rows.end).saturating_sub(image_offset.1).max(1), + ); + let input_stride = (columns.end - origin.0) as usize; + let input = core::iter::successors(Some(0.5_f32), |value| Some(value + 1.0)) + .take(input_stride * (rows.end - origin.1) as usize) + .collect::>(); + let placement = SubsampledPlacement { + input: &input, + input_stride, + component_origin: origin, + columns: columns.clone(), + rows: rows.clone(), + scale, + image_size, + image_offset, + }; + let len = (image_size.0 * image_size.1) as usize; + let (mut fast, mut reference) = (vec![-1.0; len], vec![-1.0; len]); + place_subsampled(&placement, &mut fast); + place_reference(&placement, &mut reference); + assert_eq!(fast, reference, "scale {scale:?} offset {image_offset:?}"); + } + } + } +} diff --git a/crates/j2k-native/src/j2c/decode/subband/classic.rs b/crates/j2k-native/src/j2c/decode/subband/classic.rs index e3a8a8542..5482ea1d2 100644 --- a/crates/j2k-native/src/j2c/decode/subband/classic.rs +++ b/crates/j2k-native/src/j2c/decode/subband/classic.rs @@ -262,24 +262,22 @@ pub(super) fn decode_sub_band_classic_blocks( header.strict, )?; let base_store = &mut storage.coefficients[sub_band.coefficients.clone()]; + let midpoint = tile_ctx.bit_plane_decode_context.midpoint_reconstructor( + code_block.number_of_coding_passes, + component_info.roi_shift, + ); for coefficients in tile_ctx.bit_plane_decode_context.coefficient_rows() { let out_row = &mut base_store[base_idx..]; - for (output, coefficient) in out_row.iter_mut().zip(coefficients.iter().copied()) { - *output = if irreversible_midpoint { - tile_ctx - .bit_plane_decode_context - .reconstruct_irreversible_midpoint( - coefficient, - code_block.number_of_coding_passes, - component_info.roi_shift, - ) - } else { - apply_roi_maxshift_inverse_i64( + if irreversible_midpoint { + midpoint.dequantize_row(coefficients, out_row, dequantization_step); + } else { + for (output, coefficient) in out_row.iter_mut().zip(coefficients) { + *output = apply_roi_maxshift_inverse_i64( coefficient.get_i64(), component_info.roi_shift, ) as f32 - }; - *output *= dequantization_step; + * dequantization_step; + } } base_idx += output_stride; } diff --git a/crates/j2k-native/src/j2c/ht_block_decode/magnitude.rs b/crates/j2k-native/src/j2c/ht_block_decode/magnitude.rs index 5f69cc108..54e0e484f 100644 --- a/crates/j2k-native/src/j2c/ht_block_decode/magnitude.rs +++ b/crates/j2k-native/src/j2c/ht_block_decode/magnitude.rs @@ -10,8 +10,10 @@ fn sample_mask(bit: u32) -> u32 { #[expect( clippy::cast_possible_truncation, - reason = "the widened rounding expression is bounded to the 31-bit decoded magnitude field" + clippy::inline_always, + reason = "31-bit decoded magnitudes; inlined so the `&mut` MagSgn reader stays in registers" )] +#[inline(always)] fn decode_mag_sgn_sample_with_vn( magsgn: &mut ForwardBitReader<0xFF>, inf: u32, @@ -69,9 +71,7 @@ pub(super) fn decode_magnitude_sign_phase( return None; } v_n_scratch[..v_n_width].fill(0); - let mut magsgn = ForwardBitReader::<0xFF>::new(&coded_data[..lcup - scup]); - decode_magnitude_sign_first_row_from_cleanup( &mut magsgn, scratch, diff --git a/crates/j2k-native/src/j2c/ht_block_decode/readers.rs b/crates/j2k-native/src/j2c/ht_block_decode/readers.rs index fa84706cd..f7419aba4 100644 --- a/crates/j2k-native/src/j2c/ht_block_decode/readers.rs +++ b/crates/j2k-native/src/j2c/ht_block_decode/readers.rs @@ -98,6 +98,8 @@ impl<'a> MelDecoder<'a> { } } +/// `Clone` (not `Copy`) keeps the by-value refill hand-offs explicit. +#[derive(Clone)] pub(super) struct ForwardBitReader<'a, const PAD: u8> { data: &'a [u8], pos: usize, @@ -117,7 +119,33 @@ impl<'a, const PAD: u8> ForwardBitReader<'a, PAD> { } } + /// Top the reservoir up to at least 33 bits, as the byte loop does. A byte + /// after `0xFF` carries seven bits, so four bytes with no `0xFF` before + /// or among their first three are added as one word; stuffing, an empty + /// reservoir, and the padded tail take the byte loop. + #[expect(clippy::inline_always, reason = "per-sample refill check")] + #[inline(always)] fn fill(&mut self) { + if !self.unstuff && self.bits > 0 { + if let Some(&[b0, b1, b2, b3]) = self.data.get(self.pos..self.pos + 4) { + let word = u32::from_le_bytes([b0, b1, b2, b3]); + if b0 != 0xFF && b1 != 0xFF && b2 != 0xFF { + self.tmp |= u64::from(word) << self.bits; + self.bits += 32; + self.pos += 4; + self.unstuff = b3 == 0xFF; + return; + } + } + } + *self = self.clone().fill_bytes(); + } + + /// Taking and returning the reader by value keeps the caller's local + /// reader out of memory across this cold call. + #[cold] + #[inline(never)] + fn fill_bytes(mut self) -> Self { while self.bits <= 32 { let byte = if self.pos < self.data.len() { let byte = self.data[self.pos]; @@ -134,6 +162,7 @@ impl<'a, const PAD: u8> ForwardBitReader<'a, PAD> { self.bits += valid_bits; self.unstuff = next_unstuff; } + self } #[expect(clippy::cast_possible_truncation, reason = "low reservoir word")] @@ -152,6 +181,8 @@ impl<'a, const PAD: u8> ForwardBitReader<'a, PAD> { } } +/// `Clone` (not `Copy`) keeps the by-value refill hand-offs explicit. +#[derive(Clone)] pub(super) struct ReverseBitReader<'a> { data: &'a [u8], pos: isize, @@ -190,8 +221,46 @@ impl<'a> ReverseBitReader<'a> { } } - #[expect(clippy::cast_sign_loss, reason = "nonnegative live cursor")] + /// Top the reservoir up to at least 33 bits, as the byte loop does. Read + /// backward, a byte is stuffed when it ends in seven ones after a byte + /// above 0x8F; four bytes with no such pair are added as one word, and + /// stuffing, an empty reservoir, and the padded tail take the byte loop. + #[expect( + clippy::cast_sign_loss, + clippy::inline_always, + reason = "nonnegative live cursor; per-quad refill check" + )] + #[inline(always)] fn fill(&mut self) { + if self.bits > 0 && self.remaining >= 4 && self.pos >= 3 { + let last = self.pos as usize; + if let Some(&[b3, b2, b1, b0]) = self.data.get(last - 3..=last) { + // `b0` is consumed first, so it lands in the low byte. Each + // per-byte flag below sits in that byte's bit 7. + let word = u32::from_le_bytes([b0, b1, b2, b3]); + // Above 0x8F: bit 7 set and bits 4-6 nonzero. + let above_8f = word & 0x8080_8080 & ((word & 0x7070_7070) + 0x7070_7070); + // Low seven bits all ones. + let low_ones = ((word & 0x7F7F_7F7F) + 0x0101_0101) & 0x8080_8080; + let after_above_8f = (above_8f << 8) | (u32::from(self.unstuff) << 7); + if after_above_8f & low_ones == 0 { + self.tmp |= u64::from(word) << self.bits; + self.bits += 32; + self.pos -= 4; + self.remaining -= 4; + self.unstuff = above_8f >> 31 != 0; + return; + } + } + } + *self = self.clone().fill_bytes(); + } + + /// By value for the same reason as [`ForwardBitReader::fill_bytes`]. + #[expect(clippy::cast_sign_loss, reason = "nonnegative live cursor")] + #[cold] + #[inline(never)] + fn fill_bytes(mut self) -> Self { while self.bits <= 32 { let byte = if self.remaining > 0 { let byte = self.data[self.pos as usize]; @@ -210,6 +279,7 @@ impl<'a> ReverseBitReader<'a> { self.bits += d_bits; self.unstuff = next_unstuff; } + self } #[expect(clippy::cast_possible_truncation, reason = "low reservoir word")] @@ -237,66 +307,4 @@ pub(super) fn read_u32_pair(values: &[u16], index: usize) -> u32 { } #[cfg(test)] -mod tests { - use super::{ForwardBitReader, MelDecoder, ReverseBitReader}; - - #[test] - fn reader_state_and_bit_consumption_match_pre_split_goldens() { - let data = [0xAA, 0xFF, 0x01, 0x7F, 0x80]; - let mut forward = ForwardBitReader::<0xFF>::new(&data); - assert_eq!(forward.fetch(), 0x3F81_FFAA); - forward.advance(5); - assert_eq!(forward.fetch(), 0x01FC_0FFD); - forward.advance(19); - assert_eq!(forward.fetch(), 0xFFFF_C03F); - assert_eq!( - (forward.pos, forward.bits, forward.tmp, forward.unstuff), - (5, 37, 0x0000_001F_FFFF_C03F, true) - ); - - let mut reverse = ReverseBitReader::new_mrp(&data); - assert_eq!(reverse.fetch(), 0xFF01_7F80); - assert_eq!(reverse.advance(7), 0x55FE_02FF); - assert_eq!(reverse.fetch(), 0x55FE_02FF); - assert_eq!( - ( - reverse.pos, - reverse.remaining, - reverse.bits, - reverse.tmp, - reverse.unstuff, - ), - (-1, 0, 33, 0x0000_0001_55FE_02FF, true) - ); - - let mel_data = [0x12, 0x34, 0x56, 0x78]; - let mut mel = MelDecoder::new(&mel_data, mel_data.len(), 2); - let mut runs = [0i32; 8]; - for run in &mut runs { - *run = mel.get_run().expect("MEL run"); - } - assert_eq!(runs, [1, 0, 1, 0, 0, 0, 2, 2]); - assert_eq!( - ( - mel.pos, - mel.remaining, - mel.bits_left, - mel.k, - mel.num_runs, - mel.runs, - mel.unstuff, - ), - (3, 0, 0, 5, 0, 0, false) - ); - - let vlc_data = [0x12, 0x34, 0x56, 0x78, 0x9A]; - let mut vlc = ReverseBitReader::new_vlc(&vlc_data, vlc_data.len(), 3); - assert_eq!(vlc.fetch(), 0x0000_02B7); - assert_eq!(vlc.advance(9), 0x0000_0001); - assert_eq!(vlc.fetch(), 0x0000_0001); - assert_eq!( - (vlc.pos, vlc.remaining, vlc.bits, vlc.tmp, vlc.unstuff), - (1, 0, 34, 1, false) - ); - } -} +mod tests; diff --git a/crates/j2k-native/src/j2c/ht_block_decode/readers/tests.rs b/crates/j2k-native/src/j2c/ht_block_decode/readers/tests.rs new file mode 100644 index 000000000..3d3a9b4bf --- /dev/null +++ b/crates/j2k-native/src/j2c/ht_block_decode/readers/tests.rs @@ -0,0 +1,119 @@ +// SPDX-License-Identifier: MIT OR Apache-2.0 + +use super::{ForwardBitReader, MelDecoder, ReverseBitReader}; + +#[test] +fn reader_state_and_bit_consumption_match_pre_split_goldens() { + let data = [0xAA, 0xFF, 0x01, 0x7F, 0x80]; + let mut forward = ForwardBitReader::<0xFF>::new(&data); + assert_eq!(u64::from(forward.fetch()), 0x3F81_FFAA); + forward.advance(5); + assert_eq!(u64::from(forward.fetch()), 0x01FC_0FFD); + forward.advance(19); + assert_eq!(u64::from(forward.fetch()), 0xFFFF_C03F); + assert_eq!( + (forward.pos, forward.bits, forward.tmp, forward.unstuff), + (5, 37, 0x0000_001F_FFFF_C03F, true) + ); + + let mut reverse = ReverseBitReader::new_mrp(&data); + assert_eq!(u64::from(reverse.fetch()), 0xFF01_7F80); + assert_eq!(reverse.advance(7), 0x55FE_02FF); + assert_eq!(u64::from(reverse.fetch()), 0x55FE_02FF); + assert_eq!( + ( + reverse.pos, + reverse.remaining, + reverse.bits, + reverse.tmp, + reverse.unstuff, + ), + (-1, 0, 33, 0x0000_0001_55FE_02FF, true) + ); + + let mel_data = [0x12, 0x34, 0x56, 0x78]; + let mut mel = MelDecoder::new(&mel_data, mel_data.len(), 2); + let mut runs = [0i32; 8]; + for run in &mut runs { + *run = mel.get_run().expect("MEL run"); + } + assert_eq!(runs, [1, 0, 1, 0, 0, 0, 2, 2]); + assert_eq!( + ( + mel.pos, + mel.remaining, + mel.bits_left, + mel.k, + mel.num_runs, + mel.runs, + mel.unstuff, + ), + (3, 0, 0, 5, 0, 0, false) + ); + + let vlc_data = [0x12, 0x34, 0x56, 0x78, 0x9A]; + let mut vlc = ReverseBitReader::new_vlc(&vlc_data, vlc_data.len(), 3); + assert_eq!(vlc.fetch(), 0x0000_02B7); + assert_eq!(vlc.advance(9), 0x0000_0001); + assert_eq!(vlc.fetch(), 0x0000_0001); + assert_eq!( + (vlc.pos, vlc.remaining, vlc.bits, vlc.tmp, vlc.unstuff), + (1, 0, 34, 1, false) + ); +} + +#[test] +fn word_refills_yield_the_same_bits_as_byte_refills() { + let mut seed = 0x1234_5678_u32; + let mut next = move || { + seed ^= seed << 13; + seed ^= seed >> 17; + seed ^= seed << 5; + seed + }; + // Bytes around the stuffing thresholds are over-represented. + let special = [0xFF_u8, 0x7F, 0x8F, 0x90, 0x00, 0xFE]; + for len in 0..=40_usize { + for _ in 0..20 { + let data = (0..len) + .map(|_| { + let draw = next(); + if draw % 3 == 0 { + special[(draw >> 8) as usize % special.len()] + } else { + (draw >> 16).to_le_bytes()[0] + } + }) + .collect::>(); + let advances = (0..24).map(|_| next() % 33).collect::>(); + + let mut forward = ForwardBitReader::<0xFF>::new(&data); + let mut forward_bytes = forward.clone(); + let mut reverse = ReverseBitReader::new_mrp(&data); + let mut reverse_bytes = reverse.clone(); + for &count in &advances { + if forward_bytes.bits < 32 { + forward_bytes = forward_bytes.fill_bytes(); + } + assert_eq!( + u64::from(forward.fetch()), + forward_bytes.tmp & u64::from(u32::MAX), + "forward len {len}" + ); + forward.advance(count); + forward_bytes.advance(count); + + if reverse_bytes.bits < 32 { + reverse_bytes = reverse_bytes.fill_bytes(); + } + assert_eq!( + u64::from(reverse.fetch()), + reverse_bytes.tmp & u64::from(u32::MAX), + "reverse len {len}" + ); + reverse.advance(count); + reverse_bytes.advance(count); + } + } + } +} diff --git a/crates/j2k-native/src/scalar/classic_decode.rs b/crates/j2k-native/src/scalar/classic_decode.rs index df52943a4..280015cfa 100644 --- a/crates/j2k-native/src/scalar/classic_decode.rs +++ b/crates/j2k-native/src/scalar/classic_decode.rs @@ -146,6 +146,8 @@ fn write_j2k_code_block_output( output: &mut [f32], irreversible_midpoint: bool, ) { + let midpoint = + decode_context.midpoint_reconstructor(job.number_of_coding_passes, job.roi_shift); for (row_idx, coeff_row) in decode_context .coefficient_rows() .enumerate() @@ -153,17 +155,13 @@ fn write_j2k_code_block_output( { let row_start = row_idx * job.output_stride; let output_row = &mut output[row_start..row_start + layout.stride]; + if irreversible_midpoint { + midpoint.dequantize_row(coeff_row, output_row, job.dequantization_step); + continue; + } for (coefficient, sample) in coeff_row.iter().zip(output_row.iter_mut()) { - let coefficient = if irreversible_midpoint { - decode_context.reconstruct_irreversible_midpoint( - *coefficient, - job.number_of_coding_passes, - job.roi_shift, - ) - } else { - apply_roi_maxshift_inverse_i64(coefficient.get_i64(), job.roi_shift) as f32 - }; - *sample = coefficient * job.dequantization_step; + *sample = apply_roi_maxshift_inverse_i64(coefficient.get_i64(), job.roi_shift) as f32 + * job.dequantization_step; } } } diff --git a/crates/j2k-native/tests/ht_capabilities.rs b/crates/j2k-native/tests/ht_capabilities.rs index 086c761c9..e85148edb 100644 --- a/crates/j2k-native/tests/ht_capabilities.rs +++ b/crates/j2k-native/tests/ht_capabilities.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: MIT OR Apache-2.0 use j2k_native::{ - encode_htj2k, encode_precomputed_htj2k_53, encode_typed_component_planes_53, + encode, encode_htj2k, encode_precomputed_htj2k_53, encode_typed_component_planes_53, inspect_htj2k_capabilities, DecodeError, DecodeSettings, EncodeOptions, EncodeTypedComponentPlane, Htj2kCapabilityMode, Image, J2kForwardDwt53Output, MarkerError, PrecomputedHtj2k53Component, PrecomputedHtj2k53Image, @@ -250,6 +250,59 @@ fn reserved_ht_mode_is_rejected_by_inspection_and_production_decode() { )); } +#[test] +fn cap_without_the_part15_rsiz_bit_still_decodes_as_htj2k() { + // Encoders before 2026-06-28 (including this crate's) wrote CAP with + // Pcap15 but left SIZ Rsiz at 0; DICOM pyramids made then are in use. + let conformant = encoded_fixture(); + let mut legacy = conformant.clone(); + rewrite_rsiz(&mut legacy, |value| value & !0x4000); + assert_ne!(legacy, conformant); + + let capabilities = inspect_htj2k_capabilities(&legacy) + .expect("legacy codestream inspects") + .expect("CAP still advertises Part 15"); + assert!(capabilities.default_ht_block_coding()); + let decode = |bytes: &[u8]| { + Image::new(bytes, &DecodeSettings::default()) + .and_then(|image| image.decode()) + .expect("decode") + }; + assert_eq!(decode(&legacy), decode(&conformant)); +} + +#[test] +fn part15_rsiz_bit_without_cap_is_rejected() { + let mut codestream = encode( + &[0, 1, 2, 3], + 2, + 2, + 1, + 8, + false, + &EncodeOptions { + num_decomposition_levels: 0, + ..Default::default() + }, + ) + .expect("encode classic fixture"); + rewrite_rsiz(&mut codestream, |value| value | 0x4000); + + assert!(inspect_htj2k_capabilities(&codestream).is_err()); + assert!(matches!( + Image::new(&codestream, &DecodeSettings::default()), + Err(DecodeError::Marker(MarkerError::ParseFailure( + "SIZ advertises Part 15 without Pcap15" + ))) + )); +} + +fn rewrite_rsiz(codestream: &mut [u8], rewrite: impl FnOnce(u16) -> u16) { + let siz = marker_offset(codestream, 0x51); + let value = u16::from_be_bytes([codestream[siz + 4], codestream[siz + 5]]); + codestream[siz + 4..siz + 6].copy_from_slice(&rewrite(value).to_be_bytes()); +} + fn rewrite_ccap15(codestream: &mut [u8], rewrite: impl FnOnce(u16) -> u16) { let cap = marker_offset(codestream, 0x50); let value = u16::from_be_bytes([codestream[cap + 8], codestream[cap + 9]]); From 0b7e7cba4bbdbcf1e8e37a4341553bc71d3533b5 Mon Sep 17 00:00:00 2001 From: GF Date: Sat, 26 Sep 2026 20:44:54 -0400 Subject: [PATCH 5/5] refactor(jpeg): keep bounded NEON window loads with memory helpers --- crates/j2k-jpeg/src/backend/neon.rs | 48 +------------------------ crates/j2k-jpeg/src/simd/neon_memory.rs | 46 ++++++++++++++++++++++++ 2 files changed, 47 insertions(+), 47 deletions(-) diff --git a/crates/j2k-jpeg/src/backend/neon.rs b/crates/j2k-jpeg/src/backend/neon.rs index 89c30b018..e133870d2 100644 --- a/crates/j2k-jpeg/src/backend/neon.rs +++ b/crates/j2k-jpeg/src/backend/neon.rs @@ -12,7 +12,7 @@ use super::row_pair::{normalize_simd_row_pair, normalize_ycbcr_row}; use super::{scalar, Rgb420ChromaRows, Rgb420Crop, Rgb420CroppedRowPair, Rgb420RowPair}; use crate::color::upsample::h2v2_fancy_sample_for_width; use crate::color::ycbcr::{ycbcr_to_rgb, FIX_0_34414, FIX_0_71414, FIX_1_40200, FIX_1_77200}; -use crate::simd::neon_memory; +use crate::simd::neon_memory::{self, load_head_window, load_tail_window}; pub(crate) fn fill_rgb_row_from_gray(neon: fearless_simd::Neon, gray_row: &[u8], dst: &mut [u8]) { let width = gray_row.len().min(dst.len() / 3); @@ -1194,52 +1194,6 @@ fn fill_rgb_row_pair_from_420_tail_neon_top_only( const TAIL_WINDOW: usize = LANES + 2; -fn load_tail_window(src: &[u8], start: usize, len: usize) -> [u8; UPSAMPLED_LANES] { - debug_assert!(start < src.len()); - debug_assert!(len > 0); - debug_assert!(len <= UPSAMPLED_LANES); - let mut out = [0u8; UPSAMPLED_LANES]; - let available = src.len() - start; - let copy_len = available.min(len); - out[..copy_len].copy_from_slice(&src[start..start + copy_len]); - if copy_len < len { - let pad = out[copy_len - 1]; - for value in &mut out[copy_len..len] { - *value = pad; - } - } - if len < UPSAMPLED_LANES { - let pad = out[len - 1]; - for value in &mut out[len..UPSAMPLED_LANES] { - *value = pad; - } - } - out -} - -fn load_head_window(src: &[u8], len: usize) -> [u8; UPSAMPLED_LANES] { - debug_assert!(!src.is_empty()); - debug_assert!(len > 0); - debug_assert!(len <= UPSAMPLED_LANES); - let mut out = [0u8; UPSAMPLED_LANES]; - let copy_len = src.len().min(len); - out[0] = src[0]; - out[1..=copy_len].copy_from_slice(&src[..copy_len]); - if copy_len < len { - let pad = out[copy_len]; - for value in &mut out[copy_len + 1..=len] { - *value = pad; - } - } - if len + 1 < UPSAMPLED_LANES { - let pad = out[len]; - for value in &mut out[len + 1..UPSAMPLED_LANES] { - *value = pad; - } - } - out -} - fn can_use_tail_420_chunk(chroma_width: usize, sample_offset: usize, out_len: usize) -> bool { out_len <= UPSAMPLED_LANES && sample_offset > 0 && sample_offset + LANES >= chroma_width } diff --git a/crates/j2k-jpeg/src/simd/neon_memory.rs b/crates/j2k-jpeg/src/simd/neon_memory.rs index f6052d04f..20c6255cb 100644 --- a/crates/j2k-jpeg/src/simd/neon_memory.rs +++ b/crates/j2k-jpeg/src/simd/neon_memory.rs @@ -123,3 +123,49 @@ pub(crate) fn store_rgb8x16( // - Initialization: the store initializes all forty-eight output bytes. unsafe { vst3q_u8(dst.as_mut_ptr(), uint8x16x3_t(red, green, blue)) }; } + +pub(crate) fn load_tail_window(src: &[u8], start: usize, len: usize) -> [u8; 16] { + debug_assert!(start < src.len()); + debug_assert!(len > 0); + debug_assert!(len <= 16); + let mut out = [0u8; 16]; + let available = src.len() - start; + let copy_len = available.min(len); + out[..copy_len].copy_from_slice(&src[start..start + copy_len]); + if copy_len < len { + let pad = out[copy_len - 1]; + for value in &mut out[copy_len..len] { + *value = pad; + } + } + if len < 16 { + let pad = out[len - 1]; + for value in &mut out[len..16] { + *value = pad; + } + } + out +} + +pub(crate) fn load_head_window(src: &[u8], len: usize) -> [u8; 16] { + debug_assert!(!src.is_empty()); + debug_assert!(len > 0); + debug_assert!(len <= 16); + let mut out = [0u8; 16]; + let copy_len = src.len().min(len); + out[0] = src[0]; + out[1..=copy_len].copy_from_slice(&src[..copy_len]); + if copy_len < len { + let pad = out[copy_len]; + for value in &mut out[copy_len + 1..=len] { + *value = pad; + } + } + if len + 1 < 16 { + let pad = out[len]; + for value in &mut out[len + 1..16] { + *value = pad; + } + } + out +}