hpvcd 0.3.2

/*
 * // Copyright (c) Radzivon Bartoshyk 6/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.
 */

/// intraPredAngle for modes 2..=34 (Table 8-5), indexed by mode-2.
pub(crate) static INTRA_PRED_ANGLE: [i32; 33] = [
    32, 26, 21, 17, 13, 9, 5, 2, 0, -2, -5, -9, -13, -17, -21, -26, -32, -26, -21, -17, -13, -9,
    -5, -2, 0, 2, 5, 9, 13, 17, 21, 26, 32,
];

/// invAngle for modes 11..=25 (Table 8-6), indexed by mode-11.
pub(crate) static INV_ANGLE: [i32; 15] = [
    -4096, -1638, -910, -630, -482, -390, -315, -256, -315, -390, -482, -630, -910, -1638, -4096,
];

pub(crate) const PLANAR: u8 = 0;
pub(crate) const DC: u8 = 1;

/// Pre-allocated scratch for the intra prediction pipeline (reused every TU).
pub(crate) struct IntraScratch {
    pub(crate) sub_s: Vec<u16>, // 4*N+1
    pub(crate) sub_avail: Vec<bool>,
    pub(crate) above: Vec<u16>, // 2*N+1
    pub(crate) left: Vec<u16>,
    pub(crate) fa: Vec<u16>,
    pub(crate) fl: Vec<u16>,
    pub(crate) refs_ang: Vec<i32>,          // 3*N+1
    pub(crate) pred: Vec<u16>,              // N*N
    pub(crate) raw_above: Vec<Option<u16>>, // 2*N, reused for ref gathering (no per-TU alloc)
    pub(crate) raw_left: Vec<Option<u16>>,  // 2*N
}

impl Default for IntraScratch {
    fn default() -> Self {
        Self::new()
    }
}

impl IntraScratch {
    pub(crate) fn new() -> Self {
        let max_n = 32usize;
        IntraScratch {
            sub_s: vec![0u16; 4 * max_n + 1],
            sub_avail: vec![false; 4 * max_n + 1],
            above: vec![0u16; 2 * max_n + 1],
            left: vec![0u16; 2 * max_n + 1],
            fa: vec![0u16; 2 * max_n + 1],
            fl: vec![0u16; 2 * max_n + 1],
            refs_ang: vec![0i32; 3 * max_n + 1],
            pred: vec![0u16; max_n * max_n],
            raw_above: vec![None; 2 * max_n],
            raw_left: vec![None; 2 * max_n],
        }
    }
}

/// Like `substitute_refs` but writes into caller-provided slices (length ≥ 2N+1).
#[allow(clippy::too_many_arguments)]
pub(crate) fn substitute_refs_into(
    raw_corner: Option<u16>,
    raw_above: &[Option<u16>],
    raw_left: &[Option<u16>],
    n: usize,
    neutral: u16,
    s: &mut [u16],
    avail: &mut [bool],
    above_out: &mut [u16],
    left_out: &mut [u16],
) {
    let total = 4 * n + 1;
    let s = &mut s[..total];
    let avail = &mut avail[..total];

    for v in s.iter_mut() {
        *v = 0;
    }
    for a in avail.iter_mut() {
        *a = false;
    }

    for i in 0..2 * n {
        if let Some(v) = raw_left[2 * n - 1 - i] {
            s[i] = v;
            avail[i] = true;
        }
    }
    if let Some(v) = raw_corner {
        s[2 * n] = v;
        avail[2 * n] = true;
    }
    for i in 0..2 * n {
        if let Some(v) = raw_above[i] {
            s[2 * n + 1 + i] = v;
            avail[2 * n + 1 + i] = true;
        }
    }

    if !avail.iter().any(|&a| a) {
        for v in s.iter_mut() {
            *v = neutral;
        }
    } else {
        let first = avail.iter().position(|&a| a).unwrap();
        for i in 0..first {
            s[i] = s[first];
        }
        for i in 1..total {
            if !avail[i] {
                s[i] = s[i - 1];
            }
        }
    }

    let above_out = &mut above_out[..2 * n + 1];
    let left_out = &mut left_out[..2 * n + 1];
    above_out[0] = s[2 * n];
    left_out[0] = s[2 * n];
    for i in 0..2 * n {
        above_out[1 + i] = s[2 * n + 1 + i];
        left_out[1 + i] = s[2 * n - 1 - i];
    }
}

/// Like `filter_refs` but writes into `fa_out` / `fl_out`.
#[allow(clippy::too_many_arguments)]
pub(crate) fn filter_refs_into(
    above: &[u16],
    left: &[u16],
    n: usize,
    mode: u8,
    is_luma: bool,
    intra_smoothing_enabled: bool,
    strong_intra_smoothing: bool,
    bit_depth: u8,
    fa_out: &mut [u16],
    fl_out: &mut [u16],
) {
    let len = 2 * n + 1;
    fa_out[..len].copy_from_slice(&above[..len]);
    fl_out[..len].copy_from_slice(&left[..len]);

    // RExt's intra_smoothing_disabled_flag disables the complete reference
    // sample filtering process, including both the normal [1 2 1] filter and
    // the optional strong 32x32 bilinear filter.
    if !intra_smoothing_enabled || !is_luma || n == 4 || mode == DC {
        return;
    }

    let min_dist = (mode as i32 - 26).abs().min((mode as i32 - 10).abs());
    let thresh = match n {
        8 => 7,
        16 => 1,
        32 => 0,
        _ => 8,
    };
    if !(mode == PLANAR || min_dist > thresh) {
        return;
    }

    if strong_intra_smoothing && n == 32 {
        let thr = 1 << (bit_depth - 5);
        let bl = left[2 * n] as i32;
        let tl = above[0] as i32;
        let tr = above[2 * n] as i32;
        let cl = left[n] as i32;
        let ct = above[n] as i32;
        if (bl + tl - 2 * cl).abs() < thr && (tr + tl - 2 * ct).abs() < thr {
            for ((i, fa), fl) in (1..2 * n)
                .zip(fa_out[1..2 * n].iter_mut())
                .zip(fl_out[1..2 * n].iter_mut())
            {
                *fa = ((((2 * n - i) as i32) * tl + (i as i32) * tr + n as i32) >> 6) as u16;
                *fl = ((((2 * n - i) as i32) * tl + (i as i32) * bl + n as i32) >> 6) as u16;
            }
            fa_out[0] = above[0];
            fl_out[0] = left[0];
            fa_out[2 * n] = above[2 * n];
            fl_out[2 * n] = left[2 * n];
            return;
        }
    }

    let corner = above[0] as i32;
    fa_out[0] = ((left[1] as i32 + 2 * corner + above[1] as i32 + 2) >> 2) as u16;
    fl_out[0] = fa_out[0];
    for ((i, fa), fl) in (1..2 * n)
        .zip(fa_out[1..2 * n].iter_mut())
        .zip(fl_out[1..2 * n].iter_mut())
    {
        *fa = ((above[i - 1] as i32 + 2 * above[i] as i32 + above[i + 1] as i32 + 2) >> 2) as u16;
        *fl = ((left[i - 1] as i32 + 2 * left[i] as i32 + left[i + 1] as i32 + 2) >> 2) as u16;
    }
    fa_out[2 * n] = above[2 * n];
    fl_out[2 * n] = left[2 * n];
}

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

static PREDICT: std::sync::OnceLock<PredictFn> = std::sync::OnceLock::new();

#[inline]
pub(crate) fn resolve_predict() -> PredictFn {
    *PREDICT.get_or_init(|| {
        let mut _f: PredictFn = predict_into_scalar;

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

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

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

        _f
    })
}

/// Scalar predictor fallback. Kept separate so SIMD backends can dispatch only the
/// cheap structured modes and fall back without going through OnceLock again.
#[allow(clippy::too_many_arguments)]
pub(crate) fn predict_into_scalar(
    mode: u8,
    above: &[u16],
    left: &[u16],
    n: usize,
    is_luma: bool,
    bit_depth: u8,
    out: &mut [u16],
    refs_ang: &mut [i32],
) {
    let out = &mut out[..n * n];
    match mode {
        PLANAR => {
            let tr = above[n + 1] as i32;
            let bl = left[n + 1] as i32;
            for y in 0..n {
                let out = &mut out[y * n..y * n + n];
                for (x, dst) in out.iter_mut().enumerate() {
                    let h = (n - 1 - x) as i32 * left[y + 1] as i32 + (x + 1) as i32 * tr;
                    let v = (n - 1 - y) as i32 * above[x + 1] as i32 + (y + 1) as i32 * bl;
                    *dst = ((h + v + n as i32) >> (n.trailing_zeros() as i32 + 1)) as u16;
                }
            }
        }
        DC => {
            // Mirror predict_dc exactly: N samples from above + N from left.
            let mut sum = 0i32;
            for (&above, &left) in above[1..=n].iter().zip(left[1..=n].iter()) {
                sum += above as i32 + left as i32;
            }
            let dc = (sum + n as i32) >> ((n as u32).trailing_zeros() + 1);
            out.fill(dc as u16);
            // Edge filter: luma only, and only for blocks smaller than 32×32
            if is_luma && n < 32 {
                out[0] = ((left[1] as i32 + 2 * dc + above[1] as i32 + 2) >> 2) as u16;
                for (dst, &above) in out[1..n].iter_mut().zip(above[2..n + 1].iter()) {
                    *dst = ((above as i32 + 3 * dc + 2) >> 2) as u16;
                }
                for (y, &left) in (1..n).zip(left[2..n + 1].iter()) {
                    out[y * n] = ((left as i32 + 3 * dc + 2) >> 2) as u16;
                }
            }
        }
        _ => {
            let angle = INTRA_PRED_ANGLE[mode as usize - 2];
            let vertical = mode >= 18;
            let main = if vertical { above } else { left };
            let side = if vertical { left } else { above };
            let base = n as i32;
            let refs = &mut refs_ang[..3 * n + 1];
            for (dst, &main) in refs[n..=3 * n].iter_mut().zip(main[..=2 * n].iter()) {
                *dst = main as i32;
            }
            if angle < 0 {
                let inv = INV_ANGLE[mode as usize - 11];
                let lim = (n as i32 * angle) >> 5;
                let mut k = -1i32;
                while k >= lim {
                    let idx = ((k * inv + 128) >> 8).min(2 * n as i32);
                    refs[(k + base) as usize] = side[idx.max(0) as usize] as i32;
                    k -= 1;
                }
            }
            let max_v = ((1i32 << bit_depth) - 1).max(0);
            if vertical {
                for (row, out_row) in out.chunks_exact_mut(n).enumerate() {
                    let pos = (row as i32 + 1) * angle;
                    let i_idx = pos >> 5;
                    let frac = pos & 31;
                    for (col, dst) in out_row.iter_mut().enumerate() {
                        let col = col as i32;
                        let r0 = refs[(col + i_idx + 1 + base) as usize];
                        let v = if frac == 0 {
                            r0
                        } else {
                            let r1 = refs[(col + i_idx + 2 + base) as usize];
                            ((32 - frac) * r0 + frac * r1 + 16) >> 5
                        };
                        *dst = v.clamp(0, max_v) as u16;
                    }
                }
            } else {
                for (col, out_row) in out.chunks_exact_mut(n).enumerate() {
                    let col = col as i32;
                    for (row, dst) in out_row.iter_mut().enumerate() {
                        let pos = (row as i32 + 1) * angle;
                        let i_idx = pos >> 5;
                        let frac = pos & 31;
                        let r0 = refs[(col + i_idx + 1 + base) as usize];
                        let v = if frac == 0 {
                            r0
                        } else {
                            let r1 = refs[(col + i_idx + 2 + base) as usize];
                            ((32 - frac) * r0 + frac * r1 + 16) >> 5
                        };
                        *dst = v.clamp(0, max_v) as u16;
                    }
                }
            }
            if is_luma && n < 32 {
                let max = ((1i32 << bit_depth) - 1).max(0);
                if mode == 26 {
                    for (y, &left) in left[1..n + 1].iter().enumerate() {
                        let v = above[1] as i32 + ((left as i32 - above[0] as i32) >> 1);
                        out[y * n] = v.clamp(0, max) as u16;
                    }
                } else if mode == 10 {
                    for (dst, &above_p1) in out[..n].iter_mut().zip(above[1..n + 1].iter()) {
                        let v = left[1] as i32 + ((above_p1 as i32 - above[0] as i32) >> 1);
                        *dst = v.clamp(0, max) as u16;
                    }
                }
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::{PLANAR, filter_refs_into};

    #[test]
    fn disabled_reference_smoothing_preserves_samples() {
        let n = 8;
        let above: Vec<u16> = (0..=2 * n).map(|i| 100 + (i * i) as u16).collect();
        let left: Vec<u16> = (0..=2 * n).map(|i| 100 + (3 * i * i) as u16).collect();
        let mut filtered_above = vec![0; 2 * n + 1];
        let mut filtered_left = vec![0; 2 * n + 1];

        filter_refs_into(
            &above,
            &left,
            n,
            PLANAR,
            true,
            false,
            true,
            10,
            &mut filtered_above,
            &mut filtered_left,
        );

        assert_eq!(filtered_above, above);
        assert_eq!(filtered_left, left);
    }

    #[test]
    fn disabled_reference_smoothing_skips_strong_filter() {
        let n = 32;
        let mut above = vec![777; 2 * n + 1];
        let mut left = vec![555; 2 * n + 1];
        above[0] = 100;
        above[n] = 200;
        above[2 * n] = 300;
        left[0] = 100;
        left[n] = 250;
        left[2 * n] = 400;

        let mut disabled_above = vec![0; 2 * n + 1];
        let mut disabled_left = vec![0; 2 * n + 1];
        filter_refs_into(
            &above,
            &left,
            n,
            PLANAR,
            true,
            false,
            true,
            10,
            &mut disabled_above,
            &mut disabled_left,
        );
        assert_eq!(disabled_above, above);
        assert_eq!(disabled_left, left);

        let mut enabled_above = vec![0; 2 * n + 1];
        let mut enabled_left = vec![0; 2 * n + 1];
        filter_refs_into(
            &above,
            &left,
            n,
            PLANAR,
            true,
            true,
            true,
            10,
            &mut enabled_above,
            &mut enabled_left,
        );
        assert_ne!(enabled_above, above);
        assert_ne!(enabled_left, left);
    }
}