reed_solomon_simd/rate/
decoder_work.rs

1use fixedbitset::FixedBitSet;
2
3use crate::{
4    engine::{Shards, ShardsRefMut},
5    Error,
6};
7
8// ======================================================================
9// DecoderWork - PUBLIC
10
11/// Working space for [`RateDecoder`].
12///
13/// [`RateDecoder`]: crate::rate::RateDecoder
14pub struct DecoderWork {
15    original_count: usize,
16    recovery_count: usize,
17    shard_bytes: usize,
18
19    original_base_pos: usize,
20    recovery_base_pos: usize,
21
22    original_received_count: usize,
23    recovery_received_count: usize,
24    // May contain extra zero bits.
25    received: FixedBitSet,
26    shards: Shards,
27}
28
29impl DecoderWork {
30    /// Creates new [`DecoderWork`] which initially
31    /// has no working space allocated.
32    pub fn new() -> Self {
33        Self {
34            original_count: 0,
35            recovery_count: 0,
36            shard_bytes: 0,
37
38            original_base_pos: 0,
39            recovery_base_pos: 0,
40
41            original_received_count: 0,
42            recovery_received_count: 0,
43            received: FixedBitSet::new(),
44            shards: Shards::new(),
45        }
46    }
47}
48
49// ======================================================================
50// DecoderWork - IMPL Default
51
52impl Default for DecoderWork {
53    fn default() -> Self {
54        Self::new()
55    }
56}
57
58// ======================================================================
59// DecoderWork - CRATE
60
61impl DecoderWork {
62    pub(crate) fn add_original_shard<T: AsRef<[u8]>>(
63        &mut self,
64        index: usize,
65        original_shard: T,
66    ) -> Result<(), Error> {
67        let pos = self.original_base_pos + index;
68        let original_shard = original_shard.as_ref();
69
70        if index >= self.original_count {
71            Err(Error::InvalidOriginalShardIndex {
72                original_count: self.original_count,
73                index,
74            })
75        } else if self.received[pos] {
76            Err(Error::DuplicateOriginalShardIndex { index })
77        } else if original_shard.len() != self.shard_bytes {
78            Err(Error::DifferentShardSize {
79                shard_bytes: self.shard_bytes,
80                got: original_shard.len(),
81            })
82        } else {
83            self.shards.insert(pos, original_shard);
84
85            self.original_received_count += 1;
86            self.received.set(pos, true);
87            Ok(())
88        }
89    }
90
91    pub(crate) fn add_recovery_shard<T: AsRef<[u8]>>(
92        &mut self,
93        index: usize,
94        recovery_shard: T,
95    ) -> Result<(), Error> {
96        let pos = self.recovery_base_pos + index;
97        let recovery_shard = recovery_shard.as_ref();
98
99        if index >= self.recovery_count {
100            Err(Error::InvalidRecoveryShardIndex {
101                recovery_count: self.recovery_count,
102                index,
103            })
104        } else if self.received[pos] {
105            Err(Error::DuplicateRecoveryShardIndex { index })
106        } else if recovery_shard.len() != self.shard_bytes {
107            Err(Error::DifferentShardSize {
108                shard_bytes: self.shard_bytes,
109                got: recovery_shard.len(),
110            })
111        } else {
112            self.shards.insert(pos, recovery_shard);
113
114            self.recovery_received_count += 1;
115            self.received.set(pos, true);
116            Ok(())
117        }
118    }
119
120    // Begin decode.
121    // - Returned `FixedBitSet` may contain extra zero bits.
122    pub(crate) fn decode_begin(
123        &mut self,
124    ) -> Result<Option<(ShardsRefMut, usize, usize, &FixedBitSet)>, Error> {
125        if self.original_received_count + self.recovery_received_count < self.original_count {
126            Err(Error::NotEnoughShards {
127                original_count: self.original_count,
128                original_received_count: self.original_received_count,
129                recovery_received_count: self.recovery_received_count,
130            })
131        } else if self.original_received_count == self.original_count {
132            Ok(None)
133        } else {
134            Ok(Some((
135                self.shards.as_ref_mut(),
136                self.original_count,
137                self.recovery_count,
138                &self.received,
139            )))
140        }
141    }
142
143    pub(crate) fn original_count(&self) -> usize {
144        self.original_count
145    }
146
147    pub(crate) fn reset(
148        &mut self,
149        original_count: usize,
150        recovery_count: usize,
151        shard_bytes: usize,
152
153        original_base_pos: usize,
154        recovery_base_pos: usize,
155        work_count: usize,
156    ) {
157        assert!(shard_bytes % 2 == 0);
158
159        self.original_count = original_count;
160        self.recovery_count = recovery_count;
161        self.shard_bytes = shard_bytes;
162
163        self.original_base_pos = original_base_pos;
164        self.recovery_base_pos = recovery_base_pos;
165
166        self.original_received_count = 0;
167        self.recovery_received_count = 0;
168
169        let max_received_pos = std::cmp::max(
170            original_base_pos + original_count,
171            recovery_base_pos + recovery_count,
172        );
173
174        self.received.clear();
175        if self.received.len() < max_received_pos {
176            self.received.grow(max_received_pos);
177        }
178
179        self.shards.resize(work_count, shard_bytes.div_ceil(64));
180    }
181
182    pub(crate) fn reset_received(&mut self) {
183        self.original_received_count = 0;
184        self.recovery_received_count = 0;
185        self.received.clear();
186    }
187
188    // This must only be called by `DecoderResult`.
189    pub(crate) fn restored_original(&self, index: usize) -> Option<&[u8]> {
190        let pos = self.original_base_pos + index;
191
192        if index < self.original_count && !self.received[pos] {
193            Some(&self.shards[pos].as_flattened()[..self.shard_bytes])
194        } else {
195            None
196        }
197    }
198
199    pub(crate) fn undo_last_chunk_encoding(&mut self) {
200        self.shards.undo_last_chunk_encoding(
201            self.shard_bytes,
202            self.original_base_pos..self.original_base_pos + self.original_count,
203        );
204    }
205}