diff --git a/src/decoder_result.rs b/src/decoder_result.rs index 452710f..2bfc059 100644 --- a/src/decoder_result.rs +++ b/src/decoder_result.rs @@ -54,7 +54,7 @@ impl Drop for DecoderResult<'_> { /// /// This struct is created by [`DecoderResult::restored_original_iter`]. pub struct RestoredOriginal<'a> { - ended: bool, + remaining: usize, next_index: usize, work: &'a DecoderWork, } @@ -65,30 +65,42 @@ pub struct RestoredOriginal<'a> { impl<'a> Iterator for RestoredOriginal<'a> { type Item = (usize, &'a [u8]); fn next(&mut self) -> Option<(usize, &'a [u8])> { - if self.ended { - None - } else { - let mut index = self.next_index; - while index < self.work.original_count() { - if let Some(original) = self.work.restored_original(index) { - self.next_index = index + 1; - return Some((index, original)); - } - index += 1; + if self.remaining == 0 { + return None; + } + + let mut index = self.next_index; + while index < self.work.original_count() { + if let Some(original) = self.work.restored_original(index) { + self.next_index = index + 1; + self.remaining -= 1; + return Some((index, original)); } - self.ended = true; - None + index += 1; } + + // If we reach this point, there's an inconsistency in the internal data structures. + debug_assert!(false); + None + } + + fn size_hint(&self) -> (usize, Option) { + (self.remaining, Some(self.remaining)) } } +// ====================================================================== +// RestoredOriginal - IMPL ExactSizeIterator + +impl ExactSizeIterator for RestoredOriginal<'_> {} + // ====================================================================== // RestoredOriginal - CRATE impl<'a> RestoredOriginal<'a> { pub(crate) fn new(work: &'a DecoderWork) -> Self { Self { - ended: false, + remaining: work.missing_original_count(), next_index: 0, work, } @@ -150,4 +162,39 @@ mod tests { simple_roundtrip(shard_size); } } + + #[test] + fn decoder_result_size_hint() { + let shard_size = 64; + let original = test_util::generate_original(3, shard_size, 0); + + let mut encoder = ReedSolomonEncoder::new(3, 2, shard_size).unwrap(); + let mut decoder = ReedSolomonDecoder::new(3, 2, shard_size).unwrap(); + + for original in &original { + encoder.add_original_shard(original).unwrap(); + } + + let result = encoder.encode().unwrap(); + let recovery: Vec<_> = result.recovery_iter().collect(); + + decoder.add_original_shard(1, &original[1]).unwrap(); + decoder.add_recovery_shard(0, recovery[0]).unwrap(); + decoder.add_recovery_shard(1, recovery[1]).unwrap(); + + let result: DecoderResult = decoder.decode().unwrap(); + + let mut iter: RestoredOriginal = result.restored_original_iter(); + + assert_eq!(iter.len(), 2); + + assert!(iter.next().is_some()); + assert_eq!(iter.len(), 1); + + assert!(iter.next().is_some()); + assert_eq!(iter.len(), 0); + + assert!(iter.next().is_none()); + assert_eq!(iter.len(), 0); + } } diff --git a/src/encoder_result.rs b/src/encoder_result.rs index 3c84562..8ec349d 100644 --- a/src/encoder_result.rs +++ b/src/encoder_result.rs @@ -79,8 +79,18 @@ impl<'a> Iterator for Recovery<'a> { None } } + + fn size_hint(&self) -> (usize, Option) { + let remaining = self.work.recovery_count() - self.next_index; + (remaining, Some(remaining)) + } } +// ====================================================================== +// Recovery - IMPL ExactSizeIterator + +impl ExactSizeIterator for Recovery<'_> {} + // ====================================================================== // Recovery - CRATE @@ -131,4 +141,30 @@ mod tests { assert!(iter.next().is_none()); test_util::assert_hash(all, test_util::LOW_2_3); } + + #[test] + fn encoder_result_size_hint() { + let original = test_util::generate_original(2, 1024, 123); + let mut encoder = ReedSolomonEncoder::new(2, 3, 1024).unwrap(); + + for original in &original { + encoder.add_original_shard(original).unwrap(); + } + + let result: EncoderResult = encoder.encode().unwrap(); + + let mut iter: Recovery = result.recovery_iter(); + + assert_eq!(iter.len(), 3); + + assert!(iter.next().is_some()); + assert!(iter.next().is_some()); + assert_eq!(iter.len(), 1); + + assert!(iter.next().is_some()); + assert_eq!(iter.len(), 0); + + assert!(iter.next().is_none()); + assert_eq!(iter.len(), 0); + } } diff --git a/src/rate/decoder_work.rs b/src/rate/decoder_work.rs index bd8bcf2..ec44622 100644 --- a/src/rate/decoder_work.rs +++ b/src/rate/decoder_work.rs @@ -202,4 +202,8 @@ impl DecoderWork { self.original_base_pos..self.original_base_pos + self.original_count, ); } + + pub(crate) fn missing_original_count(&self) -> usize { + self.original_count - self.original_received_count + } } diff --git a/src/rate/encoder_work.rs b/src/rate/encoder_work.rs index 8596b15..178a355 100644 --- a/src/rate/encoder_work.rs +++ b/src/rate/encoder_work.rs @@ -120,4 +120,8 @@ impl EncoderWork { self.shards .undo_last_chunk_encoding(self.shard_bytes, 0..self.recovery_count); } + + pub(crate) fn recovery_count(&self) -> usize { + self.recovery_count + } }