use crate::algebraic::{berlekamp_massey, chien_search, eval_at_power};
use crate::block::{BlockCode, container_traits, expect_len};
use comnum_core::signal::hard_decisions;
use comnum_core::{Error, Result, ensure};
use comnum_math::bits::bits_to_indices;
use comnum_math::gf2m::{Field, primitive_polynomial};
use comnum_math::{index_to_bits, ln_binomial_pmf};
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ReedSolomon {
field: Field,
t: usize,
n: usize,
first_root: usize,
root_step: usize,
generator: Vec<u16>,
}
impl ReedSolomon {
pub fn new(m: u32, t: usize) -> Result<ReedSolomon> {
let n = (1usize << m.min(16)) - 1;
ReedSolomon::shortened(m, t, n)
}
pub fn shortened(m: u32, t: usize, n: usize) -> Result<ReedSolomon> {
let poly = primitive_polynomial(m)
.ok_or_else(|| Error::Config(format!("RS: the symbol size must be 2 to 16 bits, got {m}")))?;
ReedSolomon::with_field(m, poly, 1, 1, t, n)
}
pub fn with_field(
m: u32,
poly: u32,
first_root: usize,
root_step: usize,
t: usize,
n: usize,
) -> Result<ReedSolomon> {
ensure((2..=16).contains(&m), || format!("RS: the symbol size must be 2 to 16 bits, got {m}"))?;
let field = Field::new(m, poly).map_err(|e| e.context("RS"))?;
let order = field.order();
ensure(t >= 1, || "RS: t must be at least 1".into())?;
ensure(root_step >= 1 && comnum_math::gcd(root_step, order) == 1, || {
format!("RS: the root step {root_step} must be coprime with {order}")
})?;
ensure(2 * t < n && n <= order, || {
format!("RS: the length {n} must exceed the {} parity symbols and not exceed {order}", 2 * t)
})?;
let roots = (0..2 * t).map(|i| (field.pow(root_step * ((first_root + i) % order)), 1));
let generator = linear_product(&field, roots).iter().map(|&a| a as u16).collect();
Ok(ReedSolomon { field, t, n, first_root, root_step, generator })
}
pub fn n_symbols(&self) -> usize {
self.n
}
pub fn k_symbols(&self) -> usize {
self.n - 2 * self.t
}
pub fn t(&self) -> usize {
self.t
}
pub fn m(&self) -> u32 {
self.field.m()
}
pub fn first_root(&self) -> usize {
self.first_root
}
pub fn root_step(&self) -> usize {
self.root_step
}
pub fn field(&self) -> &Field {
&self.field
}
pub fn generator(&self) -> &[u16] {
&self.generator
}
fn check_symbols(&self, symbols: &[u16]) -> Result<()> {
let max = self.field.order() as u16;
match symbols.iter().find(|&&s| s > max) {
Some(s) => Err(Error::Config(format!("RS: the symbol {s} exceeds 2^{} - 1 = {max}", self.m()))),
None => Ok(()),
}
}
pub fn encode_symbols(&self, msg: &[u16]) -> Result<Vec<u16>> {
expect_len(msg.len(), self.k_symbols(), "symbols")?;
self.check_symbols(msg)?;
let f = &self.field;
let r = 2 * self.t;
let mut reg = vec![0u32; r];
for &u in msg {
let fb = u32::from(u) ^ reg[r - 1];
for j in (1..r).rev() {
reg[j] = reg[j - 1] ^ f.mul(fb, u32::from(self.generator[j]));
}
reg[0] = f.mul(fb, u32::from(self.generator[0]));
}
let mut c = msg.to_vec();
c.extend(reg.iter().rev().map(|&a| a as u16));
Ok(c)
}
pub fn decode_symbols(&self, r: &[u16]) -> Result<(Vec<u16>, Option<usize>)> {
self.decode_erasures(r, &[])
}
pub fn decode_erasures(&self, r: &[u16], erasures: &[usize]) -> Result<(Vec<u16>, Option<usize>)> {
expect_len(r.len(), self.n, "symbols")?;
self.check_symbols(r)?;
let erased = self.erasure_mask(erasures)?;
let k = self.k_symbols();
let mut word: Vec<u32> = r.iter().map(|&v| u32::from(v)).collect();
for &p in erasures {
word[p] = 0;
}
let syn = self.syndromes(&word);
if syn.iter().all(|&v| v == 0) && erasures.is_empty() {
return Ok((r[..k].to_vec(), Some(0)));
}
let corrected = self.locate(&syn, erasures).and_then(|(psi, roots)| {
self.correct(&mut word, &syn, &psi, &roots, &erased)
.map(|()| roots.iter().filter(|&&i| word[i] != u32::from(r[i])).count())
});
Ok(match corrected {
Some(c) => (word[..k].iter().map(|&v| v as u16).collect(), Some(c)),
None => (r[..k].to_vec(), None),
})
}
fn locator_exp(&self, i: usize) -> usize {
self.root_step * (self.n - 1 - i) % self.field.order()
}
fn erasure_mask(&self, erasures: &[usize]) -> Result<Vec<bool>> {
let (n, two_t, s) = (self.n, 2 * self.t, erasures.len());
ensure(s <= two_t, || format!("RS: {s} erasures exceed the {two_t} the code can fill"))?;
let mut erased = vec![false; n];
for &p in erasures {
ensure(p < n, || format!("RS: the erasure position {p} is outside the word of {n} symbols"))?;
ensure(!erased[p], || format!("RS: the erasure position {p} is repeated"))?;
erased[p] = true;
}
Ok(erased)
}
fn syndromes(&self, word: &[u32]) -> Vec<u32> {
let f = &self.field;
let order = f.order();
let b = self.first_root;
let mut syn = vec![0u32; 2 * self.t];
for (i, &v) in word.iter().enumerate() {
if v != 0 {
let e = self.locator_exp(i);
for (j, sj) in syn.iter_mut().enumerate() {
*sj ^= f.mul(v, f.pow(e * ((b + j) % order)));
}
}
}
syn
}
fn locate(&self, syn: &[u32], erasures: &[usize]) -> Option<(Vec<u32>, Vec<usize>)> {
let f = &self.field;
let (two_t, s) = (2 * self.t, erasures.len());
let gamma = linear_product(f, erasures.iter().map(|&p| (1, f.pow(self.locator_exp(p)))));
let modified = poly_mul_mod(f, syn, &gamma, two_t);
let (lambda, l) = berlekamp_massey(f, &modified[s..]);
if 2 * l > two_t - s || lambda.len() - 1 != l {
return None;
}
let psi = poly_mul_mod(f, &lambda, &gamma, lambda.len() + gamma.len() - 1);
let roots = chien_search(f, &psi, self.n, |i| self.locator_exp(i));
(roots.len() == psi.len() - 1).then_some((psi, roots))
}
fn correct(&self, word: &mut [u32], syn: &[u32], psi: &[u32], roots: &[usize], erased: &[bool]) -> Option<()> {
let f = &self.field;
let order = f.order();
let omega = poly_mul_mod(f, syn, psi, 2 * self.t);
let exp_1mb = (order + 1 - self.first_root % order) % order;
for &i in roots {
let ze = self.locator_exp(i);
let inv = (order - ze) % order;
let num = eval_at_power(f, &omega, inv);
let den = psi
.iter()
.enumerate()
.skip(1)
.step_by(2)
.fold(0u32, |acc, (d, &a)| acc ^ f.mul(a, f.pow((d - 1) * inv)));
if den == 0 {
return None;
}
let y = f.mul(f.pow(ze * exp_1mb), f.div(num, den));
if y == 0 && !erased[i] {
return None;
}
word[i] ^= y;
}
Some(())
}
pub fn encode(&self, bits: &[u8]) -> Result<Vec<u8>> {
let m = self.m();
expect_len(bits.len(), self.k_symbols() * m as usize, "bits")?;
Ok(unpack_symbols(&self.encode_symbols(&pack_symbols(bits, m)?)?, m))
}
pub fn decode(&self, bits: &[u8]) -> Result<(Vec<u8>, Option<usize>)> {
let m = self.m();
expect_len(bits.len(), self.n * m as usize, "bits")?;
let (msg, corrected) = self.decode_symbols(&pack_symbols(bits, m)?)?;
Ok((unpack_symbols(&msg, m), corrected))
}
}
impl BlockCode for ReedSolomon {
fn k(&self) -> usize {
self.k_symbols() * self.m() as usize
}
fn n(&self) -> usize {
self.n * self.m() as usize
}
fn encode_block(&self, info: &[u8]) -> Result<Vec<u8>> {
self.encode(info)
}
fn decode_block(&self, llr: &[f64]) -> Result<Vec<u8>> {
expect_len(llr.len(), BlockCode::n(self), "LLRs")?;
Ok(self.decode(&hard_decisions(llr))?.0)
}
}
container_traits!(ReedSolomon);
pub fn pack_symbols(bits: &[u8], m: u32) -> Result<Vec<u16>> {
ensure((1..=16).contains(&m), || format!("the symbol size must be 1 to 16 bits, got {m}"))?;
Ok(bits_to_indices(bits, m as usize)?.into_iter().map(|i| i as u16).collect())
}
pub fn unpack_symbols(symbols: &[u16], m: u32) -> Vec<u8> {
let m = m.clamp(1, 16) as usize;
symbols.iter().flat_map(|&s| index_to_bits(usize::from(s), m)).collect()
}
fn linear_product(f: &Field, factors: impl Iterator<Item = (u32, u32)>) -> Vec<u32> {
let mut poly = vec![1u32];
for (c0, c1) in factors {
let mut next = vec![0u32; poly.len() + 1];
for (d, &a) in poly.iter().enumerate() {
next[d] ^= f.mul(a, c0);
next[d + 1] ^= f.mul(a, c1);
}
poly = next;
}
poly
}
fn poly_mul_mod(f: &Field, a: &[u32], b: &[u32], len: usize) -> Vec<u32> {
let mut out = vec![0u32; len];
for (i, &x) in a.iter().enumerate() {
if x == 0 {
continue;
}
for (j, &y) in b.iter().enumerate().take(len.saturating_sub(i)) {
out[i + j] ^= f.mul(x, y);
}
}
out
}
fn ensure_probability(p_s: f64) -> Result<()> {
ensure((0.0..=1.0).contains(&p_s), || format!("the symbol error probability must be in [0, 1], got {p_s}"))
}
pub fn word_error_probability(n: usize, t: usize, p_s: f64) -> Result<f64> {
ensure_probability(p_s)?;
if t >= n || p_s == 0.0 {
return Ok(0.0);
}
if p_s == 1.0 {
return Ok(1.0);
}
Ok(ln_binomial_pmf(n, p_s).iter().skip(t + 1).map(|l| l.exp()).sum())
}
pub fn symbol_error_probability(n: usize, t: usize, p_s: f64) -> Result<f64> {
ensure_probability(p_s)?;
if t >= n || p_s == 0.0 {
return Ok(0.0);
}
if p_s == 1.0 {
return Ok(1.0);
}
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)
}
pub fn bit_error_probability(n: usize, t: usize, m: u32, p_s: f64) -> Result<f64> {
ensure(m >= 1, || "a symbol needs at least one bit".into())?;
let q = 2f64.powi(m as i32);
Ok(q / 2.0 / (q - 1.0) * symbol_error_probability(n, t, p_s)?)
}
#[cfg(test)]
mod tests {
use super::*;
use comnum_core::rng::Rng;
pub(super) fn corrupt(c: &mut [u16], errors: usize, m: u32, avoid: &[usize], rng: &mut Rng) -> Vec<usize> {
let mut positions = Vec::with_capacity(errors);
while positions.len() < errors {
let p = rng.below(c.len());
if !positions.contains(&p) && !avoid.contains(&p) {
positions.push(p);
let e = 1 + rng.below((1 << m) - 1) as u16;
c[p] ^= e;
}
}
positions
}
#[test]
fn errors_and_erasures_are_corrected_up_to_2e_plus_s_equal_2t() {
let code = ReedSolomon::shortened(8, 4, 60).unwrap();
assert_eq!(code.k_symbols(), 52);
let mut rng = Rng::new(11);
for s in 0..=8 {
for e in 0..=(8 - s) / 2 {
let msg: Vec<u16> = (0..52).map(|_| rng.below(256) as u16).collect();
let mut c = code.encode_symbols(&msg).unwrap();
let erased = corrupt(&mut c, s, 8, &[], &mut rng);
corrupt(&mut c, e, 8, &erased, &mut rng);
let (m, n) = code.decode_erasures(&c, &erased).unwrap();
assert_eq!(m, msg, "s = {s}, e = {e}");
assert_eq!(n, Some(s + e), "s = {s}, e = {e}");
}
}
}
#[test]
fn invalid_inputs_are_errors_not_panics() {
assert!(ReedSolomon::new(1, 1).is_err());
assert!(ReedSolomon::new(17, 1).is_err());
assert!(ReedSolomon::new(3, 0).is_err());
assert!(ReedSolomon::shortened(3, 2, 4).is_err());
assert!(ReedSolomon::shortened(3, 1, 8).is_err());
assert!(ReedSolomon::with_field(4, 0b1_1111, 1, 1, 1, 15).is_err());
assert!(ReedSolomon::with_field(4, 0b1_0011, 1, 3, 1, 15).is_err());
let code = ReedSolomon::new(4, 2).unwrap();
assert!(code.encode_symbols(&[0; 10]).is_err());
assert!(code.encode_symbols(&[16; 11]).is_err());
assert!(code.decode_symbols(&[0; 14]).is_err());
assert!(code.decode_erasures(&[0; 15], &[15]).is_err());
assert!(code.decode_erasures(&[0; 15], &[1, 1]).is_err());
assert!(code.decode_erasures(&[0; 15], &[0, 1, 2, 3, 4]).is_err());
assert!(code.encode(&[0; 43]).is_err());
assert!(pack_symbols(&[1, 0, 1], 2).is_err());
assert_eq!(unpack_symbols(&[0b1011], 3), vec![0, 1, 1]);
assert_eq!(pack_symbols(&[1, 0, 1, 1, 1, 0], 3).unwrap(), vec![5, 6]);
}
#[test]
fn the_block_code_interface_counts_bits() {
let code = ReedSolomon::shortened(6, 3, 30).unwrap();
assert_eq!((BlockCode::k(&code), BlockCode::n(&code)), (24 * 6, 30 * 6));
let mut rng = Rng::new(3);
let info: Vec<u8> = (0..144).map(|_| rng.bit()).collect();
let c = code.encode_block(&info).unwrap();
let mut llr: Vec<f64> = c.iter().map(|&b| if b == 1 { -2.0 } else { 2.0 }).collect();
for i in [0, 1, 42, 174, 175, 176, 177, 178, 179] {
llr[i] = -llr[i];
}
assert_eq!(code.decode_block(&llr).unwrap(), info);
}
#[test]
fn the_closed_forms_have_the_right_limits() {
assert_eq!(word_error_probability(255, 16, 0.0).unwrap(), 0.0);
assert_eq!(word_error_probability(255, 16, 1.0).unwrap(), 1.0);
let p = 0.01;
assert!((word_error_probability(7, 0, p).unwrap() - (1.0 - (1.0 - p).powi(7))).abs() < 1e-15);
let out = symbol_error_probability(255, 16, 0.05).unwrap();
assert!(out > 0.0 && out < 0.05);
assert!((bit_error_probability(255, 16, 8, 0.05).unwrap() / out - 128.0 / 255.0).abs() < 1e-12);
for p in [-0.1, 1.5, f64::NAN] {
assert!(word_error_probability(255, 16, p).is_err() && symbol_error_probability(255, 16, p).is_err());
assert!(bit_error_probability(255, 16, 8, p).is_err());
}
assert_eq!(word_error_probability(7, 7, 0.5).unwrap(), 0.0);
}
#[test]
fn the_block_code_counts_bits_and_the_symbol_methods_count_symbols() {
let code = ReedSolomon::new(8, 16).unwrap();
assert_eq!((code.n_symbols(), code.k_symbols()), (255, 223));
assert_eq!((code.n(), code.k()), (255 * 8, 223 * 8));
assert_eq!((BlockCode::n(&code), BlockCode::k(&code)), (255 * 8, 223 * 8));
let bits = vec![1u8; code.k()];
assert_eq!(code.encode_block(&bits).unwrap().len(), code.n());
}
}
#[cfg(test)]
mod literature {
use super::*;
use comnum_core::rng::Rng;
#[test]
fn lin_costello_sec_7_3_generator_of_the_15_9_code() {
let code = ReedSolomon::with_field(4, 0b1_0011, 1, 1, 3, 15).unwrap();
let f = code.field();
let expected: Vec<u16> = [6, 9, 6, 4, 14, 10].iter().map(|&e| f.pow(e) as u16).chain([1]).collect();
assert_eq!(code.generator(), expected);
}
#[test]
fn lin_costello_sec_7_4_three_errors_of_the_15_9_code() {
let code = ReedSolomon::with_field(4, 0b1_0011, 1, 1, 3, 15).unwrap();
let f = code.field();
let mut r = vec![0u16; 15];
r[14 - 3] = f.pow(7) as u16;
r[14 - 6] = f.pow(3) as u16;
r[14 - 12] = f.pow(4) as u16;
let syndromes: Vec<u32> = (1..=6)
.map(|j| r.iter().enumerate().fold(0u32, |acc, (i, &v)| acc ^ f.mul(u32::from(v), f.pow(j * (14 - i)))))
.collect();
assert_eq!(syndromes, [f.pow(12), 1, f.pow(14), f.pow(10), 0, f.pow(12)]);
assert_eq!(code.decode_symbols(&r).unwrap(), (vec![0; 9], Some(3)));
}
#[test]
fn etsi_en_300_744_sec_4_3_2_generator_of_rs_204_188() {
let code = ReedSolomon::with_field(8, 0x11D, 0, 1, 8, 204).unwrap();
assert_eq!((code.n_symbols(), code.k_symbols()), (204, 188));
let f = code.field();
let exponents = [120, 104, 107, 109, 102, 161, 76, 3, 91, 191, 147, 169, 182, 194, 225, 120];
let mut expected: Vec<u16> = exponents.iter().rev().map(|&e| f.pow(e) as u16).collect();
expected.push(1);
assert_eq!(code.generator(), expected);
for j in 0..16 {
let root = f.pow(j);
let v = code
.generator()
.iter()
.enumerate()
.fold(0u32, |acc, (d, &a)| acc ^ f.mul(u32::from(a), f.pow(d * f.log(root))));
assert_eq!(v, 0, "λ^{j} is a root");
}
let full = ReedSolomon::with_field(8, 0x11D, 0, 1, 8, 255).unwrap();
let msg: Vec<u16> = (0..188).map(|i| (i * 13 % 256) as u16).collect();
let mut padded = vec![0u16; 51];
padded.extend(&msg);
assert_eq!(full.encode_symbols(&padded).unwrap()[51..], code.encode_symbols(&msg).unwrap()[..]);
}
#[test]
fn ccsds_131_0_b_sec_4_3_roots_of_rs_255_223() {
let code = ReedSolomon::with_field(8, 0x187, 112, 11, 16, 255).unwrap();
let f = code.field();
let g = code.generator();
assert_eq!(g.len(), 33);
for j in 112..=143 {
let root_exp = 11 * j;
let v = g.iter().enumerate().fold(0u32, |acc, (d, &a)| acc ^ f.mul(u32::from(a), f.pow(d * root_exp)));
assert_eq!(v, 0, "α^(11·{j}) is a root");
}
for i in 0..=32 {
assert_eq!(g[i], g[32 - i], "g is self-reciprocal");
}
let mut rng = Rng::new(5);
let msg: Vec<u16> = (0..223).map(|_| rng.below(256) as u16).collect();
let mut c = code.encode_symbols(&msg).unwrap();
super::tests::corrupt(&mut c, 16, 8, &[], &mut rng);
assert_eq!(code.decode_symbols(&c).unwrap(), (msg, Some(16)));
}
#[test]
fn rs_255_223_corrects_16_errors_and_not_17() {
let code = ReedSolomon::new(8, 16).unwrap();
let mut rng = Rng::new(255);
let (mut failures, mut miscorrections) = (0, 0);
for trial in 0..300 {
let msg: Vec<u16> = (0..223).map(|_| rng.below(256) as u16).collect();
let mut c = code.encode_symbols(&msg).unwrap();
super::tests::corrupt(&mut c, 16, 8, &[], &mut rng);
assert_eq!(code.decode_symbols(&c).unwrap(), (msg.clone(), Some(16)), "trial {trial}");
let mut c = code.encode_symbols(&msg).unwrap();
super::tests::corrupt(&mut c, 17, 8, &[], &mut rng);
let (decoded, status) = code.decode_symbols(&c).unwrap();
if status.is_none() {
failures += 1;
} else {
assert_ne!(decoded, msg, "trial {trial}: 17 errors cannot be corrected");
miscorrections += 1;
}
}
assert_eq!((failures, miscorrections), (300, 0));
}
#[test]
fn word_error_rate_matches_the_closed_form_at_p_s_0_01() {
let code = ReedSolomon::new(8, 4).unwrap();
let p_s = 0.01;
let trials = 20_000;
let mut rng = Rng::new(2024);
let mut lost = 0usize;
for _ in 0..trials {
let msg: Vec<u16> = (0..247).map(|_| rng.below(256) as u16).collect();
let mut c = code.encode_symbols(&msg).unwrap();
for v in &mut c {
if rng.uniform() < p_s {
*v ^= 1 + rng.below(255) as u16;
}
}
let (decoded, status) = code.decode_symbols(&c).unwrap();
if status.is_none() || decoded != msg {
lost += 1;
}
}
let p_w = word_error_probability(255, 4, p_s).unwrap();
let measured = lost as f64 / trials as f64;
let tol = 4.0 * (p_w * (1.0 - p_w) / trials as f64).sqrt();
assert!((measured - p_w).abs() < tol, "measured {measured}, closed form {p_w}, tolerance {tol}");
}
}