diff --git a/TESTING.md b/TESTING.md index 0e35080..cb41589 100644 --- a/TESTING.md +++ b/TESTING.md @@ -37,7 +37,9 @@ the CI toolchain, validator, or corpus forward deliberately. - pixel-for-pixel comparison of the `image` crate integration hook output (`ImageReader`/`DynamicImage::from_decoder`) against the direct Rust decode for every comparable verifier file, including exact ICC-profile equality - through the hook decoder's `ImageDecoder::icc_profile` + through the hook decoder's `ImageDecoder::icc_profile`; this reproduces + Ente's production hook shape, including its explicit guardrails, + `with_guessed_format`, `Limits::reserve`, and `set_limits` - embedded ICC colour-profile comparison against the validator's PNG output: when `heif-dec` embeds a profile, the Rust PNG must carry byte-identical profile data; a Rust-only profile is allowed (the Rust decoder synthesizes diff --git a/scripts/heic_tests.sh b/scripts/heic_tests.sh index 3d9478d..3793844 100755 --- a/scripts/heic_tests.sh +++ b/scripts/heic_tests.sh @@ -419,9 +419,9 @@ fn main() -> Result<(), Box> { RS cat > "$HELPER_DIR/src/bin/heif-image-adapter-bench.rs" <<'RS' -use heic_decoder::image_integration::register_image_decoder_hooks; -use heic_decoder::{DecodedRgbaImage, DecodedRgbaPixels, decode_path_to_rgba}; -use image::{DynamicImage, ImageReader}; +use heic_decoder::image_integration::register_image_decoder_hooks_with_guardrails; +use heic_decoder::{DecodeGuardrails, DecodedRgbaImage, DecodedRgbaPixels, decode_path_to_rgba}; +use image::{DynamicImage, ImageDecoder, ImageReader, Limits}; use std::error::Error; use std::path::Path; @@ -462,8 +462,18 @@ fn main() -> Result<(), Box> { let value = match args[1].as_str() { "direct" => direct_checksum(&decode_path_to_rgba(input)?), "adapter" => { - let _ = register_image_decoder_hooks(); - let decoded = ImageReader::open(input)?.decode()?; + let _ = register_image_decoder_hooks_with_guardrails(DecodeGuardrails { + max_input_bytes: Some(128 * 1024 * 1024), + max_pixels: Some(256_000_000), + max_temp_spool_bytes: Some(256 * 1024 * 1024), + temp_spool_directory: None, + }); + let mut decoder = ImageReader::open(input)?.with_guessed_format()?.into_decoder()?; + let _icc_profile = decoder.icc_profile()?; + let mut limits = Limits::default(); + limits.reserve(decoder.total_bytes())?; + decoder.set_limits(limits)?; + let decoded = DynamicImage::from_decoder(decoder)?; let (width, height) = (decoded.width(), decoded.height()); let pixels = match decoded { DynamicImage::ImageRgba8(buffer) => checksum(buffer.as_raw()), @@ -480,9 +490,9 @@ fn main() -> Result<(), Box> { RS cat > "$HELPER_DIR/src/bin/heif-image-hook-check.rs" <<'RS' -use heic_decoder::image_integration::register_image_decoder_hooks; -use heic_decoder::{DecodedRgbaPixels, decode_path_to_rgba}; -use image::{DynamicImage, ImageDecoder, ImageReader}; +use heic_decoder::image_integration::register_image_decoder_hooks_with_guardrails; +use heic_decoder::{DecodeGuardrails, DecodedRgbaPixels, decode_path_to_rgba}; +use image::{DynamicImage, ImageDecoder, ImageReader, Limits}; use std::error::Error; use std::path::Path; @@ -495,9 +505,17 @@ fn main() -> Result<(), Box> { let input = Path::new(&args[1]); let direct = decode_path_to_rgba(input)?; - let _ = register_image_decoder_hooks(); - let mut decoder = ImageReader::open(input)?.into_decoder()?; + let _ = register_image_decoder_hooks_with_guardrails(DecodeGuardrails { + max_input_bytes: Some(128 * 1024 * 1024), + max_pixels: Some(256_000_000), + max_temp_spool_bytes: Some(256 * 1024 * 1024), + temp_spool_directory: None, + }); + let mut decoder = ImageReader::open(input)?.with_guessed_format()?.into_decoder()?; let icc_profile = decoder.icc_profile()?; + let mut limits = Limits::default(); + limits.reserve(decoder.total_bytes())?; + decoder.set_limits(limits)?; if icc_profile != direct.icc_profile { return Err("image hook ICC profile differs from direct decode".into()); } diff --git a/src/bin/heif-image-adapter-bench.rs b/src/bin/heif-image-adapter-bench.rs index 9378e9d..bfbb11c 100644 --- a/src/bin/heif-image-adapter-bench.rs +++ b/src/bin/heif-image-adapter-bench.rs @@ -1,6 +1,6 @@ -use heic_decoder::image_integration::register_image_decoder_hooks; -use heic_decoder::{DecodedRgbaPixels, decode_path_to_rgb8, decode_path_to_rgba}; -use image::{DynamicImage, ImageReader}; +use heic_decoder::image_integration::register_image_decoder_hooks_with_guardrails; +use heic_decoder::{DecodeGuardrails, DecodedRgbaPixels, decode_path_to_rgb8, decode_path_to_rgba}; +use image::{DynamicImage, ImageDecoder, ImageReader, Limits}; use std::env; use std::error::Error; use std::fmt::{Display, Formatter}; @@ -80,10 +80,26 @@ fn bench_direct(input_path: &Path) -> Result> { Ok(((decoded.width as u64) << 32) ^ (decoded.height as u64) ^ checksum) } -fn bench_adapter(input_path: &Path) -> Result> { - let _ = register_image_decoder_hooks(); +fn decode_adapter(input_path: &Path) -> Result> { + let _ = register_image_decoder_hooks_with_guardrails(DecodeGuardrails { + max_input_bytes: Some(128 * 1024 * 1024), + max_pixels: Some(256_000_000), + max_temp_spool_bytes: Some(256 * 1024 * 1024), + temp_spool_directory: None, + }); + + let mut decoder = ImageReader::open(input_path)? + .with_guessed_format()? + .into_decoder()?; + let _icc_profile = decoder.icc_profile()?; + let mut limits = Limits::default(); + limits.reserve(decoder.total_bytes())?; + decoder.set_limits(limits)?; + Ok(DynamicImage::from_decoder(decoder)?) +} - let decoded = ImageReader::open(input_path)?.decode()?; +fn bench_adapter(input_path: &Path) -> Result> { + let decoded = decode_adapter(input_path)?; let (width, height) = (decoded.width(), decoded.height()); let checksum = match decoded { DynamicImage::ImageRgba8(buffer) => small_checksum(buffer.as_raw()), @@ -105,8 +121,7 @@ fn bench_rgb(input_path: &Path) -> Result> { } fn decode_adapter_rgb(input_path: &Path) -> Result> { - let _ = register_image_decoder_hooks(); - Ok(ImageReader::open(input_path)?.decode()?.into_rgb8()) + Ok(decode_adapter(input_path)?.into_rgb8()) } fn bench_adapter_rgb(input_path: &Path) -> Result> { diff --git a/src/heic-decoder/hevc/cabac.rs b/src/heic-decoder/hevc/cabac.rs index abc9e1b..cd39403 100644 --- a/src/heic-decoder/hevc/cabac.rs +++ b/src/heic-decoder/hevc/cabac.rs @@ -75,33 +75,55 @@ static LPS_TABLE: [[u8; 4]; 64] = [ [2, 2, 2, 2], ]; -/// Renormalization table -#[allow(dead_code)] -static RENORM_TABLE: [u8; 32] = [ - 6, 5, 4, 4, 3, 3, 3, 3, 2, 2, 2, 2, 2, 2, 2, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, -]; - /// State transition for MPS -static STATE_TRANS_MPS: [u8; 64] = [ +const STATE_TRANS_MPS: [u8; 64] = [ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 62, 63, ]; /// State transition for LPS -static STATE_TRANS_LPS: [u8; 64] = [ +const STATE_TRANS_LPS: [u8; 64] = [ 0, 0, 1, 2, 2, 4, 4, 5, 6, 7, 8, 9, 9, 11, 11, 12, 13, 13, 15, 15, 16, 16, 18, 18, 19, 19, 21, 21, 22, 22, 23, 24, 24, 25, 26, 26, 27, 27, 28, 29, 29, 30, 30, 30, 31, 32, 32, 33, 33, 33, 34, 34, 35, 35, 35, 36, 36, 36, 37, 37, 37, 38, 38, 63, ]; +const fn build_packed_mps_transitions() -> [u8; 128] { + let mut transitions = [0; 128]; + let mut packed = 0; + while packed < transitions.len() { + let state = packed >> 1; + let mps = packed & 1; + transitions[packed] = (STATE_TRANS_MPS[state] << 1) | mps as u8; + packed += 1; + } + transitions +} + +const fn build_packed_lps_transitions() -> [u8; 128] { + let mut transitions = [0; 128]; + let mut packed = 0; + while packed < transitions.len() { + let state = packed >> 1; + let mut mps = packed & 1; + if state == 0 { + mps ^= 1; + } + transitions[packed] = (STATE_TRANS_LPS[state] << 1) | mps as u8; + packed += 1; + } + transitions +} + +const PACKED_STATE_TRANS_MPS: [u8; 128] = build_packed_mps_transitions(); +const PACKED_STATE_TRANS_LPS: [u8; 128] = build_packed_lps_transitions(); + /// CABAC context model #[derive(Clone, Copy)] pub struct ContextModel { - /// State index (0-63) - state: u8, - /// Most probable symbol (0 or 1) - mps: u8, + /// State index (0-63) in bits 1-6 and the MPS value in bit 0. + state_mps: u8, } #[allow(dead_code)] @@ -122,12 +144,14 @@ impl ContextModel { ((63 - init_state) as u8, 0) }; - Self { state, mps } + Self { + state_mps: (state << 1) | mps, + } } /// Get the current context state and MPS pub fn get_state(&self) -> (u8, u8) { - (self.state, self.mps) + (self.state_mps >> 1, self.state_mps & 1) } /// Initialize context for a given slice QP @@ -140,11 +164,9 @@ impl ContextModel { let init_state = init_state.clamp(1, 126); if init_state >= 64 { - self.state = (init_state - 64) as u8; - self.mps = 1; + self.state_mps = ((init_state - 64) as u8) << 1 | 1; } else { - self.state = (63 - init_state) as u8; - self.mps = 0; + self.state_mps = ((63 - init_state) as u8) << 1; } } } @@ -164,8 +186,6 @@ pub struct CabacDecoder<'a> { value: u32, /// Bits needed before next byte read (negative means bits available) bits_needed: i32, - /// Bin counter for debug tracing - bin_counter: u32, } #[allow(dead_code)] @@ -198,7 +218,6 @@ impl<'a> CabacDecoder<'a> { range: 510, value: 0, bits_needed: 8, - bin_counter: 0, }; // Initialize value (matching libde265 exactly) @@ -223,19 +242,23 @@ impl<'a> CabacDecoder<'a> { /// Equivalent to libde265's init_CABAC_decoder_2(). pub fn reinit(&mut self) { self.range = 510; - self.bits_needed = 8; + // With fewer than two bytes left, the legacy per-bit refill consumed + // the available high byte and then normalized bits_needed to -8 on + // the first decoded bit. Represent that pending first-bit transition + // directly as -9 so batched renormalization remains exactly + // equivalent without a special case in its hot path. + self.bits_needed = -9; self.value = 0; let remaining = self.data.len() - self.byte_pos; if remaining > 0 { self.value = (self.data[self.byte_pos] as u32) << 8; self.byte_pos += 1; - self.bits_needed -= 8; } if remaining > 1 { self.value |= self.data[self.byte_pos] as u32; self.byte_pos += 1; - self.bits_needed -= 8; + self.bits_needed = -8; } } @@ -251,29 +274,14 @@ impl<'a> CabacDecoder<'a> { Ok(()) } - /// Read a single bit from the bitstream (for regular context decoding) - fn read_bit(&mut self) -> Result { - self.value <<= 1; - self.bits_needed += 1; - - if self.bits_needed >= 0 { - if self.byte_pos < self.data.len() { - self.bits_needed = -8; - self.value |= self.data[self.byte_pos] as u32; - self.byte_pos += 1; - } else { - self.bits_needed = -8; - } - } - - Ok(0) // Return value not used, just for error handling - } - /// Decode a single bin using context model - pub fn decode_bin(&mut self, ctx: &mut ContextModel) -> Result { - self.bin_counter += 1; + #[inline] + pub fn decode_bin(&mut self, ctx: &mut ContextModel) -> u8 { + let packed_state = ctx.state_mps as usize; + let state = packed_state >> 1; + let mps = (packed_state & 1) as u8; let q_range_idx = (self.range >> 6) & 3; - let lps_range = LPS_TABLE[ctx.state as usize][q_range_idx as usize] as u32; + let lps_range = LPS_TABLE[state][q_range_idx as usize] as u32; self.range -= lps_range; @@ -283,29 +291,25 @@ impl<'a> CabacDecoder<'a> { let bin_val; if self.value < scaled_range { // MPS path - bin_val = ctx.mps; - ctx.state = STATE_TRANS_MPS[ctx.state as usize]; + bin_val = mps; + ctx.state_mps = PACKED_STATE_TRANS_MPS[packed_state]; } else { // LPS path - bin_val = 1 - ctx.mps; + bin_val = mps ^ 1; self.value -= scaled_range; self.range = lps_range; - - if ctx.state == 0 { - ctx.mps = 1 - ctx.mps; - } - ctx.state = STATE_TRANS_LPS[ctx.state as usize]; + ctx.state_mps = PACKED_STATE_TRANS_LPS[packed_state]; } // Renormalize - self.renormalize()?; + self.renormalize(); - Ok(bin_val) + bin_val } /// Decode a bypass bin (equal probability) - libde265 compatible - pub fn decode_bypass(&mut self) -> Result { - self.bin_counter += 1; + #[inline] + pub fn decode_bypass(&mut self) -> u8 { self.value <<= 1; self.bits_needed += 1; @@ -322,19 +326,20 @@ impl<'a> CabacDecoder<'a> { let scaled_range = self.range << 7; if self.value >= scaled_range { self.value -= scaled_range; - Ok(1) + 1 } else { - Ok(0) + 0 } } /// Decode multiple bypass bins - pub fn decode_bypass_bits(&mut self, n: u8) -> Result { + #[inline] + pub fn decode_bypass_bits(&mut self, n: u8) -> u32 { let mut result = 0u32; for _ in 0..n { - result = (result << 1) | self.decode_bypass()? as u32; + result = (result << 1) | self.decode_bypass() as u32; } - Ok(result) + result } /// Decode Exp-Golomb coded value (EGk) using bypass bins @@ -342,7 +347,7 @@ impl<'a> CabacDecoder<'a> { let mut base = 0u32; let mut n = k; loop { - let bit = self.decode_bypass()?; + let bit = self.decode_bypass(); if bit == 0 { break; } @@ -352,40 +357,54 @@ impl<'a> CabacDecoder<'a> { return Err(HevcError::InvalidBitstream("EGk prefix too long")); } } - let suffix = self.decode_bypass_bits(n)?; + let suffix = self.decode_bypass_bits(n); Ok(base + suffix) } /// Decode a terminate bin (end of slice check) - pub fn decode_terminate(&mut self) -> Result { + pub fn decode_terminate(&mut self) -> u8 { self.range -= 2; let scaled_range = self.range << 7; if self.value >= scaled_range { - Ok(1) + 1 } else { - self.renormalize()?; - Ok(0) + self.renormalize(); + 0 } } /// Renormalize the decoder state - fn renormalize(&mut self) -> Result<()> { - while self.range < 256 { - self.range <<= 1; - // Shift value and read more bits - self.read_bit()?; + #[inline] + fn renormalize(&mut self) { + if self.range < 256 { + // `range` is non-zero and at most seven shifts are needed. Since + // `bits_needed` starts in -9..=-1, the batch crosses at most one + // byte boundary. (`-9` is the canonical short-reinit state and + // cannot cross within a seven-bit batch.) When a normal state + // crosses, place the new byte exactly where the remaining + // single-bit shifts would have moved it. + let shift = self.range.leading_zeros() - 23; + self.range <<= shift; + self.value <<= shift; + self.bits_needed += shift as i32; + if self.bits_needed >= 0 { + if self.byte_pos < self.data.len() { + self.value |= (self.data[self.byte_pos] as u32) << self.bits_needed; + self.byte_pos += 1; + } + self.bits_needed -= 8; + } } // Invariant: after renormalization, range >= 256 debug_assert!(self.range >= 256, "range {} < 256 after renorm", self.range); - Ok(()) } /// Decode unsigned Exp-Golomb code using bypass bins pub fn decode_eg(&mut self, k: u8) -> Result { // Count leading zeros let mut n = 0; - while self.decode_bypass()? != 0 { + while self.decode_bypass() != 0 { n += 1; if n > 31 { return Err(HevcError::CabacError("exp-golomb overflow")); @@ -394,7 +413,7 @@ impl<'a> CabacDecoder<'a> { let mut value = 0u32; for _ in 0..(n + k) { - value = (value << 1) | self.decode_bypass()? as u32; + value = (value << 1) | self.decode_bypass() as u32; } Ok((1 << n) - 1 + value) @@ -513,3 +532,79 @@ pub static INIT_VALUES: [u8; context::NUM_CONTEXTS] = [ 154, 154, 154, 154, 154, 154, 154, 154, // RES_SCALE_SIGN_FLAG (2) 154, 154, ]; + +#[cfg(test)] +mod tests { + use super::CabacDecoder; + + fn legacy_reinit(decoder: &mut CabacDecoder<'_>, byte_pos: usize) { + decoder.byte_pos = byte_pos; + decoder.range = 510; + decoder.bits_needed = 8; + decoder.value = 0; + + let remaining = decoder.data.len() - decoder.byte_pos; + if remaining > 0 { + decoder.value = (decoder.data[decoder.byte_pos] as u32) << 8; + decoder.byte_pos += 1; + decoder.bits_needed -= 8; + } + if remaining > 1 { + decoder.value |= decoder.data[decoder.byte_pos] as u32; + decoder.byte_pos += 1; + decoder.bits_needed -= 8; + } + } + + fn legacy_renormalize(decoder: &mut CabacDecoder<'_>) { + while decoder.range < 256 { + decoder.range <<= 1; + decoder.value <<= 1; + decoder.bits_needed += 1; + if decoder.bits_needed >= 0 { + decoder.bits_needed = -8; + if decoder.byte_pos < decoder.data.len() { + decoder.value |= decoder.data[decoder.byte_pos] as u32; + decoder.byte_pos += 1; + } + } + } + } + + #[test] + fn batched_renormalization_matches_legacy_after_short_reinit() { + let data = [0x12, 0x34, 0x56, 0x78]; + + for remaining in 0..=2 { + let byte_pos = data.len() - remaining; + for range in [128, 64, 32, 16, 8, 4, 2] { + let mut batched = CabacDecoder::new(&data).expect("CABAC test data"); + batched.seek_to(byte_pos).expect("valid byte offset"); + batched.range = range; + + let mut legacy = CabacDecoder::new(&data).expect("CABAC test data"); + legacy_reinit(&mut legacy, byte_pos); + legacy.range = range; + + batched.renormalize(); + legacy_renormalize(&mut legacy); + + assert_eq!( + ( + batched.range, + batched.value, + batched.bits_needed, + batched.byte_pos, + ), + ( + legacy.range, + legacy.value, + legacy.bits_needed, + legacy.byte_pos, + ), + "remaining={remaining}, range={range}" + ); + } + } + } +} diff --git a/src/heic-decoder/hevc/color_convert.rs b/src/heic-decoder/hevc/color_convert.rs index a035ed3..bd0944c 100644 --- a/src/heic-decoder/hevc/color_convert.rs +++ b/src/heic-decoder/hevc/color_convert.rs @@ -6,10 +6,1092 @@ use archmage::incant; use archmage::prelude::*; +#[cfg(target_arch = "aarch64")] +use core::arch::aarch64::{ + uint8x8x3_t, uint8x8x4_t, uint16x4_t, uint16x4x4_t, uint16x8_t, vaddq_f32, vaddq_s32, + vcombine_u16, vcvtq_f32_u32, vcvtq_s32_f32, vdup_n_u8, vdup_n_u16, vdupq_n_f32, vdupq_n_s32, + vfmaq_n_f32, vget_high_u16, vget_low_u16, vmaxq_s32, vminq_s32, vmlaq_n_s32, vmovl_u16, + vmulq_n_f32, vqmovn_u16, vqmovn_u32, vqmovun_s32, vreinterpretq_s32_u32, vreinterpretq_u32_s32, + vshlq_u32, vshrq_n_s32, vst3_u8, vst4_u8, vst4_u16, vsubq_f32, vsubq_s32, vzip1_u16, vzip2_u16, +}; +#[cfg(target_arch = "aarch64")] +use safe_unaligned_simd::aarch64::{vld1_u16, vld1q_u16}; + +/// Parameters for libheif's generic f32 matrix conversion. +/// +/// Full range uses offsets `(0, midpoint)` and scales `(1, 1)`. Limited +/// range uses libheif's exact `1.1689f` luma and `1.1429f` chroma scales. +#[derive(Clone, Copy)] +pub(crate) struct FloatMatrixParams { + pub(crate) y_offset: f32, + pub(crate) y_scale: f32, + pub(crate) chroma_midpoint: f32, + pub(crate) chroma_scale: f32, + pub(crate) r_cr: f32, + pub(crate) g_cb: f32, + pub(crate) g_cr: f32, + pub(crate) b_cb: f32, +} + // Explicit imports for safe SIMD load/store (can't glob-import alongside core::arch) #[cfg(target_arch = "x86_64")] use safe_unaligned_simd::x86_64::{_mm_loadu_si64, _mm_loadu_si128, _mm256_storeu_si256}; +/// Convert full-range 8-bit 4:2:0 planes with libheif's exact fp8 kernel. +/// +/// `channels` must be 3 or 4. The caller has already validated plane and +/// output lengths; keeping this kernel focused lets the same path serve both +/// RGB grid scratch buffers and the image adapter's RGBA buffer. +#[allow(clippy::too_many_arguments)] +pub(crate) fn convert_420_8bit_to_interleaved( + y_plane: &[u16], + cb_plane: &[u16], + cr_plane: &[u16], + width: usize, + height: usize, + chroma_width: usize, + r_cr: i32, + g_cb: i32, + g_cr: i32, + b_cb: i32, + channels: usize, + output: &mut [u8], +) { + convert_420_8bit_region_to_interleaved( + y_plane, + cb_plane, + cr_plane, + width, + chroma_width, + 0, + 0, + width, + height, + r_cr, + g_cb, + g_cr, + b_cb, + channels, + output, + ); +} + +/// Convert a rectangular region of full-range 8-bit 4:2:0 planes. +/// +/// Source coordinates remain in the uncropped planes, so an odd `x_start` +/// retains the original chroma phase. The result is tightly packed. +#[allow(clippy::too_many_arguments)] +pub(crate) fn convert_420_8bit_region_to_interleaved( + y_plane: &[u16], + cb_plane: &[u16], + cr_plane: &[u16], + y_stride: usize, + chroma_stride: usize, + x_start: usize, + y_start: usize, + width: usize, + height: usize, + r_cr: i32, + g_cb: i32, + g_cr: i32, + b_cb: i32, + channels: usize, + output: &mut [u8], +) { + assert!(matches!(channels, 3 | 4), "channels must be RGB or RGBA"); + let pixel_count = width + .checked_mul(height) + .expect("4:2:0 region pixel count must fit in usize"); + let output_len = pixel_count + .checked_mul(channels) + .expect("interleaved output length must fit in usize"); + let x_end = x_start + .checked_add(width) + .expect("4:2:0 region x extent must fit in usize"); + let y_end = y_start + .checked_add(height) + .expect("4:2:0 region y extent must fit in usize"); + let y_len = y_stride + .checked_mul(y_end) + .expect("4:2:0 luma extent must fit in usize"); + let chroma_rows = y_end.div_ceil(2); + let chroma_len = chroma_stride + .checked_mul(chroma_rows) + .expect("4:2:0 chroma extent must fit in usize"); + assert_eq!(output.len(), output_len, "interleaved output length"); + assert!(x_end <= y_stride, "4:2:0 region exceeds luma row"); + assert!(y_plane.len() >= y_len, "luma plane length"); + assert!( + chroma_stride >= x_end.div_ceil(2), + "4:2:0 region exceeds chroma row" + ); + assert!(cb_plane.len() >= chroma_len, "Cb plane length"); + assert!(cr_plane.len() >= chroma_len, "Cr plane length"); + incant!( + convert_420_8bit_region_to_interleaved( + y_plane, + cb_plane, + cr_plane, + y_stride, + chroma_stride, + x_start, + y_start, + width, + height, + r_cr, + g_cb, + g_cr, + b_cb, + channels, + output + ), + [neon, scalar] + ); +} + +#[allow(clippy::too_many_arguments)] +fn convert_420_8bit_region_to_interleaved_scalar( + _token: ScalarToken, + y_plane: &[u16], + cb_plane: &[u16], + cr_plane: &[u16], + y_stride: usize, + chroma_stride: usize, + x_start: usize, + y_start: usize, + width: usize, + height: usize, + r_cr: i32, + g_cb: i32, + g_cr: i32, + b_cb: i32, + channels: usize, + output: &mut [u8], +) { + for output_y in 0..height { + let source_y = y_start + output_y; + let y_row = source_y * y_stride; + let chroma_row = (source_y / 2) * chroma_stride; + for output_x in 0..width { + let source_x = x_start + output_x; + write_420_8bit_pixel( + y_plane[y_row + source_x], + cb_plane[chroma_row + source_x / 2], + cr_plane[chroma_row + source_x / 2], + r_cr, + g_cb, + g_cr, + b_cb, + channels, + &mut output[(output_y * width + output_x) * channels..], + ); + } + } +} + +#[inline(always)] +#[allow(clippy::too_many_arguments)] +fn write_420_8bit_pixel( + y_sample: u16, + cb_sample: u16, + cr_sample: u16, + r_cr: i32, + g_cb: i32, + g_cr: i32, + b_cb: i32, + channels: usize, + output: &mut [u8], +) { + let y_sample = i64::from(y_sample); + let cb = i64::from(cb_sample) - 128; + let cr = i64::from(cr_sample) - 128; + let r = y_sample + ((i64::from(r_cr) * cr + 128) >> 8); + let g = y_sample + ((i64::from(g_cb) * cb + i64::from(g_cr) * cr + 128) >> 8); + let b = y_sample + ((i64::from(b_cb) * cb + 128) >> 8); + output[0] = r.clamp(0, 255) as u8; + output[1] = g.clamp(0, 255) as u8; + output[2] = b.clamp(0, 255) as u8; + if channels == 4 { + output[3] = u8::MAX; + } +} + +#[cfg(target_arch = "aarch64")] +#[allow(clippy::too_many_arguments)] +#[arcane] +fn convert_420_8bit_region_to_interleaved_neon( + _token: NeonToken, + y_plane: &[u16], + cb_plane: &[u16], + cr_plane: &[u16], + y_stride: usize, + chroma_stride: usize, + x_start: usize, + y_start: usize, + width: usize, + height: usize, + r_cr: i32, + g_cb: i32, + g_cr: i32, + b_cb: i32, + channels: usize, + output: &mut [u8], +) { + // Plane samples are u16 even for an 8-bit stream. Bound coefficient + // magnitudes before using i32 lanes so unusual public color metadata + // retains the scalar implementation's i64 overflow behavior. + let max_delta = i64::from(u16::MAX) - 128; + let max_y = i64::from(u16::MAX); + let rounded = 128_i64; + let single_channel_fits = |coefficient: i32| { + i64::from(coefficient).abs() * max_delta + max_y + rounded <= i64::from(i32::MAX) + }; + let green_fits = (i64::from(g_cb).abs() + i64::from(g_cr).abs()) * max_delta + max_y + rounded + <= i64::from(i32::MAX); + if !single_channel_fits(r_cr) || !green_fits || !single_channel_fits(b_cb) { + convert_420_8bit_region_to_interleaved_scalar( + ScalarToken, + y_plane, + cb_plane, + cr_plane, + y_stride, + chroma_stride, + x_start, + y_start, + width, + height, + r_cr, + g_cb, + g_cr, + b_cb, + channels, + output, + ); + return; + } + + let zero = vdupq_n_s32(0); + let max_255 = vdupq_n_s32(255); + let center = vdupq_n_s32(128); + let rounding = vdupq_n_s32(128); + let opaque = vdup_n_u8(u8::MAX); + let x_end = x_start + width; + // Chroma pairs start on even luma coordinates. Peel an odd first sample + // so every SIMD group can duplicate four adjacent chroma samples. + let simd_start = x_start.next_multiple_of(2).min(x_end); + let simd_end = simd_start + (x_end - simd_start) / 8 * 8; + + for output_y in 0..height { + let source_y = y_start + output_y; + let y_row = source_y * y_stride; + let chroma_row = (source_y / 2) * chroma_stride; + + for source_x in x_start..simd_start { + let output_x = source_x - x_start; + write_420_8bit_pixel( + y_plane[y_row + source_x], + cb_plane[chroma_row + source_x / 2], + cr_plane[chroma_row + source_x / 2], + r_cr, + g_cb, + g_cr, + b_cb, + channels, + &mut output[(output_y * width + output_x) * channels..], + ); + } + + let mut source_x = simd_start; + while source_x < simd_end { + let y_values = vld1q_u16( + (&y_plane[y_row + source_x..y_row + source_x + 8]) + .try_into() + .unwrap(), + ); + let cb_values = vld1_u16( + (&cb_plane[chroma_row + source_x / 2..chroma_row + source_x / 2 + 4]) + .try_into() + .unwrap(), + ); + let cr_values = vld1_u16( + (&cr_plane[chroma_row + source_x / 2..chroma_row + source_x / 2 + 4]) + .try_into() + .unwrap(), + ); + let cb_values = vcombine_u16( + vzip1_u16(cb_values, cb_values), + vzip2_u16(cb_values, cb_values), + ); + let cr_values = vcombine_u16( + vzip1_u16(cr_values, cr_values), + vzip2_u16(cr_values, cr_values), + ); + + let y_lo = vreinterpretq_s32_u32(vmovl_u16(vget_low_u16(y_values))); + let y_hi = vreinterpretq_s32_u32(vmovl_u16(vget_high_u16(y_values))); + let cb_lo = vsubq_s32( + vreinterpretq_s32_u32(vmovl_u16(vget_low_u16(cb_values))), + center, + ); + let cb_hi = vsubq_s32( + vreinterpretq_s32_u32(vmovl_u16(vget_high_u16(cb_values))), + center, + ); + let cr_lo = vsubq_s32( + vreinterpretq_s32_u32(vmovl_u16(vget_low_u16(cr_values))), + center, + ); + let cr_hi = vsubq_s32( + vreinterpretq_s32_u32(vmovl_u16(vget_high_u16(cr_values))), + center, + ); + + let convert_half = |y_values, cb_values, cr_values| { + let r = vaddq_s32( + y_values, + vshrq_n_s32(vmlaq_n_s32(rounding, cr_values, r_cr), 8), + ); + let g_terms = vmlaq_n_s32(vmlaq_n_s32(rounding, cb_values, g_cb), cr_values, g_cr); + let g = vaddq_s32(y_values, vshrq_n_s32(g_terms, 8)); + let b = vaddq_s32( + y_values, + vshrq_n_s32(vmlaq_n_s32(rounding, cb_values, b_cb), 8), + ); + ( + vminq_s32(vmaxq_s32(r, zero), max_255), + vminq_s32(vmaxq_s32(g, zero), max_255), + vminq_s32(vmaxq_s32(b, zero), max_255), + ) + }; + 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 r = vqmovn_u16(vcombine_u16(vqmovun_s32(r_lo), vqmovun_s32(r_hi))); + let g = vqmovn_u16(vcombine_u16(vqmovun_s32(g_lo), vqmovun_s32(g_hi))); + let b = vqmovn_u16(vcombine_u16(vqmovun_s32(b_lo), vqmovun_s32(b_hi))); + let output_x = source_x - x_start; + let output_index = (output_y * width + output_x) * channels; + + // SAFETY: the caller-proven output length is `width * height * + // channels`; this iteration starts at an in-row group of eight + // pixels, and vst3/vst4 write exactly 8 * channels bytes. + unsafe { + let destination = output.as_mut_ptr().add(output_index); + if channels == 4 { + vst4_u8(destination, uint8x8x4_t(r, g, b, opaque)); + } else { + vst3_u8(destination, uint8x8x3_t(r, g, b)); + } + } + source_x += 8; + } + + for source_x in simd_end..x_end { + let output_x = source_x - x_start; + write_420_8bit_pixel( + y_plane[y_row + source_x], + cb_plane[chroma_row + source_x / 2], + cr_plane[chroma_row + source_x / 2], + r_cr, + g_cb, + g_cr, + b_cb, + channels, + &mut output[(output_y * width + output_x) * channels..], + ); + } + } +} + +/// Convert an 8-bit HEIC region through libheif's generic f32 matrix kernel. +/// +/// This covers limited-range 4:2:0 and full/limited 4:2:2 or 4:4:4, which do +/// not use libheif's specialized fixed-point 4:2:0 operation. Source +/// coordinates remain in the uncropped planes so odd crop origins retain +/// their original chroma phase. `subsample_x`/`subsample_y` must each be one +/// or two and `channels` must be three or four. +#[allow(clippy::too_many_arguments)] +pub(crate) fn convert_float_matrix_8bit_region_to_interleaved( + y_plane: &[u16], + cb_plane: &[u16], + cr_plane: &[u16], + y_stride: usize, + chroma_stride: usize, + subsample_x: usize, + subsample_y: usize, + x_start: usize, + y_start: usize, + width: usize, + height: usize, + params: FloatMatrixParams, + channels: usize, + output: &mut [u8], +) { + validate_float_matrix_region( + y_plane, + cb_plane, + cr_plane, + y_stride, + chroma_stride, + subsample_x, + subsample_y, + x_start, + y_start, + width, + height, + channels, + output.len(), + ); + incant!( + convert_float_matrix_8bit_region_to_interleaved( + y_plane, + cb_plane, + cr_plane, + y_stride, + chroma_stride, + subsample_x, + subsample_y, + x_start, + y_start, + width, + height, + params, + channels, + output + ), + [neon, scalar] + ); +} + +/// Convert a complete high-bit-depth HEIC image through libheif's generic +/// f32 matrix kernel and expand the clipped source samples to RGBA16. +#[allow(clippy::too_many_arguments)] +pub(crate) fn convert_float_matrix_region_to_rgba16( + y_plane: &[u16], + cb_plane: &[u16], + cr_plane: &[u16], + y_stride: usize, + chroma_stride: usize, + subsample_x: usize, + subsample_y: usize, + x_start: usize, + y_start: usize, + width: usize, + height: usize, + bit_depth: u8, + params: FloatMatrixParams, + output: &mut [u16], +) { + assert!((1..=16).contains(&bit_depth), "source bit depth"); + validate_float_matrix_region( + y_plane, + cb_plane, + cr_plane, + y_stride, + chroma_stride, + subsample_x, + subsample_y, + x_start, + y_start, + width, + height, + 4, + output.len(), + ); + incant!( + convert_float_matrix_region_to_rgba16( + y_plane, + cb_plane, + cr_plane, + y_stride, + chroma_stride, + subsample_x, + subsample_y, + x_start, + y_start, + width, + height, + bit_depth, + params, + output + ), + [neon, scalar] + ); +} + +#[allow(clippy::too_many_arguments)] +fn validate_float_matrix_region( + y_plane: &[u16], + cb_plane: &[u16], + cr_plane: &[u16], + y_stride: usize, + chroma_stride: usize, + subsample_x: usize, + subsample_y: usize, + x_start: usize, + y_start: usize, + width: usize, + height: usize, + channels: usize, + output_len: usize, +) { + assert!(matches!(subsample_x, 1 | 2), "horizontal subsampling"); + assert!(matches!(subsample_y, 1 | 2), "vertical subsampling"); + assert!(matches!(channels, 3 | 4), "channels must be RGB or RGBA"); + let x_end = x_start + .checked_add(width) + .expect("float-matrix region x extent must fit in usize"); + let y_end = y_start + .checked_add(height) + .expect("float-matrix region y extent must fit in usize"); + let expected_output_len = width + .checked_mul(height) + .and_then(|pixels| pixels.checked_mul(channels)) + .expect("float-matrix output length must fit in usize"); + let y_len = y_stride + .checked_mul(y_end) + .expect("float-matrix luma extent must fit in usize"); + let chroma_rows = y_end.div_ceil(subsample_y); + let chroma_len = chroma_stride + .checked_mul(chroma_rows) + .expect("float-matrix chroma extent must fit in usize"); + assert_eq!( + output_len, expected_output_len, + "float-matrix output length" + ); + assert!(x_end <= y_stride, "float-matrix region exceeds luma row"); + assert!(y_plane.len() >= y_len, "float-matrix luma plane length"); + assert!( + chroma_stride >= x_end.div_ceil(subsample_x), + "float-matrix region exceeds chroma row" + ); + assert!(cb_plane.len() >= chroma_len, "float-matrix Cb plane length"); + assert!(cr_plane.len() >= chroma_len, "float-matrix Cr plane length"); +} + +#[inline(always)] +fn float_matrix_fma(a: f32, b: f32, c: f32) -> f32 { + #[cfg(any( + target_arch = "aarch64", + all(target_arch = "x86_64", target_feature = "fma") + ))] + { + a.mul_add(b, c) + } + #[cfg(not(any( + target_arch = "aarch64", + all(target_arch = "x86_64", target_feature = "fma") + )))] + { + a * b + c + } +} + +#[inline(always)] +fn convert_float_matrix_pixel( + y_sample: u16, + cb_sample: u16, + cr_sample: u16, + bit_depth: u8, + params: FloatMatrixParams, +) -> (u16, u16, u16) { + let y = (f32::from(y_sample) - params.y_offset) * params.y_scale; + let cb = (f32::from(cb_sample) - params.chroma_midpoint) * params.chroma_scale; + let cr = (f32::from(cr_sample) - params.chroma_midpoint) * params.chroma_scale; + let r = float_matrix_fma(params.r_cr, cr, y); + let g = float_matrix_fma(params.g_cr, cr, float_matrix_fma(params.g_cb, cb, y)); + let b = float_matrix_fma(params.b_cb, cb, y); + let max = ((1_i32 << bit_depth) - 1).max(0); + let clip = |value: f32| ((value + 0.5) as i32).clamp(0, max) as u16; + (clip(r), clip(g), clip(b)) +} + +#[inline(always)] +fn scale_float_matrix_sample_to_u16(sample: u16, bit_depth: u8) -> u16 { + if bit_depth >= 16 { + return sample; + } + let shift = 16 - u32::from(bit_depth); + let value = u32::from(sample); + ((value << shift) | (value >> u32::from(bit_depth).saturating_sub(shift))) + .min(u32::from(u16::MAX)) as u16 +} + +#[allow(clippy::too_many_arguments)] +fn convert_float_matrix_8bit_region_to_interleaved_scalar( + _token: ScalarToken, + y_plane: &[u16], + cb_plane: &[u16], + cr_plane: &[u16], + y_stride: usize, + chroma_stride: usize, + subsample_x: usize, + subsample_y: usize, + x_start: usize, + y_start: usize, + width: usize, + height: usize, + params: FloatMatrixParams, + channels: usize, + output: &mut [u8], +) { + for output_y in 0..height { + let source_y = y_start + output_y; + let y_row = source_y * y_stride; + let chroma_row = (source_y / subsample_y) * chroma_stride; + for output_x in 0..width { + let source_x = x_start + output_x; + let chroma_index = chroma_row + source_x / subsample_x; + let (r, g, b) = convert_float_matrix_pixel( + y_plane[y_row + source_x], + cb_plane[chroma_index], + cr_plane[chroma_index], + 8, + params, + ); + let output_index = (output_y * width + output_x) * channels; + output[output_index] = r as u8; + output[output_index + 1] = g as u8; + output[output_index + 2] = b as u8; + if channels == 4 { + output[output_index + 3] = u8::MAX; + } + } + } +} + +#[allow(clippy::too_many_arguments)] +fn convert_float_matrix_region_to_rgba16_scalar( + _token: ScalarToken, + y_plane: &[u16], + cb_plane: &[u16], + cr_plane: &[u16], + y_stride: usize, + chroma_stride: usize, + subsample_x: usize, + subsample_y: usize, + x_start: usize, + y_start: usize, + width: usize, + height: usize, + bit_depth: u8, + params: FloatMatrixParams, + output: &mut [u16], +) { + for output_y in 0..height { + let source_y = y_start + output_y; + let y_row = source_y * y_stride; + let chroma_row = (source_y / subsample_y) * chroma_stride; + for output_x in 0..width { + let source_x = x_start + output_x; + let chroma_index = chroma_row + source_x / subsample_x; + let (r, g, b) = convert_float_matrix_pixel( + y_plane[y_row + source_x], + cb_plane[chroma_index], + cr_plane[chroma_index], + bit_depth, + params, + ); + let output_index = (output_y * width + output_x) * 4; + output[output_index] = scale_float_matrix_sample_to_u16(r, bit_depth); + output[output_index + 1] = scale_float_matrix_sample_to_u16(g, bit_depth); + output[output_index + 2] = scale_float_matrix_sample_to_u16(b, bit_depth); + output[output_index + 3] = u16::MAX; + } + } +} + +#[cfg(target_arch = "aarch64")] +#[inline] +fn float_matrix_neon_safe(params: FloatMatrixParams) -> bool { + let values = [ + params.y_offset, + params.y_scale, + params.chroma_midpoint, + params.chroma_scale, + params.r_cr, + params.g_cb, + params.g_cr, + params.b_cb, + ]; + if !values.iter().all(|value| value.is_finite()) { + return false; + } + + // FCVTZS does not have Rust's saturating float-to-int semantics for + // infinities/out-of-range values. Public matrix metadata can derive + // unusual coefficients, so conservatively retain the scalar path unless + // every possible u16 input remains in the ordinary i32 conversion range. + let input_max = f64::from(u16::MAX); + let y_bound = (input_max + f64::from(params.y_offset).abs()) * f64::from(params.y_scale).abs(); + let chroma_bound = (input_max + f64::from(params.chroma_midpoint).abs()) + * f64::from(params.chroma_scale).abs(); + let channel_bound = y_bound + + chroma_bound + * f64::from( + params + .r_cr + .abs() + .max(params.g_cb.abs() + params.g_cr.abs()) + .max(params.b_cb.abs()), + ); + channel_bound.is_finite() && channel_bound + 1.0 < f64::from(i32::MAX) +} + +#[cfg(target_arch = "aarch64")] +#[inline(always)] +fn load_float_matrix_chroma_8(plane: &[u16], index: usize, subsample_x: usize) -> uint16x8_t { + // SAFETY: only called from the runtime-dispatched NEON kernel; slices are + // bounds-checked before the unaligned vector loads. + unsafe { + if subsample_x == 1 { + return vld1q_u16((&plane[index..index + 8]).try_into().unwrap()); + } + let values = vld1_u16((&plane[index..index + 4]).try_into().unwrap()); + vcombine_u16(vzip1_u16(values, values), vzip2_u16(values, values)) + } +} + +#[cfg(target_arch = "aarch64")] +#[inline(always)] +fn load_float_matrix_chroma_4(plane: &[u16], index: usize, subsample_x: usize) -> uint16x4_t { + // SAFETY: only called from the runtime-dispatched NEON kernel; slices are + // bounds-checked before the unaligned vector loads. + unsafe { + if subsample_x == 1 { + return vld1_u16((&plane[index..index + 4]).try_into().unwrap()); + } + let duplicated = [ + plane[index], + plane[index], + plane[index + 1], + plane[index + 1], + ]; + vld1_u16(&duplicated) + } +} + +#[cfg(target_arch = "aarch64")] +#[inline(always)] +fn convert_float_matrix_neon_4( + y: uint16x4_t, + cb: uint16x4_t, + cr: uint16x4_t, + params: FloatMatrixParams, + max: core::arch::aarch64::int32x4_t, +) -> ( + core::arch::aarch64::int32x4_t, + core::arch::aarch64::int32x4_t, + core::arch::aarch64::int32x4_t, +) { + // SAFETY: only called from the runtime-dispatched NEON kernels. + unsafe { + let y = vmulq_n_f32( + vsubq_f32(vcvtq_f32_u32(vmovl_u16(y)), vdupq_n_f32(params.y_offset)), + params.y_scale, + ); + let cb = vmulq_n_f32( + vsubq_f32( + vcvtq_f32_u32(vmovl_u16(cb)), + vdupq_n_f32(params.chroma_midpoint), + ), + params.chroma_scale, + ); + let cr = vmulq_n_f32( + vsubq_f32( + vcvtq_f32_u32(vmovl_u16(cr)), + vdupq_n_f32(params.chroma_midpoint), + ), + params.chroma_scale, + ); + let r = vfmaq_n_f32(y, cr, params.r_cr); + let g = vfmaq_n_f32(vfmaq_n_f32(y, cb, params.g_cb), cr, params.g_cr); + let b = vfmaq_n_f32(y, cb, params.b_cb); + let zero = vdupq_n_s32(0); + let half = vdupq_n_f32(0.5); + let clip = |value| vminq_s32(vmaxq_s32(vcvtq_s32_f32(vaddq_f32(value, half)), zero), max); + (clip(r), clip(g), clip(b)) + } +} + +#[cfg(target_arch = "aarch64")] +#[allow(clippy::too_many_arguments)] +#[arcane] +fn convert_float_matrix_8bit_region_to_interleaved_neon( + _token: NeonToken, + y_plane: &[u16], + cb_plane: &[u16], + cr_plane: &[u16], + y_stride: usize, + chroma_stride: usize, + subsample_x: usize, + subsample_y: usize, + x_start: usize, + y_start: usize, + width: usize, + height: usize, + params: FloatMatrixParams, + channels: usize, + output: &mut [u8], +) { + if !float_matrix_neon_safe(params) { + convert_float_matrix_8bit_region_to_interleaved_scalar( + ScalarToken, + y_plane, + cb_plane, + cr_plane, + y_stride, + chroma_stride, + subsample_x, + subsample_y, + x_start, + y_start, + width, + height, + params, + channels, + output, + ); + return; + } + + let x_end = x_start + width; + let simd_start = if subsample_x == 2 { + x_start.next_multiple_of(2).min(x_end) + } else { + x_start + }; + let simd_end = simd_start + (x_end - simd_start) / 8 * 8; + let max = vdupq_n_s32(255); + let opaque = vdup_n_u8(u8::MAX); + + for output_y in 0..height { + let source_y = y_start + output_y; + let y_row = source_y * y_stride; + let chroma_row = (source_y / subsample_y) * chroma_stride; + + for source_x in x_start..simd_start { + let output_x = source_x - x_start; + let chroma_index = chroma_row + source_x / subsample_x; + let (r, g, b) = convert_float_matrix_pixel( + y_plane[y_row + source_x], + cb_plane[chroma_index], + cr_plane[chroma_index], + 8, + params, + ); + let output_index = (output_y * width + output_x) * channels; + output[output_index] = r as u8; + output[output_index + 1] = g as u8; + output[output_index + 2] = b as u8; + if channels == 4 { + output[output_index + 3] = u8::MAX; + } + } + + let mut source_x = simd_start; + while source_x < simd_end { + let y_values = vld1q_u16( + (&y_plane[y_row + source_x..y_row + source_x + 8]) + .try_into() + .unwrap(), + ); + let chroma_index = chroma_row + source_x / subsample_x; + let cb_values = load_float_matrix_chroma_8(cb_plane, chroma_index, subsample_x); + let cr_values = load_float_matrix_chroma_8(cr_plane, chroma_index, subsample_x); + let (r_lo, g_lo, b_lo) = convert_float_matrix_neon_4( + vget_low_u16(y_values), + vget_low_u16(cb_values), + vget_low_u16(cr_values), + params, + max, + ); + let (r_hi, g_hi, b_hi) = convert_float_matrix_neon_4( + vget_high_u16(y_values), + vget_high_u16(cb_values), + vget_high_u16(cr_values), + params, + max, + ); + let r = vqmovn_u16(vcombine_u16(vqmovun_s32(r_lo), vqmovun_s32(r_hi))); + let g = vqmovn_u16(vcombine_u16(vqmovun_s32(g_lo), vqmovun_s32(g_hi))); + let b = vqmovn_u16(vcombine_u16(vqmovun_s32(b_lo), vqmovun_s32(b_hi))); + let output_x = source_x - x_start; + let output_index = (output_y * width + output_x) * channels; + // SAFETY: validation proves the full tightly packed output size; + // this group is eight in-row pixels, so the stores write exactly + // eight times `channels` bytes within it. + unsafe { + let destination = output.as_mut_ptr().add(output_index); + if channels == 4 { + vst4_u8(destination, uint8x8x4_t(r, g, b, opaque)); + } else { + vst3_u8(destination, uint8x8x3_t(r, g, b)); + } + } + source_x += 8; + } + + for source_x in simd_end..x_end { + let output_x = source_x - x_start; + let chroma_index = chroma_row + source_x / subsample_x; + let (r, g, b) = convert_float_matrix_pixel( + y_plane[y_row + source_x], + cb_plane[chroma_index], + cr_plane[chroma_index], + 8, + params, + ); + let output_index = (output_y * width + output_x) * channels; + output[output_index] = r as u8; + output[output_index + 1] = g as u8; + output[output_index + 2] = b as u8; + if channels == 4 { + output[output_index + 3] = u8::MAX; + } + } + } +} + +#[cfg(target_arch = "aarch64")] +#[inline(always)] +fn scale_float_matrix_neon_to_u16( + values: core::arch::aarch64::int32x4_t, + bit_depth: u8, +) -> uint16x4_t { + // SAFETY: only called from the runtime-dispatched NEON RGBA16 kernel. + unsafe { + let values = vreinterpretq_u32_s32(values); + if bit_depth >= 16 { + return vqmovn_u32(values); + } + let left = 16_i32 - i32::from(bit_depth); + let right = i32::from(bit_depth).saturating_sub(left); + let expanded = core::arch::aarch64::vorrq_u32( + vshlq_u32(values, vdupq_n_s32(left)), + vshlq_u32(values, vdupq_n_s32(-right)), + ); + vqmovn_u32(expanded) + } +} + +#[cfg(target_arch = "aarch64")] +#[allow(clippy::too_many_arguments)] +#[arcane] +fn convert_float_matrix_region_to_rgba16_neon( + _token: NeonToken, + y_plane: &[u16], + cb_plane: &[u16], + cr_plane: &[u16], + y_stride: usize, + chroma_stride: usize, + subsample_x: usize, + subsample_y: usize, + x_start: usize, + y_start: usize, + width: usize, + height: usize, + bit_depth: u8, + params: FloatMatrixParams, + output: &mut [u16], +) { + if !float_matrix_neon_safe(params) { + convert_float_matrix_region_to_rgba16_scalar( + ScalarToken, + y_plane, + cb_plane, + cr_plane, + y_stride, + chroma_stride, + subsample_x, + subsample_y, + x_start, + y_start, + width, + height, + bit_depth, + params, + output, + ); + return; + } + + let max_sample = ((1_i32 << bit_depth) - 1).max(0); + let max = vdupq_n_s32(max_sample); + let opaque = vdup_n_u16(u16::MAX); + let x_end = x_start + width; + let simd_start = if subsample_x == 2 { + x_start.next_multiple_of(2).min(x_end) + } else { + x_start + }; + let simd_end = simd_start + (x_end - simd_start) / 4 * 4; + for output_y in 0..height { + let source_y = y_start + output_y; + let y_row = source_y * y_stride; + let chroma_row = (source_y / subsample_y) * chroma_stride; + for source_x in x_start..simd_start { + let output_x = source_x - x_start; + let chroma_index = chroma_row + source_x / subsample_x; + let (r, g, b) = convert_float_matrix_pixel( + y_plane[y_row + source_x], + cb_plane[chroma_index], + cr_plane[chroma_index], + bit_depth, + params, + ); + let output_index = (output_y * width + output_x) * 4; + output[output_index] = scale_float_matrix_sample_to_u16(r, bit_depth); + output[output_index + 1] = scale_float_matrix_sample_to_u16(g, bit_depth); + output[output_index + 2] = scale_float_matrix_sample_to_u16(b, bit_depth); + output[output_index + 3] = u16::MAX; + } + + let mut source_x = simd_start; + while source_x < simd_end { + let y_values = vld1_u16( + (&y_plane[y_row + source_x..y_row + source_x + 4]) + .try_into() + .unwrap(), + ); + let chroma_index = chroma_row + source_x / subsample_x; + let cb_values = load_float_matrix_chroma_4(cb_plane, chroma_index, subsample_x); + let cr_values = load_float_matrix_chroma_4(cr_plane, chroma_index, subsample_x); + let (r, g, b) = + convert_float_matrix_neon_4(y_values, cb_values, cr_values, params, max); + let r = scale_float_matrix_neon_to_u16(r, bit_depth); + let g = scale_float_matrix_neon_to_u16(g, bit_depth); + let b = scale_float_matrix_neon_to_u16(b, bit_depth); + let output_x = source_x - x_start; + let output_index = (output_y * width + output_x) * 4; + // SAFETY: validation proves `width * height * 4` samples; this + // group is four in-row pixels and vst4 writes exactly 16 samples. + unsafe { + vst4_u16( + output.as_mut_ptr().add(output_index), + uint16x4x4_t(r, g, b, opaque), + ); + } + source_x += 4; + } + for source_x in simd_end..x_end { + let output_x = source_x - x_start; + let chroma_index = chroma_row + source_x / subsample_x; + let (r, g, b) = convert_float_matrix_pixel( + y_plane[y_row + source_x], + cb_plane[chroma_index], + cr_plane[chroma_index], + bit_depth, + params, + ); + let output_index = (output_y * width + output_x) * 4; + output[output_index] = scale_float_matrix_sample_to_u16(r, bit_depth); + output[output_index + 1] = scale_float_matrix_sample_to_u16(g, bit_depth); + output[output_index + 2] = scale_float_matrix_sample_to_u16(b, bit_depth); + output[output_index + 3] = u16::MAX; + } + } +} + /// Get color matrix coefficients for YCbCr→RGB conversion. /// /// Returns (cr_r, cb_g, cr_g, cb_b, y_bias, y_scale, rounding, shift_bits). @@ -347,3 +1429,379 @@ fn scalar_pixel( rgb[*out_idx + 2] = b.clamp(0, 255) as u8; *out_idx += 3; } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn full_range_420_dispatch_matches_scalar_for_rgb_and_rgba() { + let mut state = 0x7a31_4c95_u32; + let mut sample = || { + state = state.wrapping_mul(1_664_525).wrapping_add(1_013_904_223); + (state >> 16) as u16 + }; + + for (width, height) in [(1_usize, 1_usize), (7, 3), (8, 2), (9, 5), (24, 4)] { + let chroma_width = width.div_ceil(2); + let chroma_height = height.div_ceil(2); + let y_plane = (0..width * height).map(|_| sample()).collect::>(); + let cb_plane = (0..chroma_width * chroma_height) + .map(|_| sample()) + .collect::>(); + let cr_plane = (0..chroma_width * chroma_height) + .map(|_| sample()) + .collect::>(); + + for channels in [3, 4] { + let mut expected = vec![0_u8; width * height * channels]; + convert_420_8bit_region_to_interleaved_scalar( + ScalarToken, + &y_plane, + &cb_plane, + &cr_plane, + width, + chroma_width, + 0, + 0, + width, + height, + 403, + -48, + -120, + 475, + channels, + &mut expected, + ); + + let mut actual = vec![0_u8; expected.len()]; + convert_420_8bit_to_interleaved( + &y_plane, + &cb_plane, + &cr_plane, + width, + height, + chroma_width, + 403, + -48, + -120, + 475, + channels, + &mut actual, + ); + assert_eq!(actual, expected, "{width}x{height}, {channels} channels"); + } + } + } + + #[test] + fn full_range_420_cropped_region_preserves_odd_chroma_phase() { + let y_stride = 19_usize; + let source_height = 7_usize; + let chroma_stride = y_stride.div_ceil(2); + let mut state = 0x6d28_1b45_u32; + let mut sample = || { + state = state.wrapping_mul(1_664_525).wrapping_add(1_013_904_223); + (state >> 16) as u16 + }; + let y_plane = (0..y_stride * source_height) + .map(|_| sample()) + .collect::>(); + let cb_plane = (0..chroma_stride * source_height.div_ceil(2)) + .map(|_| sample()) + .collect::>(); + let cr_plane = (0..chroma_stride * source_height.div_ceil(2)) + .map(|_| sample()) + .collect::>(); + let (x_start, y_start, width, height) = (1, 1, 17, 5); + + for channels in [3, 4] { + let mut expected = vec![0_u8; width * height * channels]; + convert_420_8bit_region_to_interleaved_scalar( + ScalarToken, + &y_plane, + &cb_plane, + &cr_plane, + y_stride, + chroma_stride, + x_start, + y_start, + width, + height, + 403, + -48, + -120, + 475, + channels, + &mut expected, + ); + + let mut actual = vec![0_u8; expected.len()]; + convert_420_8bit_region_to_interleaved( + &y_plane, + &cb_plane, + &cr_plane, + y_stride, + chroma_stride, + x_start, + y_start, + width, + height, + 403, + -48, + -120, + 475, + channels, + &mut actual, + ); + assert_eq!(actual, expected, "{channels} channels"); + } + } + + fn full_float_params(bit_depth: u8) -> FloatMatrixParams { + FloatMatrixParams { + y_offset: 0.0, + y_scale: 1.0, + chroma_midpoint: (1_u32 << (bit_depth - 1)) as f32, + chroma_scale: 1.0, + r_cr: 1.5748, + g_cb: -0.187_324, + g_cr: -0.468_124, + b_cb: 1.8556, + } + } + + fn limited_float_params(bit_depth: u8) -> FloatMatrixParams { + FloatMatrixParams { + y_offset: (16_u32 << bit_depth.saturating_sub(8)) as f32, + y_scale: 1.1689, + chroma_midpoint: (1_u32 << (bit_depth - 1)) as f32, + chroma_scale: 1.1429, + r_cr: 1.402, + g_cb: -0.344_136, + g_cr: -0.714_136, + b_cb: 1.772, + } + } + + #[test] + fn float_matrix_8bit_dispatch_matches_scalar_for_layouts_ranges_and_tails() { + let source_width = 23_usize; + let source_height = 9_usize; + let mut state = 0xa1d3_7e29_u32; + let mut sample = || { + state = state.wrapping_mul(1_664_525).wrapping_add(1_013_904_223); + ((state >> 16) & 255) as u16 + }; + let y_plane = (0..source_width * source_height) + .map(|_| sample()) + .collect::>(); + + for (subsample_x, subsample_y) in [(2_usize, 2_usize), (2, 1), (1, 1)] { + let chroma_stride = source_width.div_ceil(subsample_x); + let chroma_height = source_height.div_ceil(subsample_y); + let cb_plane = (0..chroma_stride * chroma_height) + .map(|_| sample()) + .collect::>(); + let cr_plane = (0..chroma_stride * chroma_height) + .map(|_| sample()) + .collect::>(); + + for params in [full_float_params(8), limited_float_params(8)] { + for width in [1_usize, 3, 4, 7, 8, 9, 15, 16, 17] { + let (x_start, y_start, height) = (1_usize, 1_usize, 7_usize); + for channels in [3_usize, 4_usize] { + let mut expected = vec![0_u8; width * height * channels]; + convert_float_matrix_8bit_region_to_interleaved_scalar( + ScalarToken, + &y_plane, + &cb_plane, + &cr_plane, + source_width, + chroma_stride, + subsample_x, + subsample_y, + x_start, + y_start, + width, + height, + params, + channels, + &mut expected, + ); + let mut actual = vec![0_u8; expected.len()]; + convert_float_matrix_8bit_region_to_interleaved( + &y_plane, + &cb_plane, + &cr_plane, + source_width, + chroma_stride, + subsample_x, + subsample_y, + x_start, + y_start, + width, + height, + params, + channels, + &mut actual, + ); + assert_eq!( + actual, expected, + "subsampling={subsample_x}x{subsample_y}, width={width}, channels={channels}" + ); + } + } + } + } + } + + #[test] + fn float_matrix_rgba16_dispatch_matches_scalar_and_exact_bit_replication() { + let source_width = 19_usize; + let source_height = 8_usize; + for bit_depth in [10_u8, 12_u8] { + let sample_max = (1_u16 << bit_depth) - 1; + let mut state = 0x59c8_0f13_u32 ^ u32::from(bit_depth); + let mut sample = || { + state = state.wrapping_mul(1_664_525).wrapping_add(1_013_904_223); + ((state >> 16) as u16) & sample_max + }; + let y_plane = (0..source_width * source_height) + .map(|_| sample()) + .collect::>(); + + for (subsample_x, subsample_y) in [(2_usize, 2_usize), (2, 1), (1, 1)] { + let chroma_stride = source_width.div_ceil(subsample_x); + let chroma_height = source_height.div_ceil(subsample_y); + let cb_plane = (0..chroma_stride * chroma_height) + .map(|_| sample()) + .collect::>(); + let cr_plane = (0..chroma_stride * chroma_height) + .map(|_| sample()) + .collect::>(); + + for params in [ + full_float_params(bit_depth), + limited_float_params(bit_depth), + ] { + for width in [1_usize, 3, 4, 5, 7, 8, 9, 13, 16, 17] { + let (x_start, y_start, height) = (1_usize, 1_usize, 6_usize); + let mut expected = vec![0_u16; width * height * 4]; + convert_float_matrix_region_to_rgba16_scalar( + ScalarToken, + &y_plane, + &cb_plane, + &cr_plane, + source_width, + chroma_stride, + subsample_x, + subsample_y, + x_start, + y_start, + width, + height, + bit_depth, + params, + &mut expected, + ); + let mut actual = vec![0_u16; expected.len()]; + convert_float_matrix_region_to_rgba16( + &y_plane, + &cb_plane, + &cr_plane, + source_width, + chroma_stride, + subsample_x, + subsample_y, + x_start, + y_start, + width, + height, + bit_depth, + params, + &mut actual, + ); + assert_eq!( + actual, expected, + "depth={bit_depth}, subsampling={subsample_x}x{subsample_y}, width={width}" + ); + assert!(actual.chunks_exact(4).all(|pixel| pixel[3] == u16::MAX)); + } + } + } + + for sample in [0_u16, 1, sample_max / 2, sample_max - 1, sample_max] { + let shift = 16 - u32::from(bit_depth); + let expected = ((u32::from(sample) << shift) + | (u32::from(sample) >> u32::from(bit_depth).saturating_sub(shift))) + as u16; + assert_eq!( + scale_float_matrix_sample_to_u16(sample, bit_depth), + expected + ); + } + } + } + + #[test] + fn float_matrix_dispatch_falls_back_for_non_finite_and_extreme_coefficients() { + let y_plane = [0_u16, 1, 16, 128, 254, 255, 65_534, 65_535]; + let cb_plane = [0_u16, 16, 128, 255, 512, 1023, 4095, 65_535]; + let cr_plane = [65_535_u16, 4095, 1023, 512, 255, 128, 16, 0]; + let mut variants = [full_float_params(8); 4]; + variants[0].r_cr = f32::NAN; + variants[1].g_cb = f32::INFINITY; + variants[2].g_cr = f32::NEG_INFINITY; + variants[3].b_cb = f32::MAX; + + for params in variants { + let mut expected = [0_u8; 8 * 4]; + convert_float_matrix_8bit_region_to_interleaved_scalar( + ScalarToken, + &y_plane, + &cb_plane, + &cr_plane, + 8, + 8, + 1, + 1, + 0, + 0, + 8, + 1, + params, + 4, + &mut expected, + ); + let mut actual = [0_u8; 8 * 4]; + convert_float_matrix_8bit_region_to_interleaved( + &y_plane, + &cb_plane, + &cr_plane, + 8, + 8, + 1, + 1, + 0, + 0, + 8, + 1, + params, + 4, + &mut actual, + ); + assert_eq!(actual, expected); + } + + #[cfg(target_arch = "aarch64")] + { + assert!(float_matrix_neon_safe(full_float_params(8))); + assert!(float_matrix_neon_safe(limited_float_params(12))); + assert!(!float_matrix_neon_safe(variants[0])); + assert!(!float_matrix_neon_safe(variants[1])); + assert!(!float_matrix_neon_safe(variants[2])); + assert!(!float_matrix_neon_safe(variants[3])); + } + } +} diff --git a/src/heic-decoder/hevc/ctu.rs b/src/heic-decoder/hevc/ctu.rs index 3e2ea6f..48b44f3 100644 --- a/src/heic-decoder/hevc/ctu.rs +++ b/src/heic-decoder/hevc/ctu.rs @@ -153,6 +153,10 @@ pub struct SliceContext<'a> { pub sao_map: SaoMap, /// Reusable residual buffer (inverse transform writes all elements, no re-zeroing needed) residual_buf: [i16; 1024], + /// Reusable sparse coefficient workspace. Only indices recorded in + /// `touched_coeffs` may be non-zero between residual decodes. + coeff_buf: [i16; 1024], + touched_coeffs: [u16; 1024], /// Reusable scaling matrix buffer scaling_buf: [u8; 1024], } @@ -291,6 +295,8 @@ impl<'a> SliceContext<'a> { current_qg_y: -1, sao_map: SaoMap::new(sps.pic_width_in_ctbs(), sps.pic_height_in_ctbs()), residual_buf: [0i16; 1024], + coeff_buf: [0i16; 1024], + touched_coeffs: [0u16; 1024], scaling_buf: [16u8; 1024], }) } @@ -362,7 +368,7 @@ impl<'a> SliceContext<'a> { } // Check for end of slice segment - let end_of_slice = self.cabac.decode_terminate()?; + let end_of_slice = self.cabac.decode_terminate(); se_trace("end_of_slice", end_of_slice as i64, &self.cabac); if end_of_slice != 0 { debug_trace!( @@ -386,7 +392,7 @@ impl<'a> SliceContext<'a> { // WPP: at row boundaries, decode end_of_sub_stream and reinit CABAC if wpp && self.ctb_y != prev_ctb_y { - let _eoss = self.cabac.decode_terminate()?; + let _eoss = self.cabac.decode_terminate(); let entry_point_index = self.ctb_y.saturating_sub(start_ctb_y + 1) as usize; if let Some(&offset) = self.header.entry_point_offsets.get(entry_point_index) { self.cabac.seek_to(offset as usize)?; @@ -451,7 +457,7 @@ impl<'a> SliceContext<'a> { let left_in_slice = ctb_addr_rs > slice_addr_rs; if left_in_slice { let ctx_idx = context::SAO_MERGE_FLAG; - sao_merge_left_flag = self.cabac.decode_bin(&mut self.ctx[ctx_idx])? != 0; + sao_merge_left_flag = self.cabac.decode_bin(&mut self.ctx[ctx_idx]) != 0; se_trace("sao_merge_left", sao_merge_left_flag as i64, &self.cabac); } } @@ -464,7 +470,7 @@ impl<'a> SliceContext<'a> { let up_in_slice = ctb_addr_rs >= pic_width_ctbs + slice_addr_rs; if up_in_slice { let ctx_idx = context::SAO_MERGE_FLAG; - sao_merge_up_flag = self.cabac.decode_bin(&mut self.ctx[ctx_idx])? != 0; + sao_merge_up_flag = self.cabac.decode_bin(&mut self.ctx[ctx_idx]) != 0; se_trace("sao_merge_up", sao_merge_up_flag as i64, &self.cabac); } } @@ -536,7 +542,7 @@ impl<'a> SliceContext<'a> { let mut signed_offsets = [0i8; 4]; for i in 0..4 { if offsets_abs[i] != 0 { - let sign = self.cabac.decode_bypass()?; + let sign = self.cabac.decode_bypass(); se_trace("sao_offset_sign", sign as i64, &self.cabac); let val = (offsets_abs[i] as i32 * offset_scale) as i8; signed_offsets[i] = if sign != 0 { -val } else { val }; @@ -544,7 +550,7 @@ impl<'a> SliceContext<'a> { } info.sao_offset_val[c_idx] = signed_offsets; - let band_pos = self.cabac.decode_bypass_bits(5)?; + let band_pos = self.cabac.decode_bypass_bits(5); se_trace("sao_band_position", band_pos as i64, &self.cabac); info.sao_band_position[c_idx] = band_pos as u8; } else { @@ -554,7 +560,7 @@ impl<'a> SliceContext<'a> { } if c_idx <= 1 { - let eo_class = self.cabac.decode_bypass_bits(2)?; + let eo_class = self.cabac.decode_bypass_bits(2); se_trace("sao_eo_class", eo_class as i64, &self.cabac); if c_idx == 0 { info.sao_eo_class[0] = eo_class as u8; @@ -578,11 +584,11 @@ impl<'a> SliceContext<'a> { /// Decode sao_type_idx: context bin + optional bypass bin fn decode_sao_type_idx(&mut self) -> Result { let ctx_idx = context::SAO_TYPE_IDX; - let bit0 = self.cabac.decode_bin(&mut self.ctx[ctx_idx])?; + let bit0 = self.cabac.decode_bin(&mut self.ctx[ctx_idx]); if bit0 == 0 { Ok(0) } else { - let bit1 = self.cabac.decode_bypass()?; + let bit1 = self.cabac.decode_bypass(); if bit1 == 0 { Ok(1) } else { Ok(2) } } } @@ -590,7 +596,7 @@ impl<'a> SliceContext<'a> { /// Decode truncated unary with bypass bins (for sao_offset_abs) fn decode_cabac_tu_bypass(&mut self, c_max: u32) -> Result { for i in 0..c_max { - let bit = self.cabac.decode_bypass()?; + let bit = self.cabac.decode_bypass(); if bit == 0 { return Ok(i); } @@ -768,7 +774,7 @@ impl<'a> SliceContext<'a> { } let ctx_idx = context::SPLIT_CU_FLAG + cond_l + cond_a; - let bin = self.cabac.decode_bin(&mut self.ctx[ctx_idx])?; + let bin = self.cabac.decode_bin(&mut self.ctx[ctx_idx]); se_trace("split_cu_flag", bin as i64, &self.cabac); Ok(bin != 0) @@ -804,7 +810,7 @@ impl<'a> SliceContext<'a> { // Decode transquant_bypass_flag if enabled self.cu_transquant_bypass_flag = if self.pps.transquant_bypass_enabled_flag { let ctx_idx = context::CU_TRANSQUANT_BYPASS_FLAG; - self.cabac.decode_bin(&mut self.ctx[ctx_idx])? != 0 + self.cabac.decode_bin(&mut self.ctx[ctx_idx]) != 0 } else { false }; @@ -1041,7 +1047,7 @@ impl<'a> SliceContext<'a> { { // Decode split_transform_flag let ctx_idx = context::SPLIT_TRANSFORM_FLAG + (5 - log2_size as usize).min(2); - let flag = self.cabac.decode_bin(&mut self.ctx[ctx_idx])? != 0; + let flag = self.cabac.decode_bin(&mut self.ctx[ctx_idx]) != 0; se_trace("split_transform", flag as i64, &self.cabac); flag } else if log2_size > log2_max_trafo_size || (intra_split_flag && trafo_depth == 0) { @@ -1071,10 +1077,10 @@ impl<'a> SliceContext<'a> { let cb = if cbf_cb_parent != 0 { let ctx_idx = context::CBF_CBCR + trafo_depth as usize; - let mut val = self.cabac.decode_bin(&mut self.ctx[ctx_idx])?; + let mut val = self.cabac.decode_bin(&mut self.ctx[ctx_idx]); se_trace("cbf_cb", val as i64, &self.cabac); if decode_extra_422_cbf { - let extra = self.cabac.decode_bin(&mut self.ctx[ctx_idx])?; + let extra = self.cabac.decode_bin(&mut self.ctx[ctx_idx]); se_trace("cbf_cb", extra as i64, &self.cabac); val |= extra << 1; } @@ -1084,10 +1090,10 @@ impl<'a> SliceContext<'a> { }; let cr = if cbf_cr_parent != 0 { let ctx_idx = context::CBF_CBCR + trafo_depth as usize; - let mut val = self.cabac.decode_bin(&mut self.ctx[ctx_idx])?; + let mut val = self.cabac.decode_bin(&mut self.ctx[ctx_idx]); se_trace("cbf_cr", val as i64, &self.cabac); if decode_extra_422_cbf { - let extra = self.cabac.decode_bin(&mut self.ctx[ctx_idx])?; + let extra = self.cabac.decode_bin(&mut self.ctx[ctx_idx]); se_trace("cbf_cr", extra as i64, &self.cabac); val |= extra << 1; } @@ -1215,7 +1221,7 @@ impl<'a> SliceContext<'a> { // Context: offset 0 if trafo_depth > 0, offset 1 if trafo_depth == 0 let ctx_offset = if trafo_depth == 0 { 1 } else { 0 }; let ctx_idx = context::CBF_LUMA + ctx_offset; - let cbf_luma = self.cabac.decode_bin(&mut self.ctx[ctx_idx])? != 0; + let cbf_luma = self.cabac.decode_bin(&mut self.ctx[ctx_idx]) != 0; se_trace("cbf_luma", cbf_luma as i64, &self.cabac); // Per H.265 7.3.8.11: decode cu_qp_delta before residuals @@ -1226,7 +1232,7 @@ impl<'a> SliceContext<'a> { { let cu_qp_delta_abs = self.decode_cu_qp_delta_abs()?; let cu_qp_delta_sign = if cu_qp_delta_abs != 0 { - self.cabac.decode_bypass()? + self.cabac.decode_bypass() } else { 0 }; @@ -1452,7 +1458,7 @@ impl<'a> SliceContext<'a> { fn decode_cu_qp_delta_abs(&mut self) -> Result { let first_bin = self .cabac - .decode_bin(&mut self.ctx[context::CU_QP_DELTA_ABS])?; + .decode_bin(&mut self.ctx[context::CU_QP_DELTA_ABS]); if first_bin == 0 { return Ok(0); } @@ -1460,7 +1466,7 @@ impl<'a> SliceContext<'a> { for _ in 0..4 { let bin = self .cabac - .decode_bin(&mut self.ctx[context::CU_QP_DELTA_ABS + 1])?; + .decode_bin(&mut self.ctx[context::CU_QP_DELTA_ABS + 1]); if bin == 0 { break; } @@ -1486,7 +1492,7 @@ impl<'a> SliceContext<'a> { frame: &mut DecodedFrame, ) -> Result<()> { // Decode coefficients via CABAC - let (mut coeff_buf, transform_skip) = residual::decode_residual( + let (num_nonzero, transform_skip) = residual::decode_residual( &mut self.cabac, &mut self.ctx, log2_size, @@ -1497,9 +1503,11 @@ impl<'a> SliceContext<'a> { self.pps.transform_skip_enabled_flag, x0, y0, + &mut self.coeff_buf, + &mut self.touched_coeffs, )?; - if coeff_buf.is_zero() { + if num_nonzero == 0 { return Ok(()); } @@ -1517,13 +1525,16 @@ impl<'a> SliceContext<'a> { // residual — no scaling, no inverse transform. if self.cu_transquant_bypass_flag { let residual = &mut self.residual_buf; - residual[..num_coeffs].copy_from_slice(&coeff_buf.coeffs[..num_coeffs]); + residual[..num_coeffs].copy_from_slice(&self.coeff_buf[..num_coeffs]); + for &idx in &self.touched_coeffs[..num_nonzero] { + self.coeff_buf[idx as usize] = 0; + } self.add_residual_to_plane(x0, y0, log2_size, c_idx, bit_depth, frame); return Ok(()); } // Dequantize coefficients in-place - let coeffs = &mut coeff_buf.coeffs; + let coeffs = &mut self.coeff_buf; let dequant_params = transform::DequantParams { qp, @@ -1590,6 +1601,12 @@ impl<'a> SliceContext<'a> { transform::inverse_transform(coeffs, residual, size, bit_depth, is_intra_4x4_luma); } + // Dequantization only changes non-zero inputs, so the dense list from + // residual parsing is sufficient to restore the workspace invariant. + for &idx in &self.touched_coeffs[..num_nonzero] { + coeffs[idx as usize] = 0; + } + self.add_residual_to_plane(x0, y0, log2_size, c_idx, bit_depth, frame); Ok(()) @@ -1644,7 +1661,7 @@ impl<'a> SliceContext<'a> { if pred_mode == PredMode::Intra { // For intra, first bin distinguishes 2Nx2N from NxN let ctx_idx = context::PART_MODE; - let bin = self.cabac.decode_bin(&mut self.ctx[ctx_idx])?; + let bin = self.cabac.decode_bin(&mut self.ctx[ctx_idx]); se_trace("part_mode", bin as i64, &self.cabac); if bin != 0 { @@ -1698,7 +1715,7 @@ impl<'a> SliceContext<'a> { /// Decode prev_intra_luma_pred_flag (context-coded bin) fn decode_prev_intra_luma_pred_flag(&mut self) -> Result { let ctx_idx = context::PREV_INTRA_LUMA_PRED_FLAG; - let val = self.cabac.decode_bin(&mut self.ctx[ctx_idx])? != 0; + let val = self.cabac.decode_bin(&mut self.ctx[ctx_idx]) != 0; se_trace("prev_intra_luma_pred", val as i64, &self.cabac); Ok(val) } @@ -1739,14 +1756,14 @@ impl<'a> SliceContext<'a> { /// - If candidate mode collides with luma mode → Angular34 fn decode_intra_chroma_mode(&mut self, luma_mode: IntraPredMode) -> Result { let ctx_idx = context::INTRA_CHROMA_PRED_MODE; - let first_bin = self.cabac.decode_bin(&mut self.ctx[ctx_idx])?; + let first_bin = self.cabac.decode_bin(&mut self.ctx[ctx_idx]); let intra_chroma_mode = if first_bin == 0 { // Mode 4: derived from luma se_trace("intra_chroma_mode", 4, &self.cabac); luma_mode } else { // Read 2 fixed-length bypass bits for modes 0-3 - let mode_idx = self.cabac.decode_bypass_bits(2)? as u8; + let mode_idx = self.cabac.decode_bypass_bits(2) as u8; se_trace("intra_chroma_mode", mode_idx as i64, &self.cabac); let candidate = match mode_idx { @@ -1921,9 +1938,9 @@ impl<'a> SliceContext<'a> { /// Decode mpm_idx (0, 1, or 2) fn decode_mpm_idx(&mut self) -> Result { // Truncated unary: 0, 10, 11 - let val = if self.cabac.decode_bypass()? == 0 { + let val = if self.cabac.decode_bypass() == 0 { 0 - } else if self.cabac.decode_bypass()? == 0 { + } else if self.cabac.decode_bypass() == 0 { 1 } else { 2 @@ -1936,7 +1953,7 @@ impl<'a> SliceContext<'a> { fn decode_rem_intra_luma_pred_mode(&mut self) -> Result { let mut val = 0u32; for _ in 0..5 { - val = (val << 1) | self.cabac.decode_bypass()? as u32; + val = (val << 1) | self.cabac.decode_bypass() as u32; } se_trace("rem_intra_luma", val as i64, &self.cabac); Ok(val) diff --git a/src/heic-decoder/hevc/residual.rs b/src/heic-decoder/hevc/residual.rs index 402c74b..7e5fae7 100644 --- a/src/heic-decoder/hevc/residual.rs +++ b/src/heic-decoder/hevc/residual.rs @@ -124,74 +124,37 @@ pub static SCAN_ORDER_4X4_VERT: [(u8, u8); 16] = [ (3, 3), ]; -/// Get scan order table for 4x4 sub-blocks -pub fn get_scan_4x4(order: ScanOrder) -> &'static [(u8, u8); 16] { - match order { - ScanOrder::Diagonal => &SCAN_ORDER_4X4_DIAG, - ScanOrder::Horizontal => &SCAN_ORDER_4X4_HORIZ, - ScanOrder::Vertical => &SCAN_ORDER_4X4_VERT, - } -} - -/// Coefficient buffer for a transform unit -#[derive(Clone)] -pub struct CoeffBuffer { - /// Coefficients for this TU - pub coeffs: [i16; MAX_COEFF], - /// Transform size (log2) - pub log2_size: u8, - /// Number of non-zero coefficients - pub num_nonzero: u16, -} +/// Raster position (`y * 4 + x`) for each position in the three 4x4 scans. +/// Residual significance decoding uses this compact form on every CABAC bin; +/// keep the coordinate-pair tables above for callers that need coordinates. +const SCAN_RASTER_4X4: [[u8; 16]; 3] = [ + [0, 4, 1, 8, 5, 2, 12, 9, 6, 3, 13, 10, 7, 14, 11, 15], + [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15], + [0, 4, 8, 12, 1, 5, 9, 13, 2, 6, 10, 14, 3, 7, 11, 15], +]; -impl Default for CoeffBuffer { - fn default() -> Self { - Self { - coeffs: [0; MAX_COEFF], - log2_size: 2, - num_nonzero: 0, +const fn invert_4x4_scans(scans: [[u8; 16]; 3]) -> [[u8; 16]; 3] { + let mut inverse = [[0u8; 16]; 3]; + let mut scan_idx = 0; + while scan_idx < scans.len() { + let mut pos = 0; + while pos < scans[scan_idx].len() { + inverse[scan_idx][scans[scan_idx][pos] as usize] = pos as u8; + pos += 1; } + scan_idx += 1; } + inverse } -impl CoeffBuffer { - /// Create a new coefficient buffer - #[inline] - pub fn new(log2_size: u8) -> Self { - Self { - coeffs: [0; MAX_COEFF], - log2_size, - num_nonzero: 0, - } - } - - /// Get the transform size - #[inline] - pub fn size(&self) -> usize { - 1 << self.log2_size - } - - /// Get coefficient at position - #[allow(dead_code)] - #[inline] - pub fn get(&self, x: usize, y: usize) -> i16 { - let stride = self.size(); - self.coeffs[y * stride + x] - } - - /// Set coefficient at position - #[inline] - pub fn set(&mut self, x: usize, y: usize, value: i16) { - let stride = self.size(); - self.coeffs[y * stride + x] = value; - if value != 0 { - self.num_nonzero = self.num_nonzero.saturating_add(1); - } - } +const INVERSE_SCAN_4X4: [[u8; 16]; 3] = invert_4x4_scans(SCAN_RASTER_4X4); - /// Check if all coefficients are zero - pub fn is_zero(&self) -> bool { - self.num_nonzero == 0 +/// Get scan order table for 4x4 sub-blocks +pub fn get_scan_4x4(order: ScanOrder) -> &'static [(u8, u8); 16] { + match order { + ScanOrder::Diagonal => &SCAN_ORDER_4X4_DIAG, + ScanOrder::Horizontal => &SCAN_ORDER_4X4_HORIZ, + ScanOrder::Vertical => &SCAN_ORDER_4X4_VERT, } } @@ -213,7 +176,9 @@ pub fn decode_residual( transform_skip_enabled: bool, _x0: u32, _y0: u32, -) -> Result<(CoeffBuffer, bool)> { + coeffs: &mut [i16; MAX_COEFF], + touched_coeffs: &mut [u16; MAX_COEFF], +) -> Result<(usize, bool)> { #[cfg(feature = "decoder-tracing")] DEBUG_RESIDUAL_COUNTER.fetch_add(1, core::sync::atomic::Ordering::Relaxed); @@ -240,8 +205,8 @@ pub fn decode_residual( let rc_trace = false; let rcp = "RCX"; - let mut buffer = CoeffBuffer::new(log2_size); let size = 1u32 << log2_size; + let mut num_nonzero = 0usize; // Decode transform_skip_flag (H.265 7.3.8.11) // Per spec: if transform_skip_enabled_flag && !cu_transquant_bypass_flag @@ -249,7 +214,7 @@ pub fn decode_residual( // Log2MaxTransformSkipSize defaults to 2 (4x4 blocks only) let transform_skip = if transform_skip_enabled && !cu_transquant_bypass && log2_size <= 2 { let ctx_idx = context::TRANSFORM_SKIP_FLAG + if c_idx > 0 { 1 } else { 0 }; - let flag = cabac.decode_bin(&mut ctx[ctx_idx])? != 0; + let flag = cabac.decode_bin(&mut ctx[ctx_idx]) != 0; if rc_trace { let (range, _, _) = cabac.get_state_extended(); let (byte_pos, _, _) = cabac.get_position(); @@ -293,25 +258,30 @@ pub fn decode_residual( // Get scan tables let scan_sub = get_scan_sub_block(log2_size, scan_order); - let scan_pos = get_scan_4x4(scan_order); - // Convert ScanOrder to scan_idx for context derivation let scan_idx = match scan_order { ScanOrder::Diagonal => 0, ScanOrder::Horizontal => 1, ScanOrder::Vertical => 2, }; + let scan_raster = &SCAN_RASTER_4X4[scan_idx]; + let inverse_scan = &INVERSE_SCAN_4X4[scan_idx]; // Find last sub-block let sb_width = (size / 4) as usize; let last_sb_x = last_x / 4; let last_sb_y = last_y / 4; - let last_sb_idx = find_scan_pos(scan_sub, last_sb_x, last_sb_y, sb_width as u32); + let last_sb_idx = sub_block_scan_pos( + last_sb_x, + last_sb_y, + sb_width as u32, + scan_order == ScanOrder::Horizontal, + ); // Find last position within sub-block let local_x = (last_x % 4) as u8; let local_y = (last_y % 4) as u8; - let last_pos_in_sb = find_scan_pos_4x4(scan_pos, local_x, local_y); + let last_pos_in_sb = inverse_scan[(local_y * 4 + local_x) as usize]; if rc_trace { rc_eprintln!( @@ -328,7 +298,7 @@ pub fn decode_residual( // Track coded_sub_block_flag for prevCsbf calculation // Max sub-block grid is 8x8 for 32x32 TU - let mut coded_sb_flags = [[false; 8]; 8]; + let mut coded_sb_flags = 0u64; // Track whether the previously-processed subblock had any greater1_flag == 1 // This is used for ctx_set derivation per H.265 section 9.3.4.2.6 @@ -342,14 +312,15 @@ pub fn decode_residual( // Calculate neighbor flags BEFORE decoding coded_sub_block_flag // These are used for both coded_sub_block_flag and sig_coeff_flag contexts // Per H.265: bit 0 = right neighbor coded, bit 1 = below neighbor coded + let sb_bit = sb_y as usize * 8 + sb_x as usize; let csbf_neighbors = { let right_coded = if (sb_x as usize + 1) < sb_width { - coded_sb_flags[sb_y as usize][sb_x as usize + 1] + coded_sb_flags & (1 << (sb_bit + 1)) != 0 } else { false }; let below_coded = if (sb_y as usize + 1) < sb_width { - coded_sb_flags[sb_y as usize + 1][sb_x as usize] + coded_sb_flags & (1 << (sb_bit + 8)) != 0 } else { false }; @@ -360,8 +331,12 @@ pub fn decode_residual( // Middle sub-blocks need coded_sub_block_flag decoded // First (i=0) and last sub-blocks are always considered coded let (sb_coded, infer_sb_dc_sig) = if sb_idx > 0 && sb_idx < last_sb_idx { - // Use proper context derivation with neighbor info - let coded = decode_coded_sub_block_flag(cabac, ctx, c_idx, csbf_neighbors)?; + // Context offset: 0-1 for luma, 2-3 for chroma. Hoist this + // directly into the sub-block loop so the hot CABAC call does not + // carry an infallible Result/helper layer. + let csbf_ctx = usize::from(csbf_neighbors != 0); + let ctx_idx = context::CODED_SUB_BLOCK_FLAG + csbf_ctx + if c_idx > 0 { 2 } else { 0 }; + let coded = cabac.decode_bin(&mut ctx[ctx_idx]) != 0; // If sub-block is coded, we may need to infer DC later (coded, coded) } else { @@ -370,7 +345,7 @@ pub fn decode_residual( // Track coded sub-block flag if sb_coded { - coded_sb_flags[sb_y as usize][sb_x as usize] = true; + coded_sb_flags |= 1 << sb_bit; } // prevCsbf for sig_coeff_flag context @@ -389,19 +364,55 @@ pub fn decode_residual( 15 }; - let mut coeff_values = [0i16; 16]; - let mut coeff_flags = [false; 16]; - let mut num_coeffs = 0u8; + // Significance is decoded in reverse scan order. Append raster indices + // directly to the persistent sparse list in that same order; it then + // doubles as the dense level/sign worklist for this sub-block. + let subblock_start = num_nonzero; + let size = size as usize; + let sb_offset = sb_y as usize * 4 * size + sb_x as usize * 4; + let mut first_sig_pos = 0u8; + let mut last_sig_pos = 0u8; let mut can_infer_dc = infer_sb_dc_sig; + // For all transforms larger than 4x4, only the compact local-position + // contribution varies per significance bin. The component, transform, + // scan, and sub-block contributions are invariant here. + let sig_ctx_base = if log2_size == 2 { + context::SIG_COEFF_FLAG + if c_idx > 0 { 27 } else { 0 } + } else { + let mut base = context::SIG_COEFF_FLAG + if c_idx > 0 { 27 } else { 0 }; + if c_idx == 0 { + if sb_x + sb_y > 0 { + base += 3; + } + base += if log2_size == 3 { + if scan_idx == 0 { 9 } else { 15 } + } else { + 21 + }; + } else { + base += if log2_size == 3 { 9 } else { 12 }; + } + base + }; + let sig_ctx_local = if log2_size == 2 { + &CTX_IDX_MAP_4X4 + } else { + &SIG_CTX_LOCAL[prev_csbf as usize] + }; + // Determine the last position to check // For last sub-block: start from last_pos_in_sb (the known last significant coeff) // For other sub-blocks: start from position 15 let last_coeff = if sb_idx == last_sb_idx { // Set the known last significant coefficient (no need to decode sig_coeff_flag) - coeff_flags[start_pos as usize] = true; - coeff_values[start_pos as usize] = 1; - num_coeffs = 1; + let raster_pos = scan_raster[start_pos as usize] as usize; + let dst_idx = sb_offset + (raster_pos >> 2) * size + (raster_pos & 3); + coeffs[dst_idx] = 1; + touched_coeffs[num_nonzero] = dst_idx as u16; + num_nonzero += 1; + first_sig_pos = start_pos; + last_sig_pos = start_pos; can_infer_dc = false; // Can't infer DC if we have other coeffs // Then check positions from start_pos-1 down to 1 start_pos.saturating_sub(1) @@ -413,17 +424,16 @@ pub fn decode_residual( // Decode significant_coeff_flags for positions last_coeff down to 1 // (DC at position 0 is handled separately for inference) for n in (1..=last_coeff).rev() { - let sig = decode_sig_coeff_flag( - cabac, ctx, c_idx, n, log2_size, scan_idx, sb_x, sb_y, prev_csbf, scan_pos, - )?; + let raster_pos = scan_raster[n as usize] as usize; + let ctx_idx = sig_ctx_base + sig_ctx_local[raster_pos] as usize; + let sig = cabac.decode_bin(&mut ctx[ctx_idx]) != 0; if rc_trace { let (range, _, _) = cabac.get_state_extended(); let (byte_pos, _, _) = cabac.get_position(); - let (x_in_sb, y_in_sb) = scan_pos[n as usize]; + let x_in_sb = (raster_pos & 3) as u8; + let y_in_sb = (raster_pos >> 2) as u8; let xc = sb_x * 4 + x_in_sb; let yc = sb_y * 4 + y_in_sb; - let ctx_idx = - calc_sig_coeff_flag_ctx(xc, yc, log2_size, c_idx, scan_idx, prev_csbf); rc_eprintln!( "{rcp}_SIG n={} pos=({},{}) ctx={} val={} range={} byte={}", n, @@ -436,9 +446,14 @@ pub fn decode_residual( ); } if sig { - coeff_flags[n as usize] = true; - coeff_values[n as usize] = 1; - num_coeffs += 1; + let dst_idx = sb_offset + (raster_pos >> 2) * size + (raster_pos & 3); + coeffs[dst_idx] = 1; + touched_coeffs[num_nonzero] = dst_idx as u16; + num_nonzero += 1; + if num_nonzero == subblock_start + 1 { + last_sig_pos = n; + } + first_sig_pos = n; can_infer_dc = false; // Found a coefficient, can't infer DC } } @@ -448,27 +463,29 @@ pub fn decode_residual( // DC is inferred to be significant (otherwise the sub-block would be all zeros) if start_pos > 0 { if can_infer_dc { - coeff_flags[0] = true; - coeff_values[0] = 1; - num_coeffs += 1; + let dst_idx = sb_offset; + coeffs[dst_idx] = 1; + touched_coeffs[num_nonzero] = dst_idx as u16; + num_nonzero += 1; + if num_nonzero == subblock_start + 1 { + last_sig_pos = 0; + } + first_sig_pos = 0; if rc_trace { rc_eprintln!("{rcp}_SIG n=0 DC_INFERRED"); } } else { - let sig = decode_sig_coeff_flag( - cabac, ctx, c_idx, 0, log2_size, scan_idx, sb_x, sb_y, prev_csbf, scan_pos, - )?; + let ctx_idx = if log2_size == 2 { + sig_ctx_base + CTX_IDX_MAP_4X4[0] as usize + } else if sb_x == 0 && sb_y == 0 { + context::SIG_COEFF_FLAG + if c_idx > 0 { 27 } else { 0 } + } else { + sig_ctx_base + SIG_CTX_LOCAL[prev_csbf as usize][0] as usize + }; + let sig = cabac.decode_bin(&mut ctx[ctx_idx]) != 0; if rc_trace { let (range, _, _) = cabac.get_state_extended(); let (byte_pos, _, _) = cabac.get_position(); - let ctx_idx = calc_sig_coeff_flag_ctx( - sb_x * 4, - sb_y * 4, - log2_size, - c_idx, - scan_idx, - prev_csbf, - ); rc_eprintln!( "{rcp}_SIG n=0 pos=({},{}) ctx={} val={} range={} byte={}", sb_x * 4, @@ -480,13 +497,19 @@ pub fn decode_residual( ); } if sig { - coeff_flags[0] = true; - coeff_values[0] = 1; - num_coeffs += 1; + let dst_idx = sb_offset; + coeffs[dst_idx] = 1; + touched_coeffs[num_nonzero] = dst_idx as u16; + num_nonzero += 1; + if num_nonzero == subblock_start + 1 { + last_sig_pos = 0; + } + first_sig_pos = 0; } } } + let num_coeffs = num_nonzero - subblock_start; if num_coeffs == 0 { continue; } @@ -516,29 +539,26 @@ pub fn decode_residual( let mut last_greater1_flag = false; // Decode greater-1 flags (up to 8) - let mut first_g1_idx: Option = None; - let mut g1_positions = [false; 16]; // Track which positions have g1=1 - let max_g1 = (num_coeffs as usize).min(8); - let mut g1_count = 0; + let mut first_g1_idx: Option = None; + let max_g1 = num_coeffs.min(8); - // Track which coefficients need remaining level decoding - let mut needs_remaining = [false; 16]; + // Track dense coefficient indices that need remaining level decoding. + let mut needs_remaining = 0u16; - for n in (0..=start_pos).rev() { - if !coeff_flags[n as usize] { - continue; - } - - if g1_count >= max_g1 { + let g1_ctx_base = context::COEFF_ABS_LEVEL_GREATER1_FLAG + + if c_idx > 0 { 16 } else { 0 } + + (ctx_set as usize) * 4; + for coeff_idx in 0..num_coeffs { + if coeff_idx >= max_g1 { // Beyond first 8: base=1, always needs remaining for values > 1 - needs_remaining[n as usize] = true; + needs_remaining |= 1u16 << coeff_idx; continue; } // Update greater1Ctx BEFORE decoding, using PREVIOUS flag // (skip for first coefficient in subblock) // Per libde265: greater1Ctx increments without clamping (clamped in ctx calc) - if g1_count > 0 && greater1_ctx > 0 { + if coeff_idx > 0 && greater1_ctx > 0 { if last_greater1_flag { greater1_ctx = 0; } else { @@ -547,7 +567,8 @@ pub fn decode_residual( } // Use ctx_set (captured at subblock start) for ALL greater1_flags - let g1 = decode_coeff_greater1_flag(cabac, ctx, c_idx, ctx_set, greater1_ctx)?; + let ctx_idx = g1_ctx_base + (greater1_ctx as usize).min(3); + let g1 = cabac.decode_bin(&mut ctx[ctx_idx]) != 0; if rc_trace { let (range, _, _) = cabac.get_state_extended(); let (byte_pos, _, _) = cabac.get_position(); @@ -557,7 +578,7 @@ pub fn decode_residual( + (greater1_ctx as usize).min(3); rc_eprintln!( "{rcp}_G1 c={} ctxSet={} g1ctx={} fullCtx={} val={} range={} byte={}", - g1_count, + coeff_idx, ctx_set, greater1_ctx, full_ctx, @@ -569,23 +590,25 @@ pub fn decode_residual( last_greater1_flag = g1; if g1 { - coeff_values[n as usize] = 2; - g1_positions[n as usize] = true; + let dst_idx = touched_coeffs[subblock_start + coeff_idx] as usize; + coeffs[dst_idx] = 2; this_subblock_had_gt1 = true; // Track for next subblock's ctx_set if first_g1_idx.is_none() { - first_g1_idx = Some(n); + first_g1_idx = Some(coeff_idx); } else { // Non-first g1=1: base=2, needs remaining for values > 2 - needs_remaining[n as usize] = true; + needs_remaining |= 1u16 << coeff_idx; } } - g1_count += 1; } // Decode greater-2 flag (only for first coefficient with greater-1) // Uses same ctx_set as greater1_flags (captured at subblock start) if let Some(g1_idx) = first_g1_idx { - let g2 = decode_coeff_greater2_flag(cabac, ctx, c_idx, ctx_set)?; + let ctx_idx = context::COEFF_ABS_LEVEL_GREATER2_FLAG + + if c_idx > 0 { 4 } else { 0 } + + ctx_set as usize; + let g2 = cabac.decode_bin(&mut ctx[ctx_idx]) != 0; if rc_trace { let (range, _, _) = cabac.get_state_extended(); let (byte_pos, _, _) = cabac.get_position(); @@ -598,24 +621,11 @@ pub fn decode_residual( ); } if g2 { - coeff_values[g1_idx as usize] = 3; - needs_remaining[g1_idx as usize] = true; - } - } - - // Find first and last significant positions in this sub-block - let mut first_sig_pos = None; - let mut last_sig_pos = None; - for n in 0..=start_pos { - if coeff_flags[n as usize] { - if first_sig_pos.is_none() { - first_sig_pos = Some(n); - } - last_sig_pos = Some(n); + let dst_idx = touched_coeffs[subblock_start + g1_idx] as usize; + coeffs[dst_idx] = 3; + needs_remaining |= 1u16 << g1_idx; } } - let first_sig_pos = first_sig_pos.unwrap_or(0); - let last_sig_pos = last_sig_pos.unwrap_or(start_pos); // Determine if sign is hidden for this sub-block // Per H.265 9.3.4.3: sign is hidden if: @@ -630,16 +640,6 @@ pub fn decode_residual( // Following libde265's approach: decode signs in coefficient order (high scan pos to low) // The LAST coefficient (at first_sig_pos) has its sign hidden when sign_hidden=true // - // Build list of significant positions in reverse scan order (high to low) - let mut sig_positions = [0u8; 16]; - let mut n_sig = 0usize; - for n in (0..=start_pos).rev() { - if coeff_flags[n as usize] { - sig_positions[n_sig] = n; - n_sig += 1; - } - } - // Decode signs following libde265 order: // - Decode signs for coefficients 0 to n_sig-2 (high scan pos to low) // - For coefficient at n_sig-1 (lowest scan pos, first in scan order): @@ -648,25 +648,20 @@ pub fn decode_residual( // // This matches the H.265 spec and libde265: the FIRST coefficient in // scanning order (lowest scan position) has its sign hidden. - let mut coeff_signs = [0u8; 16]; - for (i, sign) in coeff_signs[..n_sig.saturating_sub(1)] - .iter_mut() - .enumerate() - { - *sign = cabac.decode_bypass()?; - let _ = i; - } - if n_sig > 0 && !sign_hidden { - coeff_signs[n_sig - 1] = cabac.decode_bypass()?; - } + let num_signs = num_coeffs - usize::from(sign_hidden); + let coeff_signs = cabac.decode_bypass_bits(num_signs as u8); if rc_trace { let (range, _, _) = cabac.get_state_extended(); let (byte_pos, _, _) = cabac.get_position(); + let mut signs = [0u8; 16]; + for (coeff_idx, sign) in signs[..num_signs].iter_mut().enumerate() { + *sign = ((coeff_signs >> (num_signs - 1 - coeff_idx)) & 1) as u8; + } rc_eprintln!( "{rcp}_SIGNS n_sig={} hidden={} signs={:?} range={} byte={}", - n_sig, + num_coeffs, sign_hidden, - &coeff_signs[..n_sig], + &signs[..num_coeffs], range, byte_pos ); @@ -677,65 +672,62 @@ pub fn decode_residual( let mut rice_param = 0u8; // Decode remaining levels for coefficients that need it - for n in (0..=start_pos).rev() { - if coeff_flags[n as usize] && needs_remaining[n as usize] { - let base = coeff_values[n as usize]; - let (remaining, new_rice) = - decode_coeff_abs_level_remaining(cabac, rice_param, base)?; - if rc_trace { - let (range, _, _) = cabac.get_state_extended(); - let (byte_pos, _, _) = cabac.get_position(); - rc_eprintln!( - "{rcp}_REM n={} base={} rice={} rem={} final={} range={} byte={}", - n, - base, - rice_param, - remaining, - base + remaining, - range, - byte_pos - ); - } - rice_param = new_rice; - coeff_values[n as usize] = base + remaining; + let mut remaining_flags = needs_remaining; + while remaining_flags != 0 { + let coeff_idx = remaining_flags.trailing_zeros() as usize; + remaining_flags &= remaining_flags - 1; + let dst_idx = touched_coeffs[subblock_start + coeff_idx] as usize; + let base = coeffs[dst_idx]; + let (remaining, new_rice) = decode_coeff_abs_level_remaining(cabac, rice_param, base)?; + if rc_trace { + let (range, _, _) = cabac.get_state_extended(); + let (byte_pos, _, _) = cabac.get_position(); + let raster_x = dst_idx % size; + let raster_y = dst_idx / size; + let local_raster = + (raster_y - sb_y as usize * 4) * 4 + raster_x - sb_x as usize * 4; + let n = inverse_scan[local_raster]; + rc_eprintln!( + "{rcp}_REM n={} base={} rice={} rem={} final={} range={} byte={}", + n, + base, + rice_param, + remaining, + base + remaining, + range, + byte_pos + ); } + rice_param = new_rice; + coeffs[dst_idx] = base + remaining; } // Apply signs and compute sum for parity inference - // At this point coeff_values[] are all positive (absolute values) + // At this point all workspace values are positive absolute levels. let mut sum_abs_level = 0i32; - for (i, &pos) in sig_positions[..n_sig].iter().enumerate() { - let pos = pos as usize; - if coeff_signs[i] != 0 { - coeff_values[pos] = -coeff_values[pos]; + for coeff_idx in 0..num_coeffs { + let dst_idx = touched_coeffs[subblock_start + coeff_idx] as usize; + let coeff_value = &mut coeffs[dst_idx]; + if coeff_idx < num_signs && coeff_signs & (1 << (num_signs - 1 - coeff_idx)) != 0 { + *coeff_value = -*coeff_value; } - sum_abs_level += coeff_values[pos] as i32; + sum_abs_level += *coeff_value as i32; - // Infer hidden sign at the last coefficient (first in scan order) - // Per H.265: if sum of signed coefficients is odd, flip the hidden sign - if i == n_sig - 1 && sign_hidden && (sum_abs_level & 1) != 0 { - coeff_values[pos] = -coeff_values[pos]; + // Infer the hidden sign at the first coefficient in scan order. + if coeff_idx + 1 == num_coeffs && sign_hidden && (sum_abs_level & 1) != 0 { + *coeff_value = -*coeff_value; } } - // Store coefficients in buffer - for (n, &(px, py)) in scan_pos.iter().enumerate() { - if coeff_flags[n] { - let x = sb_x as usize * 4 + px as usize; - let y = sb_y as usize * 4 + py as usize; - - buffer.set(x, y, coeff_values[n]); - - // Track large coefficients (indicates CABAC desync) - #[cfg(feature = "decoder-tracing")] - if coeff_values[n].abs() > 500 { - let (byte_pos, _, _) = cabac.get_position(); - debug::track_large_coeff(byte_pos); - } + #[cfg(feature = "decoder-tracing")] + for &dst_idx in &touched_coeffs[subblock_start..num_nonzero] { + // Track large coefficients (indicates CABAC desync) + if coeffs[dst_idx as usize].abs() > 500 { + let (byte_pos, _, _) = cabac.get_position(); + debug::track_large_coeff(byte_pos); } } - // Update prev_subblock_had_gt1 for the next subblock (lower scan index) prev_subblock_had_gt1 = this_subblock_had_gt1; } @@ -745,7 +737,7 @@ pub fn decode_residual( let (byte_pos, _, _) = cabac.get_position(); rc_eprintln!("{rcp}_END range={} byte={}", range, byte_pos); } - Ok((buffer, transform_skip)) + Ok((num_nonzero, transform_skip)) } /// Get sub-block scan order @@ -862,24 +854,21 @@ fn get_scan_sub_block(log2_size: u8, order: ScanOrder) -> &'static [(u8, u8)] { } } -/// Find position in scan order -fn find_scan_pos(scan: &[(u8, u8)], x: u32, y: u32, _width: u32) -> u32 { - for (i, &(sx, sy)) in scan.iter().enumerate() { - if sx as u32 == x && sy as u32 == y { - return i as u32; - } +/// Inverse of the sub-block scan tables. Only 8x8 luma can use the horizontal +/// 2x2 order; all larger sub-block grids use the diagonal order. +#[inline] +fn sub_block_scan_pos(x: u32, y: u32, width: u32, horizontal: bool) -> u32 { + if horizontal && width == 2 { + return y * width + x; } - 0 -} -/// Find position in 4x4 scan order -fn find_scan_pos_4x4(scan: &[(u8, u8); 16], x: u8, y: u8) -> u8 { - for (i, &(sx, sy)) in scan.iter().enumerate() { - if sx == x && sy == y { - return i as u8; - } + let diagonal = x + y; + if diagonal < width { + diagonal * (diagonal + 1) / 2 + x + } else { + let tail = 2 * width - 1 - diagonal; + width * width - tail * (tail + 1) / 2 + x - (diagonal + 1 - width) } - 0 } /// Decode last significant coefficient position @@ -896,7 +885,7 @@ fn decode_last_sig_coeff_pos( // Decode suffix if needed let x = if x_prefix > 3 { let n_bits = (x_prefix >> 1) - 1; - let suffix = cabac.decode_bypass_bits(n_bits as u8)?; + let suffix = cabac.decode_bypass_bits(n_bits as u8); ((2 + (x_prefix & 1)) << n_bits) + suffix } else { x_prefix @@ -904,7 +893,7 @@ fn decode_last_sig_coeff_pos( let y = if y_prefix > 3 { let n_bits = (y_prefix >> 1) - 1; - let suffix = cabac.decode_bypass_bits(n_bits as u8)?; + let suffix = cabac.decode_bypass_bits(n_bits as u8); ((2 + (y_prefix & 1)) << n_bits) + suffix } else { y_prefix @@ -946,7 +935,7 @@ fn decode_last_sig_coeff_prefix( let mut prefix = 0u32; while prefix < max_prefix as u32 { let ctx_idx = ctx_base + ctx_offset + (prefix as usize >> ctx_shift as usize); - let bin = cabac.decode_bin(&mut ctx[ctx_idx])?; + let bin = cabac.decode_bin(&mut ctx[ctx_idx]); if bin == 0 { break; } @@ -956,189 +945,19 @@ fn decode_last_sig_coeff_prefix( Ok(prefix) } -/// Decode coded_sub_block_flag -/// Per H.265 section 9.3.4.2.4, context depends on neighbor coded_sub_block_flags: -/// - csbfCtx = (csbf_right | csbf_below) ? 1 : 0 -/// - ctx_idx = base + csbfCtx + c_idx_offset -fn decode_coded_sub_block_flag( - cabac: &mut CabacDecoder<'_>, - ctx: &mut [ContextModel], - c_idx: u8, - csbf_neighbors: u8, // bit 0 = right, bit 1 = below -) -> Result { - // Context offset: 0-1 for luma, 2-3 for chroma - // csbfCtx = 1 if either neighbor is coded, else 0 - let csbf_ctx = if csbf_neighbors != 0 { 1 } else { 0 }; - let ctx_idx = context::CODED_SUB_BLOCK_FLAG + csbf_ctx + if c_idx > 0 { 2 } else { 0 }; - Ok(cabac.decode_bin(&mut ctx[ctx_idx])? != 0) -} - /// Context index map for 4x4 TU sig_coeff_flag (H.265 Table 9-41) /// Maps position (y*4 + x) to context index static CTX_IDX_MAP_4X4: [u8; 16] = [0, 1, 4, 5, 2, 3, 4, 5, 6, 6, 8, 8, 7, 7, 8, 8]; -/// Calculate sig_coeff_flag context index -/// -/// Per H.265 section 9.3.4.2.5, the context depends on: -/// - Position within sub-block (xP, yP) -/// - Sub-block position within TU (xS, yS) -/// - coded_sub_block_flag of neighbors (prevCsbf) -/// - Component (luma/chroma) -/// - TU size and scan order -/// -/// Parameters: -/// - x_c, y_c: coefficient position within TU -/// - log2_size: TU size (2=4x4, 3=8x8, 4=16x16, 5=32x32) -/// - c_idx: component (0=Y, 1=Cb, 2=Cr) -/// - scan_idx: scan order (0=diagonal, 1=horizontal, 2=vertical) -/// - prev_csbf: coded_sub_block_flag of neighbors (bit0=right, bit1=below per H.265/libde265) -fn calc_sig_coeff_flag_ctx( - x_c: u8, - y_c: u8, - log2_size: u8, - c_idx: u8, - scan_idx: u8, - prev_csbf: u8, -) -> usize { - let sb_width = 1u8 << (log2_size - 2); - - let sig_ctx = if sb_width == 1 { - // 4x4 TU: use lookup table - CTX_IDX_MAP_4X4[(y_c as usize * 4 + x_c as usize).min(15)] - } else if x_c == 0 && y_c == 0 { - // DC coefficient - 0 - } else { - // Sub-block and position within sub-block - let x_s = x_c >> 2; - let y_s = y_c >> 2; - let x_p = x_c & 3; - let y_p = y_c & 3; - - // Base context from position and neighbor flags - // Per libde265: prevCsbf bit0=right, bit1=below - let mut ctx = match prev_csbf { - 0 => { - // No coded neighbors: context based on position sum - if x_p + y_p >= 3 { - 0 - } else if x_p + y_p > 0 { - 1 - } else { - 2 - } - } - 1 => { - // Right neighbor coded (bit0=1): context based on y position - if y_p == 0 { - 2 - } else if y_p == 1 { - 1 - } else { - 0 - } - } - 2 => { - // Below neighbor coded (bit1=1): context based on x position - if x_p == 0 { - 2 - } else if x_p == 1 { - 1 - } else { - 0 - } - } - _ => { - // Both neighbors coded - 2 - } - }; - - if c_idx == 0 { - // Luma - if x_s + y_s > 0 { - ctx += 3; // Not first sub-block - } - - // Size-dependent offset - if sb_width == 2 { - // 8x8 TU - ctx += if scan_idx == 0 { 9 } else { 15 }; - } else { - // 16x16 or 32x32 TU - ctx += 21; - } - } else { - // Chroma - if sb_width == 2 { - // 8x8 TU - ctx += 9; - } else { - // 16x16 or larger TU - ctx += 12; - } - } - - ctx - }; - - // Final context index - context::SIG_COEFF_FLAG + if c_idx > 0 { 27 } else { 0 } + sig_ctx as usize -} - -/// Decode significant_coeff_flag with proper context derivation -#[allow(clippy::too_many_arguments)] -fn decode_sig_coeff_flag( - cabac: &mut CabacDecoder<'_>, - ctx: &mut [ContextModel], - c_idx: u8, - pos: u8, - log2_size: u8, - scan_idx: u8, - sb_x: u8, - sb_y: u8, - prev_csbf: u8, - scan_table: &[(u8, u8); 16], -) -> Result { - // Get coefficient position within TU from scan position - let (x_in_sb, y_in_sb) = scan_table[pos as usize]; - let x_c = sb_x * 4 + x_in_sb; - let y_c = sb_y * 4 + y_in_sb; - - let ctx_idx = calc_sig_coeff_flag_ctx(x_c, y_c, log2_size, c_idx, scan_idx, prev_csbf); - - Ok(cabac.decode_bin(&mut ctx[ctx_idx])? != 0) -} - -/// Decode coeff_abs_level_greater1_flag -/// Per H.265 9.3.4.2.6: context index = ctxSet * 4 + min(greater1Ctx, 3) -/// Plus 16 for chroma (c_idx > 0) -fn decode_coeff_greater1_flag( - cabac: &mut CabacDecoder<'_>, - ctx: &mut [ContextModel], - c_idx: u8, - ctx_set: u8, - greater1_ctx: u8, -) -> Result { - let ctx_idx = context::COEFF_ABS_LEVEL_GREATER1_FLAG - + if c_idx > 0 { 16 } else { 0 } - + (ctx_set as usize) * 4 - + (greater1_ctx as usize).min(3); - Ok(cabac.decode_bin(&mut ctx[ctx_idx])? != 0) -} - -/// Decode coeff_abs_level_greater2_flag -/// Per H.265: context index = ctxSet + (c_idx > 0 ? 4 : 0) -fn decode_coeff_greater2_flag( - cabac: &mut CabacDecoder<'_>, - ctx: &mut [ContextModel], - c_idx: u8, - ctx_set: u8, -) -> Result { - let ctx_idx = - context::COEFF_ABS_LEVEL_GREATER2_FLAG + if c_idx > 0 { 4 } else { 0 } + ctx_set as usize; - Ok(cabac.decode_bin(&mut ctx[ctx_idx])? != 0) -} +/// Local sig_coeff_flag context contribution for each `prevCsbf` value and +/// raster position. The component, transform-size, scan, and sub-block terms +/// are invariant across a sub-block and are added once in `decode_residual`. +static SIG_CTX_LOCAL: [[u8; 16]; 4] = [ + [2, 1, 1, 0, 1, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0], + [2, 2, 2, 2, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0], + [2, 1, 0, 0, 2, 1, 0, 0, 2, 1, 0, 0, 2, 1, 0, 0], + [2; 16], +]; /// Decode coeff_abs_level_remaining (Golomb-Rice with adaptive rice parameter) /// Returns (value, updated_rice_param) @@ -1149,14 +968,14 @@ fn decode_coeff_abs_level_remaining( ) -> Result<(i16, u8)> { // Decode prefix (unary part) let mut prefix = 0u32; - while cabac.decode_bypass()? != 0 && prefix < 32 { + while cabac.decode_bypass() != 0 && prefix < 32 { prefix += 1; } let value = if prefix <= 3 { // TR part only: value = (prefix << rice_param) + suffix let suffix = if rice_param > 0 { - cabac.decode_bypass_bits(rice_param)? + cabac.decode_bypass_bits(rice_param) } else { 0 }; @@ -1172,7 +991,7 @@ fn decode_coeff_abs_level_remaining( } // EGk part: suffix bits = prefix - 3 + rice_param let suffix_bits = (prefix - 3 + rice_param as u32) as u8; - let suffix = cabac.decode_bypass_bits(suffix_bits)?; + let suffix = cabac.decode_bypass_bits(suffix_bits); // value = (((1 << (prefix-3)) + 3 - 1) << rice_param) + suffix let base = ((1u32 << (prefix - 3)) + 2) << rice_param; let v = base + suffix; diff --git a/src/lib.rs b/src/lib.rs index aabd03e..ae48946 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -30,8 +30,6 @@ use rav1d::{ Dav1dResult, Rav1dError, dav1d_close, dav1d_data_create, dav1d_data_unref, dav1d_default_settings, dav1d_get_picture, dav1d_open, dav1d_picture_unref, dav1d_send_data, }; -#[cfg(feature = "parallel-grid")] -use rayon::prelude::*; use scuffle_h265::NALUnitType; use source::{ FileSource, RandomAccessSource, SourceReadError, TempFileSpoolOptions, TempFileSpoolSource, @@ -4350,59 +4348,178 @@ fn estimate_heic_grid_tile_decode_bytes( .max(1)) } -/// Choose a conservative number of simultaneously decoded tiles from each -/// tile's own preflight estimate. An estimate failure stops the parallel -/// window before that tile so normal row-major decoding still selects the -/// user-visible error. -#[cfg(feature = "parallel-grid")] -fn heic_grid_tile_decode_window_from_estimates( - estimates: impl IntoIterator>, - available_threads: usize, -) -> usize { - let available_threads = available_threads.max(1); - let mut window = 0_usize; - let mut estimated_bytes = 0_u64; - for tile_bytes in estimates.into_iter().take(available_threads) { - let Some(tile_bytes) = tile_bytes else { - break; - }; - let next_estimated_bytes = estimated_bytes.saturating_add(tile_bytes.max(1)); - if window > 0 && next_estimated_bytes > GRID_TILE_DECODE_MEMORY_BUDGET { - break; - } - window += 1; - estimated_bytes = next_estimated_bytes; - if estimated_bytes >= GRID_TILE_DECODE_MEMORY_BUDGET { - break; - } - } - window.max(1) -} - -/// Choose a conservative number of simultaneously decoded tiles. Large, -/// malformed, or unusual tiles stay sequential instead of inheriting the -/// first grid tile's geometry and multiplying peak memory. +/// Choose a pool-aware worker ceiling. Per-tile estimates enforce the memory +/// bound dynamically in the ordered streaming scheduler. #[cfg(feature = "parallel-grid")] fn heic_grid_tile_decode_window(remaining_tiles: &[isobmff::HeicGridTileItemData]) -> usize { if remaining_tiles.len() <= 1 || cfg!(feature = "decoder-tracing") { return 1; } - let available_threads = rayon::current_num_threads() + rayon::current_num_threads() .min(MAX_GRID_TILE_DECODE_THREADS) - .min(remaining_tiles.len()); - heic_grid_tile_decode_window_from_estimates( - remaining_tiles - .iter() - .map(|tile| estimate_heic_grid_tile_decode_bytes(tile).ok()), - available_threads, - ) + .min(remaining_tiles.len()) +} + +#[cfg(feature = "parallel-grid")] +enum HeicGridTileDecodeCompletion { + Decoded { + tile_index: usize, + result: Result, + }, + Panicked { + tile_index: usize, + payload: Box, + }, +} + +/// Decode independent inputs concurrently while exposing their results only +/// in input order. A completed result retains its estimated memory permit +/// until `consume` returns, so active jobs plus out-of-order completions stay +/// within both bounds. An input that cannot be estimated drains earlier work +/// and is decoded on the caller thread, preserving malformed-input order. +#[cfg(feature = "parallel-grid")] +#[allow(clippy::too_many_arguments)] +fn for_each_ordered_bounded_parallel( + inputs: &[I], + first_index: usize, + max_in_flight: usize, + memory_budget: u64, + mut estimate: impl FnMut(&I) -> Option, + decode: impl Fn(&I) -> Result + Sync, + consume: &mut impl FnMut(usize, T) -> Result<(), E>, +) -> Result<(), E> +where + I: Sync, + T: Send, + D: Send, + E: From, +{ + use std::collections::{BTreeMap, VecDeque}; + use std::sync::mpsc; + + if inputs.is_empty() { + return Ok(()); + } + + let max_in_flight = max_in_flight.max(1).min(inputs.len()); + let memory_budget = memory_budget.max(1); + let (sender, receiver) = mpsc::channel::>(); + + rayon::in_place_scope_fifo(|scope| { + let mut next_to_schedule = 0_usize; + let mut next_to_consume = 0_usize; + let mut retained_estimates = VecDeque::::with_capacity(max_in_flight); + let mut retained_estimated_bytes = 0_u64; + let mut ready = BTreeMap::>::new(); + let mut next_estimate = None::>; + + while next_to_consume < inputs.len() { + // Prefer a result already buffered for the next row-major + // position before admitting more work. Results that race with + // admission may still allow bounded speculative decoding, but + // they are never exposed to the consumer out of order. + while let Ok(completion) = receiver.try_recv() { + let tile_index = match &completion { + HeicGridTileDecodeCompletion::Decoded { tile_index, .. } + | HeicGridTileDecodeCompletion::Panicked { tile_index, .. } => *tile_index, + }; + ready.insert(tile_index, completion); + } + if let Some(completion) = ready.remove(&(first_index + next_to_consume)) { + match completion { + HeicGridTileDecodeCompletion::Decoded { result, .. } => { + consume(first_index + next_to_consume, result.map_err(E::from)?)?; + } + HeicGridTileDecodeCompletion::Panicked { payload, .. } => { + std::panic::resume_unwind(payload); + } + } + let released_estimate = retained_estimates + .pop_front() + .expect("a completed grid tile must retain one memory estimate"); + retained_estimated_bytes = + retained_estimated_bytes.saturating_sub(released_estimate); + next_to_consume += 1; + continue; + } + + // Fill only permits made available by ordered consumption. This + // avoids a completed later tile allowing another allocation while + // an earlier decoded tile is still retained. + while next_to_schedule < inputs.len() && retained_estimates.len() < max_in_flight { + let estimated_bytes = + next_estimate.get_or_insert_with(|| estimate(&inputs[next_to_schedule])); + let Some(estimated_bytes) = *estimated_bytes else { + break; + }; + let estimated_bytes = estimated_bytes.max(1); + let next_estimated_bytes = retained_estimated_bytes.saturating_add(estimated_bytes); + if !retained_estimates.is_empty() && next_estimated_bytes > memory_budget { + break; + } + + let relative_index = next_to_schedule; + let tile_index = first_index + relative_index; + let input = &inputs[relative_index]; + let sender = sender.clone(); + let decode = &decode; + scope.spawn_fifo(move |_| { + let completion = + match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + decode(input) + })) { + Ok(result) => { + HeicGridTileDecodeCompletion::Decoded { tile_index, result } + } + Err(payload) => HeicGridTileDecodeCompletion::Panicked { + tile_index, + payload, + }, + }; + // A callback panic can drop the receiver while sibling + // workers finish. That is normal scope unwinding, not a + // second panic from the worker. + let _ = sender.send(completion); + }); + + retained_estimates.push_back(estimated_bytes); + retained_estimated_bytes = next_estimated_bytes; + next_to_schedule += 1; + next_estimate = None; + } + + // An unestimable tile is deliberately not admitted beside other + // work. Once all preceding permits drain, decode it synchronously + // so its original decoder error remains the next visible result. + if next_to_schedule == next_to_consume && retained_estimates.is_empty() { + let tile_index = first_index + next_to_consume; + let tile = decode(&inputs[next_to_consume]).map_err(E::from)?; + consume(tile_index, tile)?; + next_to_consume += 1; + next_to_schedule += 1; + next_estimate = None; + continue; + } + + let completion = receiver + .recv() + .expect("a scoped grid worker must report completion before exiting"); + let tile_index = match &completion { + HeicGridTileDecodeCompletion::Decoded { tile_index, .. } + | HeicGridTileDecodeCompletion::Panicked { tile_index, .. } => *tile_index, + }; + ready.insert(tile_index, completion); + } + + Ok(()) + }) } -/// Decode grid tiles in bounded batches, but deliver them to the caller in -/// row-major order. Keeping validation and paste work on the caller thread -/// preserves deterministic errors and output while independent HEVC payloads -/// use the available cores. +/// Decode grid tiles with a bounded number of in-flight jobs, but deliver them +/// to the caller in row-major order. Keeping validation and paste work on the +/// caller thread preserves deterministic errors and output while independent +/// HEVC payloads use the available cores. fn for_each_decoded_heic_grid_tile( grid_data: &isobmff::HeicGridPrimaryItemData, first_tile: DecodedHeicImage, @@ -4421,28 +4538,24 @@ where #[cfg(feature = "parallel-grid")] { - let mut next_tile_index = 1_usize; - while next_tile_index < grid_data.tiles.len() { - let undecoded_tiles = &grid_data.tiles[next_tile_index..]; - let window = heic_grid_tile_decode_window(undecoded_tiles); - if window > 1 { - let decoded = undecoded_tiles[..window] - .par_iter() - .map(decode_heic_grid_tile_to_image) - .collect::>(); - - for (offset, tile) in decoded.into_iter().enumerate() { - consume(next_tile_index + offset, tile.map_err(E::from)?)?; - } - next_tile_index += window; - } else { - consume( - next_tile_index, - decode_heic_grid_tile_to_image(&grid_data.tiles[next_tile_index]) - .map_err(E::from)?, - )?; - next_tile_index += 1; - } + let window = heic_grid_tile_decode_window(remaining_tiles); + if window > 1 { + return for_each_ordered_bounded_parallel( + remaining_tiles, + 1, + window, + GRID_TILE_DECODE_MEMORY_BUDGET, + |tile| estimate_heic_grid_tile_decode_bytes(tile).ok(), + decode_heic_grid_tile_to_image, + &mut consume, + ); + } + + for (offset, tile) in remaining_tiles.iter().enumerate() { + consume( + offset + 1, + decode_heic_grid_tile_to_image(tile).map_err(E::from)?, + )?; } Ok(()) } @@ -8285,6 +8398,12 @@ fn decode_avif_bytes_to_rgba_layout( trait RgbaSampleOutput { fn sample_len(&self) -> usize; fn write_sample(&mut self, index: usize, sample: T); + + fn write_samples(&mut self, index: usize, samples: &[T]) { + for (offset, &sample) in samples.iter().enumerate() { + self.write_sample(index + offset, sample); + } + } } #[cfg(feature = "image-integration")] @@ -8299,6 +8418,10 @@ impl RgbaSampleOutput for SliceRgbaOutput<'_, T> { fn write_sample(&mut self, index: usize, sample: T) { self.0[index] = sample; } + + fn write_samples(&mut self, index: usize, samples: &[T]) { + self.0[index..index + samples.len()].copy_from_slice(samples); + } } #[cfg(feature = "image-integration")] @@ -8383,6 +8506,26 @@ fn decode_primary_heic_grid_to_rgba_output>, auxiliary_alpha: Option<&HeicAuxiliaryAlphaPlane>, convert_tile: fn(&DecodedHeicImage, &mut Vec) -> Result<(), DecodeHeicError>, scale_alpha: fn(u16, u8) -> T, @@ -8550,6 +8695,21 @@ fn paste_heic_grid_tiles_to_transformed_rgba_slice>( + tile_pixels: &[T], + tile_width: usize, + tile_height: usize, + output: &mut O, + source_width: usize, + source_height: usize, + destination_width: usize, + x_origin: usize, + y_origin: usize, + orientation: Option<&RgbaOrientationTransform>, +) -> Result<(), DecodeError> { + if x_origin >= source_width || y_origin >= source_height { + return Ok(()); + } + + let copy_width = tile_width.min(source_width - x_origin); + let copy_height = tile_height.min(source_height - y_origin); + let copy_samples = copy_width * 4; + + if orientation.is_none() { + debug_assert_eq!(source_width, destination_width); + for tile_y in 0..copy_height { + let source_sample = tile_y * tile_width * 4; + let destination_sample = ((y_origin + tile_y) * destination_width + x_origin) * 4; + output.write_samples( + destination_sample, + &tile_pixels[source_sample..source_sample + copy_samples], + ); + } + return Ok(()); + } + + let orientation = orientation.expect("orientation was checked above"); + debug_assert_eq!(destination_width, orientation.destination_width_usize); + let destination_step = (orientation.destination_y_from_source_x * destination_width as i64 + + orientation.destination_x_from_source_x) + * 4; + + for tile_y in 0..copy_height { + let source_y = y_origin + tile_y; + let source_sample = tile_y * tile_width * 4; + let (destination_x, destination_y) = orientation.map_source_pixel(x_origin, source_y)?; + let mut destination_sample = + ((destination_y * destination_width + destination_x) * 4) as i64; + + if destination_step == 4 { + output.write_samples( + destination_sample as usize, + &tile_pixels[source_sample..source_sample + copy_samples], + ); + continue; + } + + for tile_x in 0..copy_width { + let source_sample = source_sample + tile_x * 4; + debug_assert!(destination_sample >= 0); + let destination_sample_usize = destination_sample as usize; + debug_assert!(destination_sample_usize + 4 <= output.sample_len()); + output.write_samples( + destination_sample_usize, + &tile_pixels[source_sample..source_sample + 4], + ); + destination_sample += destination_step; + } + } + + Ok(()) +} + #[cfg(feature = "image-integration")] fn decoded_heic_to_rgba8_slice( decoded: DecodedHeicImage, @@ -8606,6 +8843,72 @@ fn decoded_heic_to_rgba8_slice( auxiliary_alpha: Option<&HeicAuxiliaryAlphaPlane>, out: &mut [u8], ) -> Result<(), DecodeError> { + // Identity and clean-aperture-only coded images remain rectangular source + // regions. Route those directly through the shared contiguous converter + // so its architecture-specific 4:2:0 kernel writes the caller's buffer + // without materializing a full-frame RGBA intermediate. + if auxiliary_alpha.is_none() || transforms.is_empty() { + let source_bit_depth = heic_bit_depth_for_png_conversion(&decoded)?; + if heic_storage_bit_depth(source_bit_depth) != 8 { + return Err(DecodeError::Unsupported(format!( + "HEIC storage is RGBA{}, not RGBA8", + heic_storage_bit_depth(source_bit_depth) + ))); + } + let transform_plan = + RgbaTransformPlan::from_primary_transforms(decoded.width, decoded.height, transforms)?; + let expected = checked_rgba_sample_count( + transform_plan.destination_width, + transform_plan.destination_height, + )?; + if out.len() != expected { + return Err(DecodeError::TransformGuard( + TransformGuardError::RgbaSampleCountMismatch { + stage: "HEIC direct transformed image adapter output", + actual: out.len(), + expected, + width: transform_plan.destination_width, + height: transform_plan.destination_height, + }, + )); + } + if transform_plan + .steps + .iter() + .all(|step| matches!(step, ResolvedRgbaTransformStep::CleanAperture { .. })) + { + if let Some(alpha) = auxiliary_alpha { + validate_auxiliary_alpha_plane(alpha, decoded.width, decoded.height)?; + } + let (source_x, source_y) = transform_plan.map_destination_pixel(0, 0)?; + let destination_width = transform_dimension_to_usize( + "HEIC clean-aperture adapter", + "destination width", + transform_plan.destination_width, + )?; + let destination_height = transform_dimension_to_usize( + "HEIC clean-aperture adapter", + "destination height", + transform_plan.destination_height, + )?; + convert_heic_to_interleaved_rgb8_region_slice::<4>( + &decoded, + source_x, + source_y, + destination_width, + destination_height, + out, + scale_sample_to_u8, + "RGBA8", + ) + .map_err(DecodeError::from)?; + if let Some(alpha) = auxiliary_alpha { + apply_auxiliary_alpha_to_rgba8(out, decoded.width, decoded.height, alpha)?; + } + return Ok(()); + } + } + let mut output = SliceRgbaOutput(out); decoded_heic_to_rgba_output( decoded, @@ -8625,6 +8928,36 @@ fn decoded_heic_to_rgba16_native_endian_bytes( auxiliary_alpha: Option<&HeicAuxiliaryAlphaPlane>, out: &mut [u8], ) -> Result<(), DecodeError> { + // `image` allocates the identity decode buffer through the global + // allocator, which is naturally aligned on the mobile targets. Reuse the + // contiguous RGBA16 converter in that common case so its high-bit-depth + // AArch64 float kernel writes directly into the hook's caller buffer. + // A deliberately unaligned public slice or any primary transform keeps + // the byte-wise generic path below. For identity alpha images, color is + // converted first and the validated auxiliary plane replaces A in-place. + if transforms.is_empty() { + let source_bit_depth = heic_bit_depth_for_png_conversion(&decoded)?; + if heic_storage_bit_depth(source_bit_depth) != 16 { + return Err(DecodeError::Unsupported(format!( + "HEIC storage is RGBA{}, not RGBA16", + heic_storage_bit_depth(source_bit_depth) + ))); + } + // SAFETY: `align_to_mut` only exposes the maximally aligned middle; + // conversion is used solely when it spans the entire caller slice. + let (prefix, samples, suffix) = unsafe { out.align_to_mut::() }; + if prefix.is_empty() && suffix.is_empty() { + if let Some(alpha) = auxiliary_alpha { + validate_auxiliary_alpha_plane(alpha, decoded.width, decoded.height)?; + } + convert_heic_to_rgba16_slice(&decoded, samples).map_err(DecodeError::from)?; + if let Some(alpha) = auxiliary_alpha { + apply_auxiliary_alpha_to_rgba16(samples, decoded.width, decoded.height, alpha)?; + } + return Ok(()); + } + } + let mut output = NativeEndianRgba16Output(out); decoded_heic_to_rgba_output( decoded, @@ -11734,6 +12067,37 @@ fn convert_heic_to_interleaved_rgb8_slice( out: &mut [u8], scale_sample: fn(u16, u8) -> u8, output_label: &str, +) -> Result<(), DecodeHeicError> { + let width = + usize::try_from(decoded.width).map_err(|_| DecodeHeicError::InvalidDecodedFrame { + detail: format!("HEIC width does not fit in usize ({})", decoded.width), + })?; + let height = + usize::try_from(decoded.height).map_err(|_| DecodeHeicError::InvalidDecodedFrame { + detail: format!("HEIC height does not fit in usize ({})", decoded.height), + })?; + convert_heic_to_interleaved_rgb8_region_slice::( + decoded, + 0, + 0, + width, + height, + out, + scale_sample, + output_label, + ) +} + +#[allow(clippy::too_many_arguments)] +fn convert_heic_to_interleaved_rgb8_region_slice( + decoded: &DecodedHeicImage, + source_x_start: usize, + source_y_start: usize, + width: usize, + height: usize, + out: &mut [u8], + scale_sample: fn(u16, u8) -> u8, + output_label: &str, ) -> Result<(), DecodeHeicError> { debug_assert!(matches!(CHANNELS, 3 | 4)); let ycbcr_transform = @@ -11756,15 +12120,42 @@ fn convert_heic_to_interleaved_rgb8_slice( }); } - let width = + let source_width = usize::try_from(decoded.width).map_err(|_| DecodeHeicError::InvalidDecodedFrame { detail: format!("HEIC width does not fit in usize ({})", decoded.width), })?; - let height = + let source_height = usize::try_from(decoded.height).map_err(|_| DecodeHeicError::InvalidDecodedFrame { detail: format!("HEIC height does not fit in usize ({})", decoded.height), })?; - let output_len = checked_heic_interleaved_output_len(decoded, CHANNELS, output_label)?; + let source_x_end = + source_x_start + .checked_add(width) + .ok_or_else(|| DecodeHeicError::InvalidDecodedFrame { + detail: format!("{output_label} source region x extent overflow"), + })?; + let source_y_end = + source_y_start + .checked_add(height) + .ok_or_else(|| DecodeHeicError::InvalidDecodedFrame { + detail: format!("{output_label} source region y extent overflow"), + })?; + if source_x_end > source_width || source_y_end > source_height { + return Err(DecodeHeicError::InvalidDecodedFrame { + detail: format!( + "{output_label} source region ({source_x_start},{source_y_start}) {width}x{height} exceeds {}x{}", + decoded.width, decoded.height + ), + }); + } + let output_len = width + .checked_mul(height) + .and_then(|pixels| pixels.checked_mul(CHANNELS)) + .ok_or_else(|| DecodeHeicError::InvalidDecodedFrame { + detail: format!( + "{output_label} region output sample count overflow for {width}x{height}" + ), + })?; if out.len() != output_len { return Err(DecodeHeicError::InvalidDecodedFrame { detail: format!( @@ -11787,12 +12178,79 @@ fn convert_heic_to_interleaved_rgb8_slice( // 8-bit grayscale: libheif's Op_mono_to_RGB24_32 copies Y verbatim. let mono_verbatim = matches!(chroma, HeicChromaPlanes::Monochrome) && bit_depth == 8; - for y in 0..height { - let row_start = y * width; - let out_row_start = row_start * CHANNELS; + if let ( + HeicChromaPlanes::Color { + u_samples, + v_samples, + chroma_width, + layout: HeicPixelLayout::Yuv420, + }, + PreparedYcbcrTransform::MatrixFull { coeffs, .. }, + ) = (&chroma, converter.transform) + { + heic_decoder::hevc::color_convert::convert_420_8bit_region_to_interleaved( + &decoded.y_plane.samples, + u_samples, + v_samples, + source_width, + *chroma_width, + source_x_start, + source_y_start, + width, + height, + coeffs.r_cr_fp8, + coeffs.g_cb_fp8, + coeffs.g_cr_fp8, + coeffs.b_cb_fp8, + CHANNELS, + out, + ); + return Ok(()); + } - for x in 0..width { - let y_index = row_start + x; + // libheif routes limited-range 4:2:0 and all matrix 4:2:2/4:4:4 + // conversions through its generic f32 operation. On AArch64, preserve + // that exact operation graph in four-lane NEON; other targets retain the + // scalar implementation selected by the converter module. High-bit-depth + // RGB8 has distinct image-crate rounding, so it stays on the established + // loop below (the production RGBA16 hook is accelerated separately). + if bit_depth == 8 + && let HeicChromaPlanes::Color { + u_samples, + v_samples, + chroma_width, + layout, + } = &chroma + && let Some(params) = prepared_float_matrix_params(converter.transform) + { + let (subsample_x, subsample_y) = heic_chroma_subsampling(*layout); + heic_decoder::hevc::color_convert::convert_float_matrix_8bit_region_to_interleaved( + &decoded.y_plane.samples, + u_samples, + v_samples, + source_width, + *chroma_width, + subsample_x as usize, + subsample_y as usize, + source_x_start, + source_y_start, + width, + height, + params, + CHANNELS, + out, + ); + return Ok(()); + } + + for output_y in 0..height { + let source_y = source_y_start + output_y; + let row_start = source_y * source_width; + let out_row_start = output_y * width * CHANNELS; + + for output_x in 0..width { + let source_x = source_x_start + output_x; + let y_index = row_start + source_x; let y_sample = i32::from(decoded.y_plane.samples[y_index]); let (cb_sample, cr_sample) = match &chroma { @@ -11803,7 +12261,8 @@ fn convert_heic_to_interleaved_rgb8_slice( chroma_width, layout, } => { - let chroma_index = heic_chroma_sample_index(x, y, *chroma_width, *layout); + let chroma_index = + heic_chroma_sample_index(source_x, source_y, *chroma_width, *layout); ( i32::from(u_samples[chroma_index]), i32::from(v_samples[chroma_index]), @@ -11819,7 +12278,7 @@ fn convert_heic_to_interleaved_rgb8_slice( } else { converter.convert(y_sample, cb_sample, cr_sample) }; - let out_index = out_row_start + (x * CHANNELS); + let out_index = out_row_start + (output_x * CHANNELS); out[out_index] = scale_sample(r, bit_depth); out[out_index + 1] = scale_sample(g, bit_depth); out[out_index + 2] = scale_sample(b, bit_depth); @@ -11906,6 +12365,34 @@ fn convert_heic_to_rgba16_slice( decoded.layout == HeicPixelLayout::Yuv420, ); + if let HeicChromaPlanes::Color { + u_samples, + v_samples, + chroma_width, + layout, + } = &chroma + && let Some(params) = prepared_float_matrix_params(converter.transform) + { + let (subsample_x, subsample_y) = heic_chroma_subsampling(*layout); + heic_decoder::hevc::color_convert::convert_float_matrix_region_to_rgba16( + &decoded.y_plane.samples, + u_samples, + v_samples, + width, + *chroma_width, + subsample_x as usize, + subsample_y as usize, + 0, + 0, + width, + height, + bit_depth, + params, + out, + ); + return Ok(()); + } + for y in 0..height { let row_start = y * width; let out_row_start = row_start * 4; @@ -12375,6 +12862,41 @@ enum PreparedYcbcrTransform { }, } +fn prepared_float_matrix_params( + transform: PreparedYcbcrTransform, +) -> Option { + let (coeffs, y_offset, y_scale, chroma_midpoint, chroma_scale) = match transform { + PreparedYcbcrTransform::MatrixFullFloat { + coeffs, + chroma_midpoint, + } => (coeffs, 0.0_f32, 1.0_f32, chroma_midpoint as f32, 1.0_f32), + PreparedYcbcrTransform::MatrixLimited { + coeffs, + limited_offset, + chroma_midpoint, + } => ( + coeffs, + limited_offset, + 1.1689_f32, + chroma_midpoint, + 1.1429_f32, + ), + PreparedYcbcrTransform::IdentityFull + | PreparedYcbcrTransform::IdentityLimited { .. } + | PreparedYcbcrTransform::MatrixFull { .. } => return None, + }; + Some(heic_decoder::hevc::color_convert::FloatMatrixParams { + y_offset, + y_scale, + chroma_midpoint, + chroma_scale, + r_cr: coeffs.r_cr_f32, + g_cb: coeffs.g_cb_f32, + g_cr: coeffs.g_cr_f32, + b_cb: coeffs.b_cb_f32, + }) +} + #[derive(Clone, Copy)] struct PreparedYcbcrToRgb { bit_depth: u8, @@ -13167,43 +13689,164 @@ mod tests { #[cfg(feature = "parallel-grid")] #[test] - fn grid_parallel_window_uses_each_candidate_tile_estimate() { - const MIB: u64 = 1024 * 1024; + fn grid_parallel_stream_exposes_decode_errors_in_row_major_order() { + let inputs = [0_usize, 1, 2, 3]; + let mut consumed = Vec::new(); + let result: Result<(), usize> = super::for_each_ordered_bounded_parallel( + &inputs, + 10, + 3, + 64, + |_| Some(1), + |value| { + if *value == 1 { + Err(77_usize) + } else { + Ok(*value) + } + }, + &mut |tile_index, value| { + consumed.push((tile_index, value)); + Ok(()) + }, + ); - assert_eq!( - super::heic_grid_tile_decode_window_from_estimates( - [Some(4 * MIB), Some(80 * MIB), Some(4 * MIB)], - 8, - ), - 1, - "a later oversized tile must not inherit the first candidate's small estimate" + assert_eq!(result, Err(77)); + assert_eq!(consumed, [(10, 0)]); + } + + #[cfg(feature = "parallel-grid")] + #[test] + fn grid_parallel_stream_stops_consuming_after_callback_error() { + let inputs = [0_usize, 1, 2, 3]; + let mut consumed = Vec::new(); + let result: Result<(), usize> = super::for_each_ordered_bounded_parallel( + &inputs, + 20, + 3, + 64, + |_| Some(1), + |value| Ok::<_, usize>(*value), + &mut |tile_index, value| { + if tile_index == 21 { + return Err(88_usize); + } + consumed.push((tile_index, value)); + Ok(()) + }, ); + + assert_eq!(result, Err(88)); + assert_eq!(consumed, [(20, 0)]); + } + + #[cfg(feature = "parallel-grid")] + #[test] + fn grid_parallel_stream_retains_each_tiles_memory_permit_until_consumed() { + use std::sync::Mutex; + use std::time::Duration; + + let events = Mutex::new(Vec::new()); + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(4) + .build() + .expect("test pool should build"); + let result: Result<(), usize> = pool.install(|| { + super::for_each_ordered_bounded_parallel( + &[0_usize, 1], + 40, + 2, + 64, + |value| Some(if *value == 0 { 4 } else { 80 }), + |value| { + events + .lock() + .expect("test mutex poisoned") + .push(('d', *value)); + if *value == 0 { + // If the second 80-byte job incorrectly inherited the + // first job's small estimate, the dedicated pool has + // ample time to start it before tile 0 is consumed. + std::thread::sleep(Duration::from_millis(20)); + } + Ok::<_, usize>(*value) + }, + &mut |_, value| { + events + .lock() + .expect("test mutex poisoned") + .push(('c', value)); + Ok(()) + }, + ) + }); + + assert_eq!(result, Ok(())); assert_eq!( - super::heic_grid_tile_decode_window_from_estimates( - [Some(32 * MIB), Some(32 * MIB), Some(MIB)], - 8, - ), - 2, - "a window may fill the memory budget exactly" + *events.lock().expect("test mutex poisoned"), + [('d', 0), ('c', 0), ('d', 1), ('c', 1)] ); - assert_eq!( - super::heic_grid_tile_decode_window_from_estimates( - [Some(4 * MIB), None, Some(4 * MIB)], - 8, - ), - 1, - "a tile whose metadata cannot be preflighted must stay outside the parallel window" + } + + #[cfg(feature = "parallel-grid")] + #[test] + fn grid_parallel_stream_decodes_unestimable_input_on_caller_thread() { + use std::sync::Mutex; + + let caller_thread = std::thread::current().id(); + let unestimable_thread = Mutex::new(None); + let mut consumed = Vec::new(); + let result: Result<(), usize> = super::for_each_ordered_bounded_parallel( + &[0_usize, 1, 2], + 30, + 3, + 64, + |value| (*value != 1).then_some(1), + |value| { + if *value == 1 { + *unestimable_thread.lock().expect("test mutex poisoned") = + Some(std::thread::current().id()); + } + Ok::<_, usize>(*value) + }, + &mut |tile_index, value| { + consumed.push((tile_index, value)); + Ok(()) + }, ); + + assert_eq!(result, Ok(())); + assert_eq!(consumed, [(30, 0), (31, 1), (32, 2)]); assert_eq!( - super::heic_grid_tile_decode_window_from_estimates( - core::iter::repeat_n(Some(MIB), 16), - 8, - ), - 8, - "the thread cap must remain effective for small tiles" + *unestimable_thread.lock().expect("test mutex poisoned"), + Some(caller_thread) ); } + #[cfg(feature = "parallel-grid")] + #[test] + fn grid_parallel_stream_propagates_worker_panics_without_hanging() { + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + let mut consume = |_: usize, _: usize| Ok::<_, usize>(()); + let _: Result<(), usize> = super::for_each_ordered_bounded_parallel( + &[0_usize, 1], + 0, + 2, + 64, + |_| Some(1), + |value| { + if *value == 0 { + panic!("intentional grid worker panic"); + } + Ok::<_, usize>(*value) + }, + &mut consume, + ); + })); + + assert!(result.is_err()); + } + #[cfg(feature = "parallel-grid")] #[test] fn grid_parallel_estimate_uses_coded_sps_geometry_and_both_stream_copies() { @@ -13868,6 +14511,87 @@ mod tests { assert_eq!(actual, expected); } + #[cfg(feature = "image-integration")] + #[test] + fn identity_alpha_float_matrix_slice_paths_match_owned_output() { + let mut image8 = synthetic_yuv_image(super::HeicPixelLayout::Yuv422); + image8.ycbcr_range = super::YCbCrRange::Limited; + let alpha8 = super::HeicAuxiliaryAlphaPlane { + width: image8.width, + height: image8.height, + bit_depth: 8, + samples: (0..image8.width * image8.height) + .map(|index| ((index * 29) & 255) as u16) + .collect(), + }; + let owned8 = super::decoded_heic_to_rgba_image(image8.clone(), &[], Some(&alpha8), None) + .expect("owned RGBA8 alpha conversion should succeed"); + let expected8 = match owned8.pixels { + super::DecodedRgbaPixels::U8(pixels) => pixels, + other => panic!("expected RGBA8 pixels, got {other:?}"), + }; + let mut actual8 = vec![0_u8; expected8.len()]; + super::decoded_heic_to_rgba8_slice(image8, &[], Some(&alpha8), &mut actual8) + .expect("identity RGBA8 alpha slice conversion should succeed"); + assert_eq!(actual8, expected8); + + let mut image16 = synthetic_yuv_image(super::HeicPixelLayout::Yuv420); + image16.bit_depth_luma = 10; + image16.bit_depth_chroma = 10; + for sample in &mut image16.y_plane.samples { + *sample *= 4; + } + for sample in image16 + .u_plane + .as_mut() + .expect("U plane") + .samples + .iter_mut() + { + *sample *= 4; + } + for sample in image16 + .v_plane + .as_mut() + .expect("V plane") + .samples + .iter_mut() + { + *sample *= 4; + } + let alpha16 = super::HeicAuxiliaryAlphaPlane { + width: image16.width, + height: image16.height, + bit_depth: 10, + samples: (0..image16.width * image16.height) + .map(|index| ((index * 37) & 1023) as u16) + .collect(), + }; + let owned16 = super::decoded_heic_to_rgba_image(image16.clone(), &[], Some(&alpha16), None) + .expect("owned RGBA16 alpha conversion should succeed"); + let expected16 = match owned16.pixels { + super::DecodedRgbaPixels::U16(pixels) => pixels, + other => panic!("expected RGBA16 pixels, got {other:?}"), + }; + let mut actual16 = vec![0_u16; expected16.len()]; + // SAFETY: the byte view spans the initialized u16 allocation exactly; + // the production API writes native-endian u16 bytes. + let actual16_bytes = unsafe { + std::slice::from_raw_parts_mut( + actual16.as_mut_ptr().cast::(), + actual16.len() * std::mem::size_of::(), + ) + }; + super::decoded_heic_to_rgba16_native_endian_bytes( + image16, + &[], + Some(&alpha16), + actual16_bytes, + ) + .expect("identity RGBA16 alpha byte conversion should succeed"); + assert_eq!(actual16, expected16); + } + #[cfg(feature = "image-integration")] #[test] fn transformed_avif_rgba16_byte_output_matches_owned_path() {