correction: support erasures

src/primitives/correction.rs

@@ -76,7 +76,11 @@ pub trait CorrectableError {
  return None;
  }
 
- self.residue_error().map(|e| Corrector { residue: e.residue(), phantom: PhantomData })
+ self.residue_error().map(|e| Corrector {
+ erasures: FieldVec::new(),
+ residue: e.residue(),
+ phantom: PhantomData,
+ })
  }
 }
 
@@ -127,12 +131,40 @@ impl CorrectableError for DecodeError {
 }
 
 /// An error-correction context.
-pub struct Corrector<Ck> {
+pub struct Corrector<Ck: Checksum> {
+ erasures: FieldVec<usize>,
  residue: Polynomial<Fe32>,
  phantom: PhantomData<Ck>,
 }
 
 impl<Ck: Checksum> Corrector<Ck> {
+ /// A bound on the number of errors and erasures (errors with known location)
+ /// can be corrected by this corrector.
+ ///
+ /// Returns N such that, given E errors and X erasures, corection is possible
+ /// iff 2E + X <= N.
+ pub fn singleton_bound(&self) -> usize {
+ // d - 1, where d = [number of consecutive roots] + 2
+ Ck::ROOT_EXPONENTS.end() - Ck::ROOT_EXPONENTS.start() + 1
+ }
+
+ /// TODO
+ pub fn add_erasures(&mut self, locs: &[usize]) {
+ for loc in locs {
+ // If the user tries to add too many erasures, just ignore them. In
+ // this case error correction is guaranteed to fail anyway, because
+ // they will have exceeded the singleton bound. (Otherwise, the
+ // singleton bound, which is always <= the checksum length, must be
+ // greater than NO_ALLOC_MAX_LENGTH. So the checksum length must be
+ // greater than NO_ALLOC_MAX_LENGTH. Then correction will still fail.)
+ #[cfg(not(feature = "alloc"))]
+ if self.erasures.len() == NO_ALLOC_MAX_LENGTH {
+ break;
+ }
+ self.erasures.push(*loc);
+ }
+ }
+
  /// Returns an iterator over the errors in the string.
  ///
  /// Returns `None` if it can be determined that there are too many errors to be
@@ -145,29 +177,44 @@ impl<Ck: Checksum> Corrector<Ck> {
  /// string may not actually be the intended string.
  pub fn bch_errors(&self) -> Option<ErrorIterator<Ck>> {
  // 1. Compute all syndromes by evaluating the residue at each power of the generator.
- let syndromes: FieldVec<_> = Ck::ROOT_GENERATOR
+ let syndromes: Polynomial<_> = Ck::ROOT_GENERATOR
  .powers_range(Ck::ROOT_EXPONENTS)
  .map(|rt| self.residue.evaluate(&rt))
  .collect();
 
+ // 1a. Compute the "Forney syndrome polynomial" which is the product of the syndrome
+ // polynomial and the erasure locator. This "erases the erasures" so that B-M
+ // can find only the errors.
+ let mut erasure_locator = Polynomial::with_monic_leading_term(&[]); // 1
+ for loc in &self.erasures {
+ let factor: Polynomial<_> =
+ [Ck::CorrectionField::ONE, -Ck::ROOT_GENERATOR.powi(*loc as i64)]
+ .iter()
+ .cloned()
+ .collect(); // alpha^-ix - 1
+ erasure_locator = erasure_locator.mul_mod_x_d(&factor, usize::MAX);
+ }
+ let forney_syndromes = erasure_locator.convolution(&syndromes);
+
  // 2. Use the Berlekamp-Massey algorithm to find the connection polynomial of the
  // LFSR that generates these syndromes. For magical reasons this will be equal
  // to the error locator polynomial for the syndrome.
- let lfsr = LfsrIter::berlekamp_massey(&syndromes[..]);
+ let lfsr = LfsrIter::berlekamp_massey(&forney_syndromes.as_inner()[..]);
  let conn = lfsr.coefficient_polynomial();
 
  // 3. The connection polynomial is the error locator polynomial. Use this to get
  // the errors.
- let max_correctable_errors =
- (Ck::ROOT_EXPONENTS.end() - Ck::ROOT_EXPONENTS.start() + 1) / 2;
- if conn.degree() <= max_correctable_errors {
+ if erasure_locator.degree() + 2 * conn.degree() <= self.singleton_bound() {
+ // 3a. Compute the "errata locator" which is the product of the error locator
+ // and the erasure locator. Note that while we used the Forney syndromes
+ // when calling the BM algorithm, in all other cases we use the ordinary
+ // unmodified syndromes.
+ let errata_locator = conn.mul_mod_x_d(&erasure_locator, usize::MAX);
  Some(ErrorIterator {
- evaluator: conn.mul_mod_x_d(
- &Polynomial::from(syndromes),
- Ck::ROOT_EXPONENTS.end() - Ck::ROOT_EXPONENTS.start() + 1,
- ),
- locator_derivative: conn.formal_derivative(),
- inner: conn.find_nonzero_distinct_roots(Ck::ROOT_GENERATOR),
+ evaluator: errata_locator.mul_mod_x_d(&syndromes, self.singleton_bound()),
+ locator_derivative: errata_locator.formal_derivative(),
+ erasures: &self.erasures[..],
+ errors: conn.find_nonzero_distinct_roots(Ck::ROOT_GENERATOR),
  a: Ck::ROOT_GENERATOR,
  c: *Ck::ROOT_EXPONENTS.start(),
  })
@@ -206,32 +253,39 @@ impl<Ck: Checksum> Corrector<Ck> {
 /// caller should fix this before attempting error correction. If it is unknown,
 /// the caller cannot assume anything about the intended checksum, and should not
 /// attempt error correction.
-pub struct ErrorIterator<Ck: Checksum> {
+pub struct ErrorIterator<'c, Ck: Checksum> {
  evaluator: Polynomial<Ck::CorrectionField>,
  locator_derivative: Polynomial<Ck::CorrectionField>,
- inner: super::polynomial::RootIter<Ck::CorrectionField>,
+ erasures: &'c [usize],
+ errors: super::polynomial::RootIter<Ck::CorrectionField>,
  a: Ck::CorrectionField,
  c: usize,
 }
 
-impl<Ck: Checksum> Iterator for ErrorIterator<Ck> {
+impl<'c, Ck: Checksum> Iterator for ErrorIterator<'c, Ck> {
  type Item = (usize, Fe32);
 
  fn next(&mut self) -> Option<Self::Item> {
  // Compute -i, which is the location we will return to the user.
- let neg_i = match self.inner.next() {
- None => return None,
- Some(0) => 0,
- Some(x) => Ck::ROOT_GENERATOR.multiplicative_order() - x,
+ let neg_i = if self.erasures.is_empty() {
+ match self.errors.next() {
+ None => return None,
+ Some(0) => 0,
+ Some(x) => Ck::ROOT_GENERATOR.multiplicative_order() - x,
+ }
+ } else {
+ let pop = self.erasures[0];
+ self.erasures = &self.erasures[1..];
+ pop
  };
 
  // Forney's equation, as described in https://en.wikipedia.org/wiki/BCH_code#Forney_algorithm
  //
  // It is rendered as
  //
- // a^i evaluator(a^-i)
- // e_k = - ---------------------------------
- // a^(ci) locator_derivative(a^-i)
+ //   evaluator(a^-i)
+ // e_k = - -----------------------------------------
+ // (a^i)^(c - 1)) locator_derivative(a^-i)
  //
  // where here a is `Ck::ROOT_GENERATOR`, c is the first element of the range
  // `Ck::ROOT_EXPONENTS`, and both evalutor and locator_derivative are polynomials
@@ -240,8 +294,8 @@ impl<Ck: Checksum> Iterator for ErrorIterator<Ck> {
  let a_i = self.a.powi(neg_i as i64);
  let a_neg_i = a_i.clone().multiplicative_inverse();
 
- let num = self.evaluator.evaluate(&a_neg_i) * &a_i;
- let den = a_i.powi(self.c as i64) * self.locator_derivative.evaluate(&a_neg_i);
+ let num = self.evaluator.evaluate(&a_neg_i);
+ let den = a_i.powi(self.c as i64 - 1) * self.locator_derivative.evaluate(&a_neg_i);
  let ret = -num / den;
  match ret.try_into() {
  Ok(ret) => Some((neg_i, ret)),
@@ -263,9 +317,13 @@ mod tests {
  match SegwitHrpstring::new(s) {
  Ok(_) => panic!("{} successfully, and wrongly, parsed", s),
  Err(e) => {
- let ctx = e.correction_context::<Bech32>().unwrap();
+ let mut ctx = e.correction_context::<Bech32>().unwrap();
  let mut iter = ctx.bch_errors().unwrap();
+ assert_eq!(iter.next(), Some((0, Fe32::X)));
+ assert_eq!(iter.next(), None);
 
+ ctx.add_erasures(&[0]);
+ let mut iter = ctx.bch_errors().unwrap();
  assert_eq!(iter.next(), Some((0, Fe32::X)));
  assert_eq!(iter.next(), None);
  }
@@ -276,9 +334,13 @@ mod tests {
  match SegwitHrpstring::new(s) {
  Ok(_) => panic!("{} successfully, and wrongly, parsed", s),
  Err(e) => {
- let ctx = e.correction_context::<Bech32>().unwrap();
+ let mut ctx = e.correction_context::<Bech32>().unwrap();
  let mut iter = ctx.bch_errors().unwrap();
+ assert_eq!(iter.next(), Some((6, Fe32::T)));
+ assert_eq!(iter.next(), None);
 
+ ctx.add_erasures(&[6]);
+ let mut iter = ctx.bch_errors().unwrap();
  assert_eq!(iter.next(), Some((6, Fe32::T)));
  assert_eq!(iter.next(), None);
  }
@@ -297,13 +359,42 @@ mod tests {
  }
  }
 
- // Two errors.
- let s = "bc1qar0srrr7xfkvy5l643lydnw9re59gtzzwf5mxx";
+ // Two errors; cannot correct.
+ let s = "bc1qar0srrr7xfkvy5l64qlydnw9re59gtzzwf5mdx";
  match SegwitHrpstring::new(s) {
  Ok(_) => panic!("{} successfully, and wrongly, parsed", s),
  Err(e) => {
- let ctx = e.correction_context::<Bech32>().unwrap();
+ let mut ctx = e.correction_context::<Bech32>().unwrap();
  assert!(ctx.bch_errors().is_none());
+
+ // But we can correct it if we inform where an error is.
+ ctx.add_erasures(&[0]);
+ let mut iter = ctx.bch_errors().unwrap();
+ assert_eq!(iter.next(), Some((0, Fe32::X)));
+ assert_eq!(iter.next(), Some((20, Fe32::_3)));
+ assert_eq!(iter.next(), None);
+
+ ctx.add_erasures(&[20]);
+ let mut iter = ctx.bch_errors().unwrap();
+ assert_eq!(iter.next(), Some((0, Fe32::X)));
+ assert_eq!(iter.next(), Some((20, Fe32::_3)));
+ assert_eq!(iter.next(), None);
+ }
+ }
+
+ // In fact, if we know the locations, we can correct up to 3 errors.
+ let s = "bc1q9r0srrr7xfkvy5l64qlydnw9re59gtzzwf5mdx";
+ match SegwitHrpstring::new(s) {
+ Ok(_) => panic!("{} successfully, and wrongly, parsed", s),
+ Err(e) => {
+ let mut ctx = e.correction_context::<Bech32>().unwrap();
+ ctx.add_erasures(&[37, 0, 20]);
+ let mut iter = ctx.bch_errors().unwrap();
+
+ assert_eq!(iter.next(), Some((37, Fe32::C)));
+ assert_eq!(iter.next(), Some((0, Fe32::X)));
+ assert_eq!(iter.next(), Some((20, Fe32::_3)));
+ assert_eq!(iter.next(), None);
  }
  }
  }

src/primitives/polynomial.rs

@@ -17,9 +17,7 @@ pub struct Polynomial<F> {
 }
 
 impl<F: Field> PartialEq for Polynomial<F> {
- fn eq(&self, other: &Self) -> bool {
- self.inner[..self.degree()] == other.inner[..other.degree()]
- }
+ fn eq(&self, other: &Self) -> bool { self.coefficients() == other.coefficients() }
 }
 
 impl<F: Field> Eq for Polynomial<F> {}
@@ -58,9 +56,16 @@ impl<F: Field> Polynomial<F> {
  debug_assert_ne!(self.inner.len(), 0, "polynomials never have no terms");
  let degree_without_leading_zeros = self.inner.len() - 1;
  let leading_zeros = self.inner.iter().rev().take_while(|el| **el == F::ZERO).count();
- degree_without_leading_zeros - leading_zeros
+ degree_without_leading_zeros.saturating_sub(leading_zeros)
  }
 
+ /// Accessor for the coefficients of the polynomial, in "little endian" order.
+ ///
+ /// # Panics
+ ///
+ /// Panics if [`Self::has_data`] is false.
+ pub fn coefficients(&self) -> &[F] { &self.inner[..self.degree() + 1] }
+
  /// An iterator over the coefficients of the polynomial.
  ///
  /// Yields value in "little endian" order; that is, the constant term is returned first.
@@ -70,7 +75,7 @@ impl<F: Field> Polynomial<F> {
  /// Panics if [`Self::has_data`] is false.
  pub fn iter(&self) -> slice::Iter<F> {
  self.assert_has_data();
- self.inner[..self.degree() + 1].iter()
+ self.coefficients().iter()
  }
 
  /// The leading term of the polynomial.
@@ -143,6 +148,24 @@ impl<F: Field> Polynomial<F> {
  res
  }
 
+ /// TODO
+ pub fn convolution(&self, syndromes: &Self) -> Self {
+ let mut ret = FieldVec::new();
+ let terms = (1 + syndromes.inner.len()).saturating_sub(1 + self.degree());
+ if terms == 0 {
+ ret.push(F::ZERO);
+ return Self::from(ret);
+ }
+
+ let n = 1 + self.degree();
+ for idx in 0..terms {
+ ret.push(
+ (0..n).map(|i| self.inner[n - i - 1].clone() * &syndromes.inner[idx + i]).sum(),
+ );
+ }
+ Self::from(ret)
+ }
+
  /// Multiplies two polynomials modulo x^d, for some given `d`.
  ///
  /// Can be used to simply multiply two polynomials, by passing `usize::MAX` or