openipc_core/
fec.rs

1use std::sync::OnceLock;
2
3const GF_SIZE: usize = 255;
4const GF_BITS: usize = 8;
5const PRIMITIVE_POLY: &[u8; 9] = b"101110001";
6
7static GF_TABLES: OnceLock<GfTables> = OnceLock::new();
8
9/// Reed-Solomon FEC error.
10#[derive(Debug, Clone, PartialEq, Eq)]
11pub enum FecError {
12    /// FEC parameters are outside `0 < k <= n < 256`.
13    InvalidParameters,
14    /// Fewer than `k` usable fragments were available.
15    NotEnoughFragments,
16    /// A fragment index was outside the configured block.
17    InvalidFragmentIndex(usize),
18    /// The decode matrix could not be inverted.
19    SingularMatrix,
20    /// Recovered output did not match expected primary slots.
21    OutputSlotMismatch,
22}
23
24impl std::fmt::Display for FecError {
25    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
26        match self {
27            Self::InvalidParameters => write!(f, "invalid FEC parameters"),
28            Self::NotEnoughFragments => write!(f, "not enough fragments to recover block"),
29            Self::InvalidFragmentIndex(idx) => write!(f, "invalid FEC fragment index {idx}"),
30            Self::SingularMatrix => write!(f, "FEC decode matrix is singular"),
31            Self::OutputSlotMismatch => write!(f, "FEC output slot mismatch"),
32        }
33    }
34}
35
36impl std::error::Error for FecError {}
37
38/// Reed-Solomon FEC code used by WFB blocks.
39#[derive(Debug, Clone)]
40pub struct FecCode {
41    k: usize,
42    n: usize,
43    enc_matrix: Vec<u8>,
44}
45
46impl FecCode {
47    /// Create an FEC code with `k` primary fragments and `n-k` parity fragments.
48    pub fn new(k: usize, n: usize) -> Result<Self, FecError> {
49        if k == 0 || n == 0 || k > n || n >= 256 {
50            return Err(FecError::InvalidParameters);
51        }
52
53        let tables = tables();
54        let mut tmp = vec![0; n * k];
55        tmp[0] = 1;
56        for row in 0..(n - 1) {
57            for col in 0..k {
58                tmp[(row + 1) * k + col] = tables.gf_exp[modnn((row * col) as i32) as usize];
59            }
60        }
61
62        invert_vdm(&mut tmp[..k * k], k)?;
63
64        let mut enc_matrix = vec![0; n * k];
65        if n > k {
66            matmul(
67                &tmp[k * k..],
68                &tmp[..k * k],
69                &mut enc_matrix[k * k..],
70                n - k,
71                k,
72                k,
73            );
74        }
75        for col in 0..k {
76            enc_matrix[col * k + col] = 1;
77        }
78
79        Ok(Self { k, n, enc_matrix })
80    }
81
82    /// Return the number of primary fragments.
83    pub const fn k(&self) -> usize {
84        self.k
85    }
86
87    /// Return the total number of primary plus parity fragments.
88    pub const fn n(&self) -> usize {
89        self.n
90    }
91
92    /// Generate parity fragments for a full primary block.
93    pub fn encode(&self, primary: &[Vec<u8>], block_size: usize) -> Result<Vec<Vec<u8>>, FecError> {
94        if primary.len() != self.k || primary.iter().any(|fragment| fragment.len() < block_size) {
95            return Err(FecError::InvalidParameters);
96        }
97
98        let mut fecs = vec![vec![0; block_size]; self.n - self.k];
99        for (fec_offset, fec) in fecs.iter_mut().enumerate() {
100            let fecnum = self.k + fec_offset;
101            let matrix_row = &self.enc_matrix[fecnum * self.k..(fecnum + 1) * self.k];
102            for (src_idx, src) in primary.iter().enumerate() {
103                addmul(fec, src, matrix_row[src_idx], block_size);
104            }
105        }
106        Ok(fecs)
107    }
108
109    /// Recover missing primary fragments in-place.
110    pub fn recover_primary(
111        &self,
112        fragments: &mut [Option<Vec<u8>>],
113        block_size: usize,
114    ) -> Result<usize, FecError> {
115        if fragments.len() != self.n {
116            return Err(FecError::InvalidParameters);
117        }
118        if (0..self.k).all(|idx| fragments[idx].is_some()) {
119            return Ok(0);
120        }
121
122        let mut indexes = Vec::with_capacity(self.k);
123        let mut inputs = Vec::with_capacity(self.k);
124        let mut parity_cursor = self.k;
125
126        for primary_idx in 0..self.k {
127            if let Some(fragment) = fragments[primary_idx].as_ref() {
128                if fragment.len() < block_size {
129                    return Err(FecError::InvalidParameters);
130                }
131                indexes.push(primary_idx);
132                inputs.push(fragment.clone());
133            } else {
134                while parity_cursor < self.n && fragments[parity_cursor].is_none() {
135                    parity_cursor += 1;
136                }
137                if parity_cursor >= self.n {
138                    return Err(FecError::NotEnoughFragments);
139                }
140                let fragment = fragments[parity_cursor]
141                    .as_ref()
142                    .ok_or(FecError::NotEnoughFragments)?;
143                if fragment.len() < block_size {
144                    return Err(FecError::InvalidParameters);
145                }
146                indexes.push(parity_cursor);
147                inputs.push(fragment.clone());
148                parity_cursor += 1;
149            }
150        }
151
152        self.validate_indexes(&indexes)?;
153        let dec_matrix = self.decode_matrix(&indexes)?;
154        let mut recovered = 0usize;
155
156        for row in 0..self.k {
157            if indexes[row] >= self.k {
158                let mut out = vec![0; block_size];
159                for col in 0..self.k {
160                    addmul(
161                        &mut out,
162                        &inputs[col],
163                        dec_matrix[row * self.k + col],
164                        block_size,
165                    );
166                }
167                fragments[row] = Some(out);
168                recovered += 1;
169            }
170        }
171
172        Ok(recovered)
173    }
174
175    fn validate_indexes(&self, indexes: &[usize]) -> Result<(), FecError> {
176        if indexes.len() != self.k {
177            return Err(FecError::NotEnoughFragments);
178        }
179        for (row, &idx) in indexes.iter().enumerate() {
180            if idx >= self.n {
181                return Err(FecError::InvalidFragmentIndex(idx));
182            }
183            if idx < self.k && idx != row {
184                return Err(FecError::OutputSlotMismatch);
185            }
186        }
187        Ok(())
188    }
189
190    fn decode_matrix(&self, indexes: &[usize]) -> Result<Vec<u8>, FecError> {
191        let mut matrix = vec![0; self.k * self.k];
192        for (row, &idx) in indexes.iter().enumerate() {
193            let row_start = row * self.k;
194            if idx < self.k {
195                matrix[row_start + row] = 1;
196            } else {
197                matrix[row_start..row_start + self.k]
198                    .copy_from_slice(&self.enc_matrix[idx * self.k..(idx + 1) * self.k]);
199            }
200        }
201        invert_mat(&mut matrix, self.k)?;
202        Ok(matrix)
203    }
204}
205
206#[derive(Clone)]
207struct GfTables {
208    gf_exp: [u8; 510],
209    inverse: [u8; 256],
210    gf_mul: Box<[[u8; 256]; 256]>,
211}
212
213fn tables() -> &'static GfTables {
214    GF_TABLES.get_or_init(GfTables::new)
215}
216
217impl GfTables {
218    fn new() -> Self {
219        let mut gf_exp = [0; 510];
220        let mut gf_log = [0; 256];
221        let mut inverse = [0; 256];
222
223        let mut mask = 1u8;
224        gf_exp[GF_BITS] = 0;
225        for i in 0..GF_BITS {
226            gf_exp[i] = mask;
227            gf_log[mask as usize] = i as u16;
228            if PRIMITIVE_POLY[i] == b'1' {
229                gf_exp[GF_BITS] ^= mask;
230            }
231            mask <<= 1;
232        }
233        gf_log[gf_exp[GF_BITS] as usize] = GF_BITS as u16;
234
235        mask = 1 << (GF_BITS - 1);
236        for i in (GF_BITS + 1)..GF_SIZE {
237            gf_exp[i] = if gf_exp[i - 1] >= mask {
238                gf_exp[GF_BITS] ^ ((gf_exp[i - 1] ^ mask) << 1)
239            } else {
240                gf_exp[i - 1] << 1
241            };
242            gf_log[gf_exp[i] as usize] = i as u16;
243        }
244        gf_log[0] = GF_SIZE as u16;
245        for i in 0..GF_SIZE {
246            gf_exp[i + GF_SIZE] = gf_exp[i];
247        }
248
249        inverse[1] = 1;
250        for i in 2..=GF_SIZE {
251            inverse[i] = gf_exp[GF_SIZE - gf_log[i] as usize];
252        }
253
254        let mut gf_mul = Box::new([[0; 256]; 256]);
255        for i in 1..256 {
256            for j in 1..256 {
257                gf_mul[i][j] = gf_exp[modnn(gf_log[i] as i32 + gf_log[j] as i32) as usize];
258            }
259        }
260
261        Self {
262            gf_exp,
263            inverse,
264            gf_mul,
265        }
266    }
267}
268
269fn modnn(mut x: i32) -> u8 {
270    while x >= GF_SIZE as i32 {
271        x -= GF_SIZE as i32;
272        x = (x >> GF_BITS) + (x & GF_SIZE as i32);
273    }
274    x as u8
275}
276
277fn gf_mul(x: u8, y: u8) -> u8 {
278    tables().gf_mul[x as usize][y as usize]
279}
280
281fn addmul(dst: &mut [u8], src: &[u8], coefficient: u8, len: usize) {
282    if coefficient == 0 {
283        return;
284    }
285    let mul = &tables().gf_mul[coefficient as usize];
286    for idx in 0..len {
287        dst[idx] ^= mul[src[idx] as usize];
288    }
289}
290
291fn matmul(a: &[u8], b: &[u8], c: &mut [u8], n: usize, k: usize, m: usize) {
292    for row in 0..n {
293        for col in 0..m {
294            let mut acc = 0;
295            for i in 0..k {
296                acc ^= gf_mul(a[row * k + i], b[i * m + col]);
297            }
298            c[row * m + col] = acc;
299        }
300    }
301}
302
303fn invert_mat(src: &mut [u8], k: usize) -> Result<(), FecError> {
304    let mut indxc = vec![0; k];
305    let mut indxr = vec![0; k];
306    let mut ipiv = vec![0; k];
307    let mut id_row = vec![0; k];
308
309    for col in 0..k {
310        let mut irow = None;
311        let mut icol = None;
312
313        if ipiv[col] != 1 && src[col * k + col] != 0 {
314            irow = Some(col);
315            icol = Some(col);
316        } else {
317            'search: for row in 0..k {
318                if ipiv[row] != 1 {
319                    for ix in 0..k {
320                        if ipiv[ix] == 0 && src[row * k + ix] != 0 {
321                            irow = Some(row);
322                            icol = Some(ix);
323                            break 'search;
324                        }
325                    }
326                }
327            }
328        }
329
330        let irow = irow.ok_or(FecError::SingularMatrix)?;
331        let icol = icol.ok_or(FecError::SingularMatrix)?;
332        ipiv[icol] += 1;
333
334        if irow != icol {
335            for ix in 0..k {
336                src.swap(irow * k + ix, icol * k + ix);
337            }
338        }
339        indxr[col] = irow;
340        indxc[col] = icol;
341
342        let pivot = src[icol * k + icol];
343        if pivot == 0 {
344            return Err(FecError::SingularMatrix);
345        }
346        if pivot != 1 {
347            let inv = tables().inverse[pivot as usize];
348            src[icol * k + icol] = 1;
349            for ix in 0..k {
350                src[icol * k + ix] = gf_mul(inv, src[icol * k + ix]);
351            }
352        }
353
354        id_row[icol] = 1;
355        if src[icol * k..(icol + 1) * k] != id_row[..] {
356            let pivot_row = src[icol * k..(icol + 1) * k].to_vec();
357            for ix in 0..k {
358                if ix != icol {
359                    let coefficient = src[ix * k + icol];
360                    src[ix * k + icol] = 0;
361                    addmul(&mut src[ix * k..(ix + 1) * k], &pivot_row, coefficient, k);
362                }
363            }
364        }
365        id_row[icol] = 0;
366    }
367
368    for col in (0..k).rev() {
369        if indxr[col] != indxc[col] {
370            for row in 0..k {
371                src.swap(row * k + indxr[col], row * k + indxc[col]);
372            }
373        }
374    }
375    Ok(())
376}
377
378fn invert_vdm(src: &mut [u8], k: usize) -> Result<(), FecError> {
379    if k == 1 {
380        return Ok(());
381    }
382
383    let mut c = vec![0; k];
384    let mut b = vec![0; k];
385    let mut p = vec![0; k];
386
387    for i in 0..k {
388        p[i] = src[i * k + 1];
389    }
390
391    c[k - 1] = p[0];
392    for (i, p_i) in p.iter().copied().enumerate().take(k).skip(1) {
393        let start = k - 1 - (i - 1);
394        for j in start..(k - 1) {
395            c[j] ^= gf_mul(p_i, c[j + 1]);
396        }
397        c[k - 1] ^= p_i;
398    }
399
400    for row in 0..k {
401        let xx = p[row];
402        let mut t = 1;
403        b[k - 1] = 1;
404        for i in (1..k).rev() {
405            b[i - 1] = c[i] ^ gf_mul(xx, b[i]);
406            t = gf_mul(xx, t) ^ b[i - 1];
407        }
408        if t == 0 {
409            return Err(FecError::SingularMatrix);
410        }
411        let inv = tables().inverse[t as usize];
412        for col in 0..k {
413            src[col * k + row] = gf_mul(inv, b[col]);
414        }
415    }
416
417    Ok(())
418}
419
420#[cfg(test)]
421mod tests {
422    use super::*;
423
424    #[test]
425    fn recovers_missing_primary_fragment_from_parity() {
426        let fec = FecCode::new(3, 5).unwrap();
427        let primary = vec![b"aaaa".to_vec(), b"bbbb".to_vec(), b"cccc".to_vec()];
428        let parity = fec.encode(&primary, 4).unwrap();
429        let mut fragments = vec![
430            Some(primary[0].clone()),
431            None,
432            Some(primary[2].clone()),
433            Some(parity[0].clone()),
434            None,
435        ];
436
437        let recovered = fec.recover_primary(&mut fragments, 4).unwrap();
438        assert_eq!(recovered, 1);
439        assert_eq!(fragments[1].as_deref(), Some(&primary[1][..]));
440    }
441}
openipc_core/fec.rs

openipc_core/
fec.rs
