ExactSizeIterator for Result types

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<usize>) {
+        (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);
+    }
 }

src/encoder_result.rs

@@ -79,8 +79,18 @@ impl<'a> Iterator for Recovery<'a> {
             None
         }
     }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        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);
+    }
 }

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
+    }
 }

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
+    }
 }