j2k-native 0.13.0

// SPDX-License-Identifier: MIT OR Apache-2.0

//! `MagSgn` decoding of two quads at a time in an 8-lane vector.
//!
//! A quad's lanes follow its sample order: (x, y), (x, y + 1), (x + 1, y),
//! (x + 1, y + 1). The quad's bit lengths are prefix-summed, so all four
//! samples are gathered from one 16-byte window of the unstuffed segment
//! instead of one after another. Each lane computes exactly what
//! [`super::decode_mag_sgn_sample_with_vn`] computes for that sample.

use fearless_simd::{f32x4, f32x8, i32x8, u32x4, u32x8, u8x16, u8x32, Bytes, Simd, SimdBase};

use super::super::cleanup::{
    decode_cleanup_first_row_pair, decode_cleanup_row_pair, decode_cleanup_symbols_first_row,
    decode_cleanup_symbols_row, decode_initial_pair_symbols, decode_pair_symbols,
};
use super::super::readers::{FlatBitReader, MelDecoder, ReverseBitReader};
use super::super::stats::HtDecodeObserver;
use super::decode_magnitude_sign_pair_from_cleanup;

/// Largest per-sample bit budget (`missing_msbs + 2`) decoded here. A
/// quad's last lane then starts at most `7 + 3 * 25` bits into its window
/// (byte 10), and every lane needs at most `7 + 25` bits (four bytes), so a
/// 16-byte window always holds the whole quad. Larger budgets also imply
/// `p >= 7`, so the 31-bit `p == 0` reconstruction never occurs here.
pub(super) const MAX_SIMD_MMSBP2: u32 = 25;

#[derive(Clone, Copy)]
struct Lanes<S: Simd> {
    simd: S,
    sig_shift: u32x8<S>,
    emb_k_shift: u32x8<S>,
    emb_1_shift: u32x8<S>,
    byte_spread: u32x8<S>,
    byte_offsets: u32x8<S>,
}

#[derive(Clone, Copy)]
pub(super) struct FinalOutput {
    pub(super) coefficient_shift: u32,
    pub(super) midpoint_scale: f32,
    pub(super) dequantization_step: f32,
}

impl<S: Simd> Lanes<S> {
    #[expect(
        clippy::inline_always,
        reason = "constant setup hoisted into row loops"
    )]
    #[inline(always)]
    fn new(simd: S) -> Self {
        Self {
            simd,
            sig_shift: u32x8::from_slice(simd, &[4, 5, 6, 7, 4, 5, 6, 7]),
            emb_k_shift: u32x8::from_slice(simd, &[12, 13, 14, 15, 12, 13, 14, 15]),
            emb_1_shift: u32x8::from_slice(simd, &[8, 9, 10, 11, 8, 9, 10, 11]),
            byte_spread: u32x8::splat(simd, 0x0101_0101),
            byte_offsets: u32x8::splat(simd, 0x0302_0100),
        }
    }
}

/// Two quads' decoded samples and `v_n` values. Each 128-bit half owns
/// one quad, so prefix sums and byte gathers cannot cross quad boundaries.
#[expect(clippy::inline_always, reason = "per-quad hot path")]
#[inline(always)]
fn decode_quad_pair<S: Simd>(
    lanes: &Lanes<S>,
    magsgn: &mut FlatBitReader<'_>,
    quads: [(u32, u32); 2],
    p_shift: u32,
) -> (u32x8<S>, u32x8<S>) {
    let simd = lanes.simd;
    let pair =
        |first, second| simd.combine_u32x4(u32x4::splat(simd, first), u32x4::splat(simd, second));
    let zero = pair(0, 0);
    let [(inf0, uq0), (inf1, uq1)] = quads;
    let inf_v = pair(inf0, inf1);
    let significant = zero - ((inf_v >> lanes.sig_shift) & 1);
    let emb_k = (inf_v >> lanes.emb_k_shift) & 1;
    let emb_1 = (inf_v >> lanes.emb_1_shift) & 1;
    let m_n = (pair(uq0, uq1) - emb_k) & significant;

    // Exclusive prefix sum: each lane's bit offset within the quad.
    let pairs = m_n + zero.slide_within_blocks::<3>(m_n);
    let inclusive = pairs + zero.slide_within_blocks::<2>(pairs);
    let offsets = zero.slide_within_blocks::<3>(inclusive);
    let [_, _, _, first_bits, _, _, _, second_bits] = inclusive.to_array();
    let first = magsgn.position();
    let second = first + first_bits as usize;

    let (window0, bit0) = magsgn.window16(first);
    let (window1, bit1) = magsgn.window16(second);
    let positions = offsets + pair(bit0, bit1);
    let indices = (positions >> 3) * lanes.byte_spread + lanes.byte_offsets;
    let window = simd.combine_u8x16(
        u8x16::load_array_ref(simd, window0),
        u8x16::load_array_ref(simd, window1),
    );
    let gathered: u32x8<S> = window
        .swizzle_dyn_within_blocks(indices.bitcast::<u8x32<S>>())
        .bitcast();
    let ms_val = gathered >> (positions & 7);
    magsgn.advance_to(second + second_bits as usize);

    // The caller checks uq <= mmsbp2 <= MAX_SIMD_MMSBP2 (25), so
    // every m_n fits a single variable shift, including zero-length lanes.
    let mask = !(!zero << m_n);
    let e1 = emb_1 << m_n;
    let v_n = ((ms_val & mask) | e1 | 1) & significant;
    let magnitude = (v_n + 2) << p_shift;
    (((ms_val << 31) | magnitude) & significant, v_n)
}

/// Decode the two quads at columns `x..x + 4`, write both output rows, and
/// fold the bottom-row `v_n` values into the scratch as the scalar path does.
#[expect(
    clippy::inline_always,
    clippy::too_many_arguments,
    reason = "two-quad hot step shared by the first and later row walkers"
)]
#[inline(always)]
fn decode_two_quads<S: Simd>(
    lanes: &Lanes<S>,
    magsgn: &mut FlatBitReader<'_>,
    decoded_data: &mut [u32],
    v_n_scratch: &mut [u32],
    quads: [(u32, u32); 2],
    mmsbp2: u32,
    p_shift: u32,
    stride: usize,
    second_row_present: bool,
    dp: usize,
    vp: usize,
    carry: &mut u32,
) -> Option<()> {
    let [(inf0, uq0), (inf1, uq1)] = quads;
    if uq0 > mmsbp2 || uq1 > mmsbp2 {
        return None;
    }

    if (inf0 | inf1) & 0xF0 == 0 {
        store_empty_pair(
            decoded_data,
            v_n_scratch,
            stride,
            second_row_present,
            dp,
            vp,
            carry,
        );
        return Some(());
    }

    let (samples, v_n) = decode_quad_pair(lanes, magsgn, quads, p_shift);
    let (samples0, samples1) = lanes.simd.split_u32x8(samples);
    let (v_n0, v_n1) = lanes.simd.split_u32x8(v_n);

    samples0
        .unzip_low(samples1)
        .store_slice(&mut decoded_data[dp..dp + 4]);
    if second_row_present {
        samples0
            .unzip_high(samples1)
            .store_slice(&mut decoded_data[dp + stride..dp + stride + 4]);
    }
    let [c0, c1, c2, c3] = v_n0.unzip_high(v_n1).to_array();
    v_n_scratch[vp] = *carry | c0;
    v_n_scratch[vp + 1] = c1 | c2;
    *carry = c3;
    Some(())
}

#[expect(
    clippy::inline_always,
    reason = "per-pair final coefficient conversion"
)]
#[inline(always)]
fn final_samples<const IRREVERSIBLE_MIDPOINT: bool, S: Simd>(
    lanes: &Lanes<S>,
    samples: u32x8<S>,
    output: FinalOutput,
) -> f32x8<S> {
    let magnitude = samples & 0x7FFF_FFFF;
    let values = if IRREVERSIBLE_MIDPOINT {
        let magnitude = lanes.simd.cvt_f32_u32x8(magnitude) * output.midpoint_scale;
        let signed = magnitude.bitcast::<u32x8<S>>() | (samples & 0x8000_0000);
        signed.bitcast::<f32x8<S>>()
    } else {
        let magnitude = magnitude >> output.coefficient_shift;
        let sign = u32x8::splat(lanes.simd, 0) - (samples >> 31);
        let signed = ((magnitude ^ sign) - sign).bitcast::<i32x8<S>>();
        lanes.simd.cvt_f32_i32x8(signed)
    };
    values * output.dequantization_step
}

/// Decode a full-width pair into final samples: both output rows and the
/// bottom-row `v_n` of its four columns. This is used only by whole-image
/// cleanup-only decode, where no refinement pass needs the intermediate
/// sign-magnitude words.
#[expect(clippy::inline_always, reason = "two-quad direct-output hot step")]
#[inline(always)]
fn decode_two_quads_to_f32<const IRREVERSIBLE_MIDPOINT: bool, S: Simd>(
    lanes: &Lanes<S>,
    magsgn: &mut FlatBitReader<'_>,
    quads: [(u32, u32); 2],
    missing_msbs: u32,
    output: FinalOutput,
) -> Option<(f32x4<S>, f32x4<S>, [u32; 4])> {
    let [(inf0, uq0), (inf1, uq1)] = quads;
    if uq0 > missing_msbs + 2 || uq1 > missing_msbs + 2 {
        return None;
    }

    let simd = lanes.simd;
    if (inf0 | inf1) & 0xF0 == 0 {
        let zero = final_samples::<IRREVERSIBLE_MIDPOINT, _>(lanes, u32x8::splat(simd, 0), output);
        let (zero, _) = simd.split_f32x8(zero);
        return Some((zero, zero, [0; 4]));
    }

    let (samples, v_n) = decode_quad_pair(lanes, magsgn, quads, 29 - missing_msbs);
    let samples = final_samples::<IRREVERSIBLE_MIDPOINT, _>(lanes, samples, output);
    let (samples0, samples1) = simd.split_f32x8(samples);
    let (v_n0, v_n1) = simd.split_u32x8(v_n);
    Some((
        samples0.unzip_low(samples1),
        samples0.unzip_high(samples1),
        v_n0.unzip_high(v_n1).to_array(),
    ))
}

/// Store a quad pair with no significant sample: zero output and zero `v_n`.
#[expect(clippy::inline_always, reason = "per-pair hot path")]
#[inline(always)]
fn store_empty_pair(
    decoded_data: &mut [u32],
    v_n_scratch: &mut [u32],
    stride: usize,
    second_row_present: bool,
    dp: usize,
    vp: usize,
    carry: &mut u32,
) {
    decoded_data[dp..dp + 4].fill(0);
    if second_row_present {
        decoded_data[dp + stride..dp + stride + 4].fill(0);
    }
    v_n_scratch[vp] = *carry;
    v_n_scratch[vp + 1] = 0;
    *carry = 0;
}

/// `kappa` for a non-initial quad: `E_max` of the four bottom-row samples
/// above it when the quad has two or more significant samples, else 1.
#[expect(clippy::inline_always, reason = "per-quad hot path")]
#[inline(always)]
fn kappa(inf: u32, above: u32) -> u32 {
    let mut gamma = inf & 0xF0;
    gamma &= gamma.wrapping_sub(0x10);
    if gamma != 0 {
        (above | 2).ilog2()
    } else {
        1
    }
}

/// Add `kappa` to a non-initial quad pair's `u` values. Reads the row above
/// from `v_n_scratch[vp..vp + 3]` before this pair overwrites it.
#[expect(clippy::inline_always, reason = "per-pair hot path")]
#[inline(always)]
fn with_kappa(symbols: [u32; 4], v_n_scratch: &[u32], vp: usize) -> [(u32, u32); 2] {
    let [t0, u0, t1, u1] = symbols;
    let above = [v_n_scratch[vp], v_n_scratch[vp + 1], v_n_scratch[vp + 2]];
    [
        (t0, u0 + kappa(t0, above[0] | above[1])),
        (t1, u1 + kappa(t1, above[1] | above[2])),
    ]
}

/// Cursor over one quad row's `MagSgn` outputs.
struct RowCursor {
    second_row_present: bool,
    dp: usize,
    vp: usize,
    carry: u32,
}

/// Decode the quad pair at columns `x..x + 4`: two vector quads when all four
/// columns are inside the block, else the scalar quads that are.
#[expect(
    clippy::inline_always,
    clippy::too_many_arguments,
    reason = "per-pair hot step shared by the split and fused row walkers"
)]
#[inline(always)]
fn decode_pair<S: Simd>(
    lanes: &Lanes<S>,
    magsgn: &mut FlatBitReader<'_>,
    decoded_data: &mut [u32],
    v_n_scratch: &mut [u32],
    quads: [(u32, u32); 2],
    missing_msbs: u32,
    width: u32,
    stride: usize,
    x: u32,
    cursor: &mut RowCursor,
) -> Option<()> {
    let p = 30 - missing_msbs;
    let mmsbp2 = missing_msbs + 2;
    if x + 4 <= width {
        decode_two_quads(
            lanes,
            magsgn,
            decoded_data,
            v_n_scratch,
            quads,
            mmsbp2,
            p - 1,
            stride,
            cursor.second_row_present,
            cursor.dp,
            cursor.vp,
            &mut cursor.carry,
        )?;
        cursor.dp += 4;
        cursor.vp += 2;
        return Some(());
    }
    let mut column = x;
    for (inf, uq) in quads {
        if column >= width {
            break;
        }
        decode_magnitude_sign_pair_from_cleanup(
            magsgn,
            decoded_data,
            v_n_scratch,
            inf,
            uq,
            mmsbp2,
            p,
            width,
            stride,
            cursor.second_row_present,
            &mut column,
            &mut cursor.dp,
            &mut cursor.vp,
            &mut cursor.carry,
        )?;
    }
    Some(())
}

/// Symbols of the quad pair at `sp` in the cleanup scratch.
#[expect(clippy::inline_always, reason = "per-pair hot path")]
#[inline(always)]
fn stored_symbols(cleanup: &[u16], sp: usize) -> [u32; 4] {
    [
        u32::from(cleanup[sp]),
        u32::from(cleanup[sp + 1]),
        u32::from(cleanup[sp + 2]),
        u32::from(cleanup[sp + 3]),
    ]
}

#[expect(
    clippy::inline_always,
    clippy::too_many_arguments,
    reason = "SIMD twin of the scalar MagSgn row walkers"
)]
#[inline(always)]
fn decode_stored_row<S: Simd>(
    lanes: &Lanes<S>,
    magsgn: &mut FlatBitReader<'_>,
    cleanup: &[u16],
    decoded_data: &mut [u32],
    v_n_scratch: &mut [u32],
    missing_msbs: u32,
    width: u32,
    height: u32,
    y: u32,
    stride: u32,
    sstr: usize,
) -> Option<()> {
    let stride = stride as usize;
    let mut cursor = RowCursor {
        second_row_present: y + 1 < height,
        dp: y as usize * stride,
        vp: 0,
        carry: 0,
    };
    let mut sp = (y >> 1) as usize * sstr;
    let mut x = 0u32;
    while x < width {
        let symbols = stored_symbols(cleanup, sp);
        let quads = if y == 0 {
            [(symbols[0], symbols[1]), (symbols[2], symbols[3])]
        } else {
            with_kappa(symbols, v_n_scratch, cursor.vp)
        };
        decode_pair(
            lanes,
            magsgn,
            decoded_data,
            v_n_scratch,
            quads,
            missing_msbs,
            width,
            stride,
            x,
            &mut cursor,
        )?;
        sp += 4;
        x += 4;
    }
    v_n_scratch[cursor.vp] = cursor.carry;
    Some(())
}

/// SIMD twin of the `MagSgn` row loop in [`super::decode_magnitude_sign_phase`].
#[expect(
    clippy::inline_always,
    clippy::too_many_arguments,
    reason = "inlined into the dispatched SIMD level"
)]
#[inline(always)]
pub(super) fn decode_magnitude_sign_rows<S: Simd>(
    simd: S,
    magsgn: &mut FlatBitReader<'_>,
    cleanup: &[u16],
    decoded_data: &mut [u32],
    v_n_scratch: &mut [u32],
    missing_msbs: u32,
    width: u32,
    height: u32,
    stride: u32,
    sstr: usize,
) -> Option<()> {
    let lanes = Lanes::new(simd);
    for y in (0..height).step_by(2) {
        decode_stored_row(
            &lanes,
            magsgn,
            cleanup,
            decoded_data,
            v_n_scratch,
            missing_msbs,
            width,
            height,
            y,
            stride,
            sstr,
        )?;
    }
    Some(())
}

/// Decode the cleanup symbols and `MagSgn` samples of quad row `y` together,
/// one quad pair at a time. The VLC chain of each pair is serial; the vector
/// `MagSgn` work of the pair before it is independent of it, so the core
/// overlaps the two.
#[expect(
    clippy::inline_always,
    clippy::too_many_arguments,
    reason = "fused row walker inlined into the dispatched SIMD level"
)]
#[inline(always)]
fn decode_fused_row<S: Simd>(
    lanes: &Lanes<S>,
    mel: &mut MelDecoder<'_>,
    vlc: &mut ReverseBitReader<'_>,
    run: &mut i32,
    magsgn: &mut FlatBitReader<'_>,
    scratch: &mut [u16],
    decoded_data: &mut [u32],
    v_n_scratch: &mut [u32],
    missing_msbs: u32,
    width: u32,
    height: u32,
    y: u32,
    stride: u32,
    sstr: usize,
) -> Option<()> {
    let stride = stride as usize;
    let mut cursor = RowCursor {
        second_row_present: y + 1 < height,
        dp: y as usize * stride,
        vp: 0,
        carry: 0,
    };
    let mut row_offset = (y >> 1) as usize * sstr;
    let mut c_q = 0u32;
    let mut x = 0u32;
    while x < width {
        let quads = if y == 0 {
            let [t0, u0, t1, u1] = decode_cleanup_first_row_pair(
                mel, vlc, run, scratch, row_offset, &mut c_q, x, width,
            )?;
            [(t0, u0), (t1, u1)]
        } else {
            let symbols = decode_cleanup_row_pair(
                mel, vlc, run, scratch, row_offset, sstr, &mut c_q, x, width,
            )?;
            with_kappa(symbols, v_n_scratch, cursor.vp)
        };
        decode_pair(
            lanes,
            magsgn,
            decoded_data,
            v_n_scratch,
            quads,
            missing_msbs,
            width,
            stride,
            x,
            &mut cursor,
        )?;
        row_offset += 4;
        x += 4;
    }
    scratch[row_offset] = 0;
    scratch[row_offset + 1] = 0;
    v_n_scratch[cursor.vp] = cursor.carry;
    Some(())
}

/// SIMD twin of the row loop in
/// [`super::decode_cleanup_and_magnitude_sign_phase`]. Rows are fused unless
/// the observer times the two phases separately.
#[expect(
    clippy::inline_always,
    clippy::too_many_arguments,
    reason = "inlined into the dispatched SIMD level"
)]
#[inline(always)]
pub(super) fn decode_cleanup_and_magnitude_sign_rows<S: Simd>(
    simd: S,
    mel: &mut MelDecoder<'_>,
    vlc: &mut ReverseBitReader<'_>,
    run: &mut i32,
    magsgn: &mut FlatBitReader<'_>,
    scratch: &mut [u16],
    decoded_data: &mut [u32],
    v_n_scratch: &mut [u32],
    missing_msbs: u32,
    width: u32,
    height: u32,
    stride: u32,
    sstr: usize,
    observer: &mut impl HtDecodeObserver,
) -> Option<()> {
    let lanes = Lanes::new(simd);
    let timed = observer.times_phases();
    for y in (0..height).step_by(2) {
        if !timed {
            decode_fused_row(
                &lanes,
                mel,
                vlc,
                run,
                magsgn,
                scratch,
                decoded_data,
                v_n_scratch,
                missing_msbs,
                width,
                height,
                y,
                stride,
                sstr,
            )?;
            continue;
        }
        let phase_start = observer.phase_start();
        if y == 0 {
            decode_cleanup_symbols_first_row(mel, vlc, run, scratch, width)?;
        } else {
            decode_cleanup_symbols_row(mel, vlc, run, scratch, width, y, sstr)?;
        }
        observer.add_cleanup_us(phase_start);

        let phase_start = observer.phase_start();
        decode_stored_row(
            &lanes,
            magsgn,
            scratch,
            decoded_data,
            v_n_scratch,
            missing_msbs,
            width,
            height,
            y,
            stride,
            sstr,
        )?;
        observer.add_mag_sgn_us(phase_start);
    }
    Some(())
}

/// One quad row of direct output, sliced to the row's quad pairs (`v_n` to
/// one more), so the pair loop needs no bounds checks.
struct RowOutput<'o> {
    first: &'o mut [[f32; 4]],
    /// Empty for the last row of an odd-height block.
    second: &'o mut [[f32; 4]],
    symbols: &'o mut [[u16; 4]],
    /// The row above's `v_n` column pairs on entry, this row's on exit.
    v_n: &'o mut [[u32; 2]],
}

/// Decode a pair's two quads straight to final samples, store both rows,
/// and fold its bottom-row `v_n` into `row.v_n` as the stored-row walkers do.
#[expect(
    clippy::inline_always,
    clippy::too_many_arguments,
    reason = "per-pair direct-output hot step"
)]
#[inline(always)]
fn decode_pair_to_f32<const IRREVERSIBLE_MIDPOINT: bool, S: Simd>(
    lanes: &Lanes<S>,
    magsgn: &mut FlatBitReader<'_>,
    row: &mut RowOutput<'_>,
    index: usize,
    quads: [(u32, u32); 2],
    missing_msbs: u32,
    output: FinalOutput,
    carry: &mut u32,
) -> Option<()> {
    let (samples0, samples1, [c0, c1, c2, c3]) = decode_two_quads_to_f32::<IRREVERSIBLE_MIDPOINT, S>(
        lanes,
        magsgn,
        quads,
        missing_msbs,
        output,
    )?;
    samples0.store_array(row.first.get_mut(index)?);
    if let Some(second) = row.second.get_mut(index) {
        samples1.store_array(second);
    }
    *row.v_n.get_mut(index)? = [*carry | c0, c1 | c2];
    *carry = c3;
    Some(())
}

/// Decode the initial quad row straight to final samples and return its
/// last `v_n` carry.
#[expect(
    clippy::cast_possible_truncation,
    clippy::inline_always,
    clippy::too_many_arguments,
    reason = "16-bit symbol fields; row walker inlined into the dispatched SIMD level"
)]
#[inline(always)]
fn decode_first_row_to_f32<const IRREVERSIBLE_MIDPOINT: bool, S: Simd>(
    lanes: &Lanes<S>,
    mel: &mut MelDecoder<'_>,
    vlc: &mut ReverseBitReader<'_>,
    run: &mut i32,
    magsgn: &mut FlatBitReader<'_>,
    row: &mut RowOutput<'_>,
    missing_msbs: u32,
    output: FinalOutput,
) -> Option<u32> {
    let mut c_q = 0;
    let mut carry = 0;
    for index in 0..row.first.len() {
        // Two quads in zero context inside one MEL run of zeros code no VLC,
        // U-VLC or MagSgn bits: skip to the symbols they decode to (the
        // initial row's `u` values start at one).
        let symbols = if c_q == 0 && *run >= 4 {
            *run -= 4;
            [0, 1, 0, 1]
        } else {
            decode_initial_pair_symbols::<true>(mel, vlc, run, &mut c_q, true)?
        };
        *row.symbols.get_mut(index)? = symbols.map(|symbol| symbol as u16);
        let [t0, u0, t1, u1] = symbols;
        let quads = [(t0, u0), (t1, u1)];
        decode_pair_to_f32::<IRREVERSIBLE_MIDPOINT, S>(
            lanes,
            magsgn,
            row,
            index,
            quads,
            missing_msbs,
            output,
            &mut carry,
        )?;
    }
    Some(carry)
}

/// Decode a non-initial quad row straight to final samples and return its
/// last `v_n` carry. `above` is the previous quad row's symbols and zero
/// sentinel.
#[expect(
    clippy::cast_possible_truncation,
    clippy::inline_always,
    clippy::too_many_arguments,
    reason = "16-bit symbol fields; row walker inlined into the dispatched SIMD level"
)]
#[inline(always)]
fn decode_row_to_f32<const IRREVERSIBLE_MIDPOINT: bool, S: Simd>(
    lanes: &Lanes<S>,
    mel: &mut MelDecoder<'_>,
    vlc: &mut ReverseBitReader<'_>,
    run: &mut i32,
    magsgn: &mut FlatBitReader<'_>,
    above: &[[u16; 4]],
    row: &mut RowOutput<'_>,
    missing_msbs: u32,
    output: FinalOutput,
) -> Option<u32> {
    let mut c_q = 0;
    let mut carry = 0;
    for index in 0..row.first.len() {
        let [above0, above1] = above.get(index..)?.first_chunk::<2>()?;
        let above = |k: usize| u32::from(if k < 4 { above0[k] } else { above1[0] });
        // As in the first row; zero context also needs the bits of the row
        // above that `decode_pair_symbols` reads.
        let symbols = if c_q == 0
            && *run >= 4
            && (above0[0] | above0[2]) & 0xA0 == 0
            && above1[0] & 0x20 == 0
        {
            *run -= 4;
            [0; 4]
        } else {
            decode_pair_symbols::<true>(mel, vlc, run, &mut c_q, above, true)?
        };
        *row.symbols.get_mut(index)? = symbols.map(|symbol| symbol as u16);
        // As `with_kappa`: the row above's `v_n`, read before this pair's.
        let [v_n0, v_n1] = *row.v_n.get(index)?;
        let [v_n2, _] = *row.v_n.get(index + 1)?;
        let [t0, u0, t1, u1] = symbols;
        let quads = [
            (t0, u0 + kappa(t0, v_n0 | v_n1)),
            (t1, u1 + kappa(t1, v_n1 | v_n2)),
        ];
        decode_pair_to_f32::<IRREVERSIBLE_MIDPOINT, S>(
            lanes,
            magsgn,
            row,
            index,
            quads,
            missing_msbs,
            output,
            &mut carry,
        )?;
    }
    Some(carry)
}

/// Cleanup-only SIMD rows that publish final dequantized coefficients
/// directly. The readers are opened here, inside the dispatched SIMD level,
/// so their state stays in registers across the pair loop.
#[expect(
    clippy::inline_always,
    clippy::too_many_arguments,
    reason = "inlined into the dispatched SIMD level"
)]
#[inline(always)]
pub(super) fn decode_cleanup_and_magnitude_sign_to_f32_rows<
    const IRREVERSIBLE_MIDPOINT: bool,
    S: Simd,
>(
    simd: S,
    coded_data: &[u8],
    lcup: usize,
    scup: usize,
    flat: &[u8],
    scratch: &mut [u16],
    decoded_data: &mut [f32],
    v_n_scratch: &mut [u32],
    missing_msbs: u32,
    width: u32,
    height: u32,
    stride: u32,
    sstr: usize,
    output: FinalOutput,
) -> Option<()> {
    let mut mel = MelDecoder::new(coded_data, lcup, scup);
    let mut vlc = ReverseBitReader::new_vlc(coded_data, lcup, scup);
    let mut magsgn = FlatBitReader::new(flat)?;
    let mut run = mel.get_run_with::<true>()?;
    let lanes = Lanes::new(simd);
    // `width` is a nonzero multiple of four, the columns of `pairs` pairs.
    let (width, stride) = (width as usize, stride as usize);
    let pairs = width / 4;
    let Some(last_row) = (height as usize).checked_sub(1) else {
        return Some(());
    };
    if stride < width {
        return None;
    }
    let extent = last_row.checked_mul(stride)?.checked_add(width)?;
    // Each chunk holds output rows `y` and `y + 1`, or only `y` for the
    // last row of an odd-height block.
    let output_rows = decoded_data
        .get_mut(..extent)?
        .chunks_mut(stride.checked_mul(2)?);
    let v_n = v_n_scratch.get_mut(..2 * pairs + 2)?.as_chunks_mut::<2>().0;
    let mut prev: &[u16] = &[];
    for (y, (rows, symbols)) in output_rows.zip(scratch.chunks_exact_mut(sstr)).enumerate() {
        let (first, second) = match rows.split_at_mut_checked(stride) {
            Some((first, second)) => (first, second.get_mut(..width).unwrap_or_default()),
            None => (rows, Default::default()),
        };
        let (pair_symbols, sentinel) = symbols.split_at_mut_checked(4 * pairs)?;
        // Slicing to exactly `pairs` lets the pair loop drop its checks.
        let mut row = RowOutput {
            first: first.get_mut(..width)?.as_chunks_mut::<4>().0,
            second: second.as_chunks_mut::<4>().0,
            symbols: pair_symbols.as_chunks_mut::<4>().0.get_mut(..pairs)?,
            v_n: v_n.get_mut(..=pairs)?,
        };
        let carry = if y == 0 {
            decode_first_row_to_f32::<IRREVERSIBLE_MIDPOINT, S>(
                &lanes,
                &mut mel,
                &mut vlc,
                &mut run,
                &mut magsgn,
                &mut row,
                missing_msbs,
                output,
            )?
        } else {
            let above = prev.get(..4 * pairs + 4)?.as_chunks::<4>().0;
            decode_row_to_f32::<IRREVERSIBLE_MIDPOINT, S>(
                &lanes,
                &mut mel,
                &mut vlc,
                &mut run,
                &mut magsgn,
                above,
                &mut row,
                missing_msbs,
                output,
            )?
        };
        sentinel.get_mut(..2)?.fill(0);
        *v_n.get_mut(pairs)?.first_mut()? = carry;
        prev = symbols;
    }
    Some(())
}

#[cfg(test)]
mod tests {
    use alloc::{vec, vec::Vec};

    use fearless_simd::{dispatch, Level};

    use super::super::super::cleanup::{
        cleanup_segment_suffix_length, cleanup_symbol_stride, decode_cleanup_symbols,
    };
    use super::super::super::readers::{
        flat_magsgn_capacity, unstuff_magsgn, FlatBitReader, MelDecoder, ReverseBitReader,
    };
    use super::super::super::stats::NoHtDecodeStats;
    use super::super::super::store_coefficient_row;
    use super::super::{
        decode_magnitude_sign_first_row_from_cleanup, decode_magnitude_sign_row_from_cleanup,
    };
    use super::{
        decode_cleanup_and_magnitude_sign_rows, decode_cleanup_and_magnitude_sign_to_f32_rows,
        decode_magnitude_sign_rows, FinalOutput, MAX_SIMD_MMSBP2,
    };
    use crate::j2c::ht_block_encode::encode_code_block;

    struct Block<'a> {
        coded: &'a [u8],
        cleanup: &'a [u16],
        missing_msbs: u32,
        width: u32,
        height: u32,
        sstr: usize,
    }

    fn magsgn_reader<'b>(block: &Block<'_>, buffer: &'b mut Vec<u8>) -> Option<FlatBitReader<'b>> {
        let lcup = block.coded.len();
        let scup = cleanup_segment_suffix_length(block.coded, lcup)?;
        let samples = (block.width * block.height) as usize;
        let flat = unstuff_magsgn(&block.coded[..lcup - scup], buffer, samples * 4)?;
        FlatBitReader::new(flat)
    }

    fn scratch(block: &Block<'_>) -> (Vec<u32>, Vec<u32>, Vec<u8>) {
        let samples = (block.width * block.height) as usize;
        (
            vec![0xDEAD_BEEF; samples],
            vec![0; block.width.div_ceil(2) as usize + 2],
            Vec::with_capacity(flat_magsgn_capacity(samples).expect("MagSgn capacity")),
        )
    }

    fn scalar_rows(block: &Block<'_>) -> Option<Vec<u32>> {
        let (mut decoded, mut v_n, mut buffer) = scratch(block);
        let mut magsgn = magsgn_reader(block, &mut buffer)?;
        decode_magnitude_sign_first_row_from_cleanup(
            &mut magsgn,
            block.cleanup,
            &mut decoded,
            &mut v_n,
            block.missing_msbs,
            block.width,
            block.height,
            block.width,
        )?;
        for y in (2..block.height).step_by(2) {
            decode_magnitude_sign_row_from_cleanup(
                &mut magsgn,
                block.cleanup,
                &mut decoded,
                &mut v_n,
                block.missing_msbs,
                block.width,
                block.height,
                y,
                block.width,
                block.sstr,
            )?;
        }
        Some(decoded)
    }

    fn simd_rows(block: &Block<'_>, level: Level) -> Option<Vec<u32>> {
        let (mut decoded, mut v_n, mut buffer) = scratch(block);
        let mut magsgn = magsgn_reader(block, &mut buffer)?;
        dispatch!(level, simd => decode_magnitude_sign_rows(
            simd,
            &mut magsgn,
            block.cleanup,
            &mut decoded,
            &mut v_n,
            block.missing_msbs,
            block.width,
            block.height,
            block.width,
            block.sstr,
        ))?;
        Some(decoded)
    }

    /// Cleanup symbols and `MagSgn` samples decoded together, one quad pair at
    /// a time, as the fused cleanup-only phase does.
    fn fused_rows(
        coded: &[u8],
        missing_msbs: u32,
        width: u32,
        height: u32,
        level: Level,
    ) -> Option<Vec<u32>> {
        let lcup = coded.len();
        let scup = cleanup_segment_suffix_length(coded, lcup)?;
        let samples = (width * height) as usize;
        let sstr = cleanup_symbol_stride(width);
        let mut mel = MelDecoder::new(coded, lcup, scup);
        let mut vlc = ReverseBitReader::new_vlc(coded, lcup, scup);
        let mut buffer =
            Vec::with_capacity(flat_magsgn_capacity(samples).expect("MagSgn capacity"));
        let flat = unstuff_magsgn(&coded[..lcup - scup], &mut buffer, samples * 4)?;
        let mut magsgn = FlatBitReader::new(flat)?;
        let mut run = mel.get_run()?;
        let mut scratch = vec![0u16; sstr * (height.div_ceil(2) as usize + 1)];
        let mut decoded = vec![0xDEAD_BEEF; samples];
        let mut v_n = vec![0; width.div_ceil(2) as usize + 2];
        dispatch!(level, simd => decode_cleanup_and_magnitude_sign_rows(
            simd,
            &mut mel,
            &mut vlc,
            &mut run,
            &mut magsgn,
            &mut scratch,
            &mut decoded,
            &mut v_n,
            missing_msbs,
            width,
            height,
            width,
            sstr,
            &mut NoHtDecodeStats,
        ))?;
        Some(decoded)
    }

    /// The native SIMD level, the portable fallback and, on x86, each lower
    /// level the CPU supports, so every dispatched variant runs.
    fn test_levels() -> Vec<Level> {
        let native = Level::new();
        #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
        let lower = [
            native.as_avx2().map(Level::Avx2),
            native.as_sse4_2().map(Level::Sse4_2),
            native.as_sse2().map(Level::Sse2),
        ];
        #[cfg(not(any(target_arch = "x86", target_arch = "x86_64")))]
        let lower: [Option<Level>; 0] = [];
        [native, Level::fallback()]
            .into_iter()
            .chain(lower.into_iter().flatten())
            .collect()
    }

    /// Output samples on each side of the direct rows, which stay untouched.
    const GUARD: usize = 5;
    const UNTOUCHED: u32 = 0x7FC0_1234;

    /// `(irreversible midpoint, coded bitplanes, dequantization step)`.
    type Conversion = (bool, u8, f32);

    /// An output buffer for `height` rows of `width` samples, `stride`
    /// apart, between the guards.
    fn guarded_output(width: u32, height: u32, stride: u32) -> Vec<f32> {
        let extent = (height as usize - 1) * stride as usize + width as usize;
        vec![f32::from_bits(UNTOUCHED); GUARD + extent + GUARD]
    }

    /// The direct `f32` rows on `level`, as bits of the guarded output. The
    /// scratch starts as garbage, as the decoder's reused scratch does.
    fn direct_rows(
        coded: &[u8],
        missing_msbs: u32,
        (width, height, stride): (u32, u32, u32),
        (irreversible, coded_bitplanes, step): Conversion,
        level: Level,
    ) -> Option<Vec<u32>> {
        let lcup = coded.len();
        let scup = cleanup_segment_suffix_length(coded, lcup)?;
        let samples = (width * height) as usize;
        let sstr = cleanup_symbol_stride(width);
        let mut buffer =
            Vec::with_capacity(flat_magsgn_capacity(samples).expect("MagSgn capacity"));
        let flat = unstuff_magsgn(&coded[..lcup - scup], &mut buffer, samples * 4)?;
        let mut scratch = vec![0xA5A5; sstr * (height.div_ceil(2) as usize + 1)];
        let mut v_n = vec![0xDEAD_BEEF; width.div_ceil(2) as usize + 2];
        let mut output = guarded_output(width, height, stride);
        let len = output.len();
        let rows = &mut output[GUARD..len - GUARD];
        let conversion = FinalOutput {
            coefficient_shift: 31 - u32::from(coded_bitplanes),
            midpoint_scale: crate::math::pow2i(-(31_i32 - i32::from(coded_bitplanes))),
            dequantization_step: step,
        };
        if irreversible {
            dispatch!(level, simd => decode_cleanup_and_magnitude_sign_to_f32_rows::<true, _>(
                simd, coded, lcup, scup, flat, &mut scratch, rows, &mut v_n, missing_msbs, width,
                height, stride, sstr, conversion,
            ))?;
        } else {
            dispatch!(level, simd => decode_cleanup_and_magnitude_sign_to_f32_rows::<false, _>(
                simd, coded, lcup, scup, flat, &mut scratch, rows, &mut v_n, missing_msbs, width,
                height, stride, sstr, conversion,
            ))?;
        }
        Some(output.iter().map(|sample| sample.to_bits()).collect())
    }

    /// `coefficients` (`width` apart) converted as the per-row decode
    /// converts them, laid out as [`direct_rows`] returns them.
    fn converted_rows(
        coefficients: &[u32],
        (width, height, stride): (u32, u32, u32),
        (irreversible, coded_bitplanes, step): Conversion,
    ) -> Vec<u32> {
        let mut output = guarded_output(width, height, stride);
        for (y, row) in coefficients.chunks_exact(width as usize).enumerate() {
            let start = GUARD + y * stride as usize;
            let out = &mut output[start..start + width as usize];
            store_coefficient_row(row, out, coded_bitplanes, 0, irreversible, step);
        }
        output.iter().map(|sample| sample.to_bits()).collect()
    }

    /// Require the direct `f32` rows, at two strides and three conversions,
    /// to match `expected` (the scalar rows, or `None`) converted per row on
    /// every level. Only widths that are a multiple of four take these rows.
    fn assert_direct_rows_match(
        coded: &[u8],
        missing_msbs: u32,
        coded_bitplanes: u8,
        (width, height): (u32, u32),
        expected: Option<&[u32]>,
    ) {
        if !width.is_multiple_of(4) {
            return;
        }
        for stride in [width, width + 3] {
            let layout = (width, height, stride);
            for (irreversible, step) in [(false, 1.0), (false, 0.375), (true, -0.5)] {
                let conversion = (irreversible, coded_bitplanes, step);
                let expected = expected.map(|rows| converted_rows(rows, layout, conversion));
                for level in test_levels() {
                    assert!(
                        direct_rows(coded, missing_msbs, layout, conversion, level) == expected,
                        "{layout:?} {conversion:?} missing_msbs={missing_msbs} {level:?} coded={coded:02x?}"
                    );
                }
            }
        }
    }

    /// Require the fused rows, the SIMD rows over stored symbols and the
    /// direct `f32` rows to agree with the scalar split rows on success,
    /// failure, and every output sample, on every level of [`test_levels`].
    /// Returns whether the block decoded.
    fn assert_rows_match(
        coded: &[u8],
        missing_msbs: u32,
        coded_bitplanes: u8,
        width: u32,
        height: u32,
    ) -> bool {
        assert!(missing_msbs + 2 <= MAX_SIMD_MMSBP2);
        let lcup = coded.len();
        let Some(scup) = cleanup_segment_suffix_length(coded, lcup) else {
            return false;
        };
        let sstr = cleanup_symbol_stride(width);
        let mut cleanup = vec![0u16; sstr * (height.div_ceil(2) as usize + 1)];
        if decode_cleanup_symbols(coded, lcup, scup, width, height, sstr, &mut cleanup).is_none() {
            for level in test_levels() {
                assert_eq!(
                    fused_rows(coded, missing_msbs, width, height, level),
                    None,
                    "{width}x{height} missing_msbs={missing_msbs} coded={coded:02x?}"
                );
            }
            assert_direct_rows_match(coded, missing_msbs, coded_bitplanes, (width, height), None);
            return false;
        }
        let block = Block {
            coded,
            cleanup: &cleanup,
            missing_msbs,
            width,
            height,
            sstr,
        };
        let expected = scalar_rows(&block);
        let layout = (width, height);
        assert_direct_rows_match(
            coded,
            missing_msbs,
            coded_bitplanes,
            layout,
            expected.as_deref(),
        );
        for level in test_levels() {
            assert_eq!(
                simd_rows(&block, level),
                expected,
                "{width}x{height} missing_msbs={missing_msbs} coded={coded:02x?}"
            );
            assert_eq!(
                fused_rows(coded, missing_msbs, width, height, level),
                expected,
                "fused {width}x{height} missing_msbs={missing_msbs} coded={coded:02x?}"
            );
        }
        expected.is_some()
    }

    struct Rng(u64);

    impl Rng {
        fn next(&mut self) -> u64 {
            self.0 ^= self.0 << 13;
            self.0 ^= self.0 >> 7;
            self.0 ^= self.0 << 17;
            self.0
        }

        #[expect(clippy::cast_possible_truncation, reason = "bounded test draws")]
        fn below(&mut self, bound: u64) -> u32 {
            (self.next() % bound) as u32
        }
    }

    #[test]
    fn simd_rows_match_scalar_rows_for_encoded_blocks() {
        let mut rng = Rng(0x9E37_79B9_7F4A_7C15);
        let mut decoded_blocks = 0;
        let sizes = [
            (1, 1),
            (2, 2),
            (3, 5),
            (4, 4),
            (5, 3),
            (7, 9),
            (8, 8),
            (12, 7),
            (16, 16),
            (17, 33),
            (32, 32),
            (33, 17),
            (64, 64),
            (63, 61),
            (64, 1),
            (1, 64),
            (4, 63),
            (20, 2),
            (128, 5),
        ];
        for (width, height) in sizes {
            for total_bitplanes in [1_u8, 2, 5, 8, 12, 16, 20, 23] {
                for density in [0, 1, 3, 8] {
                    let limit = 1_i64 << (total_bitplanes - 1);
                    let original: Vec<i32> = (0..width * height)
                        .map(|_| {
                            if density != 8 && rng.below(8) >= density {
                                return 0;
                            }
                            let magnitude = i64::from(rng.below(limit.unsigned_abs()));
                            let value = if rng.below(2) == 0 {
                                magnitude
                            } else {
                                -magnitude
                            };
                            i32::try_from(value).expect("test coefficient fits i32")
                        })
                        .collect();
                    let encoded = encode_code_block(&original, width, height, total_bitplanes)
                        .expect("encode HT block");
                    let missing_msbs = u32::from(encoded.num_zero_bitplanes);
                    if encoded.data.len() < 2 || missing_msbs + 2 > MAX_SIMD_MMSBP2 {
                        continue;
                    }
                    if assert_rows_match(
                        &encoded.data,
                        missing_msbs,
                        total_bitplanes,
                        width,
                        height,
                    ) {
                        decoded_blocks += 1;
                    }
                }
            }
        }
        assert!(decoded_blocks > 300, "{decoded_blocks}");
    }

    #[test]
    fn simd_rows_match_scalar_rows_for_arbitrary_bytes() {
        let mut rng = Rng(0x2545_F491_4F6C_DD1D);
        let mut decoded_blocks = 0;
        for _ in 0..4000 {
            let width = 1 + rng.below(40);
            let height = 1 + rng.below(40);
            let missing_msbs = rng.below(u64::from(MAX_SIMD_MMSBP2 - 1));
            let len = 2 + rng.below(300) as usize;
            let mut coded: Vec<u8> = (0..len).map(|_| rng.next().to_be_bytes()[3]).collect();
            // Keep the cleanup suffix length valid so blocks reach the rows.
            let scup = 2 + rng.below(u64::try_from(len - 1).expect("len fits u64")) as usize;
            coded[len - 1] = u8::try_from(scup >> 4).expect("suffix high bits");
            coded[len - 2] = (coded[len - 2] & 0xF0) | u8::try_from(scup & 0xF).expect("nibble");
            let coded_bitplanes = u8::try_from(missing_msbs + 2 + height % 4).expect("fits u8");
            if assert_rows_match(&coded, missing_msbs, coded_bitplanes, width, height) {
                decoded_blocks += 1;
            }
        }
        assert!(decoded_blocks > 100, "{decoded_blocks}");
    }
}