1use crate::algebraic::{berlekamp_massey, chien_search, eval_at_power};
94use crate::block::{BlockCode, container_traits, expect_len};
95use comnum_core::signal::hard_decisions;
96use comnum_core::{Error, Result, ensure};
97use comnum_math::bits::bits_to_indices;
98use comnum_math::gf2m::{Field, primitive_polynomial};
99use comnum_math::{index_to_bits, ln_binomial_pmf};
100
101#[derive(Clone, Debug, PartialEq, Eq)]
103pub struct ReedSolomon {
104 field: Field,
105 t: usize,
106 n: usize,
107 first_root: usize,
108 root_step: usize,
109 generator: Vec<u16>,
111}
112
113impl ReedSolomon {
114 pub fn new(m: u32, t: usize) -> Result<ReedSolomon> {
121 let n = (1usize << m.min(16)) - 1;
122 ReedSolomon::shortened(m, t, n)
123 }
124
125 pub fn shortened(m: u32, t: usize, n: usize) -> Result<ReedSolomon> {
131 let poly = primitive_polynomial(m)
132 .ok_or_else(|| Error::Config(format!("RS: the symbol size must be 2 to 16 bits, got {m}")))?;
133 ReedSolomon::with_field(m, poly, 1, 1, t, n)
134 }
135
136 pub fn with_field(
146 m: u32,
147 poly: u32,
148 first_root: usize,
149 root_step: usize,
150 t: usize,
151 n: usize,
152 ) -> Result<ReedSolomon> {
153 ensure((2..=16).contains(&m), || format!("RS: the symbol size must be 2 to 16 bits, got {m}"))?;
154 let field = Field::new(m, poly).map_err(|e| e.context("RS"))?;
155 let order = field.order();
156 ensure(t >= 1, || "RS: t must be at least 1".into())?;
157 ensure(root_step >= 1 && comnum_math::gcd(root_step, order) == 1, || {
158 format!("RS: the root step {root_step} must be coprime with {order}")
159 })?;
160 ensure(2 * t < n && n <= order, || {
161 format!("RS: the length {n} must exceed the {} parity symbols and not exceed {order}", 2 * t)
162 })?;
163 let roots = (0..2 * t).map(|i| (field.pow(root_step * ((first_root + i) % order)), 1));
165 let generator = linear_product(&field, roots).iter().map(|&a| a as u16).collect();
166 Ok(ReedSolomon { field, t, n, first_root, root_step, generator })
167 }
168
169 pub fn n_symbols(&self) -> usize {
171 self.n
172 }
173
174 pub fn k_symbols(&self) -> usize {
176 self.n - 2 * self.t
177 }
178
179 pub fn t(&self) -> usize {
181 self.t
182 }
183
184 pub fn m(&self) -> u32 {
186 self.field.m()
187 }
188
189 pub fn first_root(&self) -> usize {
191 self.first_root
192 }
193
194 pub fn root_step(&self) -> usize {
196 self.root_step
197 }
198
199 pub fn field(&self) -> &Field {
201 &self.field
202 }
203
204 pub fn generator(&self) -> &[u16] {
206 &self.generator
207 }
208
209 fn check_symbols(&self, symbols: &[u16]) -> Result<()> {
211 let max = self.field.order() as u16;
212 match symbols.iter().find(|&&s| s > max) {
213 Some(s) => Err(Error::Config(format!("RS: the symbol {s} exceeds 2^{} - 1 = {max}", self.m()))),
214 None => Ok(()),
215 }
216 }
217
218 pub fn encode_symbols(&self, msg: &[u16]) -> Result<Vec<u16>> {
225 expect_len(msg.len(), self.k_symbols(), "symbols")?;
226 self.check_symbols(msg)?;
227 let f = &self.field;
228 let r = 2 * self.t;
229 let mut reg = vec![0u32; r];
231 for &u in msg {
232 let fb = u32::from(u) ^ reg[r - 1];
233 for j in (1..r).rev() {
234 reg[j] = reg[j - 1] ^ f.mul(fb, u32::from(self.generator[j]));
235 }
236 reg[0] = f.mul(fb, u32::from(self.generator[0]));
237 }
238 let mut c = msg.to_vec();
239 c.extend(reg.iter().rev().map(|&a| a as u16));
240 Ok(c)
241 }
242
243 pub fn decode_symbols(&self, r: &[u16]) -> Result<(Vec<u16>, Option<usize>)> {
251 self.decode_erasures(r, &[])
252 }
253
254 pub fn decode_erasures(&self, r: &[u16], erasures: &[usize]) -> Result<(Vec<u16>, Option<usize>)> {
265 expect_len(r.len(), self.n, "symbols")?;
266 self.check_symbols(r)?;
267 let erased = self.erasure_mask(erasures)?;
268 let k = self.k_symbols();
269 let mut word: Vec<u32> = r.iter().map(|&v| u32::from(v)).collect();
270 for &p in erasures {
271 word[p] = 0;
272 }
273 let syn = self.syndromes(&word);
274 if syn.iter().all(|&v| v == 0) && erasures.is_empty() {
275 return Ok((r[..k].to_vec(), Some(0)));
276 }
277 let corrected = self.locate(&syn, erasures).and_then(|(psi, roots)| {
278 self.correct(&mut word, &syn, &psi, &roots, &erased)
279 .map(|()| roots.iter().filter(|&&i| word[i] != u32::from(r[i])).count())
280 });
281 Ok(match corrected {
282 Some(c) => (word[..k].iter().map(|&v| v as u16).collect(), Some(c)),
283 None => (r[..k].to_vec(), None),
284 })
285 }
286
287 fn locator_exp(&self, i: usize) -> usize {
290 self.root_step * (self.n - 1 - i) % self.field.order()
291 }
292
293 fn erasure_mask(&self, erasures: &[usize]) -> Result<Vec<bool>> {
295 let (n, two_t, s) = (self.n, 2 * self.t, erasures.len());
296 ensure(s <= two_t, || format!("RS: {s} erasures exceed the {two_t} the code can fill"))?;
297 let mut erased = vec![false; n];
298 for &p in erasures {
299 ensure(p < n, || format!("RS: the erasure position {p} is outside the word of {n} symbols"))?;
300 ensure(!erased[p], || format!("RS: the erasure position {p} is repeated"))?;
301 erased[p] = true;
302 }
303 Ok(erased)
304 }
305
306 fn syndromes(&self, word: &[u32]) -> Vec<u32> {
308 let f = &self.field;
309 let order = f.order();
310 let b = self.first_root;
311 let mut syn = vec![0u32; 2 * self.t];
312 for (i, &v) in word.iter().enumerate() {
313 if v != 0 {
314 let e = self.locator_exp(i);
315 for (j, sj) in syn.iter_mut().enumerate() {
316 *sj ^= f.mul(v, f.pow(e * ((b + j) % order)));
317 }
318 }
319 }
320 syn
321 }
322
323 fn locate(&self, syn: &[u32], erasures: &[usize]) -> Option<(Vec<u32>, Vec<usize>)> {
328 let f = &self.field;
329 let (two_t, s) = (2 * self.t, erasures.len());
330 let gamma = linear_product(f, erasures.iter().map(|&p| (1, f.pow(self.locator_exp(p)))));
331 let modified = poly_mul_mod(f, syn, &gamma, two_t);
332 let (lambda, l) = berlekamp_massey(f, &modified[s..]);
335 if 2 * l > two_t - s || lambda.len() - 1 != l {
336 return None;
337 }
338 let psi = poly_mul_mod(f, &lambda, &gamma, lambda.len() + gamma.len() - 1);
340 let roots = chien_search(f, &psi, self.n, |i| self.locator_exp(i));
341 (roots.len() == psi.len() - 1).then_some((psi, roots))
342 }
343
344 fn correct(&self, word: &mut [u32], syn: &[u32], psi: &[u32], roots: &[usize], erased: &[bool]) -> Option<()> {
348 let f = &self.field;
349 let order = f.order();
350 let omega = poly_mul_mod(f, syn, psi, 2 * self.t);
351 let exp_1mb = (order + 1 - self.first_root % order) % order;
352 for &i in roots {
353 let ze = self.locator_exp(i);
354 let inv = (order - ze) % order;
355 let num = eval_at_power(f, &omega, inv);
356 let den = psi
357 .iter()
358 .enumerate()
359 .skip(1)
360 .step_by(2)
361 .fold(0u32, |acc, (d, &a)| acc ^ f.mul(a, f.pow((d - 1) * inv)));
362 if den == 0 {
363 return None;
364 }
365 let y = f.mul(f.pow(ze * exp_1mb), f.div(num, den));
366 if y == 0 && !erased[i] {
367 return None;
368 }
369 word[i] ^= y;
370 }
371 Some(())
372 }
373
374 pub fn encode(&self, bits: &[u8]) -> Result<Vec<u8>> {
380 let m = self.m();
381 expect_len(bits.len(), self.k_symbols() * m as usize, "bits")?;
382 Ok(unpack_symbols(&self.encode_symbols(&pack_symbols(bits, m)?)?, m))
383 }
384
385 pub fn decode(&self, bits: &[u8]) -> Result<(Vec<u8>, Option<usize>)> {
392 let m = self.m();
393 expect_len(bits.len(), self.n * m as usize, "bits")?;
394 let (msg, corrected) = self.decode_symbols(&pack_symbols(bits, m)?)?;
395 Ok((unpack_symbols(&msg, m), corrected))
396 }
397}
398
399impl BlockCode for ReedSolomon {
400 fn k(&self) -> usize {
401 self.k_symbols() * self.m() as usize
402 }
403
404 fn n(&self) -> usize {
405 self.n * self.m() as usize
406 }
407
408 fn encode_block(&self, info: &[u8]) -> Result<Vec<u8>> {
409 self.encode(info)
410 }
411
412 fn decode_block(&self, llr: &[f64]) -> Result<Vec<u8>> {
415 expect_len(llr.len(), BlockCode::n(self), "LLRs")?;
416 Ok(self.decode(&hard_decisions(llr))?.0)
417 }
418}
419
420container_traits!(ReedSolomon);
421
422pub fn pack_symbols(bits: &[u8], m: u32) -> Result<Vec<u16>> {
428 ensure((1..=16).contains(&m), || format!("the symbol size must be 1 to 16 bits, got {m}"))?;
429 Ok(bits_to_indices(bits, m as usize)?.into_iter().map(|i| i as u16).collect())
430}
431
432pub fn unpack_symbols(symbols: &[u16], m: u32) -> Vec<u8> {
435 let m = m.clamp(1, 16) as usize;
436 symbols.iter().flat_map(|&s| index_to_bits(usize::from(s), m)).collect()
437}
438
439fn linear_product(f: &Field, factors: impl Iterator<Item = (u32, u32)>) -> Vec<u32> {
442 let mut poly = vec![1u32];
443 for (c0, c1) in factors {
444 let mut next = vec![0u32; poly.len() + 1];
445 for (d, &a) in poly.iter().enumerate() {
446 next[d] ^= f.mul(a, c0);
447 next[d + 1] ^= f.mul(a, c1);
448 }
449 poly = next;
450 }
451 poly
452}
453
454fn poly_mul_mod(f: &Field, a: &[u32], b: &[u32], len: usize) -> Vec<u32> {
457 let mut out = vec![0u32; len];
458 for (i, &x) in a.iter().enumerate() {
459 if x == 0 {
460 continue;
461 }
462 for (j, &y) in b.iter().enumerate().take(len.saturating_sub(i)) {
463 out[i + j] ^= f.mul(x, y);
464 }
465 }
466 out
467}
468
469fn ensure_probability(p_s: f64) -> Result<()> {
471 ensure((0.0..=1.0).contains(&p_s), || format!("the symbol error probability must be in [0, 1], got {p_s}"))
472}
473
474pub fn word_error_probability(n: usize, t: usize, p_s: f64) -> Result<f64> {
487 ensure_probability(p_s)?;
488 if t >= n || p_s == 0.0 {
489 return Ok(0.0);
490 }
491 if p_s == 1.0 {
492 return Ok(1.0);
493 }
494 Ok(ln_binomial_pmf(n, p_s).iter().skip(t + 1).map(|l| l.exp()).sum())
495}
496
497pub fn symbol_error_probability(n: usize, t: usize, p_s: f64) -> Result<f64> {
510 ensure_probability(p_s)?;
511 if t >= n || p_s == 0.0 {
512 return Ok(0.0);
513 }
514 if p_s == 1.0 {
515 return Ok(1.0);
516 }
517 Ok(ln_binomial_pmf(n, p_s).iter().enumerate().skip(t + 1).map(|(i, l)| i as f64 * l.exp()).sum::<f64>() / n as f64)
518}
519
520pub fn bit_error_probability(n: usize, t: usize, m: u32, p_s: f64) -> Result<f64> {
531 ensure(m >= 1, || "a symbol needs at least one bit".into())?;
532 let q = 2f64.powi(m as i32);
533 Ok(q / 2.0 / (q - 1.0) * symbol_error_probability(n, t, p_s)?)
534}
535
536#[cfg(test)]
537mod tests {
538 use super::*;
539 use comnum_core::rng::Rng;
540
541 pub(super) fn corrupt(c: &mut [u16], errors: usize, m: u32, avoid: &[usize], rng: &mut Rng) -> Vec<usize> {
543 let mut positions = Vec::with_capacity(errors);
544 while positions.len() < errors {
545 let p = rng.below(c.len());
546 if !positions.contains(&p) && !avoid.contains(&p) {
547 positions.push(p);
548 let e = 1 + rng.below((1 << m) - 1) as u16;
549 c[p] ^= e;
550 }
551 }
552 positions
553 }
554
555 #[test]
556 fn errors_and_erasures_are_corrected_up_to_2e_plus_s_equal_2t() {
557 let code = ReedSolomon::shortened(8, 4, 60).unwrap();
558 assert_eq!(code.k_symbols(), 52);
559 let mut rng = Rng::new(11);
560 for s in 0..=8 {
561 for e in 0..=(8 - s) / 2 {
562 let msg: Vec<u16> = (0..52).map(|_| rng.below(256) as u16).collect();
563 let mut c = code.encode_symbols(&msg).unwrap();
564 let erased = corrupt(&mut c, s, 8, &[], &mut rng);
565 corrupt(&mut c, e, 8, &erased, &mut rng);
566 let (m, n) = code.decode_erasures(&c, &erased).unwrap();
567 assert_eq!(m, msg, "s = {s}, e = {e}");
568 assert_eq!(n, Some(s + e), "s = {s}, e = {e}");
569 }
570 }
571 }
572
573 #[test]
574 fn invalid_inputs_are_errors_not_panics() {
575 assert!(ReedSolomon::new(1, 1).is_err());
576 assert!(ReedSolomon::new(17, 1).is_err());
577 assert!(ReedSolomon::new(3, 0).is_err());
578 assert!(ReedSolomon::shortened(3, 2, 4).is_err());
579 assert!(ReedSolomon::shortened(3, 1, 8).is_err());
580 assert!(ReedSolomon::with_field(4, 0b1_1111, 1, 1, 1, 15).is_err());
581 assert!(ReedSolomon::with_field(4, 0b1_0011, 1, 3, 1, 15).is_err());
582 let code = ReedSolomon::new(4, 2).unwrap();
583 assert!(code.encode_symbols(&[0; 10]).is_err());
584 assert!(code.encode_symbols(&[16; 11]).is_err());
585 assert!(code.decode_symbols(&[0; 14]).is_err());
586 assert!(code.decode_erasures(&[0; 15], &[15]).is_err());
587 assert!(code.decode_erasures(&[0; 15], &[1, 1]).is_err());
588 assert!(code.decode_erasures(&[0; 15], &[0, 1, 2, 3, 4]).is_err());
589 assert!(code.encode(&[0; 43]).is_err());
590 assert!(pack_symbols(&[1, 0, 1], 2).is_err());
591 assert_eq!(unpack_symbols(&[0b1011], 3), vec![0, 1, 1]);
592 assert_eq!(pack_symbols(&[1, 0, 1, 1, 1, 0], 3).unwrap(), vec![5, 6]);
593 }
594
595 #[test]
596 fn the_block_code_interface_counts_bits() {
597 let code = ReedSolomon::shortened(6, 3, 30).unwrap();
598 assert_eq!((BlockCode::k(&code), BlockCode::n(&code)), (24 * 6, 30 * 6));
599 let mut rng = Rng::new(3);
600 let info: Vec<u8> = (0..144).map(|_| rng.bit()).collect();
601 let c = code.encode_block(&info).unwrap();
602 let mut llr: Vec<f64> = c.iter().map(|&b| if b == 1 { -2.0 } else { 2.0 }).collect();
603 for i in [0, 1, 42, 174, 175, 176, 177, 178, 179] {
605 llr[i] = -llr[i];
606 }
607 assert_eq!(code.decode_block(&llr).unwrap(), info);
608 }
609
610 #[test]
611 fn the_closed_forms_have_the_right_limits() {
612 assert_eq!(word_error_probability(255, 16, 0.0).unwrap(), 0.0);
613 assert_eq!(word_error_probability(255, 16, 1.0).unwrap(), 1.0);
614 let p = 0.01;
616 assert!((word_error_probability(7, 0, p).unwrap() - (1.0 - (1.0 - p).powi(7))).abs() < 1e-15);
617 let out = symbol_error_probability(255, 16, 0.05).unwrap();
619 assert!(out > 0.0 && out < 0.05);
620 assert!((bit_error_probability(255, 16, 8, 0.05).unwrap() / out - 128.0 / 255.0).abs() < 1e-12);
621 for p in [-0.1, 1.5, f64::NAN] {
623 assert!(word_error_probability(255, 16, p).is_err() && symbol_error_probability(255, 16, p).is_err());
624 assert!(bit_error_probability(255, 16, 8, p).is_err());
625 }
626 assert_eq!(word_error_probability(7, 7, 0.5).unwrap(), 0.0);
627 }
628 #[test]
629 fn the_block_code_counts_bits_and_the_symbol_methods_count_symbols() {
630 let code = ReedSolomon::new(8, 16).unwrap();
633 assert_eq!((code.n_symbols(), code.k_symbols()), (255, 223));
634 assert_eq!((code.n(), code.k()), (255 * 8, 223 * 8));
635 assert_eq!((BlockCode::n(&code), BlockCode::k(&code)), (255 * 8, 223 * 8));
636 let bits = vec![1u8; code.k()];
637 assert_eq!(code.encode_block(&bits).unwrap().len(), code.n());
638 }
639}
640
641#[cfg(test)]
642mod literature {
643 use super::*;
644 use comnum_core::rng::Rng;
645
646 #[test]
651 fn lin_costello_sec_7_3_generator_of_the_15_9_code() {
652 let code = ReedSolomon::with_field(4, 0b1_0011, 1, 1, 3, 15).unwrap();
653 let f = code.field();
654 let expected: Vec<u16> = [6, 9, 6, 4, 14, 10].iter().map(|&e| f.pow(e) as u16).chain([1]).collect();
655 assert_eq!(code.generator(), expected);
656 }
657
658 #[test]
663 fn lin_costello_sec_7_4_three_errors_of_the_15_9_code() {
664 let code = ReedSolomon::with_field(4, 0b1_0011, 1, 1, 3, 15).unwrap();
665 let f = code.field();
666 let mut r = vec![0u16; 15];
667 r[14 - 3] = f.pow(7) as u16;
669 r[14 - 6] = f.pow(3) as u16;
670 r[14 - 12] = f.pow(4) as u16;
671 let syndromes: Vec<u32> = (1..=6)
672 .map(|j| r.iter().enumerate().fold(0u32, |acc, (i, &v)| acc ^ f.mul(u32::from(v), f.pow(j * (14 - i)))))
673 .collect();
674 assert_eq!(syndromes, [f.pow(12), 1, f.pow(14), f.pow(10), 0, f.pow(12)]);
675 assert_eq!(code.decode_symbols(&r).unwrap(), (vec![0; 9], Some(3)));
676 }
677
678 #[test]
686 fn etsi_en_300_744_sec_4_3_2_generator_of_rs_204_188() {
687 let code = ReedSolomon::with_field(8, 0x11D, 0, 1, 8, 204).unwrap();
688 assert_eq!((code.n_symbols(), code.k_symbols()), (204, 188));
689 let f = code.field();
690 let exponents = [120, 104, 107, 109, 102, 161, 76, 3, 91, 191, 147, 169, 182, 194, 225, 120];
691 let mut expected: Vec<u16> = exponents.iter().rev().map(|&e| f.pow(e) as u16).collect();
692 expected.push(1);
693 assert_eq!(code.generator(), expected);
694 for j in 0..16 {
695 let root = f.pow(j);
696 let v = code
697 .generator()
698 .iter()
699 .enumerate()
700 .fold(0u32, |acc, (d, &a)| acc ^ f.mul(u32::from(a), f.pow(d * f.log(root))));
701 assert_eq!(v, 0, "λ^{j} is a root");
702 }
703 let full = ReedSolomon::with_field(8, 0x11D, 0, 1, 8, 255).unwrap();
705 let msg: Vec<u16> = (0..188).map(|i| (i * 13 % 256) as u16).collect();
706 let mut padded = vec![0u16; 51];
707 padded.extend(&msg);
708 assert_eq!(full.encode_symbols(&padded).unwrap()[51..], code.encode_symbols(&msg).unwrap()[..]);
709 }
710
711 #[test]
717 fn ccsds_131_0_b_sec_4_3_roots_of_rs_255_223() {
718 let code = ReedSolomon::with_field(8, 0x187, 112, 11, 16, 255).unwrap();
719 let f = code.field();
720 let g = code.generator();
721 assert_eq!(g.len(), 33);
722 for j in 112..=143 {
723 let root_exp = 11 * j;
724 let v = g.iter().enumerate().fold(0u32, |acc, (d, &a)| acc ^ f.mul(u32::from(a), f.pow(d * root_exp)));
725 assert_eq!(v, 0, "α^(11·{j}) is a root");
726 }
727 for i in 0..=32 {
728 assert_eq!(g[i], g[32 - i], "g is self-reciprocal");
729 }
730 let mut rng = Rng::new(5);
731 let msg: Vec<u16> = (0..223).map(|_| rng.below(256) as u16).collect();
732 let mut c = code.encode_symbols(&msg).unwrap();
733 super::tests::corrupt(&mut c, 16, 8, &[], &mut rng);
734 assert_eq!(code.decode_symbols(&c).unwrap(), (msg, Some(16)));
735 }
736
737 #[test]
741 fn rs_255_223_corrects_16_errors_and_not_17() {
742 let code = ReedSolomon::new(8, 16).unwrap();
743 let mut rng = Rng::new(255);
744 let (mut failures, mut miscorrections) = (0, 0);
745 for trial in 0..300 {
746 let msg: Vec<u16> = (0..223).map(|_| rng.below(256) as u16).collect();
747 let mut c = code.encode_symbols(&msg).unwrap();
748 super::tests::corrupt(&mut c, 16, 8, &[], &mut rng);
749 assert_eq!(code.decode_symbols(&c).unwrap(), (msg.clone(), Some(16)), "trial {trial}");
750 let mut c = code.encode_symbols(&msg).unwrap();
751 super::tests::corrupt(&mut c, 17, 8, &[], &mut rng);
752 let (decoded, status) = code.decode_symbols(&c).unwrap();
753 if status.is_none() {
754 failures += 1;
755 } else {
756 assert_ne!(decoded, msg, "trial {trial}: 17 errors cannot be corrected");
757 miscorrections += 1;
758 }
759 }
760 assert_eq!((failures, miscorrections), (300, 0));
763 }
764
765 #[test]
768 fn word_error_rate_matches_the_closed_form_at_p_s_0_01() {
769 let code = ReedSolomon::new(8, 4).unwrap();
770 let p_s = 0.01;
771 let trials = 20_000;
772 let mut rng = Rng::new(2024);
773 let mut lost = 0usize;
774 for _ in 0..trials {
775 let msg: Vec<u16> = (0..247).map(|_| rng.below(256) as u16).collect();
776 let mut c = code.encode_symbols(&msg).unwrap();
777 for v in &mut c {
778 if rng.uniform() < p_s {
779 *v ^= 1 + rng.below(255) as u16;
780 }
781 }
782 let (decoded, status) = code.decode_symbols(&c).unwrap();
783 if status.is_none() || decoded != msg {
784 lost += 1;
785 }
786 }
787 let p_w = word_error_probability(255, 4, p_s).unwrap();
788 let measured = lost as f64 / trials as f64;
789 let tol = 4.0 * (p_w * (1.0 - p_w) / trials as f64).sqrt();
793 assert!((measured - p_w).abs() < tol, "measured {measured}, closed form {p_w}, tolerance {tol}");
794 }
795}