rav1d-safe 0.6.0

#![deny(unsafe_op_in_unsafe_fn)]

use crate::include::common::attributes::clz;
use crate::include::common::intops::inv_recenter;
use crate::include::common::intops::ulog2;
use crate::src::c_arc::CArc;
use crate::src::cpu::CpuFlags;
use cfg_if::cfg_if;
use std::ffi::c_int;
use std::ffi::c_uint;
use std::mem;
use std::ops::Deref;
use std::ops::DerefMut;
#[cfg(asm_msac)]
use std::ops::Range;
#[cfg(asm_msac)]
use std::ptr;
#[cfg(asm_msac)]
use std::slice;

#[cfg(all(asm_msac, target_feature = "sse2"))]
unsafe extern "C" {
    fn dav1d_msac_decode_hi_tok_sse2(s: *mut MsacAsmContext, cdf: *mut u16) -> c_uint;
    fn dav1d_msac_decode_bool_sse2(s: *mut MsacAsmContext, f: c_uint) -> c_uint;
    fn dav1d_msac_decode_bool_equi_sse2(s: *mut MsacAsmContext) -> c_uint;
    fn dav1d_msac_decode_bool_adapt_sse2(s: *mut MsacAsmContext, cdf: *mut u16) -> c_uint;
    fn dav1d_msac_decode_symbol_adapt16_sse2(
        s: &mut MsacAsmContext,
        cdf: *mut u16,
        n_symbols: usize,
        _cdf_len: usize,
    ) -> c_uint;
    fn dav1d_msac_decode_symbol_adapt8_sse2(
        s: *mut MsacAsmContext,
        cdf: *mut u16,
        n_symbols: usize,
    ) -> c_uint;
    fn dav1d_msac_decode_symbol_adapt4_sse2(
        s: *mut MsacAsmContext,
        cdf: *mut u16,
        n_symbols: usize,
    ) -> c_uint;
}

#[cfg(all(asm_msac, target_arch = "x86_64"))]
unsafe extern "C" {
    fn dav1d_msac_decode_symbol_adapt16_avx2(
        s: &mut MsacAsmContext,
        cdf: *mut u16,
        n_symbols: usize,
        _cdf_len: usize,
    ) -> c_uint;
}

#[cfg(all(asm_msac, target_feature = "neon"))]
unsafe extern "C" {
    fn dav1d_msac_decode_hi_tok_neon(s: *mut MsacAsmContext, cdf: *mut u16) -> c_uint;
    fn dav1d_msac_decode_bool_neon(s: *mut MsacAsmContext, f: c_uint) -> c_uint;
    fn dav1d_msac_decode_bool_equi_neon(s: *mut MsacAsmContext) -> c_uint;
    fn dav1d_msac_decode_bool_adapt_neon(s: *mut MsacAsmContext, cdf: *mut u16) -> c_uint;
    fn dav1d_msac_decode_symbol_adapt16_neon(
        s: *mut MsacAsmContext,
        cdf: *mut u16,
        n_symbols: usize,
    ) -> c_uint;
    fn dav1d_msac_decode_symbol_adapt8_neon(
        s: *mut MsacAsmContext,
        cdf: *mut u16,
        n_symbols: usize,
    ) -> c_uint;
    fn dav1d_msac_decode_symbol_adapt4_neon(
        s: *mut MsacAsmContext,
        cdf: *mut u16,
        n_symbols: usize,
    ) -> c_uint;
}

pub struct Rav1dMsacDSPContext {
    #[cfg(asm_msac)]
    symbol_adapt16: unsafe extern "C" fn(
        s: &mut MsacAsmContext,
        cdf: *mut u16,
        n_symbols: usize,
        _cdf_len: usize,
    ) -> c_uint,
}

impl Rav1dMsacDSPContext {
    pub const fn default() -> Self {
        cfg_if! {
            if #[cfg(asm_msac)] {
                Self {
                    symbol_adapt16: rav1d_msac_decode_symbol_adapt_c,
                }
            } else {
                Self {}
            }
        }
    }

    #[cfg(all(asm_msac, any(target_arch = "x86", target_arch = "x86_64")))]
    #[inline(always)]
    const fn init_x86(mut self, flags: CpuFlags) -> Self {
        if !flags.contains(CpuFlags::SSE2) {
            return self;
        }

        self.symbol_adapt16 = dav1d_msac_decode_symbol_adapt16_sse2;

        #[cfg(target_arch = "x86_64")]
        {
            if !flags.contains(CpuFlags::AVX2) {
                return self;
            }

            self.symbol_adapt16 = dav1d_msac_decode_symbol_adapt16_avx2;
        }

        self
    }

    #[cfg(all(asm_msac, any(target_arch = "arm", target_arch = "aarch64")))]
    #[inline(always)]
    const fn init_arm(self, _flags: CpuFlags) -> Self {
        self
    }

    #[inline(always)]
    const fn init(self, flags: CpuFlags) -> Self {
        #[cfg(asm_msac)]
        {
            #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
            {
                return self.init_x86(flags);
            }
            #[cfg(any(target_arch = "arm", target_arch = "aarch64"))]
            {
                return self.init_arm(flags);
            }
        }

        #[allow(unreachable_code)] // Reachable on some #[cfg]s.
        {
            let _ = flags;
            self
        }
    }

    pub const fn new(flags: CpuFlags) -> Self {
        Self::default().init(flags)
    }
}

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

pub type EcWin = usize;

/// For asm builds: raw pointers into `MsacContext::data` for C FFI layout.
///
/// # Safety
///
/// [`Self`] must be the first field of [`MsacAsmContext`] for asm layout purposes,
/// and that [`MsacAsmContext`] must be a field of [`MsacContext`].
/// And [`Self::pos`] and [`Self::end`] must be either [`ptr::null`],
/// or [`Self::pos`] must point into (or the end of) [`MsacContext::data`],
/// and [`Self::end`] must point to the end of [`MsacContext::data`],
/// where [`MsacContext::data`] is part of the [`MsacContext`]
/// containing [`MsacAsmContext`] and thus also [`Self`].
#[cfg(asm_msac)]
#[repr(C)]
struct MsacAsmContextBuf {
    pos: *const u8,
    end: *const u8,
}

#[cfg(asm_msac)]
/// SAFETY: [`MsacAsmContextBuf`] is always contained in [`MsacAsmContext::buf`],
/// which is always contained in [`MsacContext::asm`], whose [`MsacContext::data`] field
/// is what is stored in [`MsacAsmContextBuf::pos`] and [`MsacAsmContextBuf::end`].
/// Since [`MsacContext::data`] is [`Send`], [`MsacAsmContextBuf`] is also [`Send`].
#[allow(unsafe_code)]
unsafe impl Send for MsacAsmContextBuf {}

#[cfg(asm_msac)]
/// SAFETY: [`MsacAsmContextBuf`] is always contained in [`MsacAsmContext::buf`],
/// which is always contained in [`MsacContext::asm`], whose [`MsacContext::data`] field
/// is what is stored in [`MsacAsmContextBuf::pos`] and [`MsacAsmContextBuf::end`].
/// Since [`MsacContext::data`] is [`Sync`], [`MsacAsmContextBuf`] is also [`Sync`].
#[allow(unsafe_code)]
unsafe impl Sync for MsacAsmContextBuf {}

#[cfg(asm_msac)]
impl Default for MsacAsmContextBuf {
    fn default() -> Self {
        Self {
            pos: ptr::null(),
            end: ptr::null(),
        }
    }
}

#[cfg(asm_msac)]
impl From<&[u8]> for MsacAsmContextBuf {
    fn from(value: &[u8]) -> Self {
        let Range { start, end } = value.as_ptr_range();
        Self { pos: start, end }
    }
}

/// For non-asm builds: byte indices into `MsacContext::data`.
///
/// Uses `usize` indices instead of raw pointers, so this type is
/// automatically `Send + Sync` without any unsafe impls.
#[cfg(not(asm_msac))]
#[derive(Default)]
struct MsacAsmContextBuf {
    /// Byte index of current position within `MsacContext::data`.
    pos: usize,
    /// Byte index of end (always `data.len()`).
    /// Maintained for layout parity with asm but not read in non-asm path.
    #[allow(dead_code)]
    end: usize,
}

#[cfg_attr(asm_msac, repr(C))]
pub struct MsacAsmContext {
    buf: MsacAsmContextBuf,
    pub dif: EcWin,
    pub rng: c_uint,
    pub cnt: c_int,
    allow_update_cdf: c_int,
    #[cfg(all(asm_msac, target_arch = "x86_64"))]
    symbol_adapt16: unsafe extern "C" fn(
        s: &mut MsacAsmContext,
        cdf: *mut u16,
        n_symbols: usize,
        _cdf_len: usize,
    ) -> c_uint,
}

// Manual impl needed: `symbol_adapt16` field (behind `cfg(asm_msac)`) is a fn ptr
// that doesn't implement Default, so #[derive(Default)] fails on asm builds.
#[allow(clippy::derivable_impls)]
impl Default for MsacAsmContext {
    fn default() -> Self {
        Self {
            buf: Default::default(),
            dif: Default::default(),
            rng: Default::default(),
            cnt: Default::default(),
            allow_update_cdf: Default::default(),

            #[cfg(all(asm_msac, target_arch = "x86_64"))]
            symbol_adapt16: Rav1dMsacDSPContext::default().symbol_adapt16,
        }
    }
}

impl MsacAsmContext {
    fn allow_update_cdf(&self) -> bool {
        self.allow_update_cdf != 0
    }
}

#[derive(Default)]
pub struct MsacContext {
    asm: MsacAsmContext,
    data: Option<CArc<[u8]>>,
}

impl Deref for MsacContext {
    type Target = MsacAsmContext;

    fn deref(&self) -> &Self::Target {
        &self.asm
    }
}

impl DerefMut for MsacContext {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.asm
    }
}

impl MsacContext {
    pub fn data(&self) -> &[u8] {
        &**self.data.as_ref().unwrap()
    }

    pub fn buf_index(&self) -> usize {
        cfg_if! {
            if #[cfg(asm_msac)] {
                // We safely subtract instead of unsafely use `ptr::offset_from`
                // as asm sets `buf_pos`, so we don't need to rely on its safety,
                // and because codegen is no less optimal this way.
                self.buf.pos as usize - self.data().as_ptr() as usize
            } else {
                // Without asm, pos is already a byte index.
                self.buf.pos
            }
        }
    }

    fn with_buf(&mut self, mut f: impl FnMut(&[u8]) -> &[u8]) {
        let data = &**self.data.as_ref().unwrap();
        let buf = &data[self.buf_index()..];
        let buf = f(buf);
        cfg_if! {
            if #[cfg(asm_msac)] {
                self.buf.pos = buf.as_ptr();
            } else {
                // Compute the new byte index from the returned sub-slice.
                self.buf.pos = buf.as_ptr() as usize - data.as_ptr() as usize;
            }
        }
        // We don't actually need to set `self.buf_end` since it has not changed.
    }
}

/// Return value uses `n` bits.
#[inline]
pub fn rav1d_msac_decode_bools(s: &mut MsacContext, n: u8) -> c_uint {
    let mut v = 0;
    for _ in 0..n {
        v = v << 1 | rav1d_msac_decode_bool_equi(s) as c_uint;
    }
    v
}

#[inline]
pub fn rav1d_msac_decode_uniform(s: &mut MsacContext, n: c_uint) -> c_int {
    assert!(n > 0);
    let l = ulog2(n) as u8 + 1;
    assert!(l > 1);
    let m = (1 << l) - n;
    let v = rav1d_msac_decode_bools(s, l - 1);
    (if v < m {
        v
    } else {
        (v << 1) - m + rav1d_msac_decode_bool_equi(s) as c_uint
    }) as c_int
}

const EC_PROB_SHIFT: c_uint = 6;
const EC_MIN_PROB: c_uint = 4;
const _: () = assert!(EC_MIN_PROB <= (1 << EC_PROB_SHIFT) / 16);

const EC_WIN_SIZE: usize = mem::size_of::<EcWin>() << 3;

/// Branchless CDF update after symbol decode.
///
/// For i < val: cdf[i] += (32768 - cdf[i]) >> rate (probability increases)
/// For i >= val: cdf[i] -= cdf[i] >> rate (probability decreases)
///
/// Uses mask-select to avoid branches on the val boundary.
#[inline(always)]
fn update_cdf(cdf: &mut [u16], n: usize, val: usize, rate: u16, count: u16) {
    #[cfg(all(not(asm_msac), target_arch = "x86_64"))]
    {
        use archmage::SimdToken as _;
        if n == 3 && rate < 16 {
            if let Some(token) = archmage::X64V1Token::summon() {
                return update_cdf3_simd(
                    token,
                    (&mut cdf[..4]).try_into().unwrap(),
                    val.min(3),
                    rate,
                    count,
                );
            }
        }
    }
    for i in 0..n {
        let mask = ((i < val) as u16).wrapping_neg(); // 0xFFFF if below val, 0 otherwise
        let delta_up = (32768u16.wrapping_sub(cdf[i])) >> rate;
        let delta_dn = cdf[i] >> rate;
        // Apply increase (delta_up) if below val, decrease (delta_dn) if at/above val
        cdf[i] = cdf[i]
            .wrapping_add(delta_up & mask)
            .wrapping_sub(delta_dn & !mask);
    }
    cdf[n] = count + (count < 32) as u16;
}

/// A leaf kernel: keep normalization/refill outside the SIMD call so no coder
/// state or vector registers must survive a refill. Each lane keeps the
/// scalar wrapping arithmetic, including zero and high u16 probabilities.
#[cfg(all(not(asm_msac), target_arch = "x86_64"))]
#[archmage::arcane]
fn update_cdf3_simd(
    _token: archmage::X64V1Token,
    cdf: &mut [u16; 4],
    val: usize,
    rate: u16,
    count: u16,
) {
    use core::arch::x86_64::*;
    use zerocopy::IntoBytes;

    let p = _mm_cvtsi64_si128(i64::from_ne_bytes(cdf.as_bytes().try_into().unwrap()));
    let below = _mm_cmpgt_epi16(
        _mm_set1_epi16(val as i16),
        _mm_setr_epi16(0, 1, 2, 3, 4, 5, 6, 7),
    );
    let shift = _mm_cvtsi32_si128(i32::from(rate));
    let up = _mm_srl_epi16(_mm_sub_epi16(_mm_set1_epi16(i16::MIN), p), shift);
    let down = _mm_srl_epi16(p, shift);
    let updated = _mm_sub_epi16(
        _mm_add_epi16(p, _mm_and_si128(below, up)),
        _mm_andnot_si128(below, down),
    );
    // The count is the fourth lane. Replacing it before the sole store also
    // leaves every byte after the four-entry CDF untouched.
    let updated = _mm_insert_epi16::<3>(updated, i32::from(count + u16::from(count < 32)));
    cdf.as_mut_bytes()
        .copy_from_slice(&_mm_cvtsi128_si64(updated).to_ne_bytes());
}

#[inline]
#[cold]
fn ctx_refill(s: &mut MsacContext) {
    let mut c = (EC_WIN_SIZE as c_int) - 24 - s.cnt;
    let mut dif = s.dif;
    s.with_buf(|mut buf| {
        // Bulk path: on 64-bit targets where `EcWin = u64`, dav1d's reference
        // ASM does this refill as `mov r, [buf]; bswap r; not r; shift; or dif, r`
        // — a single 8-byte big-endian load fills the entire 64-bit `dif`
        // window in one shot. This is the canonical hot path of the entropy
        // decoder (called for every `cnt < 0` refill, ~1-3 times per coefficient
        // group). The per-byte loop below is kept as the tail (when buf has
        // fewer than 8 bytes left) and as the EC_WIN_SIZE < 64 fallback.
        if EC_WIN_SIZE >= 64 && buf.len() >= 8 {
            // `c` at entry is in [0, EC_WIN_SIZE - 24], so for EC_WIN_SIZE=64
            // it is in [0, 40] when called from a steady-state ctx_norm
            // (cnt was in [-7, 16) right before refill), and at most 63
            // including initialization paths. Either way, 8 bytes is enough
            // to drive `c` strictly negative.
            let bytes: [u8; 8] = buf[..8].try_into().unwrap();
            // Read as big-endian: byte 0 goes to bits 56..63 of `raw`.
            let raw = u64::from_be_bytes(bytes);
            let inv = !raw; // matches `^ 0xff` per byte
            // We will consume exactly k = (c/8) + 1 bytes; bytes k..8 in `inv`
            // must not contribute to `dif`. Mask off the low (8-k) bytes.
            let k = (c as usize / 8) + 1; // 1..=8
            let keep_bits = (8 * k) as u32; // 8..=64
            let mask: u64 = if keep_bits == 64 {
                !0u64
            } else {
                (!0u64) << (64 - keep_bits)
            };
            let inv = inv & mask;
            // Byte 0 in `inv` sits at bits [56, 63]. We want it at bits
            // [c, c+7]. Shift by (c - 56): right-shift when c < 56, left
            // when c >= 56. Each subsequent byte i lines up automatically
            // since they're naturally 8 bits apart.
            let shift = (c as i32) - 56;
            let contribution: u64 = if shift >= 0 {
                inv << (shift as u32)
            } else {
                inv >> ((-shift) as u32)
            };
            dif |= contribution as EcWin;
            buf = &buf[k..];
            c -= 8 * k as c_int; // becomes negative -> done
        }
        // Tail / fallback per-byte loop.
        loop {
            if c < 0 {
                break;
            }
            if buf.is_empty() {
                // set remaining bits to 1;
                dif |= !(!(0xff as EcWin) << c);
                break;
            }
            dif |= ((buf[0] ^ 0xff) as EcWin) << c;
            buf = &buf[1..];
            c -= 8;
        }
        buf
    });
    s.dif = dif;
    s.cnt = (EC_WIN_SIZE as c_int) - 24 - c;
}

#[inline(always)]
fn ctx_norm(s: &mut MsacContext, dif: EcWin, rng: c_uint) {
    let d = 15 ^ (31 ^ clz(rng));
    let cnt = s.cnt;
    debug_assert!(rng <= 65535);
    s.dif = dif << d;
    s.rng = rng << d;
    // At EOB, `cnt - d` legitimately wraps negative; the unsigned compare
    // below detects this and triggers a refill. With debug overflow checks
    // (and especially with `panic = abort` in dev), the plain subtraction
    // panics on every EOB. Use `wrapping_sub` to match dav1d's C semantics.
    s.cnt = cnt.wrapping_sub(d);
    // unsigned compare avoids redundant refills at eob
    if (cnt as u32) < (d as u32) {
        ctx_refill(s);
    }
}

#[inline(always)]
#[cfg_attr(
    all(asm_msac, any(target_feature = "sse2", target_feature = "neon")),
    allow(dead_code)
)]
fn rav1d_msac_decode_bool_equi_rust(s: &mut MsacContext) -> bool {
    let r = s.rng;
    let mut dif = s.dif;
    debug_assert!(dif >> (EC_WIN_SIZE - 16) < r as EcWin);
    let mut v = (r >> 8 << 7) + EC_MIN_PROB;
    let vw = (v as EcWin) << (EC_WIN_SIZE - 16);
    let ret = dif >= vw;
    dif -= (ret as EcWin) * vw;
    v = v.wrapping_add((ret as c_uint) * (r.wrapping_sub(2 * v)));
    ctx_norm(s, dif, v);
    !ret
}

#[inline(always)]
#[cfg_attr(
    all(asm_msac, any(target_feature = "sse2", target_feature = "neon")),
    allow(dead_code)
)]
fn rav1d_msac_decode_bool_rust(s: &mut MsacContext, f: c_uint) -> bool {
    let r = s.rng;
    let mut dif = s.dif;
    debug_assert!(dif >> (EC_WIN_SIZE - 16) < r as EcWin);
    let mut v = ((r >> 8) * (f >> EC_PROB_SHIFT) >> (7 - EC_PROB_SHIFT)) + EC_MIN_PROB;
    let vw = (v as EcWin) << (EC_WIN_SIZE - 16);
    let ret = dif >= vw;
    dif -= (ret as EcWin) * vw;
    v = v.wrapping_add((ret as c_uint) * (r.wrapping_sub(2 * v)));
    ctx_norm(s, dif, v);
    !ret
}

pub fn rav1d_msac_decode_subexp(s: &mut MsacContext, r#ref: c_uint, n: c_uint, mut k: u8) -> c_int {
    assert!(n >> k == 8);
    let mut a = 0;
    if rav1d_msac_decode_bool_equi(s) {
        if rav1d_msac_decode_bool_equi(s) {
            k += rav1d_msac_decode_bool_equi(s) as u8 + 1;
        }
        a = 1 << k;
    }
    let v = rav1d_msac_decode_bools(s, k) + a;
    (if r#ref * 2 <= n {
        inv_recenter(r#ref, v)
    } else {
        n - 1 - inv_recenter(n - 1 - r#ref, v)
    }) as c_int
}

/// Return value is in the range `0..=n_symbols`.
///
/// `n_symbols` is in the range `0..16`, so it is really a `u4`.
fn rav1d_msac_decode_symbol_adapt_rust(s: &mut MsacContext, cdf: &mut [u16], n_symbols: u8) -> u8 {
    let c = (s.dif >> (EC_WIN_SIZE - 16)) as c_uint;
    let r = s.rng >> 8;
    let mut u;
    let mut v = s.rng;
    let mut val = 0;
    debug_assert!(n_symbols < 16);
    debug_assert!(cdf[n_symbols as usize] <= 32);
    loop {
        u = v;
        v = r * ((cdf[val as usize] >> EC_PROB_SHIFT) as c_uint);
        v >>= 7 - EC_PROB_SHIFT;
        v += EC_MIN_PROB * ((n_symbols as c_uint) - val);
        if !(c < v) {
            break;
        }
        val += 1;
    }
    debug_assert!(u <= s.rng);
    ctx_norm(
        s,
        s.dif.wrapping_sub((v as EcWin) << (EC_WIN_SIZE - 16)),
        u - v,
    );
    if s.allow_update_cdf() {
        let n_usize = n_symbols as usize;
        let count = cdf[n_usize];
        let rate = 4 + (count >> 4) + (n_symbols > 2) as u16;
        let val = val as usize;
        update_cdf(cdf, n_usize, val, rate, count);
    }
    debug_assert!(val <= n_symbols as _);
    val as u8
}

/// # Safety
///
/// Must be called through [`Rav1dMsacDSPContext::symbol_adapt16`]
/// in [`rav1d_msac_decode_symbol_adapt16`].
#[cfg(asm_msac)]
#[cfg_attr(not(all(asm_msac, target_arch = "x86_64")), allow(dead_code))]
#[deny(unsafe_op_in_unsafe_fn)]
unsafe extern "C" fn rav1d_msac_decode_symbol_adapt_c(
    s: &mut MsacAsmContext,
    cdf: *mut u16,
    n_symbols: usize,
    cdf_len: usize,
) -> c_uint {
    // SAFETY: In the `rav1d_msac_decode_symbol_adapt16` caller,
    // `&mut s.asm` is passed, so we can reverse this to get back `s`.
    // The `.sub` is safe since were are subtracting the offset of `asm` within `s`,
    // so that will stay in bounds of the `s: MsacContext` allocated object.
    let s = unsafe {
        &mut *ptr::from_mut(s)
            .sub(mem::offset_of!(MsacContext, asm))
            .cast::<MsacContext>()
    };

    // SAFETY: This is only called from [`dav1d_msac_decode_symbol_adapt16`],
    // where it comes from `cdf.len()`.
    let cdf = unsafe { slice::from_raw_parts_mut(cdf, cdf_len) };

    rav1d_msac_decode_symbol_adapt_rust(s, cdf, n_symbols as u8) as c_uint
}

#[inline(always)]
#[cfg_attr(
    all(asm_msac, any(target_feature = "sse2", target_feature = "neon")),
    allow(dead_code)
)]
fn rav1d_msac_decode_bool_adapt_rust(s: &mut MsacContext, cdf: &mut [u16; 2]) -> bool {
    let bit = rav1d_msac_decode_bool(s, cdf[0] as c_uint);
    if s.allow_update_cdf() {
        let count = cdf[1];
        let rate = 4 + (count >> 4);
        update_cdf(cdf, 1, bit as usize, rate, count);
    }
    bit
}

/// Return value is in the range `0..=15`.
#[inline(always)]
#[cfg_attr(
    all(asm_msac, any(target_feature = "sse2", target_feature = "neon")),
    allow(dead_code)
)]
fn rav1d_msac_decode_hi_tok_rust(s: &mut MsacContext, cdf: &mut [u16; 4]) -> u8 {
    let mut tok_br = rav1d_msac_decode_symbol_adapt4(s, cdf, 3);
    let mut tok = 3 + tok_br;
    if tok_br == 3 {
        tok_br = rav1d_msac_decode_symbol_adapt4(s, cdf, 3);
        tok = 6 + tok_br;
        if tok_br == 3 {
            tok_br = rav1d_msac_decode_symbol_adapt4(s, cdf, 3);
            tok = 9 + tok_br;
            if tok_br == 3 {
                tok = 12 + rav1d_msac_decode_symbol_adapt4(s, cdf, 3);
            }
        }
    }
    tok
}

// ============================================================================
// Branchless scalar implementations (used when asm is disabled)
// ============================================================================

/// Branchless implementation of symbol_adapt for n_symbols <= 3 (adapt4).
///
/// Eliminates branch misprediction from the serial comparison loop by
/// computing all v values and counting matches branchlessly.
#[cfg(not(asm_msac))]
#[inline(always)]
fn rav1d_msac_decode_symbol_adapt4_branchless(
    s: &mut MsacContext,
    cdf: &mut [u16],
    n_symbols: u8,
) -> u8 {
    debug_assert!(n_symbols > 0 && n_symbols <= 3);
    let c = (s.dif >> (EC_WIN_SIZE - 16)) as c_uint;
    let r = s.rng >> 8;
    let n = n_symbols as c_uint;

    // Compute all v values (at most 3 for adapt4)
    // NB: >> has lower precedence than + in Rust, so parens are required
    let v0 = (r * ((cdf[0] >> EC_PROB_SHIFT) as c_uint) >> (7 - EC_PROB_SHIFT)) + EC_MIN_PROB * n;
    let v1 = if n > 1 {
        (r * ((cdf[1] >> EC_PROB_SHIFT) as c_uint) >> (7 - EC_PROB_SHIFT)) + EC_MIN_PROB * (n - 1)
    } else {
        0
    };
    let v2 = if n > 2 {
        (r * ((cdf[2] >> EC_PROB_SHIFT) as c_uint) >> (7 - EC_PROB_SHIFT)) + EC_MIN_PROB * (n - 2)
    } else {
        0
    };

    // Branchless: count how many v[i] have c < v[i]
    // Since v is monotonically decreasing, this gives the first index where c >= v
    let val = (c < v0) as u32 + (c < v1) as u32 + (c < v2) as u32;
    debug_assert!(val <= n);

    // v_arr[0..3] for indexed access, with sentinel for boundary
    let v_arr = [v0, v1, v2, 0];
    let u = if val == 0 {
        s.rng
    } else {
        v_arr[val as usize - 1]
    };
    let v_val = v_arr[val as usize];

    ctx_norm(
        s,
        s.dif.wrapping_sub((v_val as EcWin) << (EC_WIN_SIZE - 16)),
        u - v_val,
    );

    if s.allow_update_cdf() {
        let n_usize = n_symbols as usize;
        let count = cdf[n_usize];
        let rate = 4 + (count >> 4) + (n_symbols > 2) as u16;
        update_cdf(cdf, n_usize, val as usize, rate, count);
    }

    val as u8
}

/// Branchless implementation of symbol_adapt for n_symbols <= 7 (adapt8).
///
/// Same strategy as adapt4 but handles up to 7 symbols.
#[cfg(not(asm_msac))]
#[inline(always)]
fn rav1d_msac_decode_symbol_adapt8_branchless(
    s: &mut MsacContext,
    cdf: &mut [u16],
    n_symbols: u8,
) -> u8 {
    debug_assert!(n_symbols > 0 && n_symbols <= 7);
    let c = (s.dif >> (EC_WIN_SIZE - 16)) as c_uint;
    let r = s.rng >> 8;
    let n = n_symbols as c_uint;

    // Compute v[i] = r * (cdf[i] >> 6) >> 1 + 4 * (n - i)
    // Use a fixed-size array, compute only valid entries
    // NB: >> has lower precedence than + in Rust, so parens are required
    let mut v = [0u32; 8]; // v[7] stays 0 as sentinel
    for i in 0..n_symbols as usize {
        v[i] = (r * ((cdf[i] >> EC_PROB_SHIFT) as c_uint) >> (7 - EC_PROB_SHIFT))
            + EC_MIN_PROB * (n - i as c_uint);
    }

    // Branchless: count how many v[i] have c < v[i]
    let mut val = 0u32;
    for i in 0..n_symbols as usize {
        val += (c < v[i]) as u32;
    }
    debug_assert!(val <= n);

    let u = if val == 0 { s.rng } else { v[val as usize - 1] };
    let v_val = v[val as usize];

    ctx_norm(
        s,
        s.dif.wrapping_sub((v_val as EcWin) << (EC_WIN_SIZE - 16)),
        u - v_val,
    );

    if s.allow_update_cdf() {
        let n_usize = n_symbols as usize;
        let count = cdf[n_usize];
        let rate = 4 + (count >> 4) + (n_symbols > 2) as u16;
        update_cdf(cdf, n_usize, val as usize, rate, count);
    }

    val as u8
}

impl MsacContext {
    pub fn new(data: CArc<[u8]>, disable_cdf_update_flag: bool, dsp: &Rav1dMsacDSPContext) -> Self {
        let buf = {
            cfg_if! {
                if #[cfg(asm_msac)] {
                    MsacAsmContextBuf::from(data.as_ref())
                } else {
                    MsacAsmContextBuf { pos: 0, end: data.as_ref().len() }
                }
            }
        };
        let asm = MsacAsmContext {
            buf,
            dif: 0,
            rng: 0x8000,
            cnt: -15,
            allow_update_cdf: (!disable_cdf_update_flag).into(),
            #[cfg(all(asm_msac, target_arch = "x86_64"))]
            symbol_adapt16: dsp.symbol_adapt16,
        };
        let mut s = Self {
            asm,
            data: Some(data),
        };
        let _ = dsp; // Silence unused warnings when asm is off.
        ctx_refill(&mut s);
        s
    }
}

/// Return value is in the range `0..=n_symbols`.
///
/// `n_symbols` is in the range `0..4`.
#[inline(always)]
pub fn rav1d_msac_decode_symbol_adapt4(s: &mut MsacContext, cdf: &mut [u16], n_symbols: u8) -> u8 {
    debug_assert!(n_symbols < 4);
    let ret;
    cfg_if! {
        if #[cfg(all(asm_msac, target_feature = "sse2"))] {
            // SAFETY: `checkasm` has verified that it is equivalent to [`dav1d_msac_decode_symbol_adapt_rust`].
            ret = unsafe {
                dav1d_msac_decode_symbol_adapt4_sse2(&mut s.asm, cdf.as_mut_ptr(), n_symbols as usize)
            };
        } else if #[cfg(all(asm_msac, target_feature = "neon"))] {
            // SAFETY: `checkasm` has verified that it is equivalent to [`dav1d_msac_decode_symbol_adapt_rust`].
            ret = unsafe {
                dav1d_msac_decode_symbol_adapt4_neon(&mut s.asm, cdf.as_mut_ptr(), n_symbols as usize)
            };
        } else if #[cfg(not(asm_msac))] {
            ret = rav1d_msac_decode_symbol_adapt4_branchless(s, cdf, n_symbols);
        } else {
            ret = rav1d_msac_decode_symbol_adapt_rust(s, cdf, n_symbols);
        }
    }
    debug_assert!(ret < 4);
    ret as u8 % 4
}

/// Return value is in the range `0..=n_symbols`.
///
/// `n_symbols` is in the range `0..8`.
#[inline(always)]
pub fn rav1d_msac_decode_symbol_adapt8(s: &mut MsacContext, cdf: &mut [u16], n_symbols: u8) -> u8 {
    debug_assert!(n_symbols < 8);
    let ret;
    cfg_if! {
        if #[cfg(all(asm_msac, target_feature = "sse2"))] {
            // SAFETY: `checkasm` has verified that it is equivalent to [`dav1d_msac_decode_symbol_adapt_rust`].
            ret = unsafe {
                dav1d_msac_decode_symbol_adapt8_sse2(&mut s.asm, cdf.as_mut_ptr(), n_symbols as usize)
            };
        } else if #[cfg(all(asm_msac, target_feature = "neon"))] {
            // SAFETY: `checkasm` has verified that it is equivalent to [`dav1d_msac_decode_symbol_adapt_rust`].
            ret = unsafe {
                dav1d_msac_decode_symbol_adapt8_neon(&mut s.asm, cdf.as_mut_ptr(), n_symbols as usize)
            };
        } else if #[cfg(not(asm_msac))] {
            ret = rav1d_msac_decode_symbol_adapt8_branchless(s, cdf, n_symbols);
        } else {
            ret = rav1d_msac_decode_symbol_adapt_rust(s, cdf, n_symbols);
        }
    }
    debug_assert!(ret < 8);
    ret as u8 % 8
}

/// Return value is in the range `0..=n_symbols`.
///
/// `n_symbols` is in the range `0..16`.
#[inline(always)]
#[cfg_attr(asm_msac, allow(unsafe_code))]
pub fn rav1d_msac_decode_symbol_adapt16(s: &mut MsacContext, cdf: &mut [u16], n_symbols: u8) -> u8 {
    debug_assert!(n_symbols < 16);
    let ret;
    cfg_if! {
        if #[cfg(all(asm_msac, target_arch = "x86_64"))] {
            // SAFETY: `checkasm` has verified that it is equivalent to [`dav1d_msac_decode_symbol_adapt_rust`].
            ret = unsafe {
                (s.symbol_adapt16)(&mut s.asm, cdf.as_mut_ptr(), n_symbols as usize, cdf.len())
            };
        } else if #[cfg(all(asm_msac, target_feature = "sse2"))] {
            // SAFETY: `checkasm` has verified that it is equivalent to [`dav1d_msac_decode_symbol_adapt_rust`].
            ret = unsafe {
                dav1d_msac_decode_symbol_adapt16_sse2(&mut s.asm, cdf.as_mut_ptr(), n_symbols as usize, cdf.len())
            };
        } else if #[cfg(all(asm_msac, target_feature = "neon"))] {
            // SAFETY: `checkasm` has verified that it is equivalent to [`dav1d_msac_decode_symbol_adapt_rust`].
            ret = unsafe {
                dav1d_msac_decode_symbol_adapt16_neon(&mut s.asm, cdf.as_mut_ptr(), n_symbols as usize)
            };
        } else if #[cfg(not(asm_msac))] {
            // Serial loop is faster than branchless for adapt16: typical AV1 distributions
            // exit early (3-5 iterations), while branchless always computes all n_symbols values.
            ret = rav1d_msac_decode_symbol_adapt_rust(s, cdf, n_symbols) as c_uint;
        } else {
            ret = rav1d_msac_decode_symbol_adapt_rust(s, cdf, n_symbols) as c_uint;
        }
    }
    debug_assert!(ret < 16);
    ret as u8 % 16
}

#[inline(always)]
pub fn rav1d_msac_decode_bool_adapt(s: &mut MsacContext, cdf: &mut [u16; 2]) -> bool {
    cfg_if! {
        if #[cfg(all(asm_msac, target_feature = "sse2"))] {
            // SAFETY: `checkasm` has verified that it is equivalent to [`dav1d_msac_decode_bool_adapt_rust`].
            unsafe {
                dav1d_msac_decode_bool_adapt_sse2(&mut s.asm, cdf.as_mut_ptr()) != 0
            }
        } else if #[cfg(all(asm_msac, target_feature = "neon"))] {
            // SAFETY: `checkasm` has verified that it is equivalent to [`dav1d_msac_decode_bool_adapt_rust`].
            unsafe {
                dav1d_msac_decode_bool_adapt_neon(&mut s.asm, cdf.as_mut_ptr()) != 0
            }
        } else {
            rav1d_msac_decode_bool_adapt_rust(s, cdf)
        }
    }
}

#[inline(always)]
pub fn rav1d_msac_decode_bool_equi(s: &mut MsacContext) -> bool {
    cfg_if! {
        if #[cfg(all(asm_msac, target_feature = "sse2"))] {
            // SAFETY: `checkasm` has verified that it is equivalent to [`dav1d_msac_decode_bool_equi_rust`].
            unsafe {
                dav1d_msac_decode_bool_equi_sse2(&mut s.asm) != 0
            }
        } else if #[cfg(all(asm_msac, target_feature = "neon"))] {
            // SAFETY: `checkasm` has verified that it is equivalent to [`dav1d_msac_decode_bool_equi_rust`].
            unsafe {
                dav1d_msac_decode_bool_equi_neon(&mut s.asm) != 0
            }
        } else {
            rav1d_msac_decode_bool_equi_rust(s)
        }
    }
}

#[inline(always)]
pub fn rav1d_msac_decode_bool(s: &mut MsacContext, f: c_uint) -> bool {
    cfg_if! {
        if #[cfg(all(asm_msac, target_feature = "sse2"))] {
            // SAFETY: `checkasm` has verified that it is equivalent to [`dav1d_msac_decode_bool_rust`].
            unsafe {
                dav1d_msac_decode_bool_sse2(&mut s.asm, f) != 0
            }
        } else if #[cfg(all(asm_msac, target_feature = "neon"))] {
            // SAFETY: `checkasm` has verified that it is equivalent to [`dav1d_msac_decode_bool_rust`].
            unsafe {
                dav1d_msac_decode_bool_neon(&mut s.asm, f) != 0
            }
        } else {
            rav1d_msac_decode_bool_rust(s, f)
        }
    }
}

/// Return value is in the range `0..16`.
#[inline(always)]
pub fn rav1d_msac_decode_hi_tok(s: &mut MsacContext, cdf: &mut [u16; 4]) -> u8 {
    let ret;
    cfg_if! {
        if #[cfg(all(asm_msac, target_feature = "sse2"))] {
            // SAFETY: `checkasm` has verified that it is equivalent to [`dav1d_msac_decode_hi_tok_rust`].
            ret = (unsafe {
                dav1d_msac_decode_hi_tok_sse2(&mut s.asm, cdf.as_mut_ptr())
            }) as u8;
        } else if #[cfg(all(asm_msac, target_feature = "neon"))] {
            // SAFETY: `checkasm` has verified that it is equivalent to [`dav1d_msac_decode_hi_tok_rust`].
            ret = unsafe {
                dav1d_msac_decode_hi_tok_neon(&mut s.asm, cdf.as_mut_ptr())
            } as u8;
        } else if #[cfg(not(asm_msac))] {
            ret = rav1d_msac_decode_hi_tok_rust(s, cdf);
        }
    }
    debug_assert!(ret < 16);
    ret % 16
}

#[cfg(all(test, not(asm_msac)))]
mod tests {
    use super::*;
    use crate::src::c_box::CBox;

    fn random(seed: &mut u64) -> u32 {
        *seed ^= *seed << 13;
        *seed ^= *seed >> 7;
        *seed ^= *seed << 17;
        (*seed >> 16) as u32
    }

    fn state(s: &MsacContext) -> (EcWin, c_uint, c_int, usize) {
        (s.dif, s.rng, s.cnt, s.buf.pos)
    }

    #[test]
    fn cdf_update_matches_directional_arithmetic_for_every_u16() {
        // Independent of the decoder reference, which also calls update_cdf.
        // Exercise every input bit pattern, including the nonstandard high
        // half, and every AV1 rate, symbol direction, and count transition.
        for probability in 0..=u16::MAX {
            for rate in 4..=7 {
                for n in [1, 3] {
                    for val in 0..=n {
                        for count in [0, 15, 16, 31, 32] {
                            let mut cdf = [0xfade; 5];
                            cdf[..n].fill(probability);
                            cdf[n] = count;
                            let mut expected = cdf;
                            let divisor = 1u32 << rate;
                            for (i, p) in expected[..n].iter_mut().enumerate() {
                                *p = if i >= val {
                                    u32::from(probability) - u32::from(probability) / divisor
                                } else {
                                    u32::from(probability)
                                        + ((32768 + 65536 - u32::from(probability)) % 65536)
                                            / divisor
                                } as u16;
                            }
                            expected[n] = count + u16::from(count < 32);
                            update_cdf(&mut cdf, n, val, rate, count);
                            assert_eq!(
                                cdf, expected,
                                "p={probability} rate={rate} n={n} val={val}"
                            );
                        }
                    }
                }
            }
        }
    }

    #[cfg(target_arch = "x86_64")]
    #[test]
    fn cdf_update_preserves_arithmetic_with_baseline_token_disabled() {
        use archmage::SimdToken as _;
        use archmage::testing::{CompileTimePolicy, for_each_token_permutation};

        let _lock = crate::src::safe_simd::token_test_lock();
        let (mut vector, mut scalar) = (false, false);
        let report = for_each_token_permutation(CompileTimePolicy::WarnStderr, |_| {
            if archmage::X64V1Token::summon().is_some() {
                vector = true;
            } else {
                scalar = true;
            }
            for val in [0, 1, 2, 3, 4, usize::MAX] {
                for rate in 0..16 {
                    for count in [0, 15, 16, 31, 32, u16::MAX] {
                        let mut cdf = [0, 32768, 65535, count, 0xfade];
                        let mut expected = cdf;
                        for (i, p) in expected[..3].iter_mut().enumerate() {
                            let q = u32::from(*p);
                            *p = if i >= val {
                                q - q / (1 << rate)
                            } else {
                                q + ((32768 + 65536 - q) % 65536) / (1 << rate)
                            } as u16;
                        }
                        expected[3] = count + u16::from(count < 32);
                        update_cdf(&mut cdf, 3, val, rate, count);
                        assert_eq!(cdf, expected);
                    }
                }
            }
        });
        assert!(report.permutations_run > 1);
        assert!(vector && scalar, "both dispatch paths must execute");
    }

    #[test]
    fn adapt4_matches_reference_through_refills_and_preserves_unused_cdf() {
        let mut seed = 0x61d4_f283_190b_7365;
        let mut symbols = [0usize; 4];
        for len in [0, 1, 2, 7, 8, 9, 15, 16, 17, 64, 511] {
            for n in 1..=3 {
                for update in [false, true] {
                    for count in [0, 15, 16, 31, 32] {
                        let bytes: Vec<_> = (0..len).map(|_| random(&mut seed) as u8).collect();
                        let data = CArc::wrap(CBox::Rust(bytes.into_boxed_slice())).unwrap();
                        let dsp = Rav1dMsacDSPContext::default();
                        let mut actual = MsacContext::new(data.clone(), !update, &dsp);
                        let mut reference = MsacContext::new(data, !update, &dsp);
                        let mut a = [0xfade; 8];
                        for probability in &mut a[..n] {
                            *probability = (random(&mut seed) & 32767) as u16;
                        }
                        a[..n].sort_unstable_by(|a, b| b.cmp(a));
                        a[n] = count;
                        let mut b = a;
                        for _ in 0..256 {
                            let val = rav1d_msac_decode_symbol_adapt4(&mut actual, &mut a, n as u8);
                            let expected = rav1d_msac_decode_symbol_adapt_rust(
                                &mut reference,
                                &mut b,
                                n as u8,
                            );
                            symbols[val as usize] += 1;
                            assert_eq!(val, expected, "len={len} n={n} update={update}");
                            assert_eq!(a, b, "all CDF entries, including untouched suffix");
                            assert_eq!(state(&actual), state(&reference));
                        }
                    }
                }
            }
        }
        assert!(
            symbols.into_iter().all(|n| n > 100),
            "exercise every symbol"
        );
    }

    #[test]
    fn adapt4_matches_reference_at_interval_boundaries_and_short_slices() {
        let mut checks = 0;
        for n in 1..=3 {
            for rng in [32768, 32769, 32831, 32896, 33024, 49151, 65534, 65535] {
                for probability in [0, 1, 31, 32, 63, 64, 65, 127, 128, 129, 16384, 32767, 32768] {
                    // Only construct states whose first threshold is within
                    // the coder range. The 32768 boundary needs low rng bits
                    // to accommodate the minimum-probability contribution.
                    if (rng >> 8) * u32::from(probability >> 6) / 2 + 4 * n as u32 > rng {
                        continue;
                    }
                    for count in [0, 15, 16, 31, 32] {
                        let mut initial = [0xfade; 8];
                        initial[..n].fill(probability);
                        initial[n] = count;
                        let r = rng >> 8;
                        for i in 0..=n {
                            let v = if i == n {
                                0
                            } else {
                                (r * u32::from(probability >> 6) >> 1) + 4 * (n - i) as u32
                            };
                            for c in [v.saturating_sub(1), v, v + 1, rng - 1] {
                                if c >= rng {
                                    continue;
                                }
                                for slice_len in [n + 1, 4, 8] {
                                    for update in [false, true] {
                                        let context = || MsacContext {
                                            asm: MsacAsmContext {
                                                dif: (c as EcWin) << (EC_WIN_SIZE - 16),
                                                rng,
                                                cnt: 64,
                                                allow_update_cdf: update.into(),
                                                ..Default::default()
                                            },
                                            ..Default::default()
                                        };
                                        let (mut a, mut b) = (initial, initial);
                                        let (mut actual, mut reference) = (context(), context());
                                        let val = rav1d_msac_decode_symbol_adapt4(
                                            &mut actual,
                                            &mut a[..slice_len],
                                            n as u8,
                                        );
                                        let expected = rav1d_msac_decode_symbol_adapt_rust(
                                            &mut reference,
                                            &mut b[..slice_len],
                                            n as u8,
                                        );
                                        assert_eq!(
                                            val, expected,
                                            "n={n} rng={rng} p={probability} c={c}"
                                        );
                                        assert_eq!(a, b);
                                        assert_eq!(state(&actual), state(&reference));
                                        checks += 1;
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }
        assert!(checks > 50_000);
    }
}