Skip to main content

commonware_cryptography/reed_solomon/
wrappers.rs

1use crate::reed_solomon::{
2    DecoderResult, EncoderResult, Error, Plan, RecoveryDecoderResult,
3    engine::DefaultEngine,
4    rate::{DefaultRate, DefaultRateDecoder, DefaultRateEncoder, Rate, RateDecoder, RateEncoder},
5};
6
7/// Reed-Solomon encoder using [`DefaultEngine`] and [`DefaultRate`].
8///
9/// [`DefaultEngine`]: crate::reed_solomon::engine::DefaultEngine
10pub struct Encoder(DefaultRateEncoder<DefaultEngine>);
11
12impl Encoder {
13    /// Adds one original shard to the encoder.
14    ///
15    /// Original shards have indexes `0..original_count` corresponding to the order in which they
16    /// are added, and these same indexes must be used when decoding.
17    ///
18    /// See [basic usage](crate::reed_solomon#basic-usage) for an example.
19    pub fn add_original_shard<T: AsRef<[u8]>>(&mut self, original_shard: T) -> Result<(), Error> {
20        self.0.add_original_shard(original_shard)
21    }
22
23    /// Encodes the added original shards, returning [`EncoderResult`] which contains the
24    /// generated recovery shards.
25    ///
26    /// When the returned [`EncoderResult`] is dropped the encoder is automatically [`reset`] and
27    /// ready for a new round of encoding. A result that is leaked instead of dropped (for example
28    /// with [`core::mem::forget`]) leaves the encoder reporting a complete round, and the next
29    /// [`encode`](Encoder::encode) would transform the previous round's buffers again. Call
30    /// [`reset`] before reusing an encoder whose result was leaked.
31    ///
32    /// See [basic usage](crate::reed_solomon#basic-usage) for an example.
33    ///
34    /// [`reset`]: Encoder::reset
35    pub fn encode(&mut self) -> Result<EncoderResult<'_>, Error> {
36        self.0.encode()
37    }
38
39    /// Creates a new encoder with the given configuration and allocates the required working
40    /// space.
41    ///
42    /// See [working space](crate::reed_solomon#working-space) for its size and the configurations
43    /// a coder accepts, and [basic usage](crate::reed_solomon#basic-usage) for an example.
44    pub fn new(
45        original_count: usize,
46        recovery_count: usize,
47        shard_bytes: usize,
48    ) -> Result<Self, Error> {
49        Ok(Self(DefaultRateEncoder::new(
50            original_count,
51            recovery_count,
52            shard_bytes,
53            DefaultEngine::new(),
54            None,
55        )?))
56    }
57
58    /// Resets the encoder to the given configuration.
59    ///
60    /// - Added original shards are forgotten.
61    /// - Existing working space is reused if it is large enough, or reallocated otherwise.
62    pub fn reset(
63        &mut self,
64        original_count: usize,
65        recovery_count: usize,
66        shard_bytes: usize,
67    ) -> Result<(), Error> {
68        self.0.reset(original_count, recovery_count, shard_bytes)
69    }
70
71    /// Returns `true` if the given `original_count` / `recovery_count` combination is supported.
72    ///
73    /// # Examples
74    ///
75    /// ```rust
76    /// use commonware_cryptography::reed_solomon::Encoder;
77    ///
78    /// assert_eq!(Encoder::supports(60_000, 4_000), true);
79    /// assert_eq!(Encoder::supports(60_000, 5_000), false);
80    /// ```
81    pub fn supports(original_count: usize, recovery_count: usize) -> bool {
82        DefaultRate::<DefaultEngine>::supports(original_count, recovery_count)
83    }
84}
85
86/// Reed-Solomon decoder using [`DefaultEngine`] and [`DefaultRate`].
87///
88/// [`DefaultEngine`]: crate::reed_solomon::engine::DefaultEngine
89pub struct Decoder(DefaultRateDecoder<DefaultEngine>);
90
91impl Decoder {
92    /// Adds one original shard to the decoder.
93    ///
94    /// - Shards can be added in any order.
95    /// - The index must be the same one used in encoding.
96    ///
97    /// See [basic usage](crate::reed_solomon#basic-usage) for an example.
98    pub fn add_original_shard<T: AsRef<[u8]>>(
99        &mut self,
100        index: usize,
101        original_shard: T,
102    ) -> Result<(), Error> {
103        self.0.add_original_shard(index, original_shard)
104    }
105
106    /// Adds one recovery shard to the decoder.
107    ///
108    /// - Shards can be added in any order.
109    /// - The index must be the same one used in encoding.
110    ///
111    /// See [basic usage](crate::reed_solomon#basic-usage) for an example.
112    pub fn add_recovery_shard<T: AsRef<[u8]>>(
113        &mut self,
114        index: usize,
115        recovery_shard: T,
116    ) -> Result<(), Error> {
117        self.0.add_recovery_shard(index, recovery_shard)
118    }
119
120    /// Decodes the added shards, returning `Some(`[`DecoderResult`]`)` with the restored original
121    /// shards, or `None` if every original shard was already provided (nothing to reconstruct).
122    ///
123    /// When the returned [`DecoderResult`] is dropped the decoder is automatically [`reset`] and
124    /// ready for a new round of decoding. When every original shard was supplied, this returns
125    /// `None` and clears the received shards immediately. A result that is leaked instead of
126    /// dropped (for example with [`core::mem::forget`]) leaves the decoder reporting a complete
127    /// round, and the next decode would transform the previous round's buffers again. Call
128    /// [`reset`] before reusing a decoder whose result was leaked.
129    ///
130    /// See [basic usage](crate::reed_solomon#basic-usage) for an example.
131    ///
132    /// [`reset`]: Decoder::reset
133    pub fn decode(&mut self) -> Result<Option<DecoderResult<'_>>, Error> {
134        self.0.decode(false)
135    }
136
137    /// Like [`decode`](Decoder::decode), but reads erasure coefficients from `plan` instead of
138    /// computing them.
139    ///
140    /// `plan` must match this decoder's shard counts and received shard indices. Shard byte
141    /// length is not part of the plan. A mismatched plan returns
142    /// [`Error::PlanMismatch`] and leaves the added shards in place, so the decoder can
143    /// retry with the matching plan or [`decode`](Decoder::decode).
144    pub fn decode_with_plan(&mut self, plan: &Plan) -> Result<Option<DecoderResult<'_>>, Error> {
145        self.0.decode_with_plan(false, plan)
146    }
147
148    /// Like [`decode`](Decoder::decode), but also reconstructs the missing recovery shards,
149    /// returning `Some(`[`RecoveryDecoderResult`]`)` that additionally exposes them via
150    /// [`RecoveryDecoderResult::recovery`] / [`recovery_iter`]. This adds a multiplication of each
151    /// missing recovery shard by a field element, so prefer [`decode`](Decoder::decode) when only
152    /// the original data is needed.
153    ///
154    /// [`recovery_iter`]: RecoveryDecoderResult::recovery_iter
155    pub fn decode_with_recovery(&mut self) -> Result<Option<RecoveryDecoderResult<'_>>, Error> {
156        Ok(self.0.decode(true)?.map(RecoveryDecoderResult::new))
157    }
158
159    /// Like [`decode_with_recovery`](Decoder::decode_with_recovery), but reads erasure
160    /// coefficients from `plan` as [`decode_with_plan`](Decoder::decode_with_plan) does.
161    pub fn decode_with_recovery_plan(
162        &mut self,
163        plan: &Plan,
164    ) -> Result<Option<RecoveryDecoderResult<'_>>, Error> {
165        Ok(self
166            .0
167            .decode_with_plan(true, plan)?
168            .map(RecoveryDecoderResult::new))
169    }
170
171    /// Creates a new decoder with the given configuration and allocates the required working
172    /// space.
173    ///
174    /// See [working space](crate::reed_solomon#working-space) for its size and the configurations
175    /// a coder accepts, and [basic usage](crate::reed_solomon#basic-usage) for an example.
176    pub fn new(
177        original_count: usize,
178        recovery_count: usize,
179        shard_bytes: usize,
180    ) -> Result<Self, Error> {
181        Ok(Self(DefaultRateDecoder::new(
182            original_count,
183            recovery_count,
184            shard_bytes,
185            DefaultEngine::new(),
186            None,
187        )?))
188    }
189
190    /// Resets the decoder to the given configuration.
191    ///
192    /// - Added shards are forgotten.
193    /// - Existing working space is reused if it is large enough, or reallocated otherwise.
194    pub fn reset(
195        &mut self,
196        original_count: usize,
197        recovery_count: usize,
198        shard_bytes: usize,
199    ) -> Result<(), Error> {
200        self.0.reset(original_count, recovery_count, shard_bytes)
201    }
202
203    /// Returns `true` if the given `original_count` / `recovery_count` combination is supported.
204    ///
205    /// # Examples
206    ///
207    /// ```rust
208    /// use commonware_cryptography::reed_solomon::Decoder;
209    ///
210    /// assert_eq!(Decoder::supports(60_000, 4_000), true);
211    /// assert_eq!(Decoder::supports(60_000, 5_000), false);
212    /// ```
213    pub fn supports(original_count: usize, recovery_count: usize) -> bool {
214        DefaultRate::<DefaultEngine>::supports(original_count, recovery_count)
215    }
216}
217
218#[cfg(test)]
219mod tests {
220    use super::*;
221    use crate::reed_solomon::test_util;
222    use fixedbitset::FixedBitSet;
223    use std::collections::BTreeMap;
224
225    /// Encodes `original_count` 1024-byte originals, checks the recovery hash, decodes from the
226    /// given original and recovery indices, and checks every original not given to the decoder.
227    fn roundtrip(
228        encoder: &mut Encoder,
229        decoder: &mut Decoder,
230        original_count: usize,
231        recovery_hash: &str,
232        decoder_original: &[usize],
233        decoder_recovery: &[usize],
234        seed: u8,
235    ) {
236        let original = test_util::generate_original(original_count, 1024, seed);
237
238        for original in &original {
239            encoder.add_original_shard(original).unwrap();
240        }
241
242        let result = encoder.encode().unwrap();
243        let recovery: Vec<_> = result.recovery_iter().collect();
244
245        test_util::assert_hash(&recovery, recovery_hash);
246
247        let mut original_received = FixedBitSet::with_capacity(original_count);
248
249        for i in decoder_original {
250            decoder.add_original_shard(*i, &original[*i]).unwrap();
251            original_received.set(*i, true);
252        }
253
254        for i in decoder_recovery {
255            decoder.add_recovery_shard(*i, recovery[*i]).unwrap();
256        }
257
258        let decoded = decoder.decode();
259        let restored: BTreeMap<_, _> = match &decoded {
260            Ok(Some(result)) => result.original_iter().collect(),
261            Ok(None) => BTreeMap::new(),
262            Err(e) => panic!("decode failed: {e:?}"),
263        };
264
265        for i in 0..original_count {
266            if !original_received[i] {
267                assert_eq!(restored[&i], original[i]);
268            }
269        }
270    }
271
272    #[test]
273    fn maximum_decoder_indices() {
274        for (original_count, recovery_count) in [(3, 2), (2, 3)] {
275            let original = test_util::generate_original(original_count, 2, 123);
276            let mut encoder = Encoder::new(original_count, recovery_count, 2).unwrap();
277            for shard in &original {
278                encoder.add_original_shard(shard).unwrap();
279            }
280            let encoded = encoder.encode().unwrap();
281            let mut decoder = Decoder::new(original_count, recovery_count, 2).unwrap();
282            assert_eq!(
283                decoder.add_original_shard(usize::MAX, [0; 2]),
284                Err(Error::InvalidOriginalShardIndex {
285                    original_count,
286                    index: usize::MAX
287                })
288            );
289            assert_eq!(
290                decoder.add_recovery_shard(usize::MAX, [0; 2]),
291                Err(Error::InvalidRecoveryShardIndex {
292                    recovery_count,
293                    index: usize::MAX
294                })
295            );
296            for (index, shard) in original.iter().enumerate().skip(1) {
297                decoder.add_original_shard(index, shard).unwrap();
298            }
299            decoder
300                .add_recovery_shard(0, encoded.recovery(0).unwrap())
301                .unwrap();
302            let decoded = decoder.decode_with_recovery().unwrap().unwrap();
303            assert_eq!(decoded.original(0), Some(original[0].as_slice()));
304            assert_eq!(decoded.original(usize::MAX), None);
305            assert_eq!(decoded.recovery(usize::MAX), None);
306        }
307    }
308
309    #[test]
310    fn unrepresentable_working_space() {
311        for (original_count, recovery_count) in [(3, 2), (2, 3)] {
312            for shard_bytes in [usize::MAX - 1, 1 << (usize::BITS - 2)] {
313                assert!(matches!(
314                    Encoder::new(original_count, recovery_count, shard_bytes),
315                    Err(Error::InvalidShardSize { .. })
316                ));
317                assert!(matches!(
318                    Decoder::new(original_count, recovery_count, shard_bytes),
319                    Err(Error::InvalidShardSize { .. })
320                ));
321            }
322        }
323    }
324
325    #[test]
326    fn roundtrip_two_rounds_reset_low_to_high() {
327        let mut encoder = Encoder::new(2, 3, 1024).unwrap();
328        let mut decoder = Decoder::new(2, 3, 1024).unwrap();
329
330        roundtrip(
331            &mut encoder,
332            &mut decoder,
333            2,
334            test_util::LOW_2_3,
335            &[],
336            &[0, 1],
337            123,
338        );
339
340        encoder.reset(3, 2, 1024).unwrap();
341        decoder.reset(3, 2, 1024).unwrap();
342
343        roundtrip(
344            &mut encoder,
345            &mut decoder,
346            3,
347            test_util::HIGH_3_2,
348            &[1],
349            &[0, 1],
350            132,
351        );
352    }
353
354    #[test]
355    fn failed_encoder_reset_preserves_state() {
356        let original = test_util::generate_original(2, 1024, 123);
357        let mut encoder = Encoder::new(2, 3, 1024).unwrap();
358        encoder.add_original_shard(&original[0]).unwrap();
359
360        for (original_count, recovery_count) in [(2, 3), (3, 2)] {
361            for shard_bytes in [0, 3, usize::MAX - 1] {
362                assert_eq!(
363                    encoder.reset(original_count, recovery_count, shard_bytes),
364                    Err(Error::InvalidShardSize { shard_bytes })
365                );
366            }
367        }
368
369        for shard in original.iter().skip(1) {
370            encoder.add_original_shard(shard).unwrap();
371        }
372        let result = encoder.encode().unwrap();
373        let recovery: Vec<_> = result.recovery_iter().collect();
374
375        test_util::assert_hash(&recovery, test_util::LOW_2_3);
376    }
377
378    #[test]
379    fn failed_decoder_reset_preserves_state() {
380        let original = test_util::generate_original(2, 1024, 123);
381        let mut encoder = Encoder::new(2, 3, 1024).unwrap();
382        for shard in &original {
383            encoder.add_original_shard(shard).unwrap();
384        }
385        let result = encoder.encode().unwrap();
386        let recovery: Vec<_> = result.recovery_iter().map(<[u8]>::to_vec).collect();
387
388        let mut decoder = Decoder::new(2, 3, 1024).unwrap();
389        decoder.add_recovery_shard(0, &recovery[0]).unwrap();
390
391        for (original_count, recovery_count) in [(2, 3), (3, 2)] {
392            for shard_bytes in [0, 3, usize::MAX - 1] {
393                assert_eq!(
394                    decoder.reset(original_count, recovery_count, shard_bytes),
395                    Err(Error::InvalidShardSize { shard_bytes })
396                );
397            }
398        }
399
400        decoder.add_recovery_shard(1, &recovery[1]).unwrap();
401        let decoded = decoder.decode();
402        let restored: BTreeMap<_, _> = match &decoded {
403            Ok(Some(result)) => result.original_iter().collect(),
404            Ok(None) => BTreeMap::new(),
405            Err(e) => panic!("decode failed: {e:?}"),
406        };
407
408        assert_eq!(restored[&0], original[0]);
409        assert_eq!(restored[&1], original[1]);
410    }
411
412    #[test]
413    fn supports() {
414        assert!(Encoder::supports(4096, 61440));
415        assert!(Encoder::supports(61440, 4096));
416
417        assert!(Decoder::supports(4096, 61440));
418        assert!(Decoder::supports(61440, 4096));
419    }
420}