innr 0.6.2

#![allow(unsafe_code)]
//! aarch64 SIMD implementations using NEON.
//!
//! NEON is always available on aarch64, so no runtime detection needed.
//! However, we still use target_feature for consistency with x86_64.
//!
//! # Performance Hierarchy
//!
//! | Implementation | Width | Throughput |
//! |----------------|-------|------------|
//! | dot_neon (4-way unrolled) | 4x4 f32 | ~4-8x vs scalar |
//! | Basic NEON | 4 f32 | ~2-4x vs scalar |

/// NEON dot product with 4-way unrolling.
///
/// Processes 16 floats per iteration (4 x 4), hiding memory latency.
///
/// # Safety
///
/// NEON is always available on aarch64, but we use `target_feature`
/// annotation for consistency and potential future optimizations.
#[cfg(target_arch = "aarch64")]
#[target_feature(enable = "neon")]
pub unsafe fn dot_neon(a: &[f32], b: &[f32]) -> f32 {
    use std::arch::aarch64::{
        float32x4_t, vaddq_f32, vaddvq_f32, vdupq_n_f32, vfmaq_f32, vld1q_f32,
    };

    let n = a.len().min(b.len());
    if n == 0 {
        return 0.0;
    }

    let a_ptr = a.as_ptr();
    let b_ptr = b.as_ptr();

    // 4-way unrolled: process 16 floats per iteration
    let chunks_16 = n / 16;
    let mut sum0: float32x4_t = vdupq_n_f32(0.0);
    let mut sum1: float32x4_t = vdupq_n_f32(0.0);
    let mut sum2: float32x4_t = vdupq_n_f32(0.0);
    let mut sum3: float32x4_t = vdupq_n_f32(0.0);

    for i in 0..chunks_16 {
        let base = i * 16;
        let va0 = vld1q_f32(a_ptr.add(base));
        let vb0 = vld1q_f32(b_ptr.add(base));
        let va1 = vld1q_f32(a_ptr.add(base + 4));
        let vb1 = vld1q_f32(b_ptr.add(base + 4));
        let va2 = vld1q_f32(a_ptr.add(base + 8));
        let vb2 = vld1q_f32(b_ptr.add(base + 8));
        let va3 = vld1q_f32(a_ptr.add(base + 12));
        let vb3 = vld1q_f32(b_ptr.add(base + 12));

        sum0 = vfmaq_f32(sum0, va0, vb0);
        sum1 = vfmaq_f32(sum1, va1, vb1);
        sum2 = vfmaq_f32(sum2, va2, vb2);
        sum3 = vfmaq_f32(sum3, va3, vb3);
    }

    // Combine accumulators
    let sum01 = vaddq_f32(sum0, sum1);
    let sum23 = vaddq_f32(sum2, sum3);
    let sum_all = vaddq_f32(sum01, sum23);
    let mut result = vaddvq_f32(sum_all);

    // Handle remaining 4-float chunks
    let remaining_start = chunks_16 * 16;
    let remaining = n - remaining_start;
    let chunks_4 = remaining / 4;

    let mut sum: float32x4_t = vdupq_n_f32(0.0);
    for i in 0..chunks_4 {
        let offset = remaining_start + i * 4;
        let va = vld1q_f32(a_ptr.add(offset));
        let vb = vld1q_f32(b_ptr.add(offset));
        sum = vfmaq_f32(sum, va, vb);
    }
    result += vaddvq_f32(sum);

    // Scalar tail
    let tail_start = remaining_start + chunks_4 * 4;
    for i in tail_start..n {
        // SAFETY: i is in tail_start..n where n = a.len().min(b.len()),
        // so i is always a valid index into both a and b.
        result += *a.get_unchecked(i) * *b.get_unchecked(i);
    }

    result
}

/// NEON MaxSim implementation.
///
/// Iterates over query tokens and computes max similarity against all doc tokens
/// using the unsafe dot_neon kernel directly.
///
/// # Safety
///
/// - NEON is always available on aarch64.
/// - `doc_tokens` must be non-empty; an empty slice causes `NEG_INFINITY` accumulation.
#[cfg(target_arch = "aarch64")]
#[target_feature(enable = "neon")]
pub unsafe fn maxsim_neon(query_tokens: &[&[f32]], doc_tokens: &[&[f32]]) -> f32 {
    let mut total_score = 0.0;

    for q in query_tokens {
        let mut max_score = f32::NEG_INFINITY;
        for d in doc_tokens {
            let score = dot_neon(q, d);
            if score > max_score {
                max_score = score;
            }
        }
        total_score += max_score;
    }

    total_score
}

/// NEON squared L2 distance: `Σ(a[i] - b[i])²`.
///
/// Single-pass: computes differences directly, avoiding catastrophic
/// cancellation from the expansion `||a||² + ||b||² - 2<a,b>`.
///
/// # Safety
///
/// NEON is always available on aarch64, but we use `target_feature`
/// annotation for consistency and potential future optimizations.
#[cfg(target_arch = "aarch64")]
#[target_feature(enable = "neon")]
pub unsafe fn l2_squared_neon(a: &[f32], b: &[f32]) -> f32 {
    use std::arch::aarch64::{
        float32x4_t, vaddq_f32, vaddvq_f32, vdupq_n_f32, vfmaq_f32, vld1q_f32, vsubq_f32,
    };

    let n = a.len().min(b.len());
    if n == 0 {
        return 0.0;
    }

    let a_ptr = a.as_ptr();
    let b_ptr = b.as_ptr();

    // 4-way unrolled: process 16 floats per iteration
    let chunks_16 = n / 16;
    let mut sum0: float32x4_t = vdupq_n_f32(0.0);
    let mut sum1: float32x4_t = vdupq_n_f32(0.0);
    let mut sum2: float32x4_t = vdupq_n_f32(0.0);
    let mut sum3: float32x4_t = vdupq_n_f32(0.0);

    for i in 0..chunks_16 {
        let base = i * 16;
        let va0 = vld1q_f32(a_ptr.add(base));
        let vb0 = vld1q_f32(b_ptr.add(base));
        let d0 = vsubq_f32(va0, vb0);
        let va1 = vld1q_f32(a_ptr.add(base + 4));
        let vb1 = vld1q_f32(b_ptr.add(base + 4));
        let d1 = vsubq_f32(va1, vb1);
        let va2 = vld1q_f32(a_ptr.add(base + 8));
        let vb2 = vld1q_f32(b_ptr.add(base + 8));
        let d2 = vsubq_f32(va2, vb2);
        let va3 = vld1q_f32(a_ptr.add(base + 12));
        let vb3 = vld1q_f32(b_ptr.add(base + 12));
        let d3 = vsubq_f32(va3, vb3);

        sum0 = vfmaq_f32(sum0, d0, d0);
        sum1 = vfmaq_f32(sum1, d1, d1);
        sum2 = vfmaq_f32(sum2, d2, d2);
        sum3 = vfmaq_f32(sum3, d3, d3);
    }

    // Combine accumulators
    let sum01 = vaddq_f32(sum0, sum1);
    let sum23 = vaddq_f32(sum2, sum3);
    let sum_all = vaddq_f32(sum01, sum23);
    let mut result = vaddvq_f32(sum_all);

    // Handle remaining 4-float chunks
    let remaining_start = chunks_16 * 16;
    let remaining = n - remaining_start;
    let chunks_4 = remaining / 4;

    let mut sum: float32x4_t = vdupq_n_f32(0.0);
    for i in 0..chunks_4 {
        let offset = remaining_start + i * 4;
        let va = vld1q_f32(a_ptr.add(offset));
        let vb = vld1q_f32(b_ptr.add(offset));
        let d = vsubq_f32(va, vb);
        sum = vfmaq_f32(sum, d, d);
    }
    result += vaddvq_f32(sum);

    // Scalar tail
    let tail_start = remaining_start + chunks_4 * 4;
    for i in tail_start..n {
        // SAFETY: i is in tail_start..n where n = a.len().min(b.len()),
        // so i is always a valid index into both a and b.
        let d = *a.get_unchecked(i) - *b.get_unchecked(i);
        result += d * d;
    }

    result
}

/// NEON L1 (Manhattan) distance: `Σ|a[i] - b[i]|`.
///
/// Uses `vabdq_f32` (absolute difference) for fused sub+abs in one instruction.
/// 4-way unrolled for pipeline efficiency.
///
/// # Safety
///
/// NEON is always available on aarch64.
#[cfg(target_arch = "aarch64")]
#[target_feature(enable = "neon")]
pub unsafe fn l1_neon(a: &[f32], b: &[f32]) -> f32 {
    use std::arch::aarch64::{
        float32x4_t, vabdq_f32, vaddq_f32, vaddvq_f32, vdupq_n_f32, vld1q_f32,
    };

    let n = a.len().min(b.len());
    if n == 0 {
        return 0.0;
    }

    let a_ptr = a.as_ptr();
    let b_ptr = b.as_ptr();

    // 4-way unrolled: process 16 floats per iteration
    let chunks_16 = n / 16;
    let mut sum0: float32x4_t = vdupq_n_f32(0.0);
    let mut sum1: float32x4_t = vdupq_n_f32(0.0);
    let mut sum2: float32x4_t = vdupq_n_f32(0.0);
    let mut sum3: float32x4_t = vdupq_n_f32(0.0);

    for i in 0..chunks_16 {
        let base = i * 16;
        let va0 = vld1q_f32(a_ptr.add(base));
        let vb0 = vld1q_f32(b_ptr.add(base));
        let va1 = vld1q_f32(a_ptr.add(base + 4));
        let vb1 = vld1q_f32(b_ptr.add(base + 4));
        let va2 = vld1q_f32(a_ptr.add(base + 8));
        let vb2 = vld1q_f32(b_ptr.add(base + 8));
        let va3 = vld1q_f32(a_ptr.add(base + 12));
        let vb3 = vld1q_f32(b_ptr.add(base + 12));

        // vabdq_f32: absolute difference in one instruction
        sum0 = vaddq_f32(sum0, vabdq_f32(va0, vb0));
        sum1 = vaddq_f32(sum1, vabdq_f32(va1, vb1));
        sum2 = vaddq_f32(sum2, vabdq_f32(va2, vb2));
        sum3 = vaddq_f32(sum3, vabdq_f32(va3, vb3));
    }

    // Combine accumulators
    let sum01 = vaddq_f32(sum0, sum1);
    let sum23 = vaddq_f32(sum2, sum3);
    let sum_all = vaddq_f32(sum01, sum23);
    let mut result = vaddvq_f32(sum_all);

    // Handle remaining 4-float chunks
    let remaining_start = chunks_16 * 16;
    let remaining = n - remaining_start;
    let chunks_4 = remaining / 4;

    let mut sum: float32x4_t = vdupq_n_f32(0.0);
    for i in 0..chunks_4 {
        let offset = remaining_start + i * 4;
        let va = vld1q_f32(a_ptr.add(offset));
        let vb = vld1q_f32(b_ptr.add(offset));
        sum = vaddq_f32(sum, vabdq_f32(va, vb));
    }
    result += vaddvq_f32(sum);

    // Scalar tail
    let tail_start = remaining_start + chunks_4 * 4;
    for i in tail_start..n {
        // SAFETY: i is in tail_start..n where n = a.len().min(b.len()),
        // so i is always a valid index into both a and b.
        result += (*a.get_unchecked(i) - *b.get_unchecked(i)).abs();
    }

    result
}

/// NEON fused cosine similarity: single-pass dot(a,b), norm(a)^2, norm(b)^2.
///
/// Accumulates all three products in one pass over memory, then uses
/// exact sqrt/div for IEEE-correct normalization. ~3x less memory bandwidth
/// than the 3-pass approach (dot + norm + norm).
///
/// # Safety
///
/// NEON is always available on aarch64.
#[cfg(target_arch = "aarch64")]
#[target_feature(enable = "neon")]
pub unsafe fn cosine_neon(a: &[f32], b: &[f32]) -> f32 {
    use std::arch::aarch64::{
        float32x4_t, vaddq_f32, vaddvq_f32, vdupq_n_f32, vfmaq_f32, vld1q_f32,
    };

    let n = a.len().min(b.len());
    if n == 0 {
        return 0.0;
    }

    let a_ptr = a.as_ptr();
    let b_ptr = b.as_ptr();

    // 4-way unrolled: process 16 floats per iteration, 3 accumulators each
    let chunks_16 = n / 16;
    let mut ab0: float32x4_t = vdupq_n_f32(0.0);
    let mut ab1: float32x4_t = vdupq_n_f32(0.0);
    let mut ab2: float32x4_t = vdupq_n_f32(0.0);
    let mut ab3: float32x4_t = vdupq_n_f32(0.0);
    let mut aa0: float32x4_t = vdupq_n_f32(0.0);
    let mut aa1: float32x4_t = vdupq_n_f32(0.0);
    let mut aa2: float32x4_t = vdupq_n_f32(0.0);
    let mut aa3: float32x4_t = vdupq_n_f32(0.0);
    let mut bb0: float32x4_t = vdupq_n_f32(0.0);
    let mut bb1: float32x4_t = vdupq_n_f32(0.0);
    let mut bb2: float32x4_t = vdupq_n_f32(0.0);
    let mut bb3: float32x4_t = vdupq_n_f32(0.0);

    for i in 0..chunks_16 {
        let base = i * 16;
        let va0 = vld1q_f32(a_ptr.add(base));
        let vb0 = vld1q_f32(b_ptr.add(base));
        let va1 = vld1q_f32(a_ptr.add(base + 4));
        let vb1 = vld1q_f32(b_ptr.add(base + 4));
        let va2 = vld1q_f32(a_ptr.add(base + 8));
        let vb2 = vld1q_f32(b_ptr.add(base + 8));
        let va3 = vld1q_f32(a_ptr.add(base + 12));
        let vb3 = vld1q_f32(b_ptr.add(base + 12));

        ab0 = vfmaq_f32(ab0, va0, vb0);
        ab1 = vfmaq_f32(ab1, va1, vb1);
        ab2 = vfmaq_f32(ab2, va2, vb2);
        ab3 = vfmaq_f32(ab3, va3, vb3);

        aa0 = vfmaq_f32(aa0, va0, va0);
        aa1 = vfmaq_f32(aa1, va1, va1);
        aa2 = vfmaq_f32(aa2, va2, va2);
        aa3 = vfmaq_f32(aa3, va3, va3);

        bb0 = vfmaq_f32(bb0, vb0, vb0);
        bb1 = vfmaq_f32(bb1, vb1, vb1);
        bb2 = vfmaq_f32(bb2, vb2, vb2);
        bb3 = vfmaq_f32(bb3, vb3, vb3);
    }

    // Combine the 4 accumulators
    let ab_sum = vaddq_f32(vaddq_f32(ab0, ab1), vaddq_f32(ab2, ab3));
    let aa_sum = vaddq_f32(vaddq_f32(aa0, aa1), vaddq_f32(aa2, aa3));
    let bb_sum = vaddq_f32(vaddq_f32(bb0, bb1), vaddq_f32(bb2, bb3));

    let mut ab = vaddvq_f32(ab_sum);
    let mut aa = vaddvq_f32(aa_sum);
    let mut bb = vaddvq_f32(bb_sum);

    // Handle remaining 4-float chunks
    let remaining_start = chunks_16 * 16;
    let remaining = n - remaining_start;
    let chunks_4 = remaining / 4;

    let mut ab_tail: float32x4_t = vdupq_n_f32(0.0);
    let mut aa_tail: float32x4_t = vdupq_n_f32(0.0);
    let mut bb_tail: float32x4_t = vdupq_n_f32(0.0);

    for i in 0..chunks_4 {
        let offset = remaining_start + i * 4;
        let va = vld1q_f32(a_ptr.add(offset));
        let vb = vld1q_f32(b_ptr.add(offset));
        ab_tail = vfmaq_f32(ab_tail, va, vb);
        aa_tail = vfmaq_f32(aa_tail, va, va);
        bb_tail = vfmaq_f32(bb_tail, vb, vb);
    }

    ab += vaddvq_f32(ab_tail);
    aa += vaddvq_f32(aa_tail);
    bb += vaddvq_f32(bb_tail);

    // Scalar tail
    let tail_start = remaining_start + chunks_4 * 4;
    for i in tail_start..n {
        // SAFETY: i is in tail_start..n where n = a.len().min(b.len()),
        // so i is always a valid index into both a and b.
        let ai = *a.get_unchecked(i);
        let bi = *b.get_unchecked(i);
        ab += ai * bi;
        aa += ai * ai;
        bb += bi * bi;
    }

    // Exact normalization (IEEE sqrt + div)
    if aa > crate::NORM_EPSILON_SQ && bb > crate::NORM_EPSILON_SQ {
        ab / (aa.sqrt() * bb.sqrt())
    } else {
        0.0
    }
}

/// NEON mixed-precision dot product: `sum(a_f32[i] * b_u8[i] as f32)`.
///
/// Widens u8 -> u16 -> u32 -> f32, then FMA with query f32 values.
/// Processes 16 u8 elements per iteration (4 FMA operations).
///
/// # Safety
///
/// NEON is always available on aarch64.
#[cfg(target_arch = "aarch64")]
#[target_feature(enable = "neon")]
pub unsafe fn dot_u8_f32_neon(a: &[f32], b: &[u8]) -> f32 {
    use std::arch::aarch64::{
        float32x4_t, vaddq_f32, vaddvq_f32, vcvtq_f32_u32, vdupq_n_f32, vfmaq_f32, vget_high_u16,
        vget_high_u8, vget_low_u16, vget_low_u8, vld1q_f32, vld1q_u8, vmovl_u16, vmovl_u8,
    };

    let n = a.len().min(b.len());
    if n == 0 {
        return 0.0;
    }

    let a_ptr = a.as_ptr();
    let b_ptr = b.as_ptr();

    // Process 16 elements per iteration (16 u8 -> 4x4 f32)
    let chunks_16 = n / 16;
    let mut sum0: float32x4_t = vdupq_n_f32(0.0);
    let mut sum1: float32x4_t = vdupq_n_f32(0.0);
    let mut sum2: float32x4_t = vdupq_n_f32(0.0);
    let mut sum3: float32x4_t = vdupq_n_f32(0.0);

    for i in 0..chunks_16 {
        let base = i * 16;

        // Load 16 x u8
        let vb = vld1q_u8(b_ptr.add(base));

        // Widen: u8x16 -> u16x8 (lo/hi halves)
        let b_lo_u16 = vmovl_u8(vget_low_u8(vb));
        let b_hi_u16 = vmovl_u8(vget_high_u8(vb));

        // Widen u16x4 -> u32x4 -> f32x4 (4 groups of 4)
        let b0_f32 = vcvtq_f32_u32(vmovl_u16(vget_low_u16(b_lo_u16)));
        let b1_f32 = vcvtq_f32_u32(vmovl_u16(vget_high_u16(b_lo_u16)));
        let b2_f32 = vcvtq_f32_u32(vmovl_u16(vget_low_u16(b_hi_u16)));
        let b3_f32 = vcvtq_f32_u32(vmovl_u16(vget_high_u16(b_hi_u16)));

        // Load 16 x f32 query values (4 loads of 4)
        let a0 = vld1q_f32(a_ptr.add(base));
        let a1 = vld1q_f32(a_ptr.add(base + 4));
        let a2 = vld1q_f32(a_ptr.add(base + 8));
        let a3 = vld1q_f32(a_ptr.add(base + 12));

        // FMA: sum += a * b_f32
        sum0 = vfmaq_f32(sum0, a0, b0_f32);
        sum1 = vfmaq_f32(sum1, a1, b1_f32);
        sum2 = vfmaq_f32(sum2, a2, b2_f32);
        sum3 = vfmaq_f32(sum3, a3, b3_f32);
    }

    // Combine accumulators
    let sum01 = vaddq_f32(sum0, sum1);
    let sum23 = vaddq_f32(sum2, sum3);
    let sum_all = vaddq_f32(sum01, sum23);
    let mut result = vaddvq_f32(sum_all);

    // Scalar tail
    let tail_start = chunks_16 * 16;
    for i in tail_start..n {
        // SAFETY: i is in tail_start..n where n = a.len().min(b.len()),
        // so i is always a valid index into both a and b.
        result += *a.get_unchecked(i) * (*b.get_unchecked(i) as f32);
    }

    result
}

/// NEON u8×u8 dot product using `vmull_u8` + `vpadalq_u16` + `vpadalq_u32`.
///
/// Processes 16 bytes per iteration: u8×u8 -> u16 (8 lanes), then pairwise
/// accumulate to u32 (4 lanes), then u64 (2 lanes) every 256 iterations to
/// prevent overflow.
///
/// Overflow analysis: each u16 can hold at most 255*255=65025. With 16 elements
/// per iteration widened to 8 x u16, and vpadalq_u16 summing adjacent pairs to
/// u32 (4 lanes), each u32 lane accumulates at most 2 * 65025 per iteration.
/// Draining to u64 every 256 iterations keeps each u32 lane below
/// 256 * 2 * 65025 ≈ 33.3M, well under u32::MAX.
///
/// # Safety
///
/// NEON is always available on aarch64.
#[cfg(target_arch = "aarch64")]
#[target_feature(enable = "neon")]
pub unsafe fn dot_u8_neon(a: &[u8], b: &[u8]) -> u32 {
    use std::arch::aarch64::{
        uint32x4_t, vaddlvq_u32, vdupq_n_u32, vget_high_u8, vget_low_u8, vld1q_u8, vmull_u8,
        vpadalq_u16,
    };

    let n = a.len().min(b.len());
    if n == 0 {
        return 0;
    }

    let a_ptr = a.as_ptr();
    let b_ptr = b.as_ptr();

    let chunks_16 = n / 16;
    let mut total: u64 = 0;
    let mut iter = 0usize;

    while iter < chunks_16 {
        let block = (chunks_16 - iter).min(256);
        let mut acc32: uint32x4_t = vdupq_n_u32(0);

        for j in 0..block {
            let base = (iter + j) * 16;
            let va = vld1q_u8(a_ptr.add(base));
            let vb = vld1q_u8(b_ptr.add(base));
            // vmull_u8: 8 x u8 * 8 x u8 -> 8 x u16 (low half)
            let prod_lo = vmull_u8(vget_low_u8(va), vget_low_u8(vb));
            let prod_hi = vmull_u8(vget_high_u8(va), vget_high_u8(vb));
            // vpadalq_u16: pairwise add adjacent u16 -> u32
            acc32 = vpadalq_u16(acc32, prod_lo);
            acc32 = vpadalq_u16(acc32, prod_hi);
        }

        // vaddlvq_u32: horizontal sum to u64
        total += vaddlvq_u32(acc32) as u64;
        iter += block;
    }

    // Scalar tail
    let tail_start = chunks_16 * 16;
    for i in tail_start..n {
        total += *a.get_unchecked(i) as u64 * *b.get_unchecked(i) as u64;
    }

    total as u32
}

/// NEON Hamming distance using `veorq_u8` + `vcntq_u8` (per-byte popcount).
///
/// `vcntq_u8` counts set bits in each byte in a single instruction.
/// Uses `vpaddlq_u8` -> `vpaddlq_u16` -> `vpaddlq_u32` to widen the per-byte
/// counts to u32 before accumulating.
///
/// # Safety
///
/// NEON is always available on aarch64.
#[cfg(target_arch = "aarch64")]
#[target_feature(enable = "neon")]
pub unsafe fn hamming_neon(a: &[u8], b: &[u8]) -> u32 {
    use std::arch::aarch64::{
        uint32x4_t, vaddlvq_u32, vaddq_u32, vcntq_u8, vdupq_n_u32, veorq_u8, vld1q_u8, vpaddlq_u16,
        vpaddlq_u8,
    };

    let n = a.len().min(b.len());
    if n == 0 {
        return 0;
    }

    let a_ptr = a.as_ptr();
    let b_ptr = b.as_ptr();

    let chunks_16 = n / 16;
    let mut acc32: uint32x4_t = vdupq_n_u32(0);

    for i in 0..chunks_16 {
        let base = i * 16;
        let va = vld1q_u8(a_ptr.add(base));
        let vb = vld1q_u8(b_ptr.add(base));
        let xored = veorq_u8(va, vb);
        // vcntq_u8: per-byte popcount (0..=8 per lane)
        let cnt8 = vcntq_u8(xored);
        // Widen: u8x16 -> u16x8 -> u32x4
        let cnt16 = vpaddlq_u8(cnt8);
        let cnt32 = vpaddlq_u16(cnt16);
        acc32 = vaddq_u32(acc32, cnt32);
    }

    let mut result = vaddlvq_u32(acc32) as u32;

    // Scalar tail
    let tail_start = chunks_16 * 16;
    for i in tail_start..n {
        result += (*a.get_unchecked(i) ^ *b.get_unchecked(i)).count_ones();
    }

    result
}

/// Integer-slot Hamming distance over `u32` slots: count of positions where
/// `a[i] != b[i]`. NEON path (4 lanes per 128-bit register).
///
/// # Safety
///
/// NEON is always available on aarch64; the `unsafe` is for the intrinsic
/// calls and unchecked tail indexing.
#[allow(unsafe_code)]
pub unsafe fn slot_hamming_u32_neon(a: &[u32], b: &[u32]) -> u32 {
    use std::arch::aarch64::{vaddvq_u32, vceqq_u32, vld1q_u32, vshrq_n_u32};

    let n = a.len().min(b.len());
    let a_ptr = a.as_ptr();
    let b_ptr = b.as_ptr();

    let chunks_4 = n / 4;
    let mut equal: u32 = 0;

    for i in 0..chunks_4 {
        let base = i * 4;
        let va = vld1q_u32(a_ptr.add(base));
        let vb = vld1q_u32(b_ptr.add(base));
        // vceqq_u32 -> 0xFFFFFFFF per equal lane; shift right 31 -> 1 per equal
        // lane; horizontal add -> equal-lane count for this chunk.
        let cmp = vceqq_u32(va, vb);
        equal += vaddvq_u32(vshrq_n_u32(cmp, 31));
    }

    let simd_slots = (chunks_4 * 4) as u32;
    let mut diff = simd_slots - equal;

    // Scalar tail
    for i in (chunks_4 * 4)..n {
        diff += u32::from(*a.get_unchecked(i) != *b.get_unchecked(i));
    }

    diff
}

// ---------------------------------------------------------------------------
// f64 reductions (2 doubles per q-register). Mirror the f32 kernels one width
// down; 4-way unrolled = 8 doubles per iteration.
// ---------------------------------------------------------------------------

/// Count differing u16 slots (8 per q-register). b-bit MinHash at b=16.
#[target_feature(enable = "neon")]
#[allow(unsafe_code)]
pub unsafe fn slot_hamming_u16_neon(a: &[u16], b: &[u16]) -> u32 {
    use std::arch::aarch64::{vaddvq_u16, vceqq_u16, vld1q_u16, vshrq_n_u16};
    let n = a.len().min(b.len());
    let (ap, bp) = (a.as_ptr(), b.as_ptr());
    let chunks8 = n / 8;
    let mut equal: u32 = 0;
    for i in 0..chunks8 {
        let base = i * 8;
        // vceqq sets all-ones per equal lane; >>15 maps to 0/1, then sum.
        let eq = vceqq_u16(vld1q_u16(ap.add(base)), vld1q_u16(bp.add(base)));
        equal += u32::from(vaddvq_u16(vshrq_n_u16::<15>(eq)));
    }
    let mut diff = (chunks8 as u32) * 8 - equal;
    for i in (chunks8 * 8)..n {
        diff += u32::from(*a.get_unchecked(i) != *b.get_unchecked(i));
    }
    diff
}

/// Count differing u64 slots (2 per q-register).
#[target_feature(enable = "neon")]
#[allow(unsafe_code)]
pub unsafe fn slot_hamming_u64_neon(a: &[u64], b: &[u64]) -> u64 {
    use std::arch::aarch64::{vaddvq_u64, vceqq_u64, vld1q_u64, vshrq_n_u64};
    let n = a.len().min(b.len());
    let (ap, bp) = (a.as_ptr(), b.as_ptr());
    let chunks2 = n / 2;
    let mut equal: u64 = 0;
    for i in 0..chunks2 {
        let base = i * 2;
        // vceqq sets all-ones per equal lane; shift to 0/1 then horizontal-add.
        let eq = vceqq_u64(vld1q_u64(ap.add(base)), vld1q_u64(bp.add(base)));
        equal += vaddvq_u64(vshrq_n_u64::<63>(eq));
    }
    let mut diff = (chunks2 as u64) * 2 - equal;
    for i in (chunks2 * 2)..n {
        diff += u64::from(*a.get_unchecked(i) != *b.get_unchecked(i));
    }
    diff
}

/// Dot product over f64. NEON is baseline on aarch64.
#[target_feature(enable = "neon")]
#[allow(unsafe_code)]
pub unsafe fn dot_f64_neon(a: &[f64], b: &[f64]) -> f64 {
    use std::arch::aarch64::{float64x2_t, vaddvq_f64, vdupq_n_f64, vfmaq_f64, vld1q_f64};
    let n = a.len().min(b.len());
    if n == 0 {
        return 0.0;
    }
    let (ap, bp) = (a.as_ptr(), b.as_ptr());
    let chunks8 = n / 8;
    let mut s0: float64x2_t = vdupq_n_f64(0.0);
    let mut s1: float64x2_t = vdupq_n_f64(0.0);
    let mut s2: float64x2_t = vdupq_n_f64(0.0);
    let mut s3: float64x2_t = vdupq_n_f64(0.0);
    for i in 0..chunks8 {
        let base = i * 8;
        s0 = vfmaq_f64(s0, vld1q_f64(ap.add(base)), vld1q_f64(bp.add(base)));
        s1 = vfmaq_f64(s1, vld1q_f64(ap.add(base + 2)), vld1q_f64(bp.add(base + 2)));
        s2 = vfmaq_f64(s2, vld1q_f64(ap.add(base + 4)), vld1q_f64(bp.add(base + 4)));
        s3 = vfmaq_f64(s3, vld1q_f64(ap.add(base + 6)), vld1q_f64(bp.add(base + 6)));
    }
    let mut result = vaddvq_f64(s0) + vaddvq_f64(s1) + vaddvq_f64(s2) + vaddvq_f64(s3);
    let mut i = chunks8 * 8;
    while i + 2 <= n {
        result += vaddvq_f64(vfmaq_f64(
            vdupq_n_f64(0.0),
            vld1q_f64(ap.add(i)),
            vld1q_f64(bp.add(i)),
        ));
        i += 2;
    }
    while i < n {
        // SAFETY: i < n = min(len), valid index in both.
        result += *a.get_unchecked(i) * *b.get_unchecked(i);
        i += 1;
    }
    result
}

/// Squared L2 distance over f64.
#[target_feature(enable = "neon")]
#[allow(unsafe_code)]
pub unsafe fn l2_squared_f64_neon(a: &[f64], b: &[f64]) -> f64 {
    use std::arch::aarch64::{
        float64x2_t, vaddvq_f64, vdupq_n_f64, vfmaq_f64, vld1q_f64, vsubq_f64,
    };
    let n = a.len().min(b.len());
    if n == 0 {
        return 0.0;
    }
    let (ap, bp) = (a.as_ptr(), b.as_ptr());
    let chunks8 = n / 8;
    let mut s0: float64x2_t = vdupq_n_f64(0.0);
    let mut s1: float64x2_t = vdupq_n_f64(0.0);
    let mut s2: float64x2_t = vdupq_n_f64(0.0);
    let mut s3: float64x2_t = vdupq_n_f64(0.0);
    for i in 0..chunks8 {
        let base = i * 8;
        let d0 = vsubq_f64(vld1q_f64(ap.add(base)), vld1q_f64(bp.add(base)));
        let d1 = vsubq_f64(vld1q_f64(ap.add(base + 2)), vld1q_f64(bp.add(base + 2)));
        let d2 = vsubq_f64(vld1q_f64(ap.add(base + 4)), vld1q_f64(bp.add(base + 4)));
        let d3 = vsubq_f64(vld1q_f64(ap.add(base + 6)), vld1q_f64(bp.add(base + 6)));
        s0 = vfmaq_f64(s0, d0, d0);
        s1 = vfmaq_f64(s1, d1, d1);
        s2 = vfmaq_f64(s2, d2, d2);
        s3 = vfmaq_f64(s3, d3, d3);
    }
    let mut result = vaddvq_f64(s0) + vaddvq_f64(s1) + vaddvq_f64(s2) + vaddvq_f64(s3);
    let mut i = chunks8 * 8;
    while i + 2 <= n {
        let d = vsubq_f64(vld1q_f64(ap.add(i)), vld1q_f64(bp.add(i)));
        result += vaddvq_f64(vfmaq_f64(vdupq_n_f64(0.0), d, d));
        i += 2;
    }
    while i < n {
        // SAFETY: i < n = min(len).
        let d = *a.get_unchecked(i) - *b.get_unchecked(i);
        result += d * d;
        i += 1;
    }
    result
}

/// L1 (Manhattan) distance over f64.
#[target_feature(enable = "neon")]
#[allow(unsafe_code)]
pub unsafe fn l1_f64_neon(a: &[f64], b: &[f64]) -> f64 {
    use std::arch::aarch64::{
        float64x2_t, vabdq_f64, vaddq_f64, vaddvq_f64, vdupq_n_f64, vld1q_f64,
    };
    let n = a.len().min(b.len());
    if n == 0 {
        return 0.0;
    }
    let (ap, bp) = (a.as_ptr(), b.as_ptr());
    let chunks8 = n / 8;
    let mut s0: float64x2_t = vdupq_n_f64(0.0);
    let mut s1: float64x2_t = vdupq_n_f64(0.0);
    let mut s2: float64x2_t = vdupq_n_f64(0.0);
    let mut s3: float64x2_t = vdupq_n_f64(0.0);
    for i in 0..chunks8 {
        let base = i * 8;
        s0 = vaddq_f64(
            s0,
            vabdq_f64(vld1q_f64(ap.add(base)), vld1q_f64(bp.add(base))),
        );
        s1 = vaddq_f64(
            s1,
            vabdq_f64(vld1q_f64(ap.add(base + 2)), vld1q_f64(bp.add(base + 2))),
        );
        s2 = vaddq_f64(
            s2,
            vabdq_f64(vld1q_f64(ap.add(base + 4)), vld1q_f64(bp.add(base + 4))),
        );
        s3 = vaddq_f64(
            s3,
            vabdq_f64(vld1q_f64(ap.add(base + 6)), vld1q_f64(bp.add(base + 6))),
        );
    }
    let mut result = vaddvq_f64(s0) + vaddvq_f64(s1) + vaddvq_f64(s2) + vaddvq_f64(s3);
    let mut i = chunks8 * 8;
    while i + 2 <= n {
        result += vaddvq_f64(vabdq_f64(vld1q_f64(ap.add(i)), vld1q_f64(bp.add(i))));
        i += 2;
    }
    while i < n {
        // SAFETY: i < n = min(len).
        result += (*a.get_unchecked(i) - *b.get_unchecked(i)).abs();
        i += 1;
    }
    result
}

#[cfg(test)]
mod tests {
    #[test]
    #[cfg(target_arch = "aarch64")]
    fn test_dot_neon_correctness() {
        use super::*;

        // Test various sizes around SIMD boundaries
        for size in [1, 3, 4, 5, 7, 8, 9, 15, 16, 17, 31, 32, 33, 64, 128, 256] {
            let a: Vec<f32> = (0..size).map(|i| (i as f32) * 0.1).collect();
            let b: Vec<f32> = (0..size).map(|i| (i as f32) * 0.2).collect();

            let expected: f32 = a.iter().zip(&b).map(|(x, y)| x * y).sum();
            let actual = unsafe { dot_neon(&a, &b) };

            let rel_error = if expected.abs() > 1e-6 {
                (actual - expected).abs() / expected.abs()
            } else {
                (actual - expected).abs()
            };
            assert!(
                rel_error < 1e-5,
                "size={}: expected={}, actual={}, rel_error={}",
                size,
                expected,
                actual,
                rel_error
            );
        }
    }

    #[test]
    #[cfg(target_arch = "aarch64")]
    fn test_l1_neon_correctness() {
        use super::*;

        for size in [1, 3, 4, 5, 7, 8, 9, 15, 16, 17, 31, 32, 33, 64, 128, 256] {
            let a: Vec<f32> = (0..size).map(|i| (i as f32) * 0.1).collect();
            let b: Vec<f32> = (0..size).map(|i| (i as f32) * 0.2).collect();

            let expected: f32 = a.iter().zip(&b).map(|(x, y)| (x - y).abs()).sum();
            let actual = unsafe { l1_neon(&a, &b) };

            let rel_error = if expected.abs() > 1e-6 {
                (actual - expected).abs() / expected.abs()
            } else {
                (actual - expected).abs()
            };
            assert!(
                rel_error < 1e-5,
                "L1 size={}: expected={}, actual={}, rel_error={}",
                size,
                expected,
                actual,
                rel_error
            );
        }
    }

    #[test]
    #[cfg(target_arch = "aarch64")]
    fn test_cosine_neon_correctness() {
        use super::*;

        for size in [1, 3, 4, 5, 7, 8, 9, 15, 16, 17, 31, 32, 33, 64, 128, 256] {
            let a: Vec<f32> = (0..size).map(|i| ((i * 7) as f32).sin()).collect();
            let b: Vec<f32> = (0..size).map(|i| ((i * 11) as f32).cos()).collect();

            // Reference: scalar cosine
            let ab: f32 = a.iter().zip(&b).map(|(x, y)| x * y).sum();
            let aa: f32 = a.iter().map(|x| x * x).sum();
            let bb: f32 = b.iter().map(|x| x * x).sum();
            let expected = if aa > crate::NORM_EPSILON_SQ && bb > crate::NORM_EPSILON_SQ {
                ab / (aa.sqrt() * bb.sqrt())
            } else {
                0.0
            };

            let actual = unsafe { cosine_neon(&a, &b) };

            let diff = (actual - expected).abs();
            assert!(
                diff < 1e-5,
                "size={}: expected={}, actual={}, diff={}",
                size,
                expected,
                actual,
                diff
            );
        }
    }

    #[test]
    #[cfg(target_arch = "aarch64")]
    fn test_l2_squared_neon_correctness() {
        use super::*;

        for size in [1, 3, 4, 5, 7, 8, 9, 15, 16, 17, 31, 32, 33, 64, 128, 256] {
            let a: Vec<f32> = (0..size).map(|i| (i as f32) * 0.1).collect();
            let b: Vec<f32> = (0..size).map(|i| (i as f32) * 0.2).collect();

            let expected: f32 = a.iter().zip(&b).map(|(x, y)| (x - y) * (x - y)).sum();
            let actual = unsafe { l2_squared_neon(&a, &b) };

            let rel_error = if expected.abs() > 1e-6 {
                (actual - expected).abs() / expected.abs()
            } else {
                (actual - expected).abs()
            };
            assert!(
                rel_error < 1e-5,
                "size={}: expected={}, actual={}, rel_error={}",
                size,
                expected,
                actual,
                rel_error
            );
        }
    }
}