sb-mesh 0.1.1

/// Galois Field GF(2^8) with Rijndael irreducible polynomial 0x11B (x^8 + x^4 + x^3 + x + 1).
/// Provides byte-level linear arithmetic for Forward Error Correction (FEC) and Random Linear Network Coding (RLNC).
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Gf256;

impl Gf256 {
    pub const POLY: u16 = 0x11B;

    #[inline]
    pub fn add(a: u8, b: u8) -> u8 {
        a ^ b
    }

    #[inline]
    pub fn sub(a: u8, b: u8) -> u8 {
        a ^ b
    }

    pub fn mul(mut a: u8, mut b: u8) -> u8 {
        let mut p: u8 = 0;
        for _ in 0..8 {
            if (b & 1) != 0 {
                p ^= a;
            }
            let carry = (a & 0x80) != 0;
            a <<= 1;
            if carry {
                a ^= (Self::POLY & 0xFF) as u8;
            }
            b >>= 1;
        }
        p
    }

    pub fn inv(a: u8) -> u8 {
        if a == 0 {
            return 0;
        }
        // Fermat's Little Theorem in GF(2^8): a^(254) = a^(-1)
        let mut res: u8 = 1;
        let mut base = a;
        let mut exp: u8 = 254;
        while exp > 0 {
            if (exp & 1) != 0 {
                res = Self::mul(res, base);
            }
            base = Self::mul(base, base);
            exp >>= 1;
        }
        res
    }
}

/// A coded packet containing linear combination coefficients and coded payload bytes.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CodedPacket {
    pub coefficients: Vec<u8>,
    pub payload: Vec<u8>,
}

impl CodedPacket {
    pub fn new(coefficients: Vec<u8>, payload: Vec<u8>) -> Self {
        Self {
            coefficients,
            payload,
        }
    }

    pub fn to_bytes(&self) -> Vec<u8> {
        let mut buf = Vec::with_capacity(1 + self.coefficients.len() + self.payload.len());
        buf.push(self.coefficients.len() as u8);
        buf.extend_from_slice(&self.coefficients);
        buf.extend_from_slice(&self.payload);
        buf
    }

    pub fn from_bytes(raw: &[u8]) -> Result<Self, String> {
        if raw.is_empty() {
            return Err("Empty buffer".to_string());
        }
        let coeff_len = raw[0] as usize;
        if raw.len() < 1 + coeff_len {
            return Err("Truncated coded packet coefficients".to_string());
        }
        let coefficients = raw[1..1 + coeff_len].to_vec();
        let payload = raw[1 + coeff_len..].to_vec();
        Ok(Self {
            coefficients,
            payload,
        })
    }
}

/// Sender-Side Forward Error Correction (FEC) / RLNC Encoder
pub struct FecEncoder {
    pub k: usize,
    pub chunk_size: usize,
    pub original_chunks: Vec<Vec<u8>>,
}

impl FecEncoder {
    /// Creates an encoder that splits `data` into `k` equal-sized source chunks.
    pub fn new(data: &[u8], k: usize) -> Self {
        assert!(k > 0, "k must be at least 1");
        let chunk_size = (data.len() + k - 1) / k;
        let mut original_chunks = Vec::with_capacity(k);

        for i in 0..k {
            let start = i * chunk_size;
            let end = (start + chunk_size).min(data.len());
            let mut chunk = Vec::with_capacity(chunk_size);
            if start < data.len() {
                chunk.extend_from_slice(&data[start..end]);
            }
            // Zero-pad last chunk to chunk_size if necessary
            chunk.resize(chunk_size, 0u8);
            original_chunks.push(chunk);
        }

        Self {
            k,
            chunk_size,
            original_chunks,
        }
    }

    /// Produces a systematic source packet (coefficient vector is unit basis vector e_i)
    pub fn systematic_packet(&self, index: usize) -> Option<CodedPacket> {
        if index >= self.k {
            return None;
        }
        let mut coeffs = vec![0u8; self.k];
        coeffs[index] = 1;
        Some(CodedPacket::new(coeffs, self.original_chunks[index].clone()))
    }

    /// Generates a parity packet using pseudo-random or provided linear coefficients
    pub fn generate_coded_packet(&self, seed: u64) -> CodedPacket {
        let mut state = seed.wrapping_mul(6364136223846793005).wrapping_add(1);
        let mut coeffs = Vec::with_capacity(self.k);
        for _ in 0..self.k {
            state = state.wrapping_mul(6364136223846793005).wrapping_add(1);
            let mut c = (state >> 33) as u8;
            if c == 0 {
                c = 1;
            }
            coeffs.push(c);
        }

        let mut payload = vec![0u8; self.chunk_size];
        for (i, c) in coeffs.iter().enumerate() {
            let chunk = &self.original_chunks[i];
            for (byte_idx, &b) in chunk.iter().enumerate() {
                payload[byte_idx] = Gf256::add(payload[byte_idx], Gf256::mul(*c, b));
            }
        }

        CodedPacket::new(coeffs, payload)
    }

    /// Generates the complete transmission set: `k` systematic packets + `m` redundant FEC parity packets
    pub fn generate_fec_block(&self, m: usize) -> Vec<CodedPacket> {
        let mut packets = Vec::with_capacity(self.k + m);
        for i in 0..self.k {
            if let Some(pkt) = self.systematic_packet(i) {
                packets.push(pkt);
            }
        }
        for j in 0..m {
            packets.push(self.generate_coded_packet((j as u64 + 1).wrapping_mul(0x9E3779B97F4A7C15)));
        }
        packets
    }
}

/// Receiver-Side FEC / RLNC Decoder using Gaussian Elimination over GF(2^8)
pub struct FecDecoder {
    pub k: usize,
    pub chunk_size: usize,
    pub original_len: usize,
    matrix: Vec<Vec<u8>>,
    payloads: Vec<Vec<u8>>,
}

impl FecDecoder {
    pub fn new(k: usize, original_len: usize) -> Self {
        let chunk_size = (original_len + k - 1) / k;
        Self {
            k,
            chunk_size,
            original_len,
            matrix: Vec::new(),
            payloads: Vec::new(),
        }
    }

    /// Returns the current rank (number of linearly independent packets received)
    pub fn rank(&self) -> usize {
        self.matrix.len()
    }

    /// Returns true if enough linearly independent packets have been received to reconstruct data
    pub fn is_complete(&self) -> bool {
        self.rank() == self.k
    }

    /// Adds a received coded packet. Performs incremental Gaussian Elimination.
    /// Returns true if the packet was linearly independent (innovative), false if redundant.
    pub fn add_packet(&mut self, packet: CodedPacket) -> bool {
        if self.is_complete() || packet.coefficients.len() != self.k {
            return false;
        }

        let mut row = packet.coefficients;
        let mut payload = packet.payload;
        if payload.len() < self.chunk_size {
            payload.resize(self.chunk_size, 0);
        }

        // Reduce against existing echelon rows
        for i in 0..self.matrix.len() {
            let lead_col = match self.matrix[i].iter().position(|&x| x != 0) {
                Some(pos) => pos,
                None => continue,
            };

            let factor = row[lead_col];
            if factor != 0 {
                // row = row ^ (factor * existing_row)
                for c in 0..self.k {
                    let term = Gf256::mul(factor, self.matrix[i][c]);
                    row[c] = Gf256::add(row[c], term);
                }
                for b in 0..self.chunk_size {
                    let term = Gf256::mul(factor, self.payloads[i][b]);
                    payload[b] = Gf256::add(payload[b], term);
                }
            }
        }

        // Check if row is non-zero (innovative)
        if let Some(pivot_col) = row.iter().position(|&x| x != 0) {
            let inv_pivot = Gf256::inv(row[pivot_col]);
            // Normalize pivot to 1
            for c in 0..self.k {
                row[c] = Gf256::mul(row[c], inv_pivot);
            }
            for b in 0..self.chunk_size {
                payload[b] = Gf256::mul(payload[b], inv_pivot);
            }

            // Back-eliminate in existing rows
            for i in 0..self.matrix.len() {
                let factor = self.matrix[i][pivot_col];
                if factor != 0 {
                    for c in 0..self.k {
                        let term = Gf256::mul(factor, row[c]);
                        self.matrix[i][c] = Gf256::add(self.matrix[i][c], term);
                    }
                    for b in 0..self.chunk_size {
                        let term = Gf256::mul(factor, payload[b]);
                        self.payloads[i][b] = Gf256::add(self.payloads[i][b], term);
                    }
                }
            }

            self.matrix.push(row);
            self.payloads.push(payload);
            true
        } else {
            false
        }
    }

    /// Reconstructs the original data if rank == k
    pub fn decode(&self) -> Option<Vec<u8>> {
        if !self.is_complete() {
            return None;
        }

        // With reduced row echelon form where diagonal is 1, rows correspond to unit chunks
        let mut ordered_chunks = vec![vec![0u8; self.chunk_size]; self.k];

        for i in 0..self.k {
            if let Some(col) = self.matrix[i].iter().position(|&x| x != 0) {
                if col < self.k {
                    ordered_chunks[col] = self.payloads[i].clone();
                }
            }
        }

        let mut result = Vec::with_capacity(self.original_len);
        for chunk in ordered_chunks {
            result.extend_from_slice(&chunk);
        }
        result.truncate(self.original_len);
        Some(result)
    }
}

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

    #[test]
    fn test_gf256_basic_arithmetic() {
        assert_eq!(Gf256::add(0x57, 0x83), 0xD4);
        assert_eq!(Gf256::mul(0x57, 0x83), 0xC1);
        let inv = Gf256::inv(0x57);
        assert_eq!(Gf256::mul(0x57, inv), 1);
    }

    #[test]
    fn test_fec_recovery_from_dropped_packets() {
        let original_data = b"Sovereign peer-to-peer payload resilient against random loss.".to_vec();
        let k = 4;
        let m = 2; // 2 redundant parity packets (total 6 packets)

        let encoder = FecEncoder::new(&original_data, k);
        let packets = encoder.generate_fec_block(m);
        assert_eq!(packets.len(), 6);

        // Simulate dropping 2 packets (e.g. packets 0 and 2 dropped)
        let surviving_packets = vec![
            packets[1].clone(),
            packets[3].clone(),
            packets[4].clone(),
            packets[5].clone(),
        ];

        let mut decoder = FecDecoder::new(k, original_data.len());
        for pkt in surviving_packets {
            decoder.add_packet(pkt);
        }

        assert!(decoder.is_complete());
        let recovered = decoder.decode().expect("Decoding should succeed");
        assert_eq!(recovered, original_data);
    }
}