sqlitegraph 2.2.2

//! SIMD-Accelerated Serialization for HNSW Persistence
//!
//! This module provides optimized serialization helpers for HNSW vector persistence,
//! including varint encoding and delta encoding with SIMD acceleration.
//!
//! # What is Varint Encoding?
//!
//! Varint (variable-length integer) encoding is a compact representation of integers
//! where smaller values use fewer bytes. Each byte uses 7 bits for data and 1 bit
//! as a continuation flag:
//!
//! - **Values < 128**: 1 byte (e.g., 42 → `0x2A`)
//! - **Values < 16384**: 2 bytes (e.g., 300 → `0xAC 0x02`)
//! - **Up to u32::MAX**: 5 bytes maximum
//!
//! This provides significant space savings for small integers common in graph data.
//!
//! # What is Delta Encoding?
//!
//! Delta encoding stores differences between consecutive values instead of absolute
//! values. For sequential data (sorted node IDs, monotonically increasing offsets),
//! deltas are typically small, which compresses well with varint encoding.
//!
//! Example:
//! - **Original**: `[100, 105, 110, 115, 120]`
//! - **Deltas**: `[100, 5, 5, 5, 5]`
//! - **Space savings**: Deltas varint-encode more efficiently
//!
//! # SIMD Acceleration
//!
//! Delta encoding benefits from SIMD because computing differences between adjacent
//! values is embarrasingly parallel:
//!
//! - **Scalar**: Loop through values one at a time
//! - **AVX2**: Process 8 u32 values per iteration (256-bit register)
//! - **Speedup**: ~3-5x for large arrays (> 100 elements)
//!
//! # Architecture
//!
//! - **Scalar fallback**: Pure Rust implementation, always available
//! - **AVX2 path**: x86_64 intrinsics for parallel delta computation
//! - **Runtime dispatch**: Automatic selection based on CPU features
//! - **Binary compatible**: SIMD and scalar produce identical output
//!
//! # Performance Characteristics
//!
//! ## Varint Encoding
//! - **Encode**: O(n) where n = value magnitude (1-5 bytes)
//! - **Decode**: O(n) with bounds checking
//! - **Space**: 1-5 bytes per u32 (average ~2 bytes for graph data)
//!
//! ## Delta Encoding
//! - **AVX2**: ~3-5x faster than scalar for arrays > 100 elements
//! - **Scalar**: Baseline implementation for all platforms
//! - **Threshold**: SIMD disabled for < 16 elements (overhead too high)
//!
//! # Examples
//!
//! ```rust
//! use sqlitegraph::hnsw::serialization::{encode_varint_scalar, decode_varint_scalar};
//!
//! let value = 300u32;
//! let mut buffer = Vec::new();
//!
//! // Encode: 300 → [0xAC, 0x02] (2 bytes)
//! encode_varint_scalar(&mut buffer, value).unwrap();
//!
//! // Decode: [0xAC, 0x02] → 300
//! let decoded = decode_varint_scalar(buffer.as_slice()).unwrap();
//! assert_eq!(decoded, 300);
//! # Ok::<(), Box<dyn std::error::Error>>(())
//! ```
//!
//! ```rust
//! use sqlitegraph::hnsw::serialization::{delta_encode, delta_decode};
//!
//! let values = vec![100, 105, 110, 115, 120];
//!
//! // Encode deltas: [100, 5, 5, 5, 5]
//! let deltas = delta_encode(&values);
//!
//! // Decode back: [100, 105, 110, 115, 120]
//! let restored = delta_decode(&deltas);
//! assert_eq!(restored, values);
//! ```

use std::io::{self, Read, Write};

//=============================================================================
// Varint Encoding (Scalar)
//=============================================================================

/// Encode a u32 as varint, writing to the provided writer.
///
/// Varint encoding uses 7 bits per byte, with the high bit set to indicate
/// continuation. This provides compact representation for small values.
///
/// # Arguments
///
/// * `writer` - Any type implementing Write (Vec<u8>, File, etc.)
/// * `value` - The u32 value to encode
///
/// # Returns
///
/// Number of bytes written (1-5)
///
/// # Examples
///
/// ```rust
/// use sqlitegraph::hnsw::serialization::encode_varint_scalar;
///
/// let mut buffer = Vec::new();
/// let bytes_written = encode_varint_scalar(&mut buffer, 300).unwrap();
/// assert_eq!(bytes_written, 2);
/// assert_eq!(buffer, vec![0xAC, 0x02]);
/// # Ok::<(), Box<dyn std::error::Error>>(())
/// ```
///
/// # Encoding Format
///
/// - **1 byte**: Values < 128 (0x00 to 0x7F)
/// - **2 bytes**: Values < 16384 (0x80 0x01 to 0x7F 0x7F)
/// - **3-5 bytes**: Larger values up to u32::MAX
pub fn encode_varint_scalar<W: Write>(mut writer: W, value: u32) -> io::Result<usize> {
    let mut bytes_written = 0;
    let mut v = value;

    loop {
        let byte = (v & 0x7F) as u8;
        v >>= 7;

        if v == 0 {
            writer.write_all(&[byte])?;
            bytes_written += 1;
            break;
        } else {
            writer.write_all(&[byte | 0x80])?;
            bytes_written += 1;
        }
    }

    Ok(bytes_written)
}

/// Decode a varint from the provided reader.
///
/// Reads bytes until the continuation bit is unset, reconstructing the original
/// u32 value.
///
/// # Arguments
///
/// * `reader` - Any type implementing Read (&[u8], File, etc.)
///
/// # Returns
///
/// The decoded u32 value
///
/// # Errors
///
/// - `InvalidData` if varint is malformed (too many continuation bytes)
/// - `UnexpectedEof` if input ends mid-varint
///
/// # Examples
///
/// ```rust
/// use sqlitegraph::hnsw::serialization::decode_varint_scalar;
///
/// let encoded = vec![0xAC, 0x02]; // 300
/// let decoded = decode_varint_scalar(encoded.as_slice()).unwrap();
/// assert_eq!(decoded, 300);
/// # Ok::<(), Box<dyn std::error::Error>>(())
/// ```
pub fn decode_varint_scalar<R: Read>(mut reader: R) -> io::Result<u32> {
    let mut result = 0;
    let mut shift = 0;

    loop {
        let mut byte = [0u8; 1];
        reader.read_exact(&mut byte)?;
        let b = byte[0];

        result |= ((b & 0x7F) as u32) << shift;

        if (b & 0x80) == 0 {
            break;
        }

        shift += 7;
        if shift >= 35 {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                "varint too large",
            ));
        }
    }

    Ok(result)
}

//=============================================================================
// Delta Encoding (Scalar)
//=============================================================================

/// Encode differences between consecutive values (delta encoding).
///
/// Scalar fallback for non-AVX2 systems or small arrays.
///
/// # Arguments
///
/// * `values` - Slice of u32 values to encode
///
/// # Returns
///
/// Vector of deltas where:
/// - First element is the original first value
/// - Subsequent elements are differences (values[i] - values[i-1])
///
/// # Examples
///
/// ```rust
/// use sqlitegraph::hnsw::serialization::delta_encode_scalar;
///
/// let values = vec![100, 105, 110, 115, 120];
/// let deltas = delta_encode_scalar(&values);
/// assert_eq!(deltas, vec![100, 5, 5, 5, 5]);
/// ```
pub fn delta_encode_scalar(values: &[u32]) -> Vec<u32> {
    if values.is_empty() {
        return Vec::new();
    }

    let mut deltas = Vec::with_capacity(values.len());
    deltas.push(values[0]); // First value stored as-is

    for i in 1..values.len() {
        // Use wrapping_sub to handle decreasing values (wraps around on underflow)
        deltas.push(values[i].wrapping_sub(values[i - 1]));
    }

    deltas
}

/// Decode deltas back to original values.
///
/// # Arguments
///
/// * `deltas` - Slice of delta-encoded values
///
/// # Returns
///
/// Reconstructed original values
///
/// # Examples
///
/// ```rust
/// use sqlitegraph::hnsw::serialization::delta_decode;
///
/// let deltas = vec![100, 5, 5, 5, 5];
/// let values = delta_decode(&deltas);
/// assert_eq!(values, vec![100, 105, 110, 115, 120]);
/// ```
pub fn delta_decode(deltas: &[u32]) -> Vec<u32> {
    let mut values = Vec::with_capacity(deltas.len());
    let mut acc = 0u32;

    for &delta in deltas {
        acc += delta;
        values.push(acc);
    }

    values
}

//=============================================================================
// Delta Encoding (AVX2)
//=============================================================================

#[cfg(target_arch = "x86_64")]
/// AVX2-accelerated batch delta encoding.
///
/// Computes differences between consecutive u32 values using 256-bit SIMD
/// registers, processing 8 values per iteration.
///
/// # Safety
///
/// This function is marked unsafe because it requires:
/// - AVX2 CPU feature support (verified by caller)
/// - Proper alignment for SIMD operations (uses unaligned loads)
///
/// The function is safe to call when AVX2 is available.
///
/// # Arguments
///
/// * `values` - Slice of u32 values (must have length >= 2)
///
/// # Returns
///
/// Vector of delta-encoded values
///
/// # Performance
///
/// - **Throughput**: 8 u32 values per iteration
/// - **Speedup**: ~3-5x for arrays > 100 elements
/// - **Overhead**: Falls back to scalar for < 16 elements
#[target_feature(enable = "avx2")]
#[inline]
pub unsafe fn delta_encode_avx2(values: &[u32]) -> Vec<u32> {
    unsafe {
        use std::arch::x86_64::*;

        if values.len() < 16 {
            // Too small for SIMD overhead to be worth it
            return delta_encode_scalar(values);
        }

        let mut deltas = Vec::with_capacity(values.len());
        deltas.push(values[0]);

        let len = values.len();
        let chunks = (len - 1) / 8;

        let mut i = 1;
        for _ in 0..chunks {
            // Load 8 consecutive u32 values using AVX2 integer load
            let v1 = _mm256_loadu_si256(values.as_ptr().add(i) as *const __m256i);
            let v2 = _mm256_loadu_si256(values.as_ptr().add(i - 1) as *const __m256i);

            // Compute differences using AVX2 integer subtraction
            let diff = _mm256_sub_epi32(v1, v2);

            // Extract results
            let mut tmp = [0u32; 8];
            _mm256_storeu_si256(tmp.as_mut_ptr() as *mut __m256i, diff);

            for &val in &tmp {
                deltas.push(val);
            }

            i += 8;
        }

        // Handle remainder
        while i < len {
            deltas.push(values[i] - values[i - 1]);
            i += 1;
        }

        deltas
    }
}

//=============================================================================
// Runtime Dispatch
//=============================================================================

/// Encode differences between consecutive values with automatic SIMD dispatch.
///
/// This function automatically selects the best implementation based on:
/// - CPU feature detection (AVX2 availability)
/// - Input size (SIMD only for >= 16 elements)
///
/// # Arguments
///
/// * `values` - Slice of u32 values to encode
///
/// # Returns
///
/// Vector of delta-encoded values
///
/// # Examples
///
/// ```rust
/// use sqlitegraph::hnsw::serialization::delta_encode;
///
/// let values = vec![100, 105, 110, 115, 120];
/// let deltas = delta_encode(&values);
/// assert_eq!(deltas, vec![100, 5, 5, 5, 5]);
/// ```
///
/// # Performance
///
/// - **AVX2 + large arrays**: ~3-5x speedup
/// - **Scalar or small arrays**: Baseline performance
/// - **Detection**: O(1) on first call, cached thereafter
pub fn delta_encode(values: &[u32]) -> Vec<u32> {
    #[cfg(target_arch = "x86_64")]
    {
        if is_x86_feature_detected!("avx2") && values.len() >= 16 {
            unsafe { delta_encode_avx2(values) }
        } else {
            delta_encode_scalar(values)
        }
    }
    #[cfg(not(target_arch = "x86_64"))]
    {
        delta_encode_scalar(values)
    }
}

//=============================================================================
// Tests
//=============================================================================

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

    //=========================================================================
    // Varint Encoding Tests
    //=========================================================================

    #[test]
    fn test_varint_roundtrip() {
        let test_values = vec![
            0u32,
            1,
            42,
            127,
            128,
            255,
            256,
            300,
            16383,
            16384,
            65535,
            65536,
            u32::MAX - 1,
            u32::MAX,
        ];

        for value in test_values {
            let mut buffer = Vec::new();
            encode_varint_scalar(&mut buffer, value).unwrap();

            let decoded = decode_varint_scalar(buffer.as_slice()).unwrap();
            assert_eq!(
                decoded, value,
                "Roundtrip failed for value {} (encoded as {:?})",
                value, buffer
            );
        }
    }

    #[test]
    fn test_varint_single_byte() {
        let mut buffer = Vec::new();
        encode_varint_scalar(&mut buffer, 42).unwrap();

        assert_eq!(buffer.len(), 1);
        assert_eq!(buffer[0], 0x2A); // 42 = 0b00101010
    }

    #[test]
    fn test_varint_two_bytes() {
        let mut buffer = Vec::new();
        encode_varint_scalar(&mut buffer, 300).unwrap();

        assert_eq!(buffer.len(), 2);
        assert_eq!(buffer[0], 0xAC); // 300 = 0b10101100 (continuation bit set)
        assert_eq!(buffer[1], 0x02); // 300 >> 7 = 0b00000010
    }

    #[test]
    fn test_varint_max_value() {
        let mut buffer = Vec::new();
        encode_varint_scalar(&mut buffer, u32::MAX).unwrap();

        assert_eq!(buffer.len(), 5);
        assert_eq!(buffer[0], 0xFF);
        assert_eq!(buffer[1], 0xFF);
        assert_eq!(buffer[2], 0xFF);
        assert_eq!(buffer[3], 0xFF);
        assert_eq!(buffer[4], 0x0F); // u32::MAX needs 5 bytes
    }

    #[test]
    fn test_varint_zero() {
        let mut buffer = Vec::new();
        encode_varint_scalar(&mut buffer, 0).unwrap();

        assert_eq!(buffer.len(), 1);
        assert_eq!(buffer[0], 0x00);
    }

    #[test]
    fn test_varint_continuation_bits() {
        // Test that continuation bits are properly set
        let test_cases = vec![
            (127, vec![0x7F]),               // No continuation
            (128, vec![0x80, 0x01]),         // Continuation on first byte
            (16383, vec![0xFF, 0x7F]),       // Continuation on first byte only
            (16384, vec![0x80, 0x80, 0x01]), // Continuation on first two bytes
        ];

        for (value, expected) in test_cases {
            let mut buffer = Vec::new();
            encode_varint_scalar(&mut buffer, value).unwrap();
            assert_eq!(buffer, expected, "Failed for value {}", value);
        }
    }

    #[test]
    fn test_varint_invalid_too_large() {
        // Malformed varint with too many continuation bytes
        let malformed = vec![0x80, 0x80, 0x80, 0x80, 0x80, 0x01]; // 6 bytes

        let result = decode_varint_scalar(malformed.as_slice());
        assert!(result.is_err());
        assert_eq!(result.unwrap_err().kind(), io::ErrorKind::InvalidData);
    }

    #[test]
    fn test_varint_incomplete() {
        // Varint ends mid-encoding (continuation bit set but no more data)
        let incomplete = vec![0x80]; // Continuation bit set but no following byte

        let result = decode_varint_scalar(incomplete.as_slice());
        assert!(result.is_err());
        assert_eq!(result.unwrap_err().kind(), io::ErrorKind::UnexpectedEof);
    }

    //=========================================================================
    // Delta Encoding Tests
    //=========================================================================

    #[test]
    fn test_delta_encode_basic() {
        let values = vec![100, 105, 110, 115, 120];
        let deltas = delta_encode_scalar(&values);

        assert_eq!(deltas, vec![100, 5, 5, 5, 5]);
    }

    #[test]
    fn test_delta_encode_empty() {
        let values: Vec<u32> = vec![];
        let deltas = delta_encode_scalar(&values);

        assert_eq!(deltas, Vec::<u32>::new());
    }

    #[test]
    fn test_delta_encode_single() {
        let values = vec![42];
        let deltas = delta_encode_scalar(&values);

        assert_eq!(deltas, vec![42]);
    }

    #[test]
    fn test_delta_encode_large_gaps() {
        let values = vec![0, 1000, 2000, 3000];
        let deltas = delta_encode_scalar(&values);

        assert_eq!(deltas, vec![0, 1000, 1000, 1000]);
    }

    #[test]
    fn test_delta_encode_decreasing() {
        // Delta encoding with decreasing values uses wrapping subtraction
        let values = vec![100u32, 95, 90, 85, 80];
        let deltas = delta_encode_scalar(&values);

        // Expected: 100, wrapping_sub(95), wrapping_sub(90), wrapping_sub(85), wrapping_sub(80)
        // 95 - 100 = 4294967291 (u32::MAX - 4)
        // 90 - 95 = 4294967291
        // etc.
        assert_eq!(
            deltas,
            vec![100, 4294967291, 4294967291, 4294967291, 4294967291]
        );
    }

    #[test]
    fn test_delta_decode_roundtrip() {
        let original = vec![100, 105, 110, 115, 120];
        let deltas = delta_encode_scalar(&original);
        let restored = delta_decode(&deltas);

        assert_eq!(restored, original);
    }

    #[test]
    fn test_delta_decode_empty() {
        let deltas: Vec<u32> = vec![];
        let values = delta_decode(&deltas);

        assert_eq!(values, Vec::<u32>::new());
    }

    #[test]
    fn test_delta_decode_single() {
        let deltas = vec![42];
        let values = delta_decode(&deltas);

        assert_eq!(values, vec![42]);
    }

    #[test]
    fn test_delta_accumulation() {
        let deltas = vec![100, 5, 5, 5, 5];
        let values = delta_decode(&deltas);

        assert_eq!(values, vec![100, 105, 110, 115, 120]);
    }

    #[test]
    fn test_delta_encode_large_array() {
        // Test with larger array to exercise SIMD path
        let values: Vec<u32> = (0..1000).map(|i| i * 10).collect();
        let deltas = delta_encode(&values);

        assert_eq!(deltas.len(), 1000);
        assert_eq!(deltas[0], 0); // First value
        assert_eq!(deltas[1], 10); // All subsequent deltas are 10
        assert_eq!(deltas[999], 10);
    }

    #[test]
    fn test_delta_encode_small_array_scalar() {
        // Small arrays should use scalar path
        let values = vec![1, 2, 3, 4, 5];
        let deltas = delta_encode(&values);

        assert_eq!(deltas, vec![1, 1, 1, 1, 1]);
    }

    #[test]
    #[cfg(target_arch = "x86_64")]
    fn test_delta_encode_simd_matches_scalar() {
        // Ensure SIMD and scalar produce identical results
        let values: Vec<u32> = (0..1000).map(|i| i * 7).collect();

        let scalar_deltas = delta_encode_scalar(&values);

        // Only test SIMD path if AVX2 is available
        if is_x86_feature_detected!("avx2") {
            let simd_deltas = unsafe { delta_encode_avx2(&values) };
            assert_eq!(
                simd_deltas, scalar_deltas,
                "SIMD and scalar produced different results"
            );
        } else {
            // Skip test if AVX2 not available
            println!("AVX2 not available, skipping SIMD comparison test");
        }
    }

    #[test]
    #[cfg(target_arch = "x86_64")]
    fn test_avx2_availability() {
        // Test that AVX2 detection doesn't panic
        let _has_avx2 = is_x86_feature_detected!("avx2");
        // If we get here without panicking, the test passes
    }

    #[test]
    fn test_varint_encode_bytes_written() {
        let mut buffer = Vec::new();

        // Single byte
        let bytes = encode_varint_scalar(&mut buffer, 42).unwrap();
        assert_eq!(bytes, 1);

        buffer.clear();

        // Two bytes
        let bytes = encode_varint_scalar(&mut buffer, 300).unwrap();
        assert_eq!(bytes, 2);

        buffer.clear();

        // Five bytes (max)
        let bytes = encode_varint_scalar(&mut buffer, u32::MAX).unwrap();
        assert_eq!(bytes, 5);
    }

    #[test]
    fn test_combined_varint_delta_workflow() {
        // Test a realistic workflow: delta encode + varint encode
        let values = vec![1000, 1050, 1100, 1150, 1200];

        // Step 1: Delta encode
        let deltas = delta_encode(&values);
        assert_eq!(deltas, vec![1000, 50, 50, 50, 50]);

        // Step 2: Varint encode each delta
        let mut encoded_data = Vec::new();
        for delta in &deltas {
            encode_varint_scalar(&mut encoded_data, *delta).unwrap();
        }

        // Step 3: Varint decode
        let mut decoded_deltas = Vec::new();
        let mut slice = encoded_data.as_slice();
        while !slice.is_empty() {
            let value = decode_varint_scalar(slice).unwrap();
            decoded_deltas.push(value);

            // Advance slice by bytes consumed
            let mut temp_buf = Vec::new();
            encode_varint_scalar(&mut temp_buf, value).unwrap();
            slice = &slice[temp_buf.len()..];
        }

        // Step 4: Delta decode
        let restored_values = delta_decode(&decoded_deltas);

        assert_eq!(restored_values, values);
    }
}