commonware_cryptography/reed_solomon/engine/
shards.rs

1use crate::reed_solomon::engine::SHARD_CHUNK_BYTES;
2#[cfg(not(feature = "std"))]
3use alloc::vec::Vec;
4use core::ops::{Bound, Index, IndexMut, Range, RangeBounds};
5
6// ======================================================================
7// Shards - CRATE
8
9pub(crate) struct Shards {
10    shard_count: usize,
11    // Shard length in `SHARD_CHUNK_BYTES` chunks.
12    shard_chunk_count: usize,
13
14    // Flat Vec of `shard_count * shard_chunk_count * SHARD_CHUNK_BYTES` bytes.
15    data: Vec<[u8; SHARD_CHUNK_BYTES]>,
16}
17
18impl Shards {
19    pub(crate) fn as_ref_mut(&mut self) -> ShardsRefMut<'_> {
20        ShardsRefMut::new(self.shard_count, self.shard_chunk_count, self.data.as_mut())
21    }
22
23    pub(crate) const fn new() -> Self {
24        Self {
25            shard_count: 0,
26            shard_chunk_count: 0,
27            data: Vec::new(),
28        }
29    }
30
31    pub(crate) fn resize(&mut self, shard_count: usize, shard_chunk_count: usize) {
32        self.shard_count = shard_count;
33        self.shard_chunk_count = shard_chunk_count;
34
35        self.data.resize(
36            self.shard_count * self.shard_chunk_count,
37            [0; SHARD_CHUNK_BYTES],
38        );
39    }
40
41    pub(crate) fn insert(&mut self, index: usize, shard: &[u8]) {
42        assert_eq!(shard.len() % 2, 0);
43
44        let whole_chunk_count = shard.len() / SHARD_CHUNK_BYTES;
45        let tail_len = shard.len() % SHARD_CHUNK_BYTES;
46
47        let (src_chunks, src_tail) = shard.split_at(shard.len() - tail_len);
48
49        let dst = &mut self[index];
50        dst[..whole_chunk_count]
51            .as_flattened_mut()
52            .copy_from_slice(src_chunks);
53
54        // Last chunk is special if shard.len() % SHARD_CHUNK_BYTES != 0.
55        // See src/algorithm.md for an explanation.
56        if tail_len > 0 {
57            let (src_lo, src_hi) = src_tail.split_at(tail_len / 2);
58            let (dst_lo, dst_hi) = dst[whole_chunk_count].split_at_mut(SHARD_CHUNK_BYTES / 2);
59            dst_lo[..src_lo.len()].copy_from_slice(src_lo);
60            dst_hi[..src_hi.len()].copy_from_slice(src_hi);
61        }
62    }
63
64    // Undoes the encoding of the last chunk for the given range of shards
65    pub(crate) fn undo_last_chunk_encoding(&mut self, shard_bytes: usize, range: Range<usize>) {
66        let whole_chunk_count = shard_bytes / SHARD_CHUNK_BYTES;
67        let tail_len = shard_bytes % SHARD_CHUNK_BYTES;
68
69        if tail_len == 0 {
70            return;
71        }
72
73        for idx in range {
74            let last_chunk = &mut self[idx][whole_chunk_count];
75            last_chunk.copy_within(
76                SHARD_CHUNK_BYTES / 2..SHARD_CHUNK_BYTES / 2 + tail_len / 2,
77                tail_len / 2,
78            );
79        }
80    }
81}
82
83// ======================================================================
84// Shards - IMPL Index
85
86impl Index<usize> for Shards {
87    type Output = [[u8; SHARD_CHUNK_BYTES]];
88    fn index(&self, index: usize) -> &Self::Output {
89        &self.data[index * self.shard_chunk_count..(index + 1) * self.shard_chunk_count]
90    }
91}
92
93// ======================================================================
94// Shards - IMPL IndexMut
95
96impl IndexMut<usize> for Shards {
97    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
98        &mut self.data[index * self.shard_chunk_count..(index + 1) * self.shard_chunk_count]
99    }
100}
101
102// ======================================================================
103// ShardsRefMut - PUBLIC
104
105/// Mutable reference to a shard array.
106pub struct ShardsRefMut<'a> {
107    shard_count: usize,
108    shard_chunk_count: usize,
109
110    data: &'a mut [[u8; SHARD_CHUNK_BYTES]],
111}
112
113type FourShardsMut<'a> = (
114    &'a mut [[u8; SHARD_CHUNK_BYTES]],
115    &'a mut [[u8; SHARD_CHUNK_BYTES]],
116    &'a mut [[u8; SHARD_CHUNK_BYTES]],
117    &'a mut [[u8; SHARD_CHUNK_BYTES]],
118);
119
120impl<'a> ShardsRefMut<'a> {
121    /// Returns mutable references to shards at `pos` and `pos + dist`.
122    ///
123    /// See source code of [`Naive::fft`] for an example.
124    ///
125    /// # Panics
126    ///
127    /// If `dist` is `0`.
128    ///
129    /// [`Naive::fft`]: crate::reed_solomon::engine::Naive#method.fft
130    pub fn dist2_mut(
131        &mut self,
132        mut pos: usize,
133        mut dist: usize,
134    ) -> (
135        &mut [[u8; SHARD_CHUNK_BYTES]],
136        &mut [[u8; SHARD_CHUNK_BYTES]],
137    ) {
138        pos *= self.shard_chunk_count;
139        dist *= self.shard_chunk_count;
140
141        let (a, b) = self.data[pos..].split_at_mut(dist);
142        (
143            &mut a[..self.shard_chunk_count],
144            &mut b[..self.shard_chunk_count],
145        )
146    }
147
148    /// Returns mutable references to shards at
149    /// `pos`, `pos + dist`, `pos + dist * 2` and `pos + dist * 3`.
150    ///
151    /// See source code of [`NoSimd::fft`] for an example
152    /// (specifically the private method `fft_butterfly_two_layers`).
153    ///
154    /// # Panics
155    ///
156    /// If `dist` is `0`.
157    ///
158    /// [`NoSimd::fft`]: crate::reed_solomon::engine::NoSimd#method.fft
159    pub fn dist4_mut(&mut self, mut pos: usize, mut dist: usize) -> FourShardsMut<'_> {
160        pos *= self.shard_chunk_count;
161        dist *= self.shard_chunk_count;
162
163        let (ab, cd) = self.data[pos..].split_at_mut(dist * 2);
164        let (a, b) = ab.split_at_mut(dist);
165        let (c, d) = cd.split_at_mut(dist);
166
167        (
168            &mut a[..self.shard_chunk_count],
169            &mut b[..self.shard_chunk_count],
170            &mut c[..self.shard_chunk_count],
171            &mut d[..self.shard_chunk_count],
172        )
173    }
174
175    /// Returns `true` if this contains no shards.
176    pub const fn is_empty(&self) -> bool {
177        self.shard_count == 0
178    }
179
180    /// Returns number of shards.
181    pub const fn len(&self) -> usize {
182        self.shard_count
183    }
184
185    /// Creates new [`ShardsRefMut`] that references given `data`.
186    ///
187    /// # Panics
188    ///
189    /// If `data.len() < shard_count * shard_chunk_count`.
190    pub fn new(
191        shard_count: usize,
192        shard_chunk_count: usize,
193        data: &'a mut [[u8; SHARD_CHUNK_BYTES]],
194    ) -> Self {
195        assert!(data.len() >= shard_count * shard_chunk_count);
196
197        Self {
198            shard_count,
199            shard_chunk_count,
200            data: &mut data[..shard_count * shard_chunk_count],
201        }
202    }
203
204    /// Splits this [`ShardsRefMut`] into two so that
205    /// first includes shards `0..mid` and second includes shards `mid..`.
206    pub fn split_at_mut(&mut self, mid: usize) -> (ShardsRefMut<'_>, ShardsRefMut<'_>) {
207        let (a, b) = self.data.split_at_mut(mid * self.shard_chunk_count);
208
209        (
210            ShardsRefMut::new(mid, self.shard_chunk_count, a),
211            ShardsRefMut::new(self.shard_count - mid, self.shard_chunk_count, b),
212        )
213    }
214
215    /// Fills the given shard-range with `0u8`:s.
216    pub fn zero<R: RangeBounds<usize>>(&mut self, range: R) {
217        let start = match range.start_bound() {
218            Bound::Included(start) => start * self.shard_chunk_count,
219            Bound::Excluded(start) => (start + 1) * self.shard_chunk_count,
220            Bound::Unbounded => 0,
221        };
222
223        let end = match range.end_bound() {
224            Bound::Included(end) => (end + 1) * self.shard_chunk_count,
225            Bound::Excluded(end) => end * self.shard_chunk_count,
226            Bound::Unbounded => self.shard_count * self.shard_chunk_count,
227        };
228
229        self.data[start..end].fill([0; SHARD_CHUNK_BYTES]);
230    }
231}
232
233// ======================================================================
234// ShardsRefMut - IMPL Index
235
236impl Index<usize> for ShardsRefMut<'_> {
237    type Output = [[u8; SHARD_CHUNK_BYTES]];
238    fn index(&self, index: usize) -> &Self::Output {
239        &self.data[index * self.shard_chunk_count..(index + 1) * self.shard_chunk_count]
240    }
241}
242
243// ======================================================================
244// ShardsRefMut - IMPL IndexMut
245
246impl IndexMut<usize> for ShardsRefMut<'_> {
247    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
248        &mut self.data[index * self.shard_chunk_count..(index + 1) * self.shard_chunk_count]
249    }
250}
251
252// ======================================================================
253// ShardsRefMut - CRATE
254
255impl ShardsRefMut<'_> {
256    pub(crate) fn copy_within(&mut self, mut src: usize, mut dest: usize, mut count: usize) {
257        src *= self.shard_chunk_count;
258        dest *= self.shard_chunk_count;
259        count *= self.shard_chunk_count;
260
261        self.data.copy_within(src..src + count, dest);
262    }
263
264    // Returns mutable references to flat-arrays of shard-ranges
265    // `x .. x + count` and `y .. y + count`. Ranges must not overlap.
266    pub(crate) fn flat2_mut(
267        &mut self,
268        mut x: usize,
269        mut y: usize,
270        mut count: usize,
271    ) -> (
272        &mut [[u8; SHARD_CHUNK_BYTES]],
273        &mut [[u8; SHARD_CHUNK_BYTES]],
274    ) {
275        x *= self.shard_chunk_count;
276        y *= self.shard_chunk_count;
277        count *= self.shard_chunk_count;
278
279        if x < y {
280            let (head, tail) = self.data.split_at_mut(y);
281            (&mut head[x..x + count], &mut tail[..count])
282        } else {
283            let (head, tail) = self.data.split_at_mut(x);
284            (&mut tail[..count], &mut head[y..y + count])
285        }
286    }
287}
commonware_cryptography/reed_solomon/engine/shards.rs

commonware_cryptography/reed_solomon/engine/
shards.rs
