vsec 0.0.1

//! SIMD-accelerated scoring using vectorized dot products.
//!
//! Converts scoring factors into a dense feature vector and weight vector,
//! then uses SIMD dot product to calculate the final score in a single
//! instruction (on supported platforms).
//!
//! # Performance
//!
//! For a typical scoring operation with 16-32 factors, this can be 2-4x
//! faster than iterating through factors individually.

/// Maximum number of features in the scoring vector.
/// Aligned to 8 for SIMD efficiency (256 bits / 32 bits per float).
pub const MAX_FEATURES: usize = 32;

/// Feature indices for the scoring vector.
/// These correspond to positions in the weight/feature vectors.
#[repr(usize)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FeatureIndex {
    // Name-based features (0-7)
    NameHasSecret = 0,
    NameHasPassword = 1,
    NameHasToken = 2,
    NameHasKey = 3,
    NameHasAuth = 4,
    NameHasCredential = 5,
    NameIsBenign = 6,
    NameHasApi = 7,

    // Value-based features (8-15)
    ValueHighEntropy = 8,
    ValueLooksLikeHex = 9,
    ValueLooksLikeBase64 = 10,
    ValueLooksLikeUuid = 11,
    ValueIsShort = 12,
    ValueIsLong = 13,
    ValueHasPlaceholder = 14,
    ValueHasSpaces = 15,

    // Context-based features (16-23)
    ContextIsTest = 16,
    ContextIsExample = 17,
    ContextIsBenchmark = 18,
    ContextHasAuthIndicators = 19,
    ContextInComparison = 20,
    ContextInFunctionArg = 21,
    ContextInReturn = 22,
    ContextInAssignment = 23,

    // RHS/Variable features (24-31)
    RhsIsCommandLike = 24,
    RhsIsAuthLike = 25,
    RhsIsInputLike = 26,
    RhsFromExternalSource = 27,
    RhsFromEnvironment = 28,
    RhsFromHeader = 29,
    RhsFromParameter = 30,
    RhsFromDatabase = 31,
}

impl FeatureIndex {
    /// Convert to array index.
    #[inline]
    pub const fn index(self) -> usize {
        self as usize
    }
}

/// A scoring feature vector with associated weights.
#[derive(Clone)]
pub struct FeatureVector {
    /// Active features (1.0 = active, 0.0 = inactive)
    features: [f32; MAX_FEATURES],
    /// Feature weights (contribution to score)
    weights: [f32; MAX_FEATURES],
}

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

impl FeatureVector {
    /// Create a new empty feature vector.
    #[inline]
    pub fn new() -> Self {
        Self {
            features: [0.0; MAX_FEATURES],
            weights: [0.0; MAX_FEATURES],
        }
    }

    /// Create a feature vector with default weights.
    pub fn with_default_weights() -> Self {
        let mut v = Self::new();
        v.set_default_weights();
        v
    }

    /// Set a feature as active with a given weight.
    #[inline]
    pub fn set(&mut self, feature: FeatureIndex, weight: f32) {
        let idx = feature.index();
        self.features[idx] = 1.0;
        self.weights[idx] = weight;
    }

    /// Set a feature as inactive.
    #[inline]
    pub fn unset(&mut self, feature: FeatureIndex) {
        let idx = feature.index();
        self.features[idx] = 0.0;
    }

    /// Check if a feature is active.
    #[inline]
    pub fn is_set(&self, feature: FeatureIndex) -> bool {
        self.features[feature.index()] != 0.0
    }

    /// Set default weights based on typical scoring patterns.
    pub fn set_default_weights(&mut self) {
        use FeatureIndex::*;

        // Name features
        self.weights[NameHasSecret.index()] = 35.0;
        self.weights[NameHasPassword.index()] = 40.0;
        self.weights[NameHasToken.index()] = 25.0;
        self.weights[NameHasKey.index()] = 20.0;
        self.weights[NameHasAuth.index()] = 25.0;
        self.weights[NameHasCredential.index()] = 35.0;
        self.weights[NameIsBenign.index()] = -100.0;
        self.weights[NameHasApi.index()] = 15.0;

        // Value features
        self.weights[ValueHighEntropy.index()] = 20.0;
        self.weights[ValueLooksLikeHex.index()] = 15.0;
        self.weights[ValueLooksLikeBase64.index()] = 15.0;
        self.weights[ValueLooksLikeUuid.index()] = 10.0;
        self.weights[ValueIsShort.index()] = -30.0;
        self.weights[ValueIsLong.index()] = 5.0;
        self.weights[ValueHasPlaceholder.index()] = -50.0;
        self.weights[ValueHasSpaces.index()] = -40.0;

        // Context features
        self.weights[ContextIsTest.index()] = -100.0;
        self.weights[ContextIsExample.index()] = -100.0;
        self.weights[ContextIsBenchmark.index()] = -100.0;
        self.weights[ContextHasAuthIndicators.index()] = 30.0;
        self.weights[ContextInComparison.index()] = 20.0;
        self.weights[ContextInFunctionArg.index()] = 15.0;
        self.weights[ContextInReturn.index()] = 10.0;
        self.weights[ContextInAssignment.index()] = 10.0;

        // RHS features
        self.weights[RhsIsCommandLike.index()] = -30.0;
        self.weights[RhsIsAuthLike.index()] = 25.0;
        self.weights[RhsIsInputLike.index()] = 15.0;
        self.weights[RhsFromExternalSource.index()] = 20.0;
        self.weights[RhsFromEnvironment.index()] = 15.0;
        self.weights[RhsFromHeader.index()] = 25.0;
        self.weights[RhsFromParameter.index()] = 10.0;
        self.weights[RhsFromDatabase.index()] = 15.0;
    }

    /// Calculate the score using SIMD dot product.
    #[inline]
    pub fn calculate_score(&self) -> f32 {
        #[cfg(target_arch = "x86_64")]
        {
            if is_x86_feature_detected!("avx") {
                return unsafe { self.dot_product_avx() };
            }
        }

        #[cfg(target_arch = "aarch64")]
        {
            if std::arch::is_aarch64_feature_detected!("neon") {
                return unsafe { self.dot_product_neon() };
            }
        }

        self.dot_product_scalar()
    }

    /// Scalar dot product fallback.
    #[inline]
    fn dot_product_scalar(&self) -> f32 {
        self.features
            .iter()
            .zip(self.weights.iter())
            .map(|(&f, &w)| f * w)
            .sum()
    }

    /// AVX dot product (8 floats at a time).
    #[cfg(target_arch = "x86_64")]
    #[target_feature(enable = "avx")]
    unsafe fn dot_product_avx(&self) -> f32 {
        use std::arch::x86_64::*;

        let mut sum = _mm256_setzero_ps();

        // Process 8 floats at a time (256 bits)
        for i in 0..4 {
            let offset = i * 8;
            let features =
                _mm256_loadu_ps(self.features.as_ptr().add(offset));
            let weights = _mm256_loadu_ps(self.weights.as_ptr().add(offset));
            let product = _mm256_mul_ps(features, weights);
            sum = _mm256_add_ps(sum, product);
        }

        // Horizontal sum of the 8 floats
        // sum = [a, b, c, d, e, f, g, h]
        let high = _mm256_extractf128_ps(sum, 1); // [e, f, g, h]
        let low = _mm256_castps256_ps128(sum); // [a, b, c, d]
        let sum128 = _mm_add_ps(low, high); // [a+e, b+f, c+g, d+h]

        let high64 = _mm_movehl_ps(sum128, sum128); // [c+g, d+h, ?, ?]
        let sum64 = _mm_add_ps(sum128, high64); // [a+e+c+g, b+f+d+h, ?, ?]

        let high32 = _mm_shuffle_ps(sum64, sum64, 1); // [b+f+d+h, ?, ?, ?]
        let sum32 = _mm_add_ss(sum64, high32);

        _mm_cvtss_f32(sum32)
    }

    /// NEON dot product (4 floats at a time).
    #[cfg(target_arch = "aarch64")]
    #[target_feature(enable = "neon")]
    unsafe fn dot_product_neon(&self) -> f32 {
        use std::arch::aarch64::*;

        let mut sum = vdupq_n_f32(0.0);

        // Process 4 floats at a time (128 bits)
        for i in 0..8 {
            let offset = i * 4;
            let features = vld1q_f32(self.features.as_ptr().add(offset));
            let weights = vld1q_f32(self.weights.as_ptr().add(offset));
            sum = vmlaq_f32(sum, features, weights); // FMA: sum += features * weights
        }

        // Horizontal sum of the 4 floats
        vaddvq_f32(sum)
    }
}

/// Quick score calculation from a list of (feature, weight) pairs.
///
/// This is a convenience function for when you don't need the full
/// feature vector infrastructure.
#[inline]
pub fn quick_score(factors: &[(FeatureIndex, f32)]) -> f32 {
    let mut vec = FeatureVector::new();
    for &(feature, weight) in factors {
        vec.set(feature, weight);
    }
    vec.calculate_score()
}

/// Calculate weighted sum using SIMD for arbitrary float slices.
#[inline]
pub fn weighted_sum(features: &[f32], weights: &[f32]) -> f32 {
    debug_assert_eq!(features.len(), weights.len());

    #[cfg(target_arch = "x86_64")]
    {
        if is_x86_feature_detected!("avx") && features.len() >= 8 {
            return unsafe { weighted_sum_avx(features, weights) };
        }
    }

    #[cfg(target_arch = "aarch64")]
    {
        if std::arch::is_aarch64_feature_detected!("neon") && features.len() >= 4 {
            return unsafe { weighted_sum_neon(features, weights) };
        }
    }

    features
        .iter()
        .zip(weights.iter())
        .map(|(&f, &w)| f * w)
        .sum()
}

#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx")]
unsafe fn weighted_sum_avx(features: &[f32], weights: &[f32]) -> f32 {
    use std::arch::x86_64::*;

    let mut sum = _mm256_setzero_ps();
    let chunks = features.len() / 8;

    for i in 0..chunks {
        let offset = i * 8;
        let f = _mm256_loadu_ps(features.as_ptr().add(offset));
        let w = _mm256_loadu_ps(weights.as_ptr().add(offset));
        sum = _mm256_add_ps(sum, _mm256_mul_ps(f, w));
    }

    // Horizontal sum
    let high = _mm256_extractf128_ps(sum, 1);
    let low = _mm256_castps256_ps128(sum);
    let sum128 = _mm_add_ps(low, high);
    let high64 = _mm_movehl_ps(sum128, sum128);
    let sum64 = _mm_add_ps(sum128, high64);
    let high32 = _mm_shuffle_ps(sum64, sum64, 1);
    let mut result = _mm_cvtss_f32(_mm_add_ss(sum64, high32));

    // Handle remainder
    for i in chunks * 8..features.len() {
        result += features[i] * weights[i];
    }

    result
}

#[cfg(target_arch = "aarch64")]
#[target_feature(enable = "neon")]
unsafe fn weighted_sum_neon(features: &[f32], weights: &[f32]) -> f32 {
    use std::arch::aarch64::*;

    let mut sum = vdupq_n_f32(0.0);
    let chunks = features.len() / 4;

    for i in 0..chunks {
        let offset = i * 4;
        let f = vld1q_f32(features.as_ptr().add(offset));
        let w = vld1q_f32(weights.as_ptr().add(offset));
        sum = vmlaq_f32(sum, f, w);
    }

    let mut result = vaddvq_f32(sum);

    // Handle remainder
    for i in chunks * 4..features.len() {
        result += features[i] * weights[i];
    }

    result
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_feature_vector_basic() {
        let mut vec = FeatureVector::with_default_weights();
        vec.set(FeatureIndex::NameHasPassword, 40.0);
        vec.set(FeatureIndex::ValueHighEntropy, 20.0);

        let score = vec.calculate_score();
        assert!((score - 60.0).abs() < 0.01);
    }

    #[test]
    fn test_benign_kills_score() {
        let mut vec = FeatureVector::with_default_weights();
        vec.set(FeatureIndex::NameHasPassword, 40.0);
        vec.set(FeatureIndex::ContextIsTest, -100.0);

        let score = vec.calculate_score();
        assert!(score < 0.0); // Should be negative
    }

    #[test]
    fn test_quick_score() {
        use FeatureIndex::*;
        let score = quick_score(&[
            (NameHasSecret, 35.0),
            (ValueHighEntropy, 20.0),
            (ContextHasAuthIndicators, 30.0),
        ]);
        assert!((score - 85.0).abs() < 0.01);
    }

    #[test]
    fn test_weighted_sum() {
        let features = vec![1.0, 0.0, 1.0, 1.0];
        let weights = vec![10.0, 20.0, 30.0, 40.0];

        let sum = weighted_sum(&features, &weights);
        assert!((sum - 80.0).abs() < 0.01); // 10 + 0 + 30 + 40
    }

    #[test]
    fn test_simd_scalar_equivalence() {
        let mut vec = FeatureVector::with_default_weights();

        // Set various features
        vec.set(FeatureIndex::NameHasPassword, 40.0);
        vec.set(FeatureIndex::ValueHighEntropy, 20.0);
        vec.set(FeatureIndex::ContextHasAuthIndicators, 30.0);
        vec.set(FeatureIndex::RhsIsAuthLike, 25.0);

        let simd_score = vec.calculate_score();
        let scalar_score = vec.dot_product_scalar();

        assert!(
            (simd_score - scalar_score).abs() < 0.01,
            "SIMD {} != scalar {}",
            simd_score,
            scalar_score
        );
    }

    #[test]
    fn test_long_weighted_sum() {
        // Long enough to trigger SIMD paths
        let features: Vec<f32> = (0..32).map(|i| if i % 2 == 0 { 1.0 } else { 0.0 }).collect();
        let weights: Vec<f32> = (0..32).map(|i| i as f32).collect();

        let sum = weighted_sum(&features, &weights);
        let expected: f32 = (0..32).filter(|i| i % 2 == 0).map(|i| i as f32).sum();

        assert!((sum - expected).abs() < 0.01);
    }
}