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