Skip to main content

commonware_cryptography/reed_solomon/
wrappers.rs

1use crate::reed_solomon::{
2    engine::DefaultEngine,
3    rate::{DefaultRate, DefaultRateDecoder, DefaultRateEncoder, Rate, RateDecoder, RateEncoder},
4    DecoderResult, EncoderResult, Error, RecoveryDecoderResult,
5};
6
7// ======================================================================
8// Encoder - PUBLIC
9
10/// Reed-Solomon encoder using [`DefaultEngine`] and [`DefaultRate`].
11///
12/// [`DefaultEngine`]: crate::reed_solomon::engine::DefaultEngine
13pub struct Encoder(DefaultRateEncoder<DefaultEngine>);
14
15impl Encoder {
16    /// Adds one original shard to the encoder.
17    ///
18    /// Original shards have indexes `0..original_count` corresponding to the order
19    /// in which they are added and these same indexes must be used when decoding.
20    ///
21    /// See [basic usage](crate::reed_solomon#basic-usage) for an example.
22    pub fn add_original_shard<T: AsRef<[u8]>>(&mut self, original_shard: T) -> Result<(), Error> {
23        self.0.add_original_shard(original_shard)
24    }
25
26    /// Encodes the added original shards returning [`EncoderResult`]
27    /// which contains the generated recovery shards.
28    ///
29    /// When returned [`EncoderResult`] is dropped the encoder is
30    /// automatically [`reset`] and ready for new round of encoding.
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 new encoder with given configuration
40    /// and allocates required working space.
41    ///
42    /// See [basic usage](crate::reed_solomon#basic-usage) for an example.
43    pub fn new(
44        original_count: usize,
45        recovery_count: usize,
46        shard_bytes: usize,
47    ) -> Result<Self, Error> {
48        Ok(Self(DefaultRateEncoder::new(
49            original_count,
50            recovery_count,
51            shard_bytes,
52            DefaultEngine::new(),
53            None,
54        )?))
55    }
56
57    /// Resets encoder to given configuration.
58    ///
59    /// - Added original shards are forgotten.
60    /// - Existing working space is re-used if it's large enough
61    ///   or re-allocated 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 given `original_count` / `recovery_count`
72    /// combination is supported.
73    ///
74    /// # Examples
75    ///
76    /// ```rust
77    /// use commonware_cryptography::reed_solomon::Encoder;
78    ///
79    /// assert_eq!(Encoder::supports(60_000, 4_000), true);
80    /// assert_eq!(Encoder::supports(60_000, 5_000), false);
81    /// ```
82    pub fn supports(original_count: usize, recovery_count: usize) -> bool {
83        DefaultRate::<DefaultEngine>::supports(original_count, recovery_count)
84    }
85}
86
87// ======================================================================
88// Decoder - PUBLIC
89
90/// Reed-Solomon decoder using [`DefaultEngine`] and [`DefaultRate`].
91///
92/// [`DefaultEngine`]: crate::reed_solomon::engine::DefaultEngine
93pub struct Decoder(DefaultRateDecoder<DefaultEngine>);
94
95impl Decoder {
96    /// Adds one original shard to the decoder.
97    ///
98    /// - Shards can be added in any order.
99    /// - Index must be the same that was used in encoding.
100    ///
101    /// See [basic usage](crate::reed_solomon#basic-usage) for an example.
102    pub fn add_original_shard<T: AsRef<[u8]>>(
103        &mut self,
104        index: usize,
105        original_shard: T,
106    ) -> Result<(), Error> {
107        self.0.add_original_shard(index, original_shard)
108    }
109
110    /// Adds one recovery shard to the decoder.
111    ///
112    /// - Shards can be added in any order.
113    /// - Index must be the same that was used in encoding.
114    ///
115    /// See [basic usage](crate::reed_solomon#basic-usage) for an example.
116    pub fn add_recovery_shard<T: AsRef<[u8]>>(
117        &mut self,
118        index: usize,
119        recovery_shard: T,
120    ) -> Result<(), Error> {
121        self.0.add_recovery_shard(index, recovery_shard)
122    }
123
124    /// Decodes the added shards, returning `Some(`[`DecoderResult`]`)` with the restored original
125    /// shards, or `None` if every original shard was already provided (nothing to reconstruct).
126    ///
127    /// When the returned [`DecoderResult`] is dropped the decoder is automatically [`reset`] and
128    /// ready for a new round of decoding.
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 also reconstructs the missing recovery shards, returning
138    /// `Some(`[`RecoveryDecoderResult`]`)` that additionally exposes them via
139    /// [`RecoveryDecoderResult::recovery`] / [`recovery_iter`]. This costs up to `recovery_count`
140    /// extra field multiplications, so prefer [`decode`](Decoder::decode) when only the original data
141    /// is needed.
142    ///
143    /// [`recovery_iter`]: RecoveryDecoderResult::recovery_iter
144    pub fn decode_with_recovery(&mut self) -> Result<Option<RecoveryDecoderResult<'_>>, Error> {
145        Ok(self.0.decode(true)?.map(RecoveryDecoderResult::new))
146    }
147
148    /// Creates new decoder with given configuration
149    /// and allocates required working space.
150    ///
151    /// See [basic usage](crate::reed_solomon#basic-usage) for an example.
152    pub fn new(
153        original_count: usize,
154        recovery_count: usize,
155        shard_bytes: usize,
156    ) -> Result<Self, Error> {
157        Ok(Self(DefaultRateDecoder::new(
158            original_count,
159            recovery_count,
160            shard_bytes,
161            DefaultEngine::new(),
162            None,
163        )?))
164    }
165
166    /// Resets decoder to given configuration.
167    ///
168    /// - Added shards are forgotten.
169    /// - Existing working space is re-used if it's large enough
170    ///   or re-allocated otherwise.
171    pub fn reset(
172        &mut self,
173        original_count: usize,
174        recovery_count: usize,
175        shard_bytes: usize,
176    ) -> Result<(), Error> {
177        self.0.reset(original_count, recovery_count, shard_bytes)
178    }
179
180    /// Returns `true` if given `original_count` / `recovery_count`
181    /// combination is supported.
182    ///
183    /// # Examples
184    ///
185    /// ```rust
186    /// use commonware_cryptography::reed_solomon::Decoder;
187    ///
188    /// assert_eq!(Decoder::supports(60_000, 4_000), true);
189    /// assert_eq!(Decoder::supports(60_000, 5_000), false);
190    /// ```
191    pub fn supports(original_count: usize, recovery_count: usize) -> bool {
192        DefaultRate::<DefaultEngine>::supports(original_count, recovery_count)
193    }
194}
195
196// ======================================================================
197// TESTS
198
199#[cfg(test)]
200mod tests {
201    use super::*;
202    use crate::reed_solomon::test_util;
203    use fixedbitset::FixedBitSet;
204    use std::collections::BTreeMap;
205
206    // ============================================================
207    // HELPERS
208
209    fn roundtrip(
210        encoder: &mut Encoder,
211        decoder: &mut Decoder,
212        original_count: usize,
213        recovery_hash: &str,
214        decoder_original: &[usize],
215        decoder_recovery: &[usize],
216        seed: u8,
217    ) {
218        let original = test_util::generate_original(original_count, 1024, seed);
219
220        for original in &original {
221            encoder.add_original_shard(original).unwrap();
222        }
223
224        let result = encoder.encode().unwrap();
225        let recovery: Vec<_> = result.recovery_iter().collect();
226
227        test_util::assert_hash(&recovery, recovery_hash);
228
229        let mut original_received = FixedBitSet::with_capacity(original_count);
230
231        for i in decoder_original {
232            decoder.add_original_shard(*i, &original[*i]).unwrap();
233            original_received.set(*i, true);
234        }
235
236        for i in decoder_recovery {
237            decoder.add_recovery_shard(*i, recovery[*i]).unwrap();
238        }
239
240        let decoded = decoder.decode();
241        let restored: BTreeMap<_, _> = match &decoded {
242            Ok(Some(result)) => result.original_iter().collect(),
243            Ok(None) => BTreeMap::new(),
244            Err(e) => panic!("decode failed: {e:?}"),
245        };
246
247        for i in 0..original_count {
248            if !original_received[i] {
249                assert_eq!(restored[&i], original[i]);
250            }
251        }
252    }
253
254    // ============================================================
255    // ROUNDTRIP - TWO ROUNDS
256
257    #[test]
258    fn roundtrip_two_rounds_reset_low_to_high() {
259        let mut encoder = Encoder::new(2, 3, 1024).unwrap();
260        let mut decoder = Decoder::new(2, 3, 1024).unwrap();
261
262        roundtrip(
263            &mut encoder,
264            &mut decoder,
265            2,
266            test_util::LOW_2_3,
267            &[],
268            &[0, 1],
269            123,
270        );
271
272        encoder.reset(3, 2, 1024).unwrap();
273        decoder.reset(3, 2, 1024).unwrap();
274
275        roundtrip(
276            &mut encoder,
277            &mut decoder,
278            3,
279            test_util::HIGH_3_2,
280            &[1],
281            &[0, 1],
282            132,
283        );
284    }
285
286    #[test]
287    fn failed_encoder_reset_preserves_state() {
288        let original = test_util::generate_original(2, 1024, 123);
289        let mut encoder = Encoder::new(2, 3, 1024).unwrap();
290
291        assert_eq!(
292            encoder.reset(3, 2, 3),
293            Err(Error::InvalidShardSize { shard_bytes: 3 })
294        );
295
296        for shard in &original {
297            encoder.add_original_shard(shard).unwrap();
298        }
299        let result = encoder.encode().unwrap();
300        let recovery: Vec<_> = result.recovery_iter().collect();
301
302        test_util::assert_hash(&recovery, test_util::LOW_2_3);
303    }
304
305    #[test]
306    fn failed_decoder_reset_preserves_state() {
307        let original = test_util::generate_original(2, 1024, 123);
308        let mut encoder = Encoder::new(2, 3, 1024).unwrap();
309        for shard in &original {
310            encoder.add_original_shard(shard).unwrap();
311        }
312        let result = encoder.encode().unwrap();
313        let recovery: Vec<_> = result.recovery_iter().map(<[u8]>::to_vec).collect();
314
315        let mut decoder = Decoder::new(2, 3, 1024).unwrap();
316
317        assert_eq!(
318            decoder.reset(3, 2, 3),
319            Err(Error::InvalidShardSize { shard_bytes: 3 })
320        );
321
322        decoder.add_recovery_shard(0, &recovery[0]).unwrap();
323        decoder.add_recovery_shard(1, &recovery[1]).unwrap();
324        let decoded = decoder.decode();
325        let restored: BTreeMap<_, _> = match &decoded {
326            Ok(Some(result)) => result.original_iter().collect(),
327            Ok(None) => BTreeMap::new(),
328            Err(e) => panic!("decode failed: {e:?}"),
329        };
330
331        assert_eq!(restored[&0], original[0]);
332        assert_eq!(restored[&1], original[1]);
333    }
334
335    // ==================================================
336    // supports
337
338    #[test]
339    fn supports() {
340        assert!(Encoder::supports(4096, 61440));
341        assert!(Encoder::supports(61440, 4096));
342
343        assert!(Decoder::supports(4096, 61440));
344        assert!(Decoder::supports(61440, 4096));
345    }
346}