1use crate::reed_solomon::{
2 DecoderResult, EncoderResult, Error, Plan, RecoveryDecoderResult,
3 engine::DefaultEngine,
4 rate::{DefaultRate, DefaultRateDecoder, DefaultRateEncoder, Rate, RateDecoder, RateEncoder},
5};
6
7pub struct Encoder(DefaultRateEncoder<DefaultEngine>);
11
12impl Encoder {
13 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 pub fn encode(&mut self) -> Result<EncoderResult<'_>, Error> {
36 self.0.encode()
37 }
38
39 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 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 pub fn supports(original_count: usize, recovery_count: usize) -> bool {
82 DefaultRate::<DefaultEngine>::supports(original_count, recovery_count)
83 }
84}
85
86pub struct Decoder(DefaultRateDecoder<DefaultEngine>);
90
91impl Decoder {
92 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 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 pub fn decode(&mut self) -> Result<Option<DecoderResult<'_>>, Error> {
134 self.0.decode(false)
135 }
136
137 pub fn decode_with_plan(&mut self, plan: &Plan) -> Result<Option<DecoderResult<'_>>, Error> {
145 self.0.decode_with_plan(false, plan)
146 }
147
148 pub fn decode_with_recovery(&mut self) -> Result<Option<RecoveryDecoderResult<'_>>, Error> {
156 Ok(self.0.decode(true)?.map(RecoveryDecoderResult::new))
157 }
158
159 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 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 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 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 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}