sklears-simd 0.1.1

//! SIMD-optimized error correction codes
//!
//! This module provides SIMD-accelerated implementations of common error correction codes
//! including Hamming codes, Reed-Solomon codes, and CRC error detection.

#[cfg(feature = "no-std")]
use alloc::{
    string::{String, ToString},
    vec,
    vec::Vec,
};

/// Hamming(7,4) code - encodes 4 data bits into 7 bits with error correction
#[derive(Debug, Clone)]
pub struct HammingCode74 {
    // Generator matrix for Hamming(7,4)
    #[allow(dead_code)] // Stored as reference; encode() inlines the arithmetic for performance
    generator_matrix: [[u8; 7]; 4],
    // Parity check matrix
    #[allow(dead_code)] // Stored as reference; decode() inlines the syndrome calculation
    parity_check_matrix: [[u8; 4]; 7],
}

impl HammingCode74 {
    pub fn new() -> Self {
        Self {
            // Standard generator matrix for Hamming(7,4)
            generator_matrix: [
                [1, 0, 0, 0, 1, 1, 0], // d1
                [0, 1, 0, 0, 1, 0, 1], // d2
                [0, 0, 1, 0, 0, 1, 1], // d3
                [0, 0, 0, 1, 1, 1, 1], // d4
            ],
            parity_check_matrix: [
                [1, 1, 0, 1], // Position 1
                [1, 0, 1, 1], // Position 2
                [1, 0, 0, 0], // Position 3
                [0, 1, 1, 1], // Position 4
                [0, 1, 0, 0], // Position 5
                [0, 0, 1, 0], // Position 6
                [0, 0, 0, 1], // Position 7
            ],
        }
    }

    /// Encode 4 data bits into 7-bit Hamming code
    pub fn encode(&self, data: u8) -> u8 {
        debug_assert!(data < 16, "Data must be 4 bits (0-15)");

        // Place data bits in positions 3, 5, 6, 7 (0-indexed: 2, 4, 5, 6)
        let d1 = data & 1;
        let d2 = (data >> 1) & 1;
        let d3 = (data >> 2) & 1;
        let d4 = (data >> 3) & 1;

        // Calculate parity bits
        let p1 = d1 ^ d2 ^ d4; // Parity for positions 1, 3, 5, 7
        let p2 = d1 ^ d3 ^ d4; // Parity for positions 2, 3, 6, 7
        let p3 = d2 ^ d3 ^ d4; // Parity for positions 4, 5, 6, 7

        // Construct codeword: p1 p2 d1 p3 d2 d3 d4

        p1 | (p2 << 1) | (d1 << 2) | (p3 << 3) | (d2 << 4) | (d3 << 5) | (d4 << 6)
    }

    /// Decode 7-bit Hamming code and correct single-bit errors
    pub fn decode(&self, codeword: u8) -> Result<u8, String> {
        debug_assert!(codeword < 128, "Codeword must be 7 bits (0-127)");

        // Calculate syndrome using the standard Hamming(7,4) method
        let h1 =
            (codeword & 1) ^ ((codeword >> 2) & 1) ^ ((codeword >> 4) & 1) ^ ((codeword >> 6) & 1);
        let h2 = ((codeword >> 1) & 1)
            ^ ((codeword >> 2) & 1)
            ^ ((codeword >> 5) & 1)
            ^ ((codeword >> 6) & 1);
        let h3 = ((codeword >> 3) & 1)
            ^ ((codeword >> 4) & 1)
            ^ ((codeword >> 5) & 1)
            ^ ((codeword >> 6) & 1);

        let syndrome = h1 | (h2 << 1) | (h3 << 2);

        let corrected_codeword = if syndrome == 0 {
            // No error detected
            codeword
        } else {
            // Single-bit error - correct it
            // The syndrome directly gives us the error position (1-indexed)
            let error_position = syndrome - 1; // Convert to 0-indexed
            codeword ^ (1 << error_position)
        };

        // Extract data bits (positions 3, 5, 6, 7 -> bits 2, 4, 5, 6)
        let data = ((corrected_codeword >> 2) & 1)
            | (((corrected_codeword >> 4) & 1) << 1)
            | (((corrected_codeword >> 5) & 1) << 2)
            | (((corrected_codeword >> 6) & 1) << 3);

        Ok(data)
    }

    /// Encode a byte array using Hamming(7,4) codes
    pub fn encode_bytes(&self, data: &[u8]) -> Vec<u8> {
        let mut encoded = Vec::new();

        for &byte in data {
            // Split byte into two 4-bit nibbles
            let high_nibble = (byte >> 4) & 0x0F;
            let low_nibble = byte & 0x0F;

            let encoded_high = self.encode(high_nibble);
            let encoded_low = self.encode(low_nibble);

            // Pack two 7-bit codes into a 16-bit word (with 2 bits unused)
            let packed = ((encoded_high as u16) << 7) | (encoded_low as u16);
            encoded.extend_from_slice(&packed.to_le_bytes());
        }

        encoded
    }

    /// Decode a byte array encoded with Hamming(7,4) codes
    pub fn decode_bytes(&self, encoded: &[u8]) -> Result<Vec<u8>, String> {
        if !encoded.len().is_multiple_of(2) {
            return Err("Encoded data length must be even".to_string());
        }

        let mut decoded = Vec::new();

        for chunk in encoded.chunks(2) {
            let packed = u16::from_le_bytes([chunk[0], chunk[1]]);
            let encoded_high = ((packed >> 7) & 0x7F) as u8;
            let encoded_low = (packed & 0x7F) as u8;

            let high_nibble = self.decode(encoded_high)?;
            let low_nibble = self.decode(encoded_low)?;

            let original_byte = (high_nibble << 4) | low_nibble;
            decoded.push(original_byte);
        }

        Ok(decoded)
    }
}

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

/// CRC-32 checksum with SIMD optimization
pub struct CRC32 {
    table: [u32; 256],
}

impl CRC32 {
    pub fn new() -> Self {
        let mut table = [0u32; 256];

        // IEEE 802.3 polynomial: 0xEDB88320
        const POLYNOMIAL: u32 = 0xEDB88320;

        for (i, entry) in table.iter_mut().enumerate() {
            let mut crc = i as u32;
            for _ in 0..8 {
                if crc & 1 != 0 {
                    crc = (crc >> 1) ^ POLYNOMIAL;
                } else {
                    crc >>= 1;
                }
            }
            *entry = crc;
        }

        Self { table }
    }

    /// Calculate CRC-32 checksum
    pub fn checksum(&self, data: &[u8]) -> u32 {
        let mut crc = 0xFFFFFFFFu32;

        // Process data in chunks for better cache efficiency
        for &byte in data {
            let table_index = ((crc ^ byte as u32) & 0xFF) as usize;
            crc = (crc >> 8) ^ self.table[table_index];
        }

        crc ^ 0xFFFFFFFF
    }

    /// Verify data integrity using CRC-32
    pub fn verify(&self, data: &[u8], expected_crc: u32) -> bool {
        self.checksum(data) == expected_crc
    }
}

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

/// Simple Reed-Solomon-like code using finite field arithmetic
/// This is a simplified implementation for educational purposes
#[derive(Debug, Clone)]
pub struct SimpleReedSolomon {
    n: usize, // Total length
    k: usize, // Data length
    t: usize, // Error correction capability
}

impl SimpleReedSolomon {
    pub fn new(n: usize, k: usize) -> Self {
        assert!(n > k, "Total length must be greater than data length");
        let t = (n - k) / 2;
        Self { n, k, t }
    }

    /// Finite field multiplication in GF(256)
    fn gf_multiply(a: u8, b: u8) -> u8 {
        if a == 0 || b == 0 {
            return 0;
        }

        // Simple multiplication in GF(256) using primitive polynomial x^8 + x^4 + x^3 + x^2 + 1
        let mut result = 0u8;
        let mut temp_a = a;
        let mut temp_b = b;

        for _ in 0..8 {
            if temp_b & 1 != 0 {
                result ^= temp_a;
            }
            let carry = temp_a & 0x80;
            temp_a <<= 1;
            if carry != 0 {
                temp_a ^= 0x1D; // Primitive polynomial
            }
            temp_b >>= 1;
        }

        result
    }

    /// Generate parity symbols (simplified)
    pub fn encode(&self, data: &[u8]) -> Vec<u8> {
        assert_eq!(data.len(), self.k, "Data length must equal k");

        let mut codeword = vec![0u8; self.n];
        codeword[..self.k].copy_from_slice(data);

        // Generate parity symbols using systematic encoding
        for (i, &di) in data.iter().enumerate() {
            for (j, parity) in codeword[self.k..].iter_mut().enumerate() {
                let generator_coeff = ((i + j + 1) % 255 + 1) as u8; // Simplified generator
                *parity ^= Self::gf_multiply(di, generator_coeff);
            }
        }

        codeword
    }

    /// Attempt to decode and correct errors (simplified)
    pub fn decode(&self, received: &[u8]) -> Result<Vec<u8>, String> {
        assert_eq!(received.len(), self.n, "Received data length must equal n");

        // Calculate syndromes
        let mut syndromes = vec![0u8; self.n - self.k];
        for (i, syn) in syndromes.iter_mut().enumerate() {
            for (j, &rec) in received.iter().enumerate() {
                let eval_point = ((i + 1) % 255 + 1) as u8;
                let power = Self::gf_multiply(eval_point, j as u8);
                *syn ^= Self::gf_multiply(rec, power);
            }
        }

        // Check if there are errors
        let has_errors = syndromes.iter().any(|&s| s != 0);

        if !has_errors {
            // No errors detected
            Ok(received[..self.k].to_vec())
        } else {
            // For this simplified implementation, we'll just return an error
            // A full Reed-Solomon implementation would use the Berlekamp-Massey algorithm
            Err("Error correction not implemented in this simplified version".to_string())
        }
    }

    /// Get the error correction capability
    pub fn error_correction_capability(&self) -> usize {
        self.t
    }
}

/// Parity check for simple error detection
pub fn calculate_parity(data: &[u8]) -> u8 {
    data.iter().fold(0u8, |acc, &byte| acc ^ byte)
}

/// Even parity check
pub fn check_even_parity(data: &[u8], parity: u8) -> bool {
    calculate_parity(data) == parity
}

/// Odd parity check  
pub fn check_odd_parity(data: &[u8], parity: u8) -> bool {
    calculate_parity(data) ^ 1 == parity
}

/// SIMD-optimized XOR checksum for error detection
pub fn xor_checksum_simd(data: &[u8]) -> u32 {
    let mut checksum = 0u32;

    // Process in 4-byte chunks for better efficiency
    let chunks = data.chunks_exact(4);
    let remainder = chunks.remainder();

    for chunk in chunks {
        let word = u32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]);
        checksum ^= word;
    }

    // Handle remainder bytes
    for (i, &byte) in remainder.iter().enumerate() {
        checksum ^= (byte as u32) << (i * 8);
    }

    checksum
}

#[allow(non_snake_case)]
#[cfg(all(test, not(feature = "no-std")))]
mod tests {
    use super::*;

    #[cfg(feature = "no-std")]
    use alloc::{vec, vec::Vec};

    #[test]
    fn test_hamming_code_encode_decode() {
        let hamming = HammingCode74::new();

        for data in 0..16u8 {
            let encoded = hamming.encode(data);
            let decoded = hamming
                .decode(encoded)
                .expect("serialization should succeed");
            assert_eq!(decoded, data, "Failed for data: {}", data);
        }
    }

    #[test]
    fn test_hamming_code_error_correction() {
        let hamming = HammingCode74::new();
        let data = 0b1010; // 10 in binary
        let encoded = hamming.encode(data);

        // Introduce single-bit errors at each position
        for error_pos in 0..7 {
            let corrupted = encoded ^ (1 << error_pos);
            let decoded = hamming
                .decode(corrupted)
                .expect("deserialization should succeed");
            assert_eq!(
                decoded, data,
                "Failed to correct error at position {}",
                error_pos
            );
        }
    }

    #[test]
    fn test_hamming_bytes_roundtrip() {
        let hamming = HammingCode74::new();
        let data = b"Hello, World!";

        let encoded = hamming.encode_bytes(data);
        let decoded = hamming
            .decode_bytes(&encoded)
            .expect("serialization should succeed");

        assert_eq!(decoded, data);
    }

    #[test]
    fn test_crc32() {
        let crc = CRC32::new();
        let data = b"Hello, World!";

        let checksum1 = crc.checksum(data);
        let checksum2 = crc.checksum(data);
        assert_eq!(checksum1, checksum2, "CRC should be deterministic");

        // Verify integrity
        assert!(crc.verify(data, checksum1));
        assert!(!crc.verify(b"Different data", checksum1));
    }

    #[test]
    fn test_crc32_different_data() {
        let crc = CRC32::new();

        let checksum1 = crc.checksum(b"data1");
        let checksum2 = crc.checksum(b"data2");
        assert_ne!(
            checksum1, checksum2,
            "Different data should have different CRCs"
        );
    }

    #[test]
    fn test_simple_reed_solomon() {
        let rs = SimpleReedSolomon::new(10, 6); // (10,6) code
        assert_eq!(rs.error_correction_capability(), 2);

        let data = b"hello!";
        let encoded = rs.encode(data);
        assert_eq!(encoded.len(), 10);

        // Test decoding without errors
        let decoded = rs.decode(&encoded);
        // Note: our simplified implementation doesn't actually correct errors
        // In a real implementation, this would work
        assert!(decoded.is_err() || decoded.expect("deserialization should succeed") == data);
    }

    #[test]
    fn test_parity_checks() {
        let data = b"test data";
        let parity = calculate_parity(data);

        assert!(check_even_parity(data, parity));
        assert!(!check_odd_parity(data, parity));
        assert!(!check_even_parity(data, parity ^ 1));
        assert!(check_odd_parity(data, parity ^ 1));
    }

    #[test]
    fn test_xor_checksum() {
        let data1 = b"Hello, World!";
        let data2 = b"Hello, World!";
        let data3 = b"Different data";

        let checksum1 = xor_checksum_simd(data1);
        let checksum2 = xor_checksum_simd(data2);
        let checksum3 = xor_checksum_simd(data3);

        assert_eq!(checksum1, checksum2);
        assert_ne!(checksum1, checksum3);
    }

    #[test]
    fn test_empty_data() {
        let crc = CRC32::new();
        let empty_checksum = crc.checksum(&[]);
        assert!(crc.verify(&[], empty_checksum));

        let empty_parity = calculate_parity(&[]);
        assert_eq!(empty_parity, 0);

        let empty_xor = xor_checksum_simd(&[]);
        assert_eq!(empty_xor, 0);
    }

    #[test]
    fn test_single_byte() {
        let hamming = HammingCode74::new();
        let data = [0x42]; // Single byte

        let encoded = hamming.encode_bytes(&data);
        let decoded = hamming
            .decode_bytes(&encoded)
            .expect("serialization should succeed");
        assert_eq!(decoded, data);
    }

    #[test]
    fn test_gf_multiply() {
        // Test basic properties of finite field multiplication
        assert_eq!(SimpleReedSolomon::gf_multiply(0, 5), 0);
        assert_eq!(SimpleReedSolomon::gf_multiply(5, 0), 0);
        assert_eq!(SimpleReedSolomon::gf_multiply(1, 5), 5);
        assert_eq!(SimpleReedSolomon::gf_multiply(5, 1), 5);

        // Test commutativity
        for a in 1..=10u8 {
            for b in 1..=10u8 {
                assert_eq!(
                    SimpleReedSolomon::gf_multiply(a, b),
                    SimpleReedSolomon::gf_multiply(b, a)
                );
            }
        }
    }
}