Skip to main content

commonware_cryptography/reed_solomon/rate/
encoder_work.rs

1use crate::reed_solomon::{
2    Error,
3    engine::{SHARD_CHUNK_BYTES, Shards, ShardsRefMut},
4};
5
6/// Working space for [`RateEncoder`].
7///
8/// [`RateEncoder`]: crate::reed_solomon::rate::RateEncoder
9pub struct EncoderWork {
10    original_count: usize,
11    recovery_count: usize,
12
13    pub(crate) shard_bytes: usize,
14
15    /// Number of originals added, which is also the position of the next one.
16    original_received_count: usize,
17
18    /// Encoding workspace. Originals are stored at `0..original_count`, and an encode leaves the
19    /// recovery shards at `0..recovery_count`.
20    shards: Shards,
21}
22
23impl EncoderWork {
24    /// Creates a new [`EncoderWork`] with no working space allocated.
25    pub const fn new() -> Self {
26        Self {
27            original_count: 0,
28            recovery_count: 0,
29            shard_bytes: 0,
30
31            original_received_count: 0,
32            shards: Shards::new(),
33        }
34    }
35}
36
37impl Default for EncoderWork {
38    fn default() -> Self {
39        Self::new()
40    }
41}
42
43impl EncoderWork {
44    /// Stores the next original shard at position `original_received_count`.
45    ///
46    /// Returns an error if `original_count` shards were already added or if the shard is not
47    /// `shard_bytes` long.
48    pub(crate) fn add_original_shard<T: AsRef<[u8]>>(
49        &mut self,
50        original_shard: T,
51    ) -> Result<(), Error> {
52        let original_shard = original_shard.as_ref();
53
54        if self.original_received_count == self.original_count {
55            Err(Error::TooManyOriginalShards {
56                original_count: self.original_count,
57            })
58        } else if original_shard.len() != self.shard_bytes {
59            Err(Error::DifferentShardSize {
60                shard_bytes: self.shard_bytes,
61                got: original_shard.len(),
62            })
63        } else {
64            self.shards
65                .insert(self.original_received_count, original_shard);
66
67            self.original_received_count += 1;
68            Ok(())
69        }
70    }
71
72    /// Returns the work shards and the original and recovery counts for an encode.
73    ///
74    /// Returns [`Error::TooFewOriginalShards`] unless all `original_count` originals were added.
75    pub(crate) fn encode_begin(&mut self) -> Result<(ShardsRefMut<'_>, usize, usize), Error> {
76        if self.original_received_count == self.original_count {
77            Ok((
78                self.shards.as_ref_mut(),
79                self.original_count,
80                self.recovery_count,
81            ))
82        } else {
83            Err(Error::TooFewOriginalShards {
84                original_count: self.original_count,
85                original_received_count: self.original_received_count,
86            })
87        }
88    }
89
90    /// Returns recovery shard `index`, or `None` if `index >= recovery_count`.
91    ///
92    /// This must only be called by `EncoderResult`.
93    pub(crate) fn recovery(&self, index: usize) -> Option<&[u8]> {
94        if index < self.recovery_count {
95            Some(&self.shards[index].as_flattened()[..self.shard_bytes])
96        } else {
97            None
98        }
99    }
100
101    /// Configures this work for new shard counts, clears the received originals, and resizes
102    /// `shards` to `work_count` shards.
103    ///
104    /// # Panics
105    ///
106    /// Panics if `shard_bytes` is odd.
107    pub(crate) fn reset(
108        &mut self,
109        original_count: usize,
110        recovery_count: usize,
111        shard_bytes: usize,
112        work_count: usize,
113    ) {
114        assert!(shard_bytes.is_multiple_of(2));
115
116        self.original_count = original_count;
117        self.recovery_count = recovery_count;
118        self.shard_bytes = shard_bytes;
119
120        self.original_received_count = 0;
121        self.shards
122            .resize(work_count, shard_bytes.div_ceil(SHARD_CHUNK_BYTES));
123    }
124
125    /// Clears the received originals so new ones can be added under the same configuration.
126    pub(crate) const fn reset_received(&mut self) {
127        self.original_received_count = 0;
128    }
129
130    /// Undoes the last-chunk encoding of the recovery shards at `0..recovery_count`.
131    pub(crate) fn undo_last_chunk_encoding(&mut self) {
132        self.shards
133            .undo_last_chunk_encoding(self.shard_bytes, 0..self.recovery_count);
134    }
135
136    /// Returns the number of recovery shards.
137    pub(crate) const fn recovery_count(&self) -> usize {
138        self.recovery_count
139    }
140}