maroontree 0.1.2

AV1 & AV2 tiny still-image (AVIF) encoder
Documentation
/*
 * 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 std::arch::aarch64::*;

#[inline]
#[target_feature(enable = "neon")]
fn neon_log2p1_f32(x: float32x4_t) -> float32x4_t {
    let y = vaddq_f32(x, vdupq_n_f32(1.0));

    // Decompose y = m * 2^e, m in [1, 2).
    let bits = vreinterpretq_u32_f32(y);
    let exp_u = vshrq_n_u32::<23>(bits);
    let exp_i = vsubq_s32(vreinterpretq_s32_u32(exp_u), vdupq_n_s32(127));
    let e = vcvtq_f32_s32(exp_i);

    let mant_bits = vorrq_u32(
        vandq_u32(bits, vdupq_n_u32(0x007f_ffff)),
        vdupq_n_u32(0x3f80_0000),
    );
    let m = vreinterpretq_f32_u32(mant_bits);
    let t = vsubq_f32(m, vdupq_n_f32(1.0));

    // Generated by Sollya.
    let c0 = vdupq_n_f32(1.4426934719085693359375);
    let c1 = vdupq_n_f32(-0.721179187297821044921875);
    let c2 = vdupq_n_f32(0.477900087833404541015625);
    let c3 = vdupq_n_f32(-0.340080082416534423828125);
    let c4 = vdupq_n_f32(0.21719777584075927734375);
    let c5 = vdupq_n_f32(-9.749893844127655029296875e-2);
    let c6 = vdupq_n_f32(2.096841670572757720947265625e-2);

    let mut p = c6;
    p = vfmaq_f32(c5, t, p);
    p = vfmaq_f32(c4, t, p);
    p = vfmaq_f32(c3, t, p);
    p = vfmaq_f32(c2, t, p);
    p = vfmaq_f32(c1, t, p);
    p = vfmaq_f32(c0, t, p);

    vaddq_f32(e, vmulq_f32(t, p))
}

#[target_feature(enable = "neon")]
pub(crate) fn coeff_rate_f32_neon(lev: &[f32]) -> f32 {
    let mut sum = vdupq_n_f32(0.0);
    let two = vdupq_n_f32(2.0);
    let zero = vdupq_n_f32(0.0);

    let chunks = lev.as_chunks::<4>();
    let rem = chunks.1;

    for chunk in chunks.0.iter() {
        let v = unsafe { vld1q_f32(chunk.as_ptr()) };
        let a = vabsq_f32(v);
        let nz = vcgtq_f32(a, zero);
        let cost = vfmaq_f32(two, two, neon_log2p1_f32(a));
        sum = vaddq_f32(sum, vbslq_f32(nz, cost, zero));
    }

    let mut out = vaddvq_f32(sum);
    for &v in rem {
        if v != 0.0 {
            out += 2.0 + 2.0 * crate::av2::helpers::log2p1_approx_f32(v.abs());
        }
    }
    out
}

#[target_feature(enable = "neon")]
pub(crate) fn coeff_abs_rate_f32_neon(lev: &[f32]) -> f32 {
    let mut sum = vdupq_n_f32(0.0);

    let chunks = lev.as_chunks::<4>();
    let rem = chunks.1;

    for chunk in chunks.0.iter() {
        let v = unsafe { vld1q_f32(chunk.as_ptr()) };
        sum = vaddq_f32(sum, vabsq_f32(v));
    }

    let mut out = vaddvq_f32(sum);
    for &v in rem {
        out += v.abs();
    }
    out
}