velesdb-core 7.0.0

//! Distance computation engines for native HNSW.
//!
//! Provides trait abstraction for different distance computation backends:
//! - CPU scalar (baseline)
//! - CPU SIMD (AVX2/AVX-512/NEON)
//! - GPU (future: CUDA/Vulkan compute)

use crate::distance::DistanceMetric;
use smallvec::SmallVec;

/// Trait for distance computation engines.
///
/// This abstraction allows swapping between CPU, SIMD, and GPU backends
/// without changing the HNSW algorithm implementation.
pub trait DistanceEngine: Send + Sync {
    /// Computes distance between two vectors.
    fn distance(&self, a: &[f32], b: &[f32]) -> f32;

    /// Batch distance computation (one query vs many candidates).
    ///
    /// Returns distances in the same order as candidates.
    /// Default implementation calls `distance()` in a loop.
    fn batch_distance(&self, query: &[f32], candidates: &[&[f32]]) -> Vec<f32> {
        candidates.iter().map(|c| self.distance(query, c)).collect()
    }

    /// Returns the metric type for this engine.
    fn metric(&self) -> DistanceMetric;

    /// Returns whether the engine expects cosine inputs to be pre-normalized.
    #[must_use]
    fn is_pre_normalized(&self) -> bool {
        false
    }
}

// =============================================================================
// Shared SIMD distance helpers (RF-DEDUP: eliminates 3x copy-paste)
// =============================================================================

/// Batch distance with CPU prefetch hints to hide memory latency.
///
/// Returns `SmallVec<[f32; 64]>` to avoid heap allocation for real batch
/// sizes: layer 0 runs at `M0 = max_connections * 2` (48 or 64 with the
/// default params), so a 32-slot inline buffer spilled to the heap on
/// every dense candidate expansion. 64 slots (256 bytes) stay on the stack.
///
/// Used by `CachedSimdDistance` and called directly from the HNSW search
/// hot loop to bypass the
/// `DistanceEngine::batch_distance` trait method (which returns `Vec<f32>`).
#[inline]
pub(crate) fn batch_distance_with_prefetch(
    engine: &impl DistanceEngine,
    query: &[f32],
    candidates: &[&[f32]],
) -> SmallVec<[f32; 64]> {
    let prefetch_distance = crate::simd_native::calculate_prefetch_distance(query.len());
    let mut results = SmallVec::with_capacity(candidates.len());

    for (i, candidate) in candidates.iter().enumerate() {
        if i + prefetch_distance < candidates.len() {
            crate::simd_native::prefetch_vector(candidates[i + prefetch_distance]);
        }
        results.push(engine.distance(query, candidate));
    }

    results
}

/// CPU scalar distance computation (baseline, no SIMD).
pub struct CpuDistance {
    metric: DistanceMetric,
}

impl CpuDistance {
    /// Creates a new CPU distance engine with the given metric.
    #[must_use]
    pub fn new(metric: DistanceMetric) -> Self {
        Self { metric }
    }
}

impl DistanceEngine for CpuDistance {
    fn distance(&self, a: &[f32], b: &[f32]) -> f32 {
        #[cfg(feature = "internal-bench")]
        crate::index::hnsw::eval_count::record_eval();
        match self.metric {
            DistanceMetric::Cosine => cosine_distance_scalar(a, b),
            DistanceMetric::Euclidean => euclidean_distance_scalar(a, b),
            DistanceMetric::DotProduct => dot_product_scalar(a, b),
            DistanceMetric::Hamming => hamming_distance_scalar(a, b),
            DistanceMetric::Jaccard => jaccard_distance_scalar(a, b),
        }
    }

    fn metric(&self) -> DistanceMetric {
        self.metric
    }
}

/// SIMD distance with fully cached kernel resolution for HNSW hot loops.
///
/// All 5 distance metrics use pre-resolved fn pointers via `simd_native::DistanceEngine`.
/// The only per-call branch is `match self.metric` for the distance-to-similarity
/// transform (1-cosine, -dot), which is perfectly predicted by the branch predictor
/// since a given HNSW index always uses the same metric.
pub struct CachedSimdDistance {
    metric: DistanceMetric,
    engine: crate::simd_native::DistanceEngine,
    pre_normalized: bool,
}

impl CachedSimdDistance {
    /// Creates a cached SIMD distance engine optimized for the given metric and dimension.
    #[must_use]
    pub fn new(metric: DistanceMetric, dimension: usize) -> Self {
        Self {
            metric,
            engine: crate::simd_native::DistanceEngine::new(dimension),
            pre_normalized: false,
        }
    }

    /// Creates a cached SIMD engine that expects cosine vectors to be pre-normalized.
    ///
    /// For non-cosine metrics, this is equivalent to [`Self::new`].
    #[must_use]
    pub fn new_prenormalized(metric: DistanceMetric, dimension: usize) -> Self {
        Self {
            metric,
            engine: crate::simd_native::DistanceEngine::new(dimension),
            pre_normalized: metric == DistanceMetric::Cosine,
        }
    }
}

impl DistanceEngine for CachedSimdDistance {
    #[allow(clippy::inline_always)] // Reason: HNSW hot loop --- single branch + fn pointer call
    #[inline(always)]
    fn distance(&self, a: &[f32], b: &[f32]) -> f32 {
        #[cfg(feature = "internal-bench")]
        crate::index::hnsw::eval_count::record_eval();
        match self.metric {
            DistanceMetric::Cosine if self.pre_normalized => {
                // Both sides are unit-norm (enforced at insert and load), so
                // cosine reduces to a single dot-product chain: no norm
                // accumulators, no sqrt, no divide. The clamp absorbs FP
                // drift so the distance stays in [0, 2] like the exact arm.
                1.0 - self.engine.dot_product(a, b).clamp(-1.0, 1.0)
            }
            DistanceMetric::Cosine => 1.0 - self.engine.cosine_similarity(a, b),
            // Reason: Returns squared L2 (no sqrt) because HNSW traversal only
            // needs ordering and sqrt is monotone. The sqrt is deferred to
            // `transform_score()` which is applied to the final k results only.
            // This saves one f32::sqrt() per distance computation in the hot loop.
            DistanceMetric::Euclidean => self.engine.euclidean_squared(a, b),
            DistanceMetric::DotProduct => -self.engine.dot_product(a, b),
            DistanceMetric::Hamming => self.engine.hamming(a, b),
            DistanceMetric::Jaccard => 1.0 - self.engine.jaccard(a, b),
        }
    }

    fn batch_distance(&self, query: &[f32], candidates: &[&[f32]]) -> Vec<f32> {
        batch_distance_with_prefetch(self, query, candidates).into_vec()
    }

    fn metric(&self) -> DistanceMetric {
        self.metric
    }

    fn is_pre_normalized(&self) -> bool {
        self.pre_normalized
    }
}

// =============================================================================
// Scalar implementations (baseline for comparison)
// =============================================================================

/// Cosine distance, derived from the same similarity kernel as production.
///
/// `cosine_scalar` clamps to `[-1, 1]` and scores a zero-norm pair 0.0, so the
/// distance stays in `[0, 2]` and a zero vector is maximally distant — the
/// same contract `CachedSimdDistance` enforces on its pre-normalized arm. The
/// hand-rolled copy this replaced omitted the clamp, so two identical vectors
/// could come back at a distance slightly below zero.
#[inline]
fn cosine_distance_scalar(a: &[f32], b: &[f32]) -> f32 {
    1.0 - crate::simd_native::scalar::cosine_scalar(a, b)
}

#[inline]
fn euclidean_distance_scalar(a: &[f32], b: &[f32]) -> f32 {
    a.iter()
        .zip(b.iter())
        .map(|(x, y)| (x - y).powi(2))
        .sum::<f32>()
        .sqrt()
}

#[inline]
fn dot_product_scalar(a: &[f32], b: &[f32]) -> f32 {
    // Return negative because we want distance (lower = better)
    -a.iter().zip(b.iter()).map(|(x, y)| x * y).sum::<f32>()
}

/// Hamming distance, delegating to the kernel production falls back on.
///
/// A baseline that computes a *different* formula is worse than no baseline:
/// a parity test between it and the SIMD path reports a divergence as "the
/// SIMD kernel is wrong". This engine's job is to be the same metric without
/// the vectorization, so the formula has exactly one definition —
/// [`crate::simd_native::scalar::hamming_scalar`], which buckets each
/// component at the 0.5 binary threshold.
#[inline]
fn hamming_distance_scalar(a: &[f32], b: &[f32]) -> f32 {
    crate::simd_native::scalar::hamming_scalar(a, b)
}

/// Jaccard distance, derived from the same similarity kernel as production.
///
/// `jaccard_scalar` returns a *similarity*, and an empty union scores 1.0
/// (two all-zero vectors are identical), so the distance is 0.0 there.
#[inline]
fn jaccard_distance_scalar(a: &[f32], b: &[f32]) -> f32 {
    1.0 - crate::simd_native::scalar::jaccard_scalar(a, b)
}