Skip to main content

commonware_cryptography/reed_solomon/
encoder_result.rs

1use crate::reed_solomon::rate::EncoderWork;
2
3/// Result of encoding. Contains the generated recovery shards.
4///
5/// This struct is created by [`Encoder::encode`] and [`RateEncoder::encode`].
6///
7/// [`RateEncoder::encode`]: crate::reed_solomon::rate::RateEncoder::encode
8/// [`Encoder::encode`]: crate::reed_solomon::Encoder::encode
9pub struct EncoderResult<'a> {
10    work: &'a mut EncoderWork,
11}
12
13impl EncoderResult<'_> {
14    /// Returns the recovery shard with the given `index`, or `None` if `index >= recovery_count`.
15    ///
16    /// Recovery shards have indexes `0..recovery_count`, and these same indexes must be used when
17    /// decoding.
18    pub fn recovery(&self, index: usize) -> Option<&[u8]> {
19        self.work.recovery(index)
20    }
21
22    /// Returns an iterator over all recovery shards, ordered by index.
23    ///
24    /// Recovery shards have indexes `0..recovery_count`, and these same indexes must be used when
25    /// decoding.
26    pub const fn recovery_iter(&self) -> Recovery<'_> {
27        Recovery::new(self.work)
28    }
29}
30
31impl<'a> EncoderResult<'a> {
32    /// Wraps `work` after an encode that wrote the recovery shards.
33    ///
34    /// Dropping the result clears the received original count in `work`.
35    pub(crate) const fn new(work: &'a mut EncoderWork) -> Self {
36        Self { work }
37    }
38}
39
40impl Drop for EncoderResult<'_> {
41    fn drop(&mut self) {
42        self.work.reset_received();
43    }
44}
45
46/// Iterator over generated recovery shards.
47///
48/// This struct is created by [`EncoderResult::recovery_iter`].
49pub struct Recovery<'a> {
50    /// Set once `next` has returned `None`.
51    ended: bool,
52    /// Index of the next recovery shard to yield.
53    next_index: usize,
54    work: &'a EncoderWork,
55}
56
57impl<'a> Iterator for Recovery<'a> {
58    type Item = &'a [u8];
59    fn next(&mut self) -> Option<&'a [u8]> {
60        if self.ended {
61            None
62        } else if let Some(next) = self.work.recovery(self.next_index) {
63            self.next_index += 1;
64            Some(next)
65        } else {
66            self.ended = true;
67            None
68        }
69    }
70
71    fn size_hint(&self) -> (usize, Option<usize>) {
72        let remaining = self.work.recovery_count() - self.next_index;
73        (remaining, Some(remaining))
74    }
75}
76
77impl ExactSizeIterator for Recovery<'_> {}
78
79impl<'a> Recovery<'a> {
80    /// Creates an iterator over all recovery shards in `work`, starting at index 0.
81    pub(crate) const fn new(work: &'a EncoderWork) -> Self {
82        Self {
83            ended: false,
84            next_index: 0,
85            work,
86        }
87    }
88}
89
90#[cfg(test)]
91mod tests {
92    use super::*;
93    use crate::reed_solomon::{Encoder, test_util};
94
95    /// Covers `EncoderResult::recovery`, `EncoderResult::recovery_iter`, and `Recovery`.
96    #[test]
97    fn encoder_result() {
98        let original = test_util::generate_original(2, 1024, 123);
99        let mut encoder = Encoder::new(2, 3, 1024).unwrap();
100
101        for original in &original {
102            encoder.add_original_shard(original).unwrap();
103        }
104
105        let result: EncoderResult<'_> = encoder.encode().unwrap();
106
107        let all = vec![
108            result.recovery(0).unwrap(),
109            result.recovery(1).unwrap(),
110            result.recovery(2).unwrap(),
111        ];
112        assert!(result.recovery(3).is_none());
113        test_util::assert_hash(all, test_util::LOW_2_3);
114
115        let mut iter: Recovery<'_> = result.recovery_iter();
116        let all = vec![
117            iter.next().unwrap(),
118            iter.next().unwrap(),
119            iter.next().unwrap(),
120        ];
121        assert!(iter.next().is_none());
122        test_util::assert_hash(all, test_util::LOW_2_3);
123    }
124
125    #[test]
126    fn encoder_result_size_hint() {
127        let original = test_util::generate_original(2, 1024, 123);
128        let mut encoder = Encoder::new(2, 3, 1024).unwrap();
129
130        for original in &original {
131            encoder.add_original_shard(original).unwrap();
132        }
133
134        let result: EncoderResult<'_> = encoder.encode().unwrap();
135
136        let mut iter: Recovery<'_> = result.recovery_iter();
137
138        assert_eq!(iter.len(), 3);
139
140        assert!(iter.next().is_some());
141        assert!(iter.next().is_some());
142        assert_eq!(iter.len(), 1);
143
144        assert!(iter.next().is_some());
145        assert_eq!(iter.len(), 0);
146
147        assert!(iter.next().is_none());
148        assert_eq!(iter.len(), 0);
149    }
150}