hpvcd 0.3.2

/*
 * // Copyright (c) Radzivon Bartoshyk 7/2026. All rights reserved.
 * //
 * // Redistribution and use in source and binary forms, with or without modification,
 * // are permitted provided that the following conditions are met:
 * //
 * // 1.  Redistributions of source code must retain the above copyright notice, this
 * // list of conditions and the following disclaimer.
 * //
 * // 2.  Redistributions in binary form must reproduce the above copyright notice,
 * // this list of conditions and the following disclaimer in the documentation
 * // and/or other materials provided with the distribution.
 * //
 * // 3.  Neither the name of the copyright holder nor the names of its
 * // contributors may be used to endorse or promote products derived from
 * // this software without specific prior written permission.
 * //
 * // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * // AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * // IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * // DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
 * // FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * // DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
 * // SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
 * // CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
 * // OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

use crate::exec::ExecContext;

pub(crate) type SaoPlaneFn =
    fn(&mut [u16], &[u16], usize, usize, usize, usize, usize, usize, u8, &[i32; 4], u8, u8, u8);

pub(crate) type SaoPlaneBandedFn = fn(
    &mut [u16],
    &[u16],
    usize,
    usize,
    usize,
    usize,
    usize,
    usize,
    usize,
    u8,
    &[i32; 4],
    u8,
    u8,
    u8,
);

pub(crate) type SaoBandOffsetInplaceFn =
    fn(&mut [u16], usize, usize, usize, usize, usize, &[i32; 4], u8, u8);

pub(crate) type SaoBandOffsetBandedInplaceFn =
    fn(&mut [u16], usize, usize, usize, usize, usize, usize, &[i32; 4], u8, u8);

static APPLY_SAO_PLANE: std::sync::OnceLock<SaoPlaneFn> = std::sync::OnceLock::new();
static APPLY_SAO_PLANE_BANDED: std::sync::OnceLock<SaoPlaneBandedFn> = std::sync::OnceLock::new();
static APPLY_SAO_BAND_OFFSET_INPLACE: std::sync::OnceLock<SaoBandOffsetInplaceFn> =
    std::sync::OnceLock::new();
static APPLY_SAO_BAND_OFFSET_BANDED_INPLACE: std::sync::OnceLock<SaoBandOffsetBandedInplaceFn> =
    std::sync::OnceLock::new();

#[inline]
pub(crate) fn resolve_apply_sao_plane() -> SaoPlaneFn {
    *APPLY_SAO_PLANE.get_or_init(|| {
        let mut _f: SaoPlaneFn = apply_sao_plane_scalar;

        #[cfg(all(feature = "neon", target_arch = "aarch64"))]
        {
            _f = crate::neon::apply_sao_plane_neon;
        }

        #[cfg(all(feature = "sse", any(target_arch = "x86", target_arch = "x86_64")))]
        {
            if std::is_x86_feature_detected!("sse4.1") {
                _f = crate::sse::apply_sao_plane_sse41;
            }
        }

        #[cfg(all(feature = "avx", target_arch = "x86_64"))]
        {
            if std::is_x86_feature_detected!("avx2") {
                _f = crate::avx::apply_sao_plane_avx2;
            }
        }

        _f
    })
}

#[inline]
pub(crate) fn resolve_apply_sao_plane_banded() -> SaoPlaneBandedFn {
    *APPLY_SAO_PLANE_BANDED.get_or_init(|| {
        let mut _f: SaoPlaneBandedFn = apply_sao_plane_banded_scalar;

        #[cfg(all(feature = "neon", target_arch = "aarch64"))]
        {
            _f = crate::neon::apply_sao_plane_banded_neon;
        }

        #[cfg(all(feature = "sse", any(target_arch = "x86", target_arch = "x86_64")))]
        {
            if std::is_x86_feature_detected!("sse4.1") {
                _f = crate::sse::apply_sao_plane_banded_sse41;
            }
        }

        #[cfg(all(feature = "avx", target_arch = "x86_64"))]
        {
            if std::is_x86_feature_detected!("avx2") {
                _f = crate::avx::apply_sao_plane_banded_avx2;
            }
        }

        _f
    })
}

#[inline]
pub(crate) fn resolve_apply_sao_band_offset_inplace() -> SaoBandOffsetInplaceFn {
    *APPLY_SAO_BAND_OFFSET_INPLACE.get_or_init(|| {
        let mut _f: SaoBandOffsetInplaceFn = apply_sao_band_offset_inplace_scalar;

        #[cfg(all(feature = "neon", target_arch = "aarch64"))]
        {
            _f = crate::neon::apply_sao_band_offset_inplace_neon;
        }

        #[cfg(all(feature = "sse", any(target_arch = "x86", target_arch = "x86_64")))]
        {
            if std::is_x86_feature_detected!("sse4.1") {
                _f = crate::sse::apply_sao_band_offset_inplace_sse41;
            }
        }

        #[cfg(all(feature = "avx", target_arch = "x86_64"))]
        {
            if std::is_x86_feature_detected!("avx2") {
                _f = crate::avx::apply_sao_band_offset_inplace_avx2;
            }
        }

        _f
    })
}

#[inline]
pub(crate) fn resolve_apply_sao_band_offset_banded_inplace() -> SaoBandOffsetBandedInplaceFn {
    *APPLY_SAO_BAND_OFFSET_BANDED_INPLACE.get_or_init(|| {
        let mut _f: SaoBandOffsetBandedInplaceFn = apply_sao_band_offset_banded_inplace_scalar;

        #[cfg(all(feature = "neon", target_arch = "aarch64"))]
        {
            _f = crate::neon::apply_sao_band_offset_banded_inplace_neon;
        }

        #[cfg(all(feature = "sse", any(target_arch = "x86", target_arch = "x86_64")))]
        {
            if std::is_x86_feature_detected!("sse4.1") {
                _f = crate::sse::apply_sao_band_offset_banded_inplace_sse41;
            }
        }

        #[cfg(all(feature = "avx", target_arch = "x86_64"))]
        {
            if std::is_x86_feature_detected!("avx2") {
                _f = crate::avx::apply_sao_band_offset_banded_inplace_avx2;
            }
        }

        _f
    })
}

#[allow(clippy::too_many_arguments)]
#[inline]
pub(crate) fn apply_sao_plane_scalar(
    dst: &mut [u16],
    src: &[u16],
    w: usize,
    h: usize,
    x0: usize,
    y0: usize,
    x_end: usize,
    y_end: usize,
    type_idx: u8,
    offsets: &[i32; 4],
    band_pos: u8,
    eo_class: u8,
    bd: u8,
) {
    match type_idx {
        1 => apply_sao_band_offset_scalar(dst, src, w, x0, y0, x_end, y_end, offsets, band_pos, bd),
        2 => apply_sao_edge_offset_scalar(
            dst, src, w, h, x0, y0, x_end, y_end, offsets, eo_class, bd,
        ),
        _ => {}
    }
}

#[allow(clippy::too_many_arguments)]
#[inline]
pub(crate) fn apply_sao_band_offset_scalar(
    dst: &mut [u16],
    src: &[u16],
    w: usize,
    x0: usize,
    y0: usize,
    x_end: usize,
    y_end: usize,
    offsets: &[i32; 4],
    band_pos: u8,
    bd: u8,
) {
    let max_val = ((1u32 << bd) - 1) as i32;
    let shift = bd - 5;

    for y in y0..y_end {
        let row = y * w;
        for x in x0..x_end {
            let s = src[row + x] as i32;
            let band = (s >> shift) as u8;
            let rel = band.wrapping_sub(band_pos);
            if rel < 4 {
                let v = (s + offsets[rel as usize]).clamp(0, max_val);
                dst[row + x] = v as u16;
            }
        }
    }
}

#[allow(clippy::too_many_arguments)]
#[inline]
pub(crate) fn apply_sao_band_offset_inplace_scalar(
    dst: &mut [u16],
    w: usize,
    x0: usize,
    y0: usize,
    x_end: usize,
    y_end: usize,
    offsets: &[i32; 4],
    band_pos: u8,
    bd: u8,
) {
    let max_val = ((1u32 << bd) - 1) as i32;
    let shift = bd - 5;

    for y in y0..y_end {
        let row = y * w;
        for x in x0..x_end {
            let dst = &mut dst[row + x];
            let s = *dst as i32;
            let band = (s >> shift) as u8;
            let rel = band.wrapping_sub(band_pos);
            if rel < 4 {
                *dst = (s + offsets[rel as usize]).clamp(0, max_val) as u16;
            }
        }
    }
}

#[allow(clippy::too_many_arguments)]
#[inline]
fn apply_sao_edge_offset_scalar(
    dst: &mut [u16],
    src: &[u16],
    w: usize,
    h: usize,
    x0: usize,
    y0: usize,
    x_end: usize,
    y_end: usize,
    offsets: &[i32; 4],
    eo_class: u8,
    bd: u8,
) {
    let max_val = ((1u32 << bd) - 1) as i32;

    // Direction vectors for the two neighbors (§8.7.3.2.4).
    let (dx, dy): (i32, i32) = match eo_class {
        0 => (1, 0),  // horizontal
        1 => (0, 1),  // vertical
        2 => (1, 1),  // 135°
        _ => (1, -1), // 45°
    };

    for y in y0..y_end {
        for x in x0..x_end {
            let s = src[y * w + x] as i32;

            let x1 = x as i32 + dx;
            let y1 = y as i32 + dy;
            let x2 = x as i32 - dx;
            let y2 = y as i32 - dy;

            let inb = |xx: i32, yy: i32| -> bool {
                xx >= 0 && yy >= 0 && (xx as usize) < w && (yy as usize) < h
            };
            // §8.7.3.2: if either neighbor lies outside the picture the sample
            // is left unchanged (edgeIdx not applied). Substituting `s` is wrong
            // — the other neighbor can still yield a nonzero edgeIdx.
            if !inb(x1, y1) || !inb(x2, y2) {
                continue;
            }
            let n1 = src[y1 as usize * w + x1 as usize] as i32;
            let n2 = src[y2 as usize * w + x2 as usize] as i32;

            let sign1 = (s > n1) as i32 - (s < n1) as i32;
            let sign2 = (s > n2) as i32 - (s < n2) as i32;
            let edge_idx = sign1 + sign2 + 2;

            let offset = match edge_idx {
                0 => offsets[0],
                1 => offsets[1],
                3 => offsets[2],
                4 => offsets[3],
                _ => 0,
            };
            if offset != 0 {
                dst[y * w + x] = (s + offset).clamp(0, max_val) as u16;
            }
        }
    }
}

/// Per-CTU SAO parameters, flattened out of the decoder's private `SaoCtb` so
/// this driver has no dependency on `decode.rs` internals. One entry per CTB in
/// raster order (`ry * ctb_cols + rx`). Component order is [luma, Cb, Cr].
#[derive(Clone, Copy, Default)]
pub(crate) struct SaoCtbParams {
    pub type_idx: [u8; 3],
    pub offsets: [[i32; 4]; 3],
    pub band_pos: [u8; 3],
    pub eo_class: [u8; 3],
}

/// Immutable geometry + toggles shared by every SAO band worker.
pub(crate) struct SaoPlanesCtx<'a> {
    pub exec: ExecContext,
    pub params: &'a [SaoCtbParams],
    pub ctb_cols: usize,
    pub ctb_rows: usize,
    pub log2_ctb: u32,
    pub w: usize,
    pub h: usize,
    pub cw: usize,
    pub ch: usize,
    pub sub_w: usize,
    pub sub_h: usize,
    pub bd: u8,
    pub bd_c: u8,
    pub sao_luma: bool,
    pub sao_chroma: bool,
}

/// Per-4×4-luma boundary maps + flags used to decide SAO edge-offset neighbor
/// availability (§8.7.3.2) on the restricted (serial) path.
pub(crate) struct SaoBoundary<'a> {
    pub gw: usize,
    pub log2_ctb: u32,
    pub sub_w: usize,
    pub sub_h: usize,
    pub slice_idx: &'a [u16],
    pub tqb: &'a [bool],
    pub pcm: &'a [bool],
    pub loop_filter_across_slices: bool,
    pub loop_filter_across_tiles: bool,
    pub pcm_loop_filter_disabled: bool,
    pub tile_grid: Option<&'a crate::tiles::TileGrid>,
}

impl SaoBoundary<'_> {
    /// Whether luma sample `(x, y)` is a valid SAO edge-offset neighbor of the
    /// sample at `(cx, cy)` (luma coords). A cross slice/tile boundary with
    /// filtering disabled, or a PCM/TQB exemption on either endpoint, makes the
    /// neighbor unavailable.
    #[inline]
    pub(crate) fn luma_neighbor_ok(&self, cx: usize, cy: usize, x: usize, y: usize) -> bool {
        if self.gw == 0 {
            return true;
        }
        let gc = (cy / 4) * self.gw + (cx / 4);
        let gn = (y / 4) * self.gw + (x / 4);
        // The current sample gets no SAO when it is transquant-bypass or in an
        // I_PCM CU with pcm_loop_filter_disabled (§8.7.3.1); reporting the
        // neighbor as unavailable is how the gated edge-offset kernel skips it. A
        // *neighbor* that is exempt is still available for edge classification —
        // exemption never removes a sample from §6.4.1 availability — so only the
        // center `(cx, cy)` is tested here. (The final substitution of exempt
        // samples back to their pre-SAO values is done in `apply_sao`.)
        if self.tqb.get(gc).copied().unwrap_or(false) {
            return false;
        }
        if self.pcm_loop_filter_disabled && self.pcm.get(gc).copied().unwrap_or(false) {
            return false;
        }
        if !self.loop_filter_across_slices
            && self.slice_idx.get(gc).copied().unwrap_or(0)
                != self.slice_idx.get(gn).copied().unwrap_or(0)
        {
            return false;
        }
        if !self.loop_filter_across_tiles
            && let Some(g) = self.tile_grid
        {
            let c = self.log2_ctb;
            if g.tile_id_at(cx >> c, cy >> c) != g.tile_id_at(x >> c, y >> c) {
                return false;
            }
        }
        true
    }

    /// Chroma analogue: map chroma coords back to luma to consult the maps.
    #[inline]
    pub(crate) fn chroma_neighbor_ok(&self, ccx: usize, ccy: usize, cx: usize, cy: usize) -> bool {
        self.luma_neighbor_ok(
            ccx * self.sub_w,
            ccy * self.sub_h,
            cx * self.sub_w,
            cy * self.sub_h,
        )
    }
}

/// Edge-offset SAO with per-neighbor boundary gating (§8.7.3.2). Used only on
/// the restricted path; when either neighbor is unavailable the sample is left
/// unchanged (edgeIdx forced to 0). `neighbor_ok(cx, cy, nx, ny)` reports
/// whether neighbor `(nx, ny)` may be used for the sample at `(cx, cy)`.
#[allow(clippy::too_many_arguments)]
pub(crate) fn apply_sao_edge_offset_gated(
    dst: &mut [u16],
    src: &[u16],
    w: usize,
    h: usize,
    x0: usize,
    y0: usize,
    x_end: usize,
    y_end: usize,
    offsets: &[i32; 4],
    eo_class: u8,
    bd: u8,
    neighbor_ok: &dyn Fn(usize, usize, usize, usize) -> bool,
) {
    let max_val = ((1u32 << bd) - 1) as i32;
    let (dx, dy): (i32, i32) = match eo_class {
        0 => (1, 0),
        1 => (0, 1),
        2 => (1, 1),
        _ => (1, -1),
    };
    for y in y0..y_end {
        for x in x0..x_end {
            let s = src[y * w + x] as i32;
            let x1 = x as i32 + dx;
            let y1 = y as i32 + dy;
            let x2 = x as i32 - dx;
            let y2 = y as i32 - dy;
            let ok1 = x1 >= 0
                && y1 >= 0
                && (x1 as usize) < w
                && (y1 as usize) < h
                && neighbor_ok(x, y, x1 as usize, y1 as usize);
            let ok2 = x2 >= 0
                && y2 >= 0
                && (x2 as usize) < w
                && (y2 as usize) < h
                && neighbor_ok(x, y, x2 as usize, y2 as usize);
            // §8.7.3.2: if either neighbor is unavailable, the sample is not
            // modified.
            if !ok1 || !ok2 {
                continue;
            }
            let n1 = src[y1 as usize * w + x1 as usize] as i32;
            let n2 = src[y2 as usize * w + x2 as usize] as i32;
            let sign1 = (s > n1) as i32 - (s < n1) as i32;
            let sign2 = (s > n2) as i32 - (s < n2) as i32;
            let edge_idx = sign1 + sign2 + 2;
            let offset = match edge_idx {
                0 => offsets[0],
                1 => offsets[1],
                3 => offsets[2],
                4 => offsets[3],
                _ => 0,
            };
            if offset != 0 {
                dst[y * w + x] = (s + offset).clamp(0, max_val) as u16;
            }
        }
    }
}

/// Apply SAO to one CTB-row's worth of output. `*_dst` are band slices.
/// Edge-offset SAO reads from optional full-plane snapshots; band-offset SAO is
/// pointwise and runs in-place, so BO-only pictures avoid full-plane clones.
#[allow(clippy::too_many_arguments)]
fn apply_sao_ctb_row(
    ctx: &SaoPlanesCtx<'_>,
    ry: usize,
    luma_dst: &mut [u16],
    luma_src: Option<&[u16]>,
    band_y0: usize,
    cb_dst: &mut [u16],
    cr_dst: &mut [u16],
    cb_src: Option<&[u16]>,
    cr_src: Option<&[u16]>,
    band_cy0: usize,
) {
    let ctb = 1usize << ctx.log2_ctb;
    for rx in 0..ctx.ctb_cols {
        let p = &ctx.params[ry * ctx.ctb_cols + rx];
        let x0 = rx * ctb;
        let y0 = ry * ctb;

        if ctx.sao_luma && p.type_idx[0] != 0 {
            let x_end = (x0 + ctb).min(ctx.w);
            let y_end = (y0 + ctb).min(ctx.h);
            match p.type_idx[0] {
                1 => (ctx.exec.sao_band_offset_banded_inplace)(
                    luma_dst,
                    ctx.w,
                    band_y0,
                    x0,
                    y0,
                    x_end,
                    y_end,
                    &p.offsets[0],
                    p.band_pos[0],
                    ctx.bd,
                ),
                2 => (ctx.exec.sao_plane_banded)(
                    luma_dst,
                    luma_src.expect("SAO EO requires luma snapshot"),
                    ctx.w,
                    ctx.h,
                    band_y0,
                    x0,
                    y0,
                    x_end,
                    y_end,
                    2,
                    &p.offsets[0],
                    p.band_pos[0],
                    p.eo_class[0],
                    ctx.bd,
                ),
                _ => {}
            }
        }

        if ctx.sao_chroma {
            let cx0 = x0 / ctx.sub_w;
            let cy0 = y0 / ctx.sub_h;
            let cx_end = ((x0 + ctb) / ctx.sub_w).min(ctx.cw);
            let cy_end = ((y0 + ctb) / ctx.sub_h).min(ctx.ch);
            match p.type_idx[1] {
                1 => (ctx.exec.sao_band_offset_banded_inplace)(
                    cb_dst,
                    ctx.cw,
                    band_cy0,
                    cx0,
                    cy0,
                    cx_end,
                    cy_end,
                    &p.offsets[1],
                    p.band_pos[1],
                    ctx.bd_c,
                ),
                2 => (ctx.exec.sao_plane_banded)(
                    cb_dst,
                    cb_src.expect("SAO EO requires Cb snapshot"),
                    ctx.cw,
                    ctx.ch,
                    band_cy0,
                    cx0,
                    cy0,
                    cx_end,
                    cy_end,
                    2,
                    &p.offsets[1],
                    p.band_pos[1],
                    p.eo_class[1],
                    ctx.bd_c,
                ),
                _ => {}
            }
            match p.type_idx[2] {
                1 => (ctx.exec.sao_band_offset_banded_inplace)(
                    cr_dst,
                    ctx.cw,
                    band_cy0,
                    cx0,
                    cy0,
                    cx_end,
                    cy_end,
                    &p.offsets[2],
                    p.band_pos[2],
                    ctx.bd_c,
                ),
                2 => (ctx.exec.sao_plane_banded)(
                    cr_dst,
                    cr_src.expect("SAO EO requires Cr snapshot"),
                    ctx.cw,
                    ctx.ch,
                    band_cy0,
                    cx0,
                    cy0,
                    cx_end,
                    cy_end,
                    2,
                    &p.offsets[2],
                    p.band_pos[2],
                    p.eo_class[2],
                    ctx.bd_c,
                ),
                _ => {}
            }
        }
    }
}

/// Banded SAO kernel: reads the *full* source plane at global coordinates
/// (so edge-offset neighbor lookups across the band's top/bottom edge stay
/// correct) and writes to a *band* destination whose local row 0 is global row
/// `band_y0` — i.e. `dst_band[(y - band_y0) * w + x]`. This is the parallel
/// analogue of the serial `apply_sao_plane`, split so disjoint bands can run on
/// separate threads. Bit-exact with the serial path because both read only from
/// the untouched clone and write each output pixel exactly once.
#[allow(clippy::too_many_arguments)]
#[inline]
fn apply_sao_band_offset_banded_inplace_scalar(
    dst_band: &mut [u16],
    w: usize,
    band_y0: usize,
    x0: usize,
    y0: usize,
    x_end: usize,
    y_end: usize,
    offsets: &[i32; 4],
    band_pos: u8,
    bd: u8,
) {
    let max_val = ((1u32 << bd) - 1) as i32;
    let shift = bd - 5;

    for y in y0..y_end {
        if y < band_y0 {
            continue;
        }
        let dst_row = (y - band_y0) * w;
        let Some(dst_row) = dst_band.get_mut(dst_row + x0..dst_row + x_end) else {
            continue;
        };
        for dst in dst_row.iter_mut() {
            let s = *dst as i32;
            let band = (s >> shift) as u8;
            let rel = band.wrapping_sub(band_pos);
            if rel < 4 {
                *dst = (s + offsets[rel as usize]).clamp(0, max_val) as u16;
            }
        }
    }
}

#[allow(clippy::too_many_arguments)]
pub(crate) fn apply_sao_plane_banded_scalar(
    dst_band: &mut [u16],
    src_full: &[u16],
    w: usize,
    h: usize,
    band_y0: usize,
    x0: usize,
    y0: usize,
    x_end: usize,
    y_end: usize,
    type_idx: u8,
    offsets: &[i32; 4],
    band_pos: u8,
    eo_class: u8,
    bd: u8,
) {
    if w == 0 || x_end <= x0 || y_end <= y0 || y_end <= band_y0 {
        return;
    }

    let Some(max_val) = (1u32)
        .checked_shl(bd as u32)
        .map(|v| v.saturating_sub(1) as i32)
    else {
        return;
    };

    match type_idx {
        // Band offset: purely pointwise, no neighbor reads.
        1 => {
            let shift = bd.saturating_sub(5);
            for y in y0..y_end {
                let src_row = y * w;
                let dst_row = (y - band_y0) * w;
                let src_range = src_row + x0..src_row + x_end;
                let dst_range = dst_row + x0..dst_row + x_end;
                let (Some(src_row), Some(dst_row)) =
                    (src_full.get(src_range), dst_band.get_mut(dst_range))
                else {
                    continue;
                };
                for (s, dst) in src_row.iter().copied().zip(dst_row.iter_mut()) {
                    let s = s as i32;
                    let band = (s >> shift) as u8;
                    let rel = band.wrapping_sub(band_pos);
                    if rel < 4 {
                        *dst = (s + offsets[rel as usize]).clamp(0, max_val) as u16;
                    }
                }
            }
        }
        // Edge offset: reads two neighbors from the full source plane.
        2 => {
            let (dx, dy): (i32, i32) = match eo_class {
                0 => (1, 0),
                1 => (0, 1),
                2 => (1, 1),
                _ => (1, -1),
            };
            let inb = |xx: i32, yy: i32| -> bool {
                xx >= 0 && yy >= 0 && (xx as usize) < w && (yy as usize) < h
            };
            for y in y0..y_end {
                let src_base = y * w;
                let dst_base = (y - band_y0) * w;
                for x in x0..x_end {
                    let Some(&s0) = src_full.get(src_base + x) else {
                        continue;
                    };
                    let s = s0 as i32;
                    let (x1, y1) = (x as i32 + dx, y as i32 + dy);
                    let (x2, y2) = (x as i32 - dx, y as i32 - dy);
                    // §8.7.3.2: an out-of-picture neighbor makes the sample
                    // unavailable → leave it unchanged (edgeIdx not applied).
                    if !inb(x1, y1) || !inb(x2, y2) {
                        continue;
                    }
                    let n1 = src_full
                        .get(y1 as usize * w + x1 as usize)
                        .copied()
                        .unwrap_or(s0) as i32;
                    let n2 = src_full
                        .get(y2 as usize * w + x2 as usize)
                        .copied()
                        .unwrap_or(s0) as i32;
                    let sign1 = (s > n1) as i32 - (s < n1) as i32;
                    let sign2 = (s > n2) as i32 - (s < n2) as i32;
                    let offset = match sign1 + sign2 + 2 {
                        0 => offsets[0],
                        1 => offsets[1],
                        3 => offsets[2],
                        4 => offsets[3],
                        _ => 0,
                    };
                    if offset != 0
                        && let Some(dst) = dst_band.get_mut(dst_base + x)
                    {
                        *dst = (s + offset).clamp(0, max_val) as u16;
                    }
                }
            }
        }
        _ => {}
    }
}

#[derive(Clone, Copy, Default)]
struct SaoUsage {
    active: [bool; 3],
    needs_src: [bool; 3],
}

fn sao_usage(ctx: &SaoPlanesCtx<'_>) -> SaoUsage {
    let mut usage = SaoUsage::default();
    for p in ctx.params.iter() {
        if ctx.sao_luma {
            usage.active[0] |= p.type_idx[0] != 0;
            usage.needs_src[0] |= p.type_idx[0] == 2;
        }
        if ctx.sao_chroma && ctx.cw != 0 && ctx.ch != 0 {
            usage.active[1] |= p.type_idx[1] != 0;
            usage.active[2] |= p.type_idx[2] != 0;
            usage.needs_src[1] |= p.type_idx[1] == 2;
            usage.needs_src[2] |= p.type_idx[2] == 2;
        }
    }
    usage
}

/// Parallel SAO over the whole picture. Splits each plane into CTB-row bands
/// (`ctb` output rows each, clipped at the picture edge), hands every band an
/// exclusive `&mut` region via [`DisjointMut`], and applies that CTB-row's SAO
/// on a pool worker. BO-only planes run in-place; only EO planes allocate an
/// untouched source snapshot for neighbor reads.
///
/// `dst_*` are consumed and the filtered planes returned (matching the serial
/// path's in-place semantics from the caller's view). When the pool is
/// single-threaded or there is only one band, it runs serially with no dispatch.
#[allow(clippy::type_complexity)]
pub(crate) fn apply_sao_parallel(
    pool: &crate::threadpool::ThreadPool,
    ctx: &SaoPlanesCtx<'_>,
    mut y: Vec<u16>,
    mut cb: Vec<u16>,
    mut cr: Vec<u16>,
) -> Result<(Vec<u16>, Vec<u16>, Vec<u16>), crate::DecodeError> {
    let ctb = 1usize << ctx.log2_ctb;
    let rows = ctx.ctb_rows;
    let usage = sao_usage(ctx);

    if !usage.active.iter().any(|&x| x) {
        return Ok((y, cb, cr));
    }

    let snapshot = |src: &[u16], what: &'static str| -> Result<Vec<u16>, crate::DecodeError> {
        let mut out = try_vec![0u16; src.len(), what];
        out.copy_from_slice(src);
        Ok(out)
    };
    let src_y = usage.needs_src[0]
        .then(|| snapshot(&y, "parallel SAO luma snapshot"))
        .transpose()?;
    let src_cb = usage.needs_src[1]
        .then(|| snapshot(&cb, "parallel SAO Cb snapshot"))
        .transpose()?;
    let src_cr = usage.needs_src[2]
        .then(|| snapshot(&cr, "parallel SAO Cr snapshot"))
        .transpose()?;

    // Serial fallback: no benefit from dispatch.
    if pool.threads() <= 1 || rows <= 1 {
        for ry in 0..rows {
            let y_lo = ry * ctb * ctx.w;
            let cy0 = (ry * ctb) / ctx.sub_h;
            let c_lo = cy0 * ctx.cw;
            apply_sao_ctb_row(
                ctx,
                ry,
                &mut y[y_lo..],
                src_y.as_deref(),
                ry * ctb,
                &mut cb[c_lo..],
                &mut cr[c_lo..],
                src_cb.as_deref(),
                src_cr.as_deref(),
                cy0,
            );
        }
        return Ok((y, cb, cr));
    }

    let y_dm = crate::threadpool::DisjointMut::new(y);
    let cb_dm = crate::threadpool::DisjointMut::new(cb);
    let cr_dm = crate::threadpool::DisjointMut::new(cr);

    crate::threadpool::parallel_for(pool, rows, |ry| {
        // Luma band: rows [ry*ctb, (ry+1)*ctb) clipped to h.
        let ly0 = ry * ctb;
        let ly1 = ((ry + 1) * ctb).min(ctx.h);
        if ly0 >= ly1 {
            return;
        }
        let y_lo = ly0 * ctx.w;
        let y_hi = ly1 * ctx.w;

        // Chroma band: the chroma rows this CTB-row covers.
        let cy0 = ly0 / ctx.sub_h;
        let cy1 = ((ly0 + ctb) / ctx.sub_h).min(ctx.ch);
        let (c_lo, c_hi) = (cy0 * ctx.cw, cy1.max(cy0) * ctx.cw);

        let mut y_band = y_dm.slice_mut(y_lo..y_hi);
        if ctx.sao_chroma && ctx.cw > 0 && c_lo < c_hi {
            let mut cb_band = cb_dm.slice_mut(c_lo..c_hi);
            let mut cr_band = cr_dm.slice_mut(c_lo..c_hi);
            apply_sao_ctb_row(
                ctx,
                ry,
                &mut y_band,
                src_y.as_deref(),
                ly0,
                &mut cb_band,
                &mut cr_band,
                src_cb.as_deref(),
                src_cr.as_deref(),
                cy0,
            );
        } else {
            apply_sao_ctb_row(
                ctx,
                ry,
                &mut y_band,
                src_y.as_deref(),
                ly0,
                &mut [],
                &mut [],
                src_cb.as_deref(),
                src_cr.as_deref(),
                cy0,
            );
        }
    });

    Ok((y_dm.into_inner(), cb_dm.into_inner(), cr_dm.into_inner()))
}

#[cfg(test)]
mod tests {
    use super::*;

    // Build a boundary where the vertical line x=4 splits slice 0 (left) from
    // slice 1 (right), on a 8×2 luma grid (gw=2, one 4×4 row).
    fn split_slice_bnd(slice_idx: &[u16]) -> SaoBoundary<'_> {
        SaoBoundary {
            gw: 2,
            log2_ctb: 6,
            sub_w: 2,
            sub_h: 2,
            slice_idx,
            tqb: &[false, false],
            pcm: &[false, false],
            loop_filter_across_slices: false,
            loop_filter_across_tiles: true,
            pcm_loop_filter_disabled: false,
            tile_grid: None,
        }
    }

    #[test]
    fn cross_slice_neighbor_unavailable() {
        // grid cell 0 = slice 0 (x 0..4), cell 1 = slice 1 (x 4..8).
        let slices = [0u16, 1];
        let b = split_slice_bnd(&slices);
        // Sample at x=3 (slice 0), neighbor at x=4 (slice 1) → unavailable.
        assert!(!b.luma_neighbor_ok(3, 0, 4, 0));
        // Same-slice neighbor is available.
        assert!(b.luma_neighbor_ok(3, 0, 2, 0));
        assert!(b.luma_neighbor_ok(5, 0, 6, 0));
    }

    #[test]
    fn across_slices_enabled_allows_neighbor() {
        let slices = [0u16, 1];
        let mut b = split_slice_bnd(&slices);
        b.loop_filter_across_slices = true;
        assert!(b.luma_neighbor_ok(3, 0, 4, 0));
    }

    #[test]
    fn cross_tile_neighbor_respects_loop_filter_flag() {
        let grid = crate::tiles::TileGrid {
            cols: 2,
            rows: 1,
            col_bd: vec![0, 1],
            row_bd: vec![0],
            col_width: vec![1, 1],
            row_height: vec![1],
            rs_to_ts: vec![0, 1],
            ts_to_rs: vec![0, 1],
            tile_id: vec![0, 1],
            ctb_cols: 2,
            ctb_rows: 1,
            loop_filter_across_tiles: false,
        };
        let slices = [0u16; 32];
        let flags = [false; 32];
        let mut b = SaoBoundary {
            gw: 32,
            log2_ctb: 6,
            sub_w: 2,
            sub_h: 2,
            slice_idx: &slices,
            tqb: &flags,
            pcm: &flags,
            loop_filter_across_slices: true,
            loop_filter_across_tiles: false,
            pcm_loop_filter_disabled: false,
            tile_grid: Some(&grid),
        };

        assert!(!b.luma_neighbor_ok(63, 0, 64, 0));
        b.loop_filter_across_tiles = true;
        assert!(b.luma_neighbor_ok(63, 0, 64, 0));
    }

    #[test]
    fn gated_eo_skips_when_neighbor_unavailable() {
        // 4×1 luma row: values chosen so the middle sample would get an offset
        // if both neighbors were used, but the right neighbor is "unavailable".
        // Horizontal EO (eo_class 0): neighbors are x-1 and x+1.
        let src = [10u16, 5, 10, 5];
        let mut dst = src;
        let offsets = [7i32, 0, 0, 0]; // edge_idx 0 (local min) → +7
        // Sample x=1 (val 5) has neighbors 10 and 10 → local min → edge_idx 0.
        // Make the x+1 neighbor unavailable → sample must be left unchanged.
        let ok = |_cx: usize, _cy: usize, nx: usize, _ny: usize| nx != 2;
        apply_sao_edge_offset_gated(&mut dst, &src, 4, 1, 0, 0, 4, 1, &offsets, 0, 8, &ok);
        // x=1 unchanged because its right neighbor (x=2) was unavailable.
        assert_eq!(dst[1], 5);
    }

    #[test]
    fn banded_eo_matches_serial_scalar_for_all_classes() {
        let w = 19usize;
        let h = 11usize;
        let bd = 10u8;
        let max_v = (1u16 << bd) - 1;
        let src: Vec<u16> = (0..w * h)
            .map(|i| {
                let x = i % w;
                let y = i / w;
                ((x * 73 + y * 191 + (x * y * 17)) & max_v as usize) as u16
            })
            .collect();
        let offsets = [3i32, -2, 4, -5];
        let x0 = 1usize;
        let x_end = w - 1;
        let y0 = 2usize;
        let y_end = h - 2;

        for eo_class in 0..4u8 {
            let mut serial = src.clone();
            apply_sao_plane_scalar(
                &mut serial,
                &src,
                w,
                h,
                x0,
                y0,
                x_end,
                y_end,
                2,
                &offsets,
                0,
                eo_class,
                bd,
            );

            let band_y0 = y0;
            let mut band = src[band_y0 * w..y_end * w].to_vec();
            apply_sao_plane_banded_scalar(
                &mut band, &src, w, h, band_y0, x0, y0, x_end, y_end, 2, &offsets, 0, eo_class, bd,
            );

            assert_eq!(&band[..], &serial[band_y0 * w..y_end * w]);
        }
    }

    #[test]
    fn gated_eo_applies_when_all_available() {
        let src = [10u16, 5, 10, 5];
        let mut dst = src;
        let offsets = [7i32, 0, 0, 0];
        let ok = |_cx: usize, _cy: usize, _nx: usize, _ny: usize| true;
        apply_sao_edge_offset_gated(&mut dst, &src, 4, 1, 0, 0, 4, 1, &offsets, 0, 8, &ok);
        // x=1 is a local minimum (5 between 10 and 10) → edge_idx 0 → +7 = 12.
        assert_eq!(dst[1], 12);
    }
}