1use crate::reed_solomon::{
2 DecoderResult, EncoderResult, Error, Plan,
3 engine::{self, Engine, GF_MODULUS, GF_ORDER, GfElement, SHARD_CHUNK_BYTES, tables},
4 rate::{DecoderWork, EncoderWork, Rate, RateDecoder, RateEncoder},
5};
6#[cfg(not(feature = "std"))]
7use alloc::vec;
8use core::marker::PhantomData;
9use fixedbitset::FixedBitSet;
10
11pub(crate) const DIRECT_EVALUATION_LIMIT: usize = 32;
14
15pub(crate) fn with_erasures<E: Engine, T>(
18 original_count: usize,
19 recovery_count: usize,
20 received: &FixedBitSet,
21 f: impl FnOnce(&[GfElement]) -> T,
22) -> T {
23 let end = original_count.next_power_of_two() + recovery_count;
24 if end <= DIRECT_EVALUATION_LIMIT {
25 let mut erasures = [0; DIRECT_EVALUATION_LIMIT];
26 eval_direct(&mut erasures[..end], original_count, received);
27 f(&erasures[..end])
28 } else {
29 let mut erasures = vec![0; end.next_power_of_two()];
30 eval_walsh::<E>(&mut erasures, original_count, recovery_count, received);
31 f(&erasures[..end])
32 }
33}
34
35fn eval_direct(erasures: &mut [GfElement], original_count: usize, received: &FixedBitSet) {
41 let chunk_size = original_count.next_power_of_two();
42 let mut known = [0; DIRECT_EVALUATION_LIMIT];
43 let mut count = 0;
44 for i in 0..erasures.len() {
45 if received[i] || (original_count..chunk_size).contains(&i) {
47 known[count] = i;
48 count += 1;
49 }
50 }
51
52 let log = &tables::get_exp_log().log;
55 for (i, erasure) in erasures.iter_mut().enumerate() {
56 let mut sum = 0u32;
58 for &j in &known[..count] {
59 sum += u32::from(log[i ^ j]);
61 }
62 *erasure = GF_MODULUS - (sum % u32::from(GF_MODULUS)) as GfElement;
63 }
64}
65
66pub(crate) fn eval_walsh<E: Engine>(
76 erasures: &mut [GfElement],
77 original_count: usize,
78 recovery_count: usize,
79 received: &FixedBitSet,
80) {
81 let chunk_size = original_count.next_power_of_two();
82 let end = chunk_size + recovery_count;
83
84 for i in 0..end {
88 if received[i] || (original_count..chunk_size).contains(&i) {
89 erasures[i] = 1;
90 }
91 }
92 super::eval_locator::<E>(erasures, end);
93 for value in &mut erasures[..end] {
94 *value = GF_MODULUS - *value;
95 }
96}
97
98pub struct LowRate<E: Engine>(PhantomData<E>);
100
101impl<E: Engine> Rate<E> for LowRate<E> {
102 type RateEncoder = LowRateEncoder<E>;
103 type RateDecoder = LowRateDecoder<E>;
104
105 fn supports(original_count: usize, recovery_count: usize) -> bool {
106 original_count > 0
107 && recovery_count > 0
108 && original_count < GF_ORDER
109 && recovery_count < GF_ORDER
110 && original_count.next_power_of_two() + recovery_count <= GF_ORDER
111 }
112}
113
114pub struct LowRateEncoder<E: Engine> {
116 engine: E,
117 work: EncoderWork,
120}
121
122impl<E: Engine> RateEncoder<E> for LowRateEncoder<E> {
123 type Rate = LowRate<E>;
124
125 fn validate(
126 original_count: usize,
127 recovery_count: usize,
128 shard_bytes: usize,
129 ) -> Result<(), Error> {
130 Self::Rate::validate(original_count, recovery_count, shard_bytes)?;
131 super::validate_work_size(
132 shard_bytes,
133 Self::work_count(original_count, recovery_count),
134 )
135 }
136
137 fn add_original_shard<T: AsRef<[u8]>>(&mut self, original_shard: T) -> Result<(), Error> {
138 self.work.add_original_shard(original_shard)
139 }
140
141 fn encode(&mut self) -> Result<EncoderResult<'_>, Error> {
142 let (mut work, original_count, recovery_count) = self.work.encode_begin()?;
143 let chunk_size = original_count.next_power_of_two();
144 let engine = &self.engine;
145
146 work.zero(original_count..chunk_size);
148 engine.ifft(&mut work, 0, chunk_size, original_count, 0);
149
150 let mut chunk_start = chunk_size;
152 while chunk_start < recovery_count {
153 work.copy_within(0, chunk_start, chunk_size);
154 chunk_start += chunk_size;
155 }
156
157 let mut chunk_start = 0;
159 while chunk_start + chunk_size <= recovery_count {
160 engine::fft_skew_end(engine, &mut work, chunk_start, chunk_size, chunk_size);
161 chunk_start += chunk_size;
162 }
163
164 let last_count = recovery_count % chunk_size;
166 if last_count > 0 {
167 engine::fft_skew_end(engine, &mut work, chunk_start, chunk_size, last_count);
168 }
169
170 self.work.undo_last_chunk_encoding();
171
172 Ok(EncoderResult::new(&mut self.work))
173 }
174
175 fn into_parts(self) -> (E, EncoderWork) {
176 (self.engine, self.work)
177 }
178
179 fn new(
180 original_count: usize,
181 recovery_count: usize,
182 shard_bytes: usize,
183 engine: E,
184 work: Option<EncoderWork>,
185 ) -> Result<Self, Error> {
186 let mut work = work.unwrap_or_default();
187 Self::reset_work(original_count, recovery_count, shard_bytes, &mut work)?;
188 Ok(Self { engine, work })
189 }
190
191 fn reset(
192 &mut self,
193 original_count: usize,
194 recovery_count: usize,
195 shard_bytes: usize,
196 ) -> Result<(), Error> {
197 Self::reset_work(original_count, recovery_count, shard_bytes, &mut self.work)
198 }
199}
200
201impl<E: Engine> LowRateEncoder<E> {
202 fn reset_work(
206 original_count: usize,
207 recovery_count: usize,
208 shard_bytes: usize,
209 work: &mut EncoderWork,
210 ) -> Result<(), Error> {
211 Self::validate(original_count, recovery_count, shard_bytes)?;
212 work.reset(
213 original_count,
214 recovery_count,
215 shard_bytes,
216 Self::work_count(original_count, recovery_count),
217 );
218 Ok(())
219 }
220
221 fn work_count(original_count: usize, recovery_count: usize) -> usize {
232 assert!(Self::supports(original_count, recovery_count));
233
234 let chunk_size = original_count.next_power_of_two();
235
236 recovery_count.next_multiple_of(chunk_size)
237 }
238}
239
240pub struct LowRateDecoder<E: Engine> {
242 engine: E,
243 work: DecoderWork,
246}
247
248impl<E: Engine> RateDecoder<E> for LowRateDecoder<E> {
249 type Rate = LowRate<E>;
250
251 fn validate(
252 original_count: usize,
253 recovery_count: usize,
254 shard_bytes: usize,
255 ) -> Result<(), Error> {
256 Self::Rate::validate(original_count, recovery_count, shard_bytes)?;
257 super::validate_work_size(
258 shard_bytes,
259 Self::work_count(original_count, recovery_count),
260 )
261 }
262
263 fn add_original_shard<T: AsRef<[u8]>>(
264 &mut self,
265 index: usize,
266 original_shard: T,
267 ) -> Result<(), Error> {
268 self.work.add_original_shard(index, original_shard)
269 }
270
271 fn add_recovery_shard<T: AsRef<[u8]>>(
272 &mut self,
273 index: usize,
274 recovery_shard: T,
275 ) -> Result<(), Error> {
276 self.work.add_recovery_shard(index, recovery_shard)
277 }
278
279 fn decode(&mut self, compute_recovery: bool) -> Result<Option<DecoderResult<'_>>, Error> {
280 self.decode_impl(compute_recovery, None)
281 }
282
283 fn into_parts(self) -> (E, DecoderWork) {
284 (self.engine, self.work)
285 }
286
287 fn new(
288 original_count: usize,
289 recovery_count: usize,
290 shard_bytes: usize,
291 engine: E,
292 work: Option<DecoderWork>,
293 ) -> Result<Self, Error> {
294 let mut work = work.unwrap_or_default();
295 Self::reset_work(original_count, recovery_count, shard_bytes, &mut work)?;
296 Ok(Self { engine, work })
297 }
298
299 fn reset(
300 &mut self,
301 original_count: usize,
302 recovery_count: usize,
303 shard_bytes: usize,
304 ) -> Result<(), Error> {
305 Self::reset_work(original_count, recovery_count, shard_bytes, &mut self.work)
306 }
307}
308
309impl<E: Engine> LowRateDecoder<E> {
310 pub(crate) fn decode_with_plan(
316 &mut self,
317 compute_recovery: bool,
318 plan: &Plan,
319 ) -> Result<Option<DecoderResult<'_>>, Error> {
320 self.decode_impl(compute_recovery, Some(plan))
321 }
322}
323
324impl<E: Engine> LowRateDecoder<E> {
325 fn decode_impl(
330 &mut self,
331 compute_recovery: bool,
332 plan: Option<&Plan>,
333 ) -> Result<Option<DecoderResult<'_>>, Error> {
334 if let Some(plan) = plan {
335 self.work.validate_plan(plan, false)?;
336 }
337 let Some((mut work, original_count, recovery_count, received)) =
338 self.work.decode_begin()?
339 else {
340 self.work.reset_received();
343 return Ok(None);
344 };
345
346 let chunk_size = original_count.next_power_of_two();
347 let recovery_end = chunk_size + recovery_count;
348 let work_count = work.len();
349
350 let mut decode = |erasures: &[GfElement]| {
354 for i in 0..original_count {
361 if received[i] {
362 self.engine.mul(&mut work[i], erasures[i]);
363 } else {
364 work[i].fill([0; SHARD_CHUNK_BYTES]);
365 }
366 }
367 work.zero(original_count..chunk_size);
368 for i in chunk_size..recovery_end {
369 if received[i] {
370 self.engine.mul(&mut work[i], erasures[i]);
371 } else {
372 work[i].fill([0; SHARD_CHUNK_BYTES]);
373 }
374 }
375 work.zero(recovery_end..);
376
377 self.engine.ifft(&mut work, 0, work_count, recovery_end, 0);
379 engine::formal_derivative(&mut work);
380 self.engine.fft(&mut work, 0, work_count, recovery_end, 0);
381
382 for i in 0..original_count {
384 if !received[i] {
385 self.engine.mul(&mut work[i], GF_MODULUS - erasures[i]);
386 }
387 }
388
389 if compute_recovery {
395 for i in chunk_size..recovery_end {
396 if !received[i] {
397 self.engine.mul(&mut work[i], GF_MODULUS - erasures[i]);
398 }
399 }
400 }
401 };
402 match plan {
403 Some(plan) => decode(plan.coefficients()),
404 None => with_erasures::<E, _>(original_count, recovery_count, received, decode),
405 }
406
407 self.work.undo_last_chunk_encoding();
409 if compute_recovery {
410 self.work.undo_last_chunk_encoding_recovery();
411 }
412
413 Ok(Some(DecoderResult::new(&mut self.work)))
414 }
415
416 fn reset_work(
420 original_count: usize,
421 recovery_count: usize,
422 shard_bytes: usize,
423 work: &mut DecoderWork,
424 ) -> Result<(), Error> {
425 Self::validate(original_count, recovery_count, shard_bytes)?;
426
427 work.reset(
430 original_count,
431 recovery_count,
432 shard_bytes,
433 0,
434 original_count.next_power_of_two(),
435 Self::work_count(original_count, recovery_count),
436 );
437
438 Ok(())
439 }
440
441 fn work_count(original_count: usize, recovery_count: usize) -> usize {
450 assert!(Self::supports(original_count, recovery_count));
451
452 (original_count.next_power_of_two() + recovery_count).next_power_of_two()
453 }
454}
455
456#[cfg(test)]
457mod tests {
458 use super::*;
459 use crate::reed_solomon::{engine::Scalar, test_util};
460 use commonware_macros::test_group;
461 use commonware_utils::test_rng;
462 use rand::RngExt as _;
463
464 #[test]
465 fn direct_matches_transform() {
466 let mut rng = test_rng();
467
468 let quarter = DIRECT_EVALUATION_LIMIT / 4;
471 let chunk = quarter.next_power_of_two();
472 let recovery = DIRECT_EVALUATION_LIMIT - chunk;
473 for (original_count, recovery_count) in [
474 (1usize, 1),
475 (3, 5),
476 (7, 13),
477 (quarter / 2, quarter),
478 (chunk - 1, recovery - 1),
479 (chunk, recovery),
480 ] {
481 let chunk_size = original_count.next_power_of_two();
482 let end = chunk_size + recovery_count;
483 for pattern in 0..18 {
484 let mut received = FixedBitSet::with_capacity(end.next_power_of_two());
485 let mut expected = [1; GF_ORDER];
486 expected[original_count..chunk_size].fill(0);
487 for i in (0..original_count).chain(chunk_size..end) {
488 if pattern == 0 || (pattern > 1 && rng.random()) {
490 received.insert(i);
491 expected[i] = 0;
492 }
493 }
494 Scalar::eval_poly(&mut expected, GF_ORDER);
495 let mut actual = [0; DIRECT_EVALUATION_LIMIT];
496 eval_direct(&mut actual[..end], original_count, &received);
497 for i in 0..end {
498 assert_eq!(
500 actual[i] % GF_MODULUS,
501 expected[i] % GF_MODULUS,
502 "originals={original_count} recoveries={recovery_count} pattern={pattern} i={i}"
503 );
504 }
505 }
506 }
507 }
508
509 #[test]
510 fn roundtrip_all_originals_missing() {
511 roundtrip_single!(
512 LowRate,
513 3,
514 3,
515 1024,
516 test_util::EITHER_3_3,
517 &[],
518 &[test_util::range(0, 3)],
519 133
520 );
521 }
522
523 #[test]
524 fn roundtrip_no_originals_missing() {
525 roundtrip_single!(
526 LowRate,
527 2,
528 3,
529 1024,
530 test_util::LOW_2_3,
531 &[test_util::index(0), test_util::index(1)],
532 &[],
533 123
534 );
535 }
536
537 #[test]
538 fn roundtrips_tiny() {
539 for (original_count, recovery_count, seed, recovery_hash) in test_util::LOW_TINY {
540 roundtrip_single!(
541 LowRate,
542 *original_count,
543 *recovery_count,
544 1024,
545 recovery_hash,
546 &[test_util::range(*recovery_count, *original_count)],
547 &[test_util::range(
548 0,
549 core::cmp::min(*original_count, *recovery_count)
550 )],
551 *seed,
552 );
553 }
554 }
555
556 #[test_group("slow")]
557 #[test]
558 fn roundtrip_3000_60000() {
559 roundtrip_single!(
560 LowRate,
561 3000,
562 60000,
563 crate::reed_solomon::SHARD_CHUNK_BYTES,
564 test_util::LOW_3000_60000_13,
565 &[],
566 &[test_util::range(0, 3000)],
567 13,
568 );
569 }
570
571 #[test_group("slow")]
572 #[test]
573 fn roundtrip_30000_3000() {
574 roundtrip_single!(
575 LowRate,
576 30000,
577 3000,
578 crate::reed_solomon::SHARD_CHUNK_BYTES,
579 test_util::LOW_30000_3000_15,
580 &[test_util::range(3000, 30000)],
581 &[test_util::range(0, 3000)],
582 15,
583 );
584 }
585
586 #[test_group("slow")]
587 #[test]
588 fn roundtrip_32768_32768() {
589 roundtrip_single!(
590 LowRate,
591 32768,
592 32768,
593 crate::reed_solomon::SHARD_CHUNK_BYTES,
594 test_util::EITHER_32768_32768_11,
595 &[],
596 &[test_util::range(0, 32768)],
597 11,
598 );
599 }
600
601 #[test]
602 fn roundtrip_2000_34000_shard_size_8() {
603 roundtrip_single!(
604 LowRate,
605 2000,
606 34000,
607 8,
608 test_util::LOW_2000_34000_123_8,
609 &[test_util::range(0, 2000)],
610 &[test_util::range(0, 32000)],
611 123
612 );
613 }
614
615 #[test]
616 fn two_rounds_implicit_reset() {
617 roundtrip_two_rounds!(
618 LowRate,
619 false,
620 (
621 2,
622 3,
623 1024,
624 test_util::LOW_2_3,
625 &[],
626 &[test_util::index(0), test_util::index(2)],
627 123
628 ),
629 (
630 2,
631 3,
632 1024,
633 test_util::LOW_2_3_223,
634 &[],
635 &[test_util::index(1), test_util::index(2)],
636 223
637 ),
638 );
639 }
640
641 #[test]
642 fn two_rounds_explicit_reset() {
643 roundtrip_two_rounds!(
644 LowRate,
645 true,
646 (
647 2,
648 3,
649 1024,
650 test_util::LOW_2_3,
651 &[],
652 &[test_util::index(0), test_util::index(2)],
653 123
654 ),
655 (
656 2,
657 5,
658 1024,
659 test_util::LOW_2_5,
660 &[],
661 &[test_util::index(0), test_util::index(4)],
662 125
663 ),
664 );
665 }
666
667 mod low_rate {
668 use crate::reed_solomon::{
669 Error, SHARD_CHUNK_BYTES,
670 engine::Scalar,
671 rate::{LowRate, Rate},
672 };
673
674 #[test]
675 fn decoder() {
676 assert!(
677 LowRate::<Scalar>::decoder(4096, 61440, SHARD_CHUNK_BYTES, Scalar::new(), None)
678 .is_ok()
679 );
680
681 assert_eq!(
682 LowRate::<Scalar>::decoder(61440, 4096, SHARD_CHUNK_BYTES, Scalar::new(), None)
683 .err(),
684 Some(Error::UnsupportedShardCount {
685 original_count: 61440,
686 recovery_count: 4096,
687 })
688 );
689 }
690
691 #[test]
692 fn encoder() {
693 assert!(
694 LowRate::<Scalar>::encoder(4096, 61440, SHARD_CHUNK_BYTES, Scalar::new(), None)
695 .is_ok()
696 );
697
698 assert_eq!(
699 LowRate::<Scalar>::encoder(61440, 4096, SHARD_CHUNK_BYTES, Scalar::new(), None)
700 .err(),
701 Some(Error::UnsupportedShardCount {
702 original_count: 61440,
703 recovery_count: 4096,
704 })
705 );
706 }
707
708 #[test]
709 fn supports() {
710 assert!(!LowRate::<Scalar>::supports(0, 1));
711 assert!(!LowRate::<Scalar>::supports(1, 0));
712
713 assert!(LowRate::<Scalar>::supports(4096, 61440));
714 assert!(!LowRate::<Scalar>::supports(4096, 61441));
715 assert!(!LowRate::<Scalar>::supports(4097, 61440));
716
717 assert!(!LowRate::<Scalar>::supports(61440, 4096));
718
719 assert!(!LowRate::<Scalar>::supports(usize::MAX, usize::MAX));
720 }
721
722 #[test]
723 fn validate() {
724 assert_eq!(
725 LowRate::<Scalar>::validate(1, 1, 123).err(),
726 Some(Error::InvalidShardSize { shard_bytes: 123 })
727 );
728
729 assert!(LowRate::<Scalar>::validate(4096, 61440, SHARD_CHUNK_BYTES).is_ok());
730
731 assert_eq!(
732 LowRate::<Scalar>::validate(61440, 4096, SHARD_CHUNK_BYTES).err(),
733 Some(Error::UnsupportedShardCount {
734 original_count: 61440,
735 recovery_count: 4096,
736 })
737 );
738 }
739 }
740
741 mod low_rate_encoder {
742 use crate::reed_solomon::{
743 Error, SHARD_CHUNK_BYTES,
744 engine::Scalar,
745 rate::{LowRateEncoder, RateEncoder},
746 };
747
748 test_rate_encoder_errors! {LowRateEncoder}
749
750 #[test]
751 fn supports() {
752 assert!(LowRateEncoder::<Scalar>::supports(4096, 61440));
753 assert!(!LowRateEncoder::<Scalar>::supports(61440, 4096));
754 }
755
756 #[test]
757 fn validate() {
758 assert_eq!(
759 LowRateEncoder::<Scalar>::validate(1, 1, 123).err(),
760 Some(Error::InvalidShardSize { shard_bytes: 123 })
761 );
762
763 assert!(LowRateEncoder::<Scalar>::validate(4096, 61440, SHARD_CHUNK_BYTES).is_ok());
764
765 assert_eq!(
766 LowRateEncoder::<Scalar>::validate(61440, 4096, SHARD_CHUNK_BYTES).err(),
767 Some(Error::UnsupportedShardCount {
768 original_count: 61440,
769 recovery_count: 4096,
770 })
771 );
772 }
773
774 #[test]
775 fn work_count() {
776 assert_eq!(LowRateEncoder::<Scalar>::work_count(1, 1), 1);
777 assert_eq!(LowRateEncoder::<Scalar>::work_count(1024, 4096), 4096);
778 assert_eq!(LowRateEncoder::<Scalar>::work_count(1024, 4097), 5120);
779 assert_eq!(LowRateEncoder::<Scalar>::work_count(1025, 4097), 6144);
780 assert_eq!(LowRateEncoder::<Scalar>::work_count(32768, 32768), 32768);
781 }
782 }
783
784 mod low_rate_decoder {
785 use crate::reed_solomon::{
786 Error, SHARD_CHUNK_BYTES,
787 engine::Scalar,
788 rate::{LowRateDecoder, RateDecoder},
789 };
790
791 test_rate_decoder_errors! {LowRateDecoder}
792
793 #[test]
794 fn supports() {
795 assert!(LowRateDecoder::<Scalar>::supports(4096, 61440));
796 assert!(!LowRateDecoder::<Scalar>::supports(61440, 4096));
797 }
798
799 #[test]
800 fn validate() {
801 assert_eq!(
802 LowRateDecoder::<Scalar>::validate(1, 1, 123).err(),
803 Some(Error::InvalidShardSize { shard_bytes: 123 })
804 );
805
806 assert!(LowRateDecoder::<Scalar>::validate(4096, 61440, SHARD_CHUNK_BYTES).is_ok());
807
808 assert_eq!(
809 LowRateDecoder::<Scalar>::validate(61440, 4096, SHARD_CHUNK_BYTES).err(),
810 Some(Error::UnsupportedShardCount {
811 original_count: 61440,
812 recovery_count: 4096,
813 })
814 );
815 }
816
817 #[test]
818 fn work_count() {
819 assert_eq!(LowRateDecoder::<Scalar>::work_count(1, 1), 2);
820 assert_eq!(LowRateDecoder::<Scalar>::work_count(1024, 3072), 4096);
821 assert_eq!(LowRateDecoder::<Scalar>::work_count(1024, 3073), 8192);
822 assert_eq!(LowRateDecoder::<Scalar>::work_count(1025, 2048), 4096);
823 assert_eq!(LowRateDecoder::<Scalar>::work_count(1025, 2049), 8192);
824 assert_eq!(LowRateDecoder::<Scalar>::work_count(32768, 32768), 65536);
825 }
826 }
827}