velocidb 0.1.0

// Vectorized query execution using SIMD instructions
// Leverages CPU vector units (AVX, SSE) for data parallelism

use crate::types::{Result, Value, VelociError};

#[cfg(target_arch = "x86_64")]
use std::arch::x86_64::*;
#[cfg(target_arch = "aarch64")]
use std::arch::aarch64::*;

/// Vectorized batch size - process this many elements at once
pub const VECTOR_BATCH_SIZE: usize = 256;

/// Vectorized filter operation for WHERE clauses
/// Processes multiple rows in parallel using SIMD
pub struct VectorizedFilter;

impl VectorizedFilter {
    /// Apply a comparison predicate to a batch of integer values
    /// Returns a bitmask indicating which elements passed the filter
    #[cfg(target_arch = "x86_64")]
    pub fn filter_integers_greater_than(values: &[i64], threshold: i64) -> Vec<bool> {
        let mut result = vec![false; values.len()];
        
        // Check if AVX2 is available
        if is_x86_feature_detected!("avx2") {
            unsafe {
                Self::filter_integers_greater_than_avx2(values, threshold, &mut result);
            }
        } else {
            // Fallback to scalar implementation
            Self::filter_integers_greater_than_scalar(values, threshold, &mut result);
        }
        
        result
    }

    /// AVX2-optimized integer comparison
    #[cfg(target_arch = "x86_64")]
    #[target_feature(enable = "avx2")]
    unsafe fn filter_integers_greater_than_avx2(values: &[i64], threshold: i64, result: &mut [bool]) {
        let threshold_vec = _mm256_set1_epi64x(threshold);
        let len = values.len();
        let chunk_size = 4; // AVX2 processes 4 x i64 at a time

        // Process in chunks of 4
        for i in (0..len).step_by(chunk_size) {
            if i + chunk_size > len {
                // Handle remainder with scalar code
                for j in i..len {
                    result[j] = values[j] > threshold;
                }
                break;
            }

            // Load 4 integers
            let values_vec = _mm256_loadu_si256(values[i..].as_ptr() as *const __m256i);
            
            // Compare: values > threshold
            let cmp_result = _mm256_cmpgt_epi64(values_vec, threshold_vec);
            
            // Extract comparison results
            let mask = _mm256_movemask_pd(_mm256_castsi256_pd(cmp_result));
            
            result[i] = (mask & 1) != 0;
            result[i + 1] = (mask & 2) != 0;
            result[i + 2] = (mask & 4) != 0;
            result[i + 3] = (mask & 8) != 0;
        }
    }

    /// Scalar fallback for integer comparison
    fn filter_integers_greater_than_scalar(values: &[i64], threshold: i64, result: &mut [bool]) {
        for (i, &value) in values.iter().enumerate() {
            result[i] = value > threshold;
        }
    }
    
    /// Generic fallback for non-x86_64 architectures
    #[cfg(not(target_arch = "x86_64"))]
    pub fn filter_integers_greater_than(values: &[i64], threshold: i64) -> Vec<bool> {
        let mut result = vec![false; values.len()];
        Self::filter_integers_greater_than_scalar(values, threshold, &mut result);
        result
    }

    /// Filter for equality comparison
    pub fn filter_integers_equal(values: &[i64], target: i64) -> Vec<bool> {
        let mut result = vec![false; values.len()];
        
        #[cfg(target_arch = "x86_64")]
        {
            if is_x86_feature_detected!("avx2") {
                unsafe {
                    Self::filter_integers_equal_avx2(values, target, &mut result);
                }
                return result;
            }
        }
        
        // Scalar fallback
        for (i, &value) in values.iter().enumerate() {
            result[i] = value == target;
        }
        
        result
    }

    #[cfg(target_arch = "x86_64")]
    #[target_feature(enable = "avx2")]
    unsafe fn filter_integers_equal_avx2(values: &[i64], target: i64, result: &mut [bool]) {
        let target_vec = _mm256_set1_epi64x(target);
        let len = values.len();
        let chunk_size = 4;

        for i in (0..len).step_by(chunk_size) {
            if i + chunk_size > len {
                for j in i..len {
                    result[j] = values[j] == target;
                }
                break;
            }

            let values_vec = _mm256_loadu_si256(values[i..].as_ptr() as *const __m256i);
            let cmp_result = _mm256_cmpeq_epi64(values_vec, target_vec);
            let mask = _mm256_movemask_pd(_mm256_castsi256_pd(cmp_result));
            
            result[i] = (mask & 1) != 0;
            result[i + 1] = (mask & 2) != 0;
            result[i + 2] = (mask & 4) != 0;
            result[i + 3] = (mask & 8) != 0;
        }
    }
}

/// Vectorized aggregation operations
pub struct VectorizedAggregation;

impl VectorizedAggregation {
    /// Sum an array of integers using SIMD
    pub fn sum_integers(values: &[i64]) -> i64 {
        #[cfg(target_arch = "x86_64")]
        {
            if is_x86_feature_detected!("avx2") {
                return unsafe { Self::sum_integers_avx2(values) };
            }
        }
        
        // Scalar fallback
        values.iter().sum()
    }

    #[cfg(target_arch = "x86_64")]
    #[target_feature(enable = "avx2")]
    unsafe fn sum_integers_avx2(values: &[i64]) -> i64 {
        let mut sum_vec = _mm256_setzero_si256();
        let len = values.len();
        let chunk_size = 4;

        // Process in chunks of 4
        for i in (0..len).step_by(chunk_size) {
            if i + chunk_size > len {
                // Handle remainder
                let mut scalar_sum = 0i64;
                for j in i..len {
                    scalar_sum += values[j];
                }
                
                // Extract vector sum
                let mut temp = [0i64; 4];
                _mm256_storeu_si256(temp.as_mut_ptr() as *mut __m256i, sum_vec);
                return temp.iter().sum::<i64>() + scalar_sum;
            }

            let values_vec = _mm256_loadu_si256(values[i..].as_ptr() as *const __m256i);
            sum_vec = _mm256_add_epi64(sum_vec, values_vec);
        }

        // Extract and sum the 4 i64 values from the vector
        let mut temp = [0i64; 4];
        _mm256_storeu_si256(temp.as_mut_ptr() as *mut __m256i, sum_vec);
        temp.iter().sum()
    }

    /// Count non-null values using SIMD
    pub fn count_non_null(values: &[Option<i64>]) -> usize {
        values.iter().filter(|v| v.is_some()).count()
    }

    /// Calculate average using vectorized sum
    pub fn average_integers(values: &[i64]) -> f64 {
        if values.is_empty() {
            return 0.0;
        }
        
        let sum = Self::sum_integers(values);
        sum as f64 / values.len() as f64
    }

    /// Find minimum value using SIMD
    pub fn min_integers(values: &[i64]) -> Option<i64> {
        if values.is_empty() {
            return None;
        }

        #[cfg(target_arch = "x86_64")]
        {
            if is_x86_feature_detected!("avx2") {
                return Some(unsafe { Self::min_integers_avx2(values) });
            }
        }

        values.iter().copied().min()
    }

    #[cfg(target_arch = "x86_64")]
    #[target_feature(enable = "avx2")]
    unsafe fn min_integers_avx2(values: &[i64]) -> i64 {
        if values.is_empty() {
            return i64::MAX;
        }

        let mut min_vec = _mm256_set1_epi64x(i64::MAX);
        let len = values.len();
        let chunk_size = 4;

        for i in (0..len).step_by(chunk_size) {
            if i + chunk_size > len {
                let mut temp = [0i64; 4];
                _mm256_storeu_si256(temp.as_mut_ptr() as *mut __m256i, min_vec);
                let mut scalar_min = temp.iter().copied().min().unwrap();
                
                for j in i..len {
                    scalar_min = scalar_min.min(values[j]);
                }
                return scalar_min;
            }

            let values_vec = _mm256_loadu_si256(values[i..].as_ptr() as *const __m256i);
            min_vec = _mm256_min_epi64(min_vec, values_vec);
        }

        let mut temp = [0i64; 4];
        _mm256_storeu_si256(temp.as_mut_ptr() as *mut __m256i, min_vec);
        temp.iter().copied().min().unwrap()
    }

    /// Find maximum value using SIMD
    pub fn max_integers(values: &[i64]) -> Option<i64> {
        if values.is_empty() {
            return None;
        }

        #[cfg(target_arch = "x86_64")]
        {
            if is_x86_feature_detected!("avx2") {
                return Some(unsafe { Self::max_integers_avx2(values) });
            }
        }

        values.iter().copied().max()
    }

    #[cfg(target_arch = "x86_64")]
    #[target_feature(enable = "avx2")]
    unsafe fn max_integers_avx2(values: &[i64]) -> i64 {
        if values.is_empty() {
            return i64::MIN;
        }

        let mut max_vec = _mm256_set1_epi64x(i64::MIN);
        let len = values.len();
        let chunk_size = 4;

        for i in (0..len).step_by(chunk_size) {
            if i + chunk_size > len {
                let mut temp = [0i64; 4];
                _mm256_storeu_si256(temp.as_mut_ptr() as *mut __m256i, max_vec);
                let mut scalar_max = temp.iter().copied().max().unwrap();
                
                for j in i..len {
                    scalar_max = scalar_max.max(values[j]);
                }
                return scalar_max;
            }

            let values_vec = _mm256_loadu_si256(values[i..].as_ptr() as *const __m256i);
            max_vec = _mm256_max_epi64(max_vec, values_vec);
        }

        let mut temp = [0i64; 4];
        _mm256_storeu_si256(temp.as_mut_ptr() as *mut __m256i, max_vec);
        temp.iter().copied().max().unwrap()
    }
}

/// Vectorized batch processor for query execution
pub struct VectorBatch {
    /// Column data stored in columnar format for SIMD processing
    columns: Vec<VectorColumn>,
    /// Number of rows in this batch
    row_count: usize,
}

/// A column of vectorized data
#[derive(Debug, Clone)]
pub enum VectorColumn {
    Integer(Vec<i64>),
    Real(Vec<f64>),
    Text(Vec<String>),
    Null(usize), // Just track count for null columns
}

impl VectorBatch {
    pub fn new() -> Self {
        Self {
            columns: Vec::new(),
            row_count: 0,
        }
    }

    /// Add a column to the batch
    pub fn add_column(&mut self, column: VectorColumn) {
        self.columns.push(column);
    }

    /// Set the row count
    pub fn set_row_count(&mut self, count: usize) {
        self.row_count = count;
    }

    /// Get row count
    pub fn row_count(&self) -> usize {
        self.row_count
    }

    /// Apply a vectorized filter to a column
    pub fn filter_column(&self, column_idx: usize, predicate: FilterPredicate) -> Result<Vec<bool>> {
        if column_idx >= self.columns.len() {
            return Err(VelociError::NotFound(format!("Column {} not found", column_idx)));
        }

        match (&self.columns[column_idx], predicate) {
            (VectorColumn::Integer(values), FilterPredicate::GreaterThan(Value::Integer(threshold))) => {
                Ok(VectorizedFilter::filter_integers_greater_than(values, threshold))
            }
            (VectorColumn::Integer(values), FilterPredicate::Equal(Value::Integer(target))) => {
                Ok(VectorizedFilter::filter_integers_equal(values, target))
            }
            _ => {
                // Scalar fallback for unsupported operations
                Err(VelociError::NotImplemented(
                    "Predicate type not supported for vectorized execution".to_string()
                ))
            }
        }
    }

    /// Aggregate a column
    pub fn aggregate_column(&self, column_idx: usize, agg_func: AggregateFunction) -> Result<Value> {
        if column_idx >= self.columns.len() {
            return Err(VelociError::NotFound(format!("Column {} not found", column_idx)));
        }

        match (&self.columns[column_idx], agg_func) {
            (VectorColumn::Integer(values), AggregateFunction::Sum) => {
                Ok(Value::Integer(VectorizedAggregation::sum_integers(values)))
            }
            (VectorColumn::Integer(values), AggregateFunction::Count) => {
                Ok(Value::Integer(values.len() as i64))
            }
            (VectorColumn::Integer(values), AggregateFunction::Avg) => {
                Ok(Value::Real(VectorizedAggregation::average_integers(values)))
            }
            (VectorColumn::Integer(values), AggregateFunction::Min) => {
                Ok(Value::Integer(VectorizedAggregation::min_integers(values).unwrap_or(0)))
            }
            (VectorColumn::Integer(values), AggregateFunction::Max) => {
                Ok(Value::Integer(VectorizedAggregation::max_integers(values).unwrap_or(0)))
            }
            _ => Err(VelociError::NotImplemented(
                "Aggregate function not supported".to_string()
            )),
        }
    }
}

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

/// Filter predicate types
#[derive(Debug, Clone)]
pub enum FilterPredicate {
    Equal(Value),
    NotEqual(Value),
    GreaterThan(Value),
    LessThan(Value),
    GreaterThanOrEqual(Value),
    LessThanOrEqual(Value),
}

/// Aggregate function types
#[derive(Debug, Clone, Copy)]
pub enum AggregateFunction {
    Sum,
    Count,
    Avg,
    Min,
    Max,
}

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

    #[test]
    fn test_vectorized_filter_greater_than() {
        let values = vec![1, 5, 10, 15, 20, 25];
        let result = VectorizedFilter::filter_integers_greater_than(&values, 10);
        
        assert_eq!(result, vec![false, false, false, true, true, true]);
    }

    #[test]
    fn test_vectorized_filter_equal() {
        let values = vec![1, 5, 10, 15, 10, 25];
        let result = VectorizedFilter::filter_integers_equal(&values, 10);
        
        assert_eq!(result, vec![false, false, true, false, true, false]);
    }

    #[test]
    fn test_vectorized_sum() {
        let values = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
        let sum = VectorizedAggregation::sum_integers(&values);
        
        assert_eq!(sum, 55);
    }

    #[test]
    fn test_vectorized_average() {
        let values = vec![10, 20, 30, 40, 50];
        let avg = VectorizedAggregation::average_integers(&values);
        
        assert_eq!(avg, 30.0);
    }

    #[test]
    fn test_vectorized_min_max() {
        let values = vec![5, 2, 8, 1, 9, 3];
        
        let min = VectorizedAggregation::min_integers(&values);
        let max = VectorizedAggregation::max_integers(&values);
        
        assert_eq!(min, Some(1));
        assert_eq!(max, Some(9));
    }

    #[test]
    fn test_vector_batch() {
        let mut batch = VectorBatch::new();
        
        let column = VectorColumn::Integer(vec![1, 5, 10, 15, 20]);
        batch.add_column(column);
        batch.set_row_count(5);
        
        // Test aggregation
        let sum = batch.aggregate_column(0, AggregateFunction::Sum).unwrap();
        assert_eq!(sum, Value::Integer(51));
        
        let avg = batch.aggregate_column(0, AggregateFunction::Avg).unwrap();
        assert!(matches!(avg, Value::Real(x) if (x - 10.2).abs() < 0.01));
    }

    #[test]
    fn test_large_vectorized_sum() {
        // Test with a larger dataset to ensure SIMD paths are exercised
        let values: Vec<i64> = (1..=1000).collect();
        let sum = VectorizedAggregation::sum_integers(&values);
        
        // Sum of 1 to 1000 = 1000 * 1001 / 2 = 500500
        assert_eq!(sum, 500500);
    }
}