lcpfs 2026.1.102

// Copyright 2025 LunaOS Contributors
// SPDX-License-Identifier: Apache-2.0

//! # Galois Field GF(2^8) Linear Solver
//!
//! This module implements Gaussian elimination over GF(2^8) for RAID-Z2/Z3
//! multi-parity reconstruction. When multiple disks fail, we need to solve
//! a system of linear equations in the Galois field to recover data.
//!
//! ## Mathematical Background
//!
//! GF(2^8) is a finite field with 256 elements (0-255). Operations:
//! - Addition: XOR
//! - Multiplication: Polynomial multiplication mod primitive polynomial 0x11D
//! - Division: a/b = a * b^(-1) = a * b^254 (Fermat's little theorem)
//!
//! ## RAID-Z Reconstruction
//!
//! For RAID-Z2 (dual parity), we use a Vandermonde matrix:
//! ```text
//! P = D0 + D1 + D2 + ...              (XOR of all data)
//! Q = 2^0*D0 + 2^1*D1 + 2^2*D2 + ...  (Reed-Solomon syndrome)
//! ```
//!
//! For RAID-Z3 (triple parity), we add a third syndrome:
//! ```text
//! R = 4^0*D0 + 4^1*D1 + 4^2*D2 + ...
//! ```

use crate::ml::gfalgo::GfAlgo;
use alloc::vec;
use alloc::vec::Vec;

/// GF(2^8) inverse lookup table.
///
/// Generated using: inv(a) = a^254 (Fermat's little theorem for prime fields).
/// For GF(2^8), the multiplicative order is 255, so a^255 = 1, thus a^(-1) = a^254.
static GF_INV: [u8; 256] = generate_inverse_table();

/// Generate the GF(2^8) inverse lookup table at compile time.
const fn generate_inverse_table() -> [u8; 256] {
    let mut table = [0u8; 256];
    // 0 has no inverse (undefined), leave as 0
    table[0] = 0;
    table[1] = 1; // 1 is its own inverse

    // For other elements, compute a^254 using repeated squaring
    let mut i = 2usize;
    while i < 256 {
        table[i] = gf_pow_const(i as u8, 254);
        i += 1;
    }
    table
}

/// Compute a^n in GF(2^8) using repeated squaring (const-compatible).
const fn gf_pow_const(a: u8, mut n: u8) -> u8 {
    let mut result = 1u8;
    let mut base = a;

    while n > 0 {
        if n & 1 != 0 {
            result = gf_mul_const(result, base);
        }
        base = gf_mul_const(base, base);
        n >>= 1;
    }
    result
}

/// GF(2^8) multiplication with primitive polynomial 0x11D (const-compatible).
const fn gf_mul_const(a: u8, b: u8) -> u8 {
    let mut p: u16 = 0;
    let mut a_copy = a as u16;
    let mut b_copy = b as u16;

    let mut i = 0;
    while i < 8 {
        if (a_copy & 1) != 0 {
            p ^= b_copy;
        }
        a_copy >>= 1;
        let carry = (b_copy & 0x80) != 0;
        b_copy <<= 1;
        if carry {
            b_copy ^= 0x11D; // Primitive polynomial for GF(2^8)
        }
        i += 1;
    }
    p as u8
}

/// Galois Field linear equation solver for RAID reconstruction
pub struct GfSolver;

impl GfSolver {
    /// Get the multiplicative inverse of a in GF(2^8).
    ///
    /// Returns 0 for input 0 (undefined, but safe for our use cases).
    #[inline]
    pub fn inverse(a: u8) -> u8 {
        GF_INV[a as usize]
    }

    /// Solve a system of linear equations over GF(2^8) using Gaussian elimination.
    ///
    /// The input matrix is an augmented matrix [A|b] where:
    /// - First `cols` columns are the coefficient matrix A
    /// - Remaining columns are the right-hand side b
    ///
    /// Returns the solution vectors if the system is solvable.
    ///
    /// # Arguments
    ///
    /// * `matrix` - Augmented matrix [A|b], modified in place
    /// * `rows` - Number of equations (rows in A)
    /// * `cols` - Number of unknowns (columns in A)
    ///
    /// # Returns
    ///
    /// * `Ok(solutions)` - Vector of solution vectors
    /// * `Err(msg)` - If the system is underdetermined (not enough parity)
    pub fn solve(
        matrix: &mut [Vec<u8>],
        rows: usize,
        cols: usize,
    ) -> Result<Vec<Vec<u8>>, &'static str> {
        crate::lcpfs_println!(
            "[ GFSOLVER] Starting GF(2^8) Matrix Solution ({} eq, {} unknowns)...",
            rows,
            cols
        );

        if rows < cols {
            return Err("Underdetermined system: need more parity disks");
        }

        // ═══════════════════════════════════════════════════════════════════════
        // PHASE 1: FORWARD ELIMINATION (Row Echelon Form)
        // ═══════════════════════════════════════════════════════════════════════
        for col in 0..cols {
            // Find pivot (first non-zero element in column)
            let mut pivot_row = col;
            while pivot_row < rows && matrix[pivot_row][col] == 0 {
                pivot_row += 1;
            }

            if pivot_row == rows {
                return Err("Unrecoverable: singular matrix (linearly dependent parities)");
            }

            // Swap rows if needed
            if pivot_row != col {
                matrix.swap(col, pivot_row);
            }

            // Scale pivot row so pivot element becomes 1
            let pivot = matrix[col][col];
            let inv_pivot = Self::inverse(pivot);

            if inv_pivot == 0 && pivot != 0 {
                return Err("Unrecoverable: inverse calculation failed");
            }

            // Scale pivot row elements
            for elem in matrix[col].iter_mut().skip(col) {
                *elem = GfAlgo::multiply(*elem, inv_pivot);
            }

            // Eliminate below pivot - need to copy pivot row values first
            let pivot_row_vals: Vec<u8> = matrix[col].iter().skip(col).copied().collect();

            for matrix_row in matrix.iter_mut().take(rows).skip(col + 1) {
                let factor = matrix_row[col];
                if factor != 0 {
                    for (idx, &pivot_val) in pivot_row_vals.iter().enumerate() {
                        let term = GfAlgo::multiply(factor, pivot_val);
                        matrix_row[col + idx] ^= term;
                    }
                }
            }
        }

        // ═══════════════════════════════════════════════════════════════════════
        // PHASE 2: BACK SUBSTITUTION (Reduced Row Echelon Form)
        // ═══════════════════════════════════════════════════════════════════════
        for col in (1..cols).rev() {
            // Copy the col-th row's relevant values first
            let col_row_vals: Vec<u8> = matrix[col].iter().skip(col).copied().collect();

            for matrix_row in matrix.iter_mut().take(col) {
                let factor = matrix_row[col];
                if factor != 0 {
                    for (idx, &col_val) in col_row_vals.iter().enumerate() {
                        let term = GfAlgo::multiply(factor, col_val);
                        matrix_row[col + idx] ^= term;
                    }
                }
            }
        }

        // ═══════════════════════════════════════════════════════════════════════
        // PHASE 3: EXTRACT SOLUTIONS
        // ═══════════════════════════════════════════════════════════════════════
        let num_rhs = matrix[0].len() - cols; // Number of right-hand side vectors
        let mut solutions = Vec::with_capacity(num_rhs);

        for rhs_col in 0..num_rhs {
            let solution: Vec<u8> = matrix
                .iter()
                .take(cols)
                .map(|row| row[cols + rhs_col])
                .collect();
            solutions.push(solution);
        }

        crate::lcpfs_println!("[ GFSOLVER] Solution complete. {} unknowns solved.", cols);
        Ok(solutions)
    }

    /// Reconstruct a single failed disk from P parity (RAID-Z1).
    ///
    /// P = D0 ⊕ D1 ⊕ ... ⊕ Dn
    /// D_failed = P ⊕ (all other D)
    pub fn reconstruct_z1(surviving_data: &[&[u8]], parity_p: &[u8], block_size: usize) -> Vec<u8> {
        let mut recovered = vec![0u8; block_size];

        // Start with parity
        recovered.copy_from_slice(parity_p);

        // XOR all surviving data blocks
        for data in surviving_data {
            for (i, &byte) in data.iter().enumerate() {
                if i < block_size {
                    recovered[i] ^= byte;
                }
            }
        }

        recovered
    }

    /// Reconstruct one or two failed disks from P and Q parities (RAID-Z2).
    ///
    /// P = D0 ⊕ D1 ⊕ D2 ⊕ ...
    /// Q = 2^0*D0 ⊕ 2^1*D1 ⊕ 2^2*D2 ⊕ ...
    ///
    /// For single failure: use P like RAID-Z1
    /// For double failure at positions x, y:
    ///   Pxy = P ⊕ (surviving data)     = Dx ⊕ Dy
    ///   Qxy = Q ⊕ (weighted surviving) = 2^x*Dx ⊕ 2^y*Dy
    ///
    /// Solve: Dx = (2^y * Pxy ⊕ Qxy) / (2^y ⊕ 2^x)
    ///        Dy = Pxy ⊕ Dx
    pub fn reconstruct_z2(
        failed_indices: &[usize],
        surviving_data: &[(&[u8], usize)], // (data, original_index)
        parity_p: &[u8],
        parity_q: &[u8],
        block_size: usize,
    ) -> Result<Vec<Vec<u8>>, &'static str> {
        match failed_indices.len() {
            0 => Ok(Vec::new()),

            1 => {
                // Single failure - use P parity
                let recovered = Self::reconstruct_z1(
                    &surviving_data.iter().map(|(d, _)| *d).collect::<Vec<_>>(),
                    parity_p,
                    block_size,
                );
                Ok(vec![recovered])
            }

            2 => {
                let x = failed_indices[0];
                let y = failed_indices[1];

                // Compute Pxy = P ⊕ (all surviving data)
                let mut pxy = vec![0u8; block_size];
                pxy.copy_from_slice(parity_p);
                for (data, _) in surviving_data {
                    for (i, &byte) in data.iter().enumerate() {
                        if i < block_size {
                            pxy[i] ^= byte;
                        }
                    }
                }

                // Compute Qxy = Q ⊕ (weighted surviving data)
                let mut qxy = vec![0u8; block_size];
                qxy.copy_from_slice(parity_q);
                for (data, idx) in surviving_data {
                    let coeff = GfAlgo::multiply(1, Self::gf_pow_2(*idx)); // 2^idx
                    for (i, &byte) in data.iter().enumerate() {
                        if i < block_size {
                            qxy[i] ^= GfAlgo::multiply(byte, coeff);
                        }
                    }
                }

                // Coefficients: 2^x, 2^y
                let coeff_x = Self::gf_pow_2(x);
                let coeff_y = Self::gf_pow_2(y);

                // Divisor: 2^y ⊕ 2^x
                let divisor = coeff_x ^ coeff_y;
                if divisor == 0 {
                    return Err("Cannot reconstruct: coefficient collision");
                }
                let inv_divisor = Self::inverse(divisor);

                // Dx = (2^y * Pxy ⊕ Qxy) / (2^y ⊕ 2^x)
                let mut dx = vec![0u8; block_size];
                for i in 0..block_size {
                    let numerator = GfAlgo::multiply(pxy[i], coeff_y) ^ qxy[i];
                    dx[i] = GfAlgo::multiply(numerator, inv_divisor);
                }

                // Dy = Pxy ⊕ Dx
                let mut dy = vec![0u8; block_size];
                for i in 0..block_size {
                    dy[i] = pxy[i] ^ dx[i];
                }

                Ok(vec![dx, dy])
            }

            _ => Err("RAID-Z2 can only recover up to 2 disk failures"),
        }
    }

    /// Reconstruct one, two, or three failed disks (RAID-Z3).
    ///
    /// P = D0 ⊕ D1 ⊕ D2 ⊕ ...
    /// Q = 2^0*D0 ⊕ 2^1*D1 ⊕ 2^2*D2 ⊕ ...
    /// R = 4^0*D0 ⊕ 4^1*D1 ⊕ 4^2*D2 ⊕ ...
    pub fn reconstruct_z3(
        failed_indices: &[usize],
        surviving_data: &[(&[u8], usize)],
        parity_p: &[u8],
        parity_q: &[u8],
        parity_r: &[u8],
        block_size: usize,
    ) -> Result<Vec<Vec<u8>>, &'static str> {
        match failed_indices.len() {
            0 => Ok(Vec::new()),

            1 => {
                let recovered = Self::reconstruct_z1(
                    &surviving_data.iter().map(|(d, _)| *d).collect::<Vec<_>>(),
                    parity_p,
                    block_size,
                );
                Ok(vec![recovered])
            }

            2 => Self::reconstruct_z2(
                failed_indices,
                surviving_data,
                parity_p,
                parity_q,
                block_size,
            ),

            3 => {
                let x = failed_indices[0];
                let y = failed_indices[1];
                let z = failed_indices[2];

                // Compute Pxyz, Qxyz, Rxyz
                let mut pxyz = vec![0u8; block_size];
                let mut qxyz = vec![0u8; block_size];
                let mut rxyz = vec![0u8; block_size];

                pxyz.copy_from_slice(parity_p);
                qxyz.copy_from_slice(parity_q);
                rxyz.copy_from_slice(parity_r);

                for (data, idx) in surviving_data {
                    let coeff_q = Self::gf_pow_2(*idx); // 2^idx
                    let coeff_r = Self::gf_pow_4(*idx); // 4^idx

                    for (i, &byte) in data.iter().enumerate() {
                        if i < block_size {
                            pxyz[i] ^= byte;
                            qxyz[i] ^= GfAlgo::multiply(byte, coeff_q);
                            rxyz[i] ^= GfAlgo::multiply(byte, coeff_r);
                        }
                    }
                }

                // Set up 3x3 system over GF(2^8)
                // [1       1       1     ] [Dx]   [Pxyz]
                // [2^x     2^y     2^z   ] [Dy] = [Qxyz]
                // [4^x     4^y     4^z   ] [Dz]   [Rxyz]

                let mut dx = vec![0u8; block_size];
                let mut dy = vec![0u8; block_size];
                let mut dz = vec![0u8; block_size];

                // Precompute coefficients
                let a11 = 1u8;
                let a12 = 1u8;
                let a13 = 1u8;
                let a21 = Self::gf_pow_2(x);
                let a22 = Self::gf_pow_2(y);
                let a23 = Self::gf_pow_2(z);
                let a31 = Self::gf_pow_4(x);
                let a32 = Self::gf_pow_4(y);
                let a33 = Self::gf_pow_4(z);

                // Compute determinant
                let det =
                    GfAlgo::multiply(a11, GfAlgo::multiply(a22, a33) ^ GfAlgo::multiply(a23, a32))
                        ^ GfAlgo::multiply(
                            a12,
                            GfAlgo::multiply(a21, a33) ^ GfAlgo::multiply(a23, a31),
                        )
                        ^ GfAlgo::multiply(
                            a13,
                            GfAlgo::multiply(a21, a32) ^ GfAlgo::multiply(a22, a31),
                        );

                if det == 0 {
                    return Err("Singular matrix in RAID-Z3 reconstruction");
                }
                let inv_det = Self::inverse(det);

                // Compute adjugate matrix elements (cofactors transposed)
                let c11 = GfAlgo::multiply(a22, a33) ^ GfAlgo::multiply(a23, a32);
                let c12 = GfAlgo::multiply(a13, a32) ^ GfAlgo::multiply(a12, a33);
                let c13 = GfAlgo::multiply(a12, a23) ^ GfAlgo::multiply(a13, a22);
                let c21 = GfAlgo::multiply(a23, a31) ^ GfAlgo::multiply(a21, a33);
                let c22 = GfAlgo::multiply(a11, a33) ^ GfAlgo::multiply(a13, a31);
                let c23 = GfAlgo::multiply(a13, a21) ^ GfAlgo::multiply(a11, a23);
                let c31 = GfAlgo::multiply(a21, a32) ^ GfAlgo::multiply(a22, a31);
                let c32 = GfAlgo::multiply(a12, a31) ^ GfAlgo::multiply(a11, a32);
                let c33 = GfAlgo::multiply(a11, a22) ^ GfAlgo::multiply(a12, a21);

                // Solve for each byte position
                for i in 0..block_size {
                    let b1 = pxyz[i];
                    let b2 = qxyz[i];
                    let b3 = rxyz[i];

                    // x = A^(-1) * b = adj(A) * b / det
                    let dx_num = GfAlgo::multiply(c11, b1)
                        ^ GfAlgo::multiply(c12, b2)
                        ^ GfAlgo::multiply(c13, b3);
                    let dy_num = GfAlgo::multiply(c21, b1)
                        ^ GfAlgo::multiply(c22, b2)
                        ^ GfAlgo::multiply(c23, b3);
                    let dz_num = GfAlgo::multiply(c31, b1)
                        ^ GfAlgo::multiply(c32, b2)
                        ^ GfAlgo::multiply(c33, b3);

                    dx[i] = GfAlgo::multiply(dx_num, inv_det);
                    dy[i] = GfAlgo::multiply(dy_num, inv_det);
                    dz[i] = GfAlgo::multiply(dz_num, inv_det);
                }

                Ok(vec![dx, dy, dz])
            }

            _ => Err("RAID-Z3 can only recover up to 3 disk failures"),
        }
    }

    /// Compute 2^n in GF(2^8).
    fn gf_pow_2(n: usize) -> u8 {
        let mut result = 1u8;
        for _ in 0..n {
            result = GfAlgo::multiply(result, 2);
        }
        result
    }

    /// Compute 4^n in GF(2^8).
    fn gf_pow_4(n: usize) -> u8 {
        let mut result = 1u8;
        let four = GfAlgo::multiply(2, 2);
        for _ in 0..n {
            result = GfAlgo::multiply(result, four);
        }
        result
    }
}

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

    #[test]
    fn test_gf_inverse() {
        // Test a * a^(-1) = 1 for various values
        for a in [1, 2, 3, 127, 255].iter() {
            let inv = GfSolver::inverse(*a);
            let product = GfAlgo::multiply(*a, inv);
            assert_eq!(
                product, 1,
                "Inverse of {} failed: {} * {} = {}",
                a, a, inv, product
            );
        }
    }

    #[test]
    fn test_gf_inverse_table_consistency() {
        // Verify entire table
        for a in 1..=255u8 {
            let inv = GfSolver::inverse(a);
            let product = GfAlgo::multiply(a, inv);
            assert_eq!(
                product, 1,
                "Inverse of {} = {}, product = {}",
                a, inv, product
            );
        }
    }

    #[test]
    fn test_z1_reconstruction() {
        let d0 = vec![0x11u8; 512];
        let d1 = vec![0x22u8; 512];

        // P = D0 XOR D1
        let parity: Vec<u8> = d0.iter().zip(&d1).map(|(a, b)| a ^ b).collect();

        // Lose D0, reconstruct from D1 + P
        let recovered = GfSolver::reconstruct_z1(&[&d1], &parity, 512);
        assert_eq!(recovered, d0, "Z1 reconstruction failed");
    }
}