ab_erasure_coding/
lib.rs

1#![expect(incomplete_features, reason = "generic_const_*")]
2#![feature(
3    gca_const_items,
4    gca_macroless_args,
5    gca_min_const_items,
6    generic_const_items
7)]
8#![no_std]
9
10use core::{array, fmt, slice};
11use reed_solomon_simd::Error;
12use reed_solomon_simd::engine::DefaultEngine;
13use reed_solomon_simd::rate::{HighRateDecoder, HighRateEncoder, RateDecoder, RateEncoder};
14
15/// Error that occurs when erasure coding data
16#[derive(Debug, Clone, PartialEq, thiserror::Error)]
17pub enum ErasureCodingError {
18    /// Decoder error
19    #[error("Decoder error: {0}")]
20    DecoderError(#[from] Error),
21}
22
23/// Erasure coding processes shards in blocks of this many bytes independently from each other
24const BLOCK_BYTES: usize = 64;
25/// Size of source shards to process at once when shards are large.
26///
27/// Small enough to fit into L2 cache of CPUs since Skylake and Zen 1 (on most of them even with
28/// both SMT threads of a core busy).
29const TARGET_COLUMNS_BYTES: usize = 256 * 1024;
30/// Narrower columns do not amortize fixed per-call overhead of erasure coding well enough
31const MIN_COLUMN_BYTES: usize = 1024;
32/// Max number of threads used by parallel methods.
33///
34/// Erasure coding of a set of shards is too little work to benefit from many threads, so only a few
35/// are used to reduce latency without wasting CPU on orchestration.
36#[cfg(feature = "parallel")]
37const MAX_PARALLELISM: usize = 4;
38
39/// Number of `u64` words needed for a bit per shard
40const NUM_WORDS<const NUM_SHARDS: usize>: usize = NUM_SHARDS.div_ceil(u64::BITS as usize);
41
42/// A bit per shard, typically saying whether that shard is present.
43///
44/// Bit `index % 64` of word `index / 64` corresponds to shard `index`, which is the same layout
45/// `reed_solomon_simd::rate::ReceivedShards` uses.
46#[derive(Copy, Clone, Eq, PartialEq)]
47pub struct ShardsBitmap<const NUM_SHARDS: usize> {
48    words: [u64; NUM_WORDS::<NUM_SHARDS>],
49}
50
51impl<const NUM_SHARDS: usize> fmt::Debug for ShardsBitmap<NUM_SHARDS> {
52    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
53        f.debug_struct("ShardsBitmap")
54            .field("num_shards", &NUM_SHARDS)
55            .field("num_set", &self.count())
56            .finish()
57    }
58}
59
60impl<const NUM_SHARDS: usize> Default for ShardsBitmap<NUM_SHARDS> {
61    #[inline(always)]
62    fn default() -> Self {
63        Self::none()
64    }
65}
66
67impl<const NUM_SHARDS: usize> ShardsBitmap<NUM_SHARDS> {
68    /// Bitmap with no shards set
69    #[inline(always)]
70    pub const fn none() -> Self {
71        Self {
72            words: [0; NUM_WORDS::<NUM_SHARDS>],
73        }
74    }
75
76    /// Bitmap with all shards set
77    #[inline(always)]
78    pub const fn all() -> Self {
79        let mut this = Self {
80            words: [u64::MAX; NUM_WORDS::<NUM_SHARDS>],
81        };
82
83        // Bits past the last shard must not be set
84        let unused_bits = NUM_WORDS::<NUM_SHARDS> * u64::BITS as usize - NUM_SHARDS;
85        if unused_bits > 0 {
86            this.words[NUM_WORDS::<NUM_SHARDS> - 1] >>= unused_bits;
87        }
88
89        this
90    }
91
92    /// Whether the shard at `index` is set, `false` if `index` is out of bounds
93    #[inline(always)]
94    pub const fn get(&self, index: usize) -> bool {
95        if index >= NUM_SHARDS {
96            return false;
97        }
98
99        (self.words[index / u64::BITS as usize] >> (index % u64::BITS as usize)) & 1 == 1
100    }
101
102    /// Sets the shard at `index`, does nothing if `index` is out of bounds
103    #[inline(always)]
104    pub const fn set(&mut self, index: usize) {
105        if index < NUM_SHARDS {
106            self.words[index / u64::BITS as usize] |= 1 << (index % u64::BITS as usize);
107        }
108    }
109
110    /// Unsets the shard at `index`, does nothing if `index` is out of bounds
111    #[inline(always)]
112    pub const fn unset(&mut self, index: usize) {
113        if index < NUM_SHARDS {
114            self.words[index / u64::BITS as usize] &= !(1 << (index % u64::BITS as usize));
115        }
116    }
117
118    /// Number of shards that are set
119    #[inline]
120    pub const fn count(&self) -> usize {
121        let mut count = 0;
122        let mut word_index = 0;
123        while word_index < NUM_WORDS::<NUM_SHARDS> {
124            count += self.words[word_index].count_ones() as usize;
125            word_index += 1;
126        }
127
128        count
129    }
130}
131
132/// Which source and parity shards are present
133#[derive(Debug, Default, Copy, Clone, Eq, PartialEq)]
134pub struct ShardsPresent<const NUM_SHARDS: usize> {
135    /// Which source shards are present
136    pub source: ShardsBitmap<NUM_SHARDS>,
137    /// Which parity shards are present
138    pub parity: ShardsBitmap<NUM_SHARDS>,
139}
140
141impl<const NUM_SHARDS: usize> ShardsPresent<NUM_SHARDS> {
142    /// Nothing is present
143    #[inline(always)]
144    pub const fn none() -> Self {
145        Self {
146            source: ShardsBitmap::none(),
147            parity: ShardsBitmap::none(),
148        }
149    }
150
151    /// Everything is present
152    #[inline(always)]
153    pub const fn all() -> Self {
154        Self {
155            source: ShardsBitmap::all(),
156            parity: ShardsBitmap::all(),
157        }
158    }
159}
160
161/// Erasure coding abstraction.
162///
163/// Supports creation of parity shards and recovery of missing data.
164///
165/// Erasure coding processes each 64-byte block of shards independently from other blocks, so large
166/// shards are split into narrower columns that are processed one by one, which keeps the working
167/// set in CPU cache and is much faster than processing whole shards at once. With `parallel`
168/// feature enabled, `par_*` methods additionally process columns concurrently.
169#[derive(Debug, Clone)]
170pub struct ErasureCoding;
171
172impl Default for ErasureCoding {
173    fn default() -> Self {
174        Self::new()
175    }
176}
177
178impl ErasureCoding {
179    /// Create new erasure coding instance
180    pub fn new() -> Self {
181        Self {}
182    }
183
184    /// Extend contiguously stored source shards with parity shards
185    #[inline]
186    pub fn extend<const NUM_SHARDS: usize, const SHARD_BYTES: usize>(
187        &self,
188        source: &[[u8; SHARD_BYTES]; NUM_SHARDS],
189        parity: &mut [[u8; SHARD_BYTES]; NUM_SHARDS],
190    ) -> Result<(), ErasureCodingError> {
191        Ok(extend::<NUM_SHARDS, SHARD_BYTES, false, _, _>(
192            source, parity,
193        )?)
194    }
195
196    /// Parallel version of [`Self::extend()`].
197    ///
198    /// Columns of large shards are processed using up to 4 threads, which is beneficial when
199    /// erasure coding is the only thing running. Prefer [`Self::extend()`] when calling from
200    /// code that is already parallel to avoid contention.
201    ///
202    /// Small shards (like many record chunks) are not split into columns and are processed by a
203    /// single thread, the same way as [`Self::extend()`] does.
204    #[cfg(feature = "parallel")]
205    #[inline]
206    pub fn par_extend<const NUM_SHARDS: usize, const SHARD_BYTES: usize>(
207        &self,
208        source: &[[u8; SHARD_BYTES]; NUM_SHARDS],
209        parity: &mut [[u8; SHARD_BYTES]; NUM_SHARDS],
210    ) -> Result<(), ErasureCodingError> {
211        Ok(extend::<NUM_SHARDS, SHARD_BYTES, true, _, _>(
212            source, parity,
213        )?)
214    }
215
216    /// Extend source shards with parity shards, where shards are not stored contiguously
217    #[inline]
218    pub fn extend_scattered<
219        const NUM_SHARDS: usize,
220        const SHARD_BYTES: usize,
221        SourceShard,
222        ParityShard,
223    >(
224        &self,
225        source: [SourceShard; NUM_SHARDS],
226        parity: [ParityShard; NUM_SHARDS],
227    ) -> Result<(), ErasureCodingError>
228    where
229        SourceShard: AsRef<[u8; SHARD_BYTES]>,
230        ParityShard: AsMut<[u8; SHARD_BYTES]>,
231    {
232        let mut parity = parity;
233
234        Ok(extend::<NUM_SHARDS, SHARD_BYTES, false, _, _>(
235            &source.each_ref().map(|shard| shard.as_ref().as_slice()),
236            &mut parity.each_mut().map(|shard| shard.as_mut().as_mut_slice()),
237        )?)
238    }
239
240    /// Parallel version of [`Self::extend_scattered()`].
241    ///
242    /// Columns of large shards are processed using up to 4 threads, which is beneficial when
243    /// erasure coding is the only thing running. Prefer [`Self::extend_scattered()`] when calling
244    /// from code that is already parallel to avoid contention.
245    ///
246    /// Small shards (like many record chunks) are not split into columns and are processed by a
247    /// single thread, the same way as [`Self::extend_scattered()`] does.
248    #[cfg(feature = "parallel")]
249    #[inline]
250    pub fn par_extend_scattered<
251        const NUM_SHARDS: usize,
252        const SHARD_BYTES: usize,
253        SourceShard,
254        ParityShard,
255    >(
256        &self,
257        source: [SourceShard; NUM_SHARDS],
258        parity: [ParityShard; NUM_SHARDS],
259    ) -> Result<(), ErasureCodingError>
260    where
261        SourceShard: AsRef<[u8; SHARD_BYTES]>,
262        ParityShard: AsMut<[u8; SHARD_BYTES]>,
263    {
264        let mut parity = parity;
265
266        Ok(extend::<NUM_SHARDS, SHARD_BYTES, true, _, _>(
267            &source.each_ref().map(|shard| shard.as_ref().as_slice()),
268            &mut parity.each_mut().map(|shard| shard.as_mut().as_mut_slice()),
269        )?)
270    }
271
272    /// Recover missing source shards in place, with everything stored contiguously.
273    ///
274    /// Parity shards are inputs only, missing ones are simply not used. Prefer this over
275    /// [`Self::recover_all()`] when parity shards are not needed.
276    #[inline]
277    pub fn recover_source<const NUM_SHARDS: usize, const SHARD_BYTES: usize>(
278        &self,
279        source: &mut [[u8; SHARD_BYTES]; NUM_SHARDS],
280        parity: &[[u8; SHARD_BYTES]; NUM_SHARDS],
281        present: &ShardsPresent<NUM_SHARDS>,
282    ) -> Result<(), ErasureCodingError> {
283        Ok(recover_source::<NUM_SHARDS, SHARD_BYTES, false, _, _>(
284            source,
285            &present.source,
286            parity
287                .iter()
288                .enumerate()
289                .map(|(index, shard)| present.parity.get(index).then_some(shard.as_slice())),
290        )?)
291    }
292
293    /// Parallel version of [`Self::recover_source()`].
294    ///
295    /// Columns of large shards are processed using up to 4 threads, which is beneficial when
296    /// erasure coding is the only thing running. Prefer [`Self::recover_source()`] when calling
297    /// from code that is already parallel to avoid contention.
298    ///
299    /// Small shards (like many record chunks) are not split into columns and are processed by a
300    /// single thread, the same way as [`Self::recover_source()`] does.
301    #[cfg(feature = "parallel")]
302    #[inline]
303    pub fn par_recover_source<const NUM_SHARDS: usize, const SHARD_BYTES: usize>(
304        &self,
305        source: &mut [[u8; SHARD_BYTES]; NUM_SHARDS],
306        parity: &[[u8; SHARD_BYTES]; NUM_SHARDS],
307        present: &ShardsPresent<NUM_SHARDS>,
308    ) -> Result<(), ErasureCodingError> {
309        Ok(recover_source::<NUM_SHARDS, SHARD_BYTES, true, _, _>(
310            source,
311            &present.source,
312            parity
313                .iter()
314                .enumerate()
315                .map(|(index, shard)| present.parity.get(index).then_some(shard.as_slice())),
316        )?)
317    }
318
319    /// Recover missing source shards in place, where shards are not stored contiguously.
320    ///
321    /// Parity shards are inputs only, so missing ones simply have no memory.
322    #[inline]
323    pub fn recover_source_scattered<
324        const NUM_SHARDS: usize,
325        const SHARD_BYTES: usize,
326        SourceShard,
327        ParityShard,
328    >(
329        &self,
330        source: [SourceShard; NUM_SHARDS],
331        source_present: &ShardsBitmap<NUM_SHARDS>,
332        parity: [Option<ParityShard>; NUM_SHARDS],
333    ) -> Result<(), ErasureCodingError>
334    where
335        SourceShard: AsMut<[u8; SHARD_BYTES]>,
336        ParityShard: AsRef<[u8; SHARD_BYTES]>,
337    {
338        let mut source = source;
339
340        Ok(recover_source::<NUM_SHARDS, SHARD_BYTES, false, _, _>(
341            &mut source.each_mut().map(|shard| shard.as_mut().as_mut_slice()),
342            source_present,
343            parity
344                .iter()
345                .map(|maybe_shard| maybe_shard.as_ref().map(|shard| shard.as_ref().as_slice())),
346        )?)
347    }
348
349    /// Parallel version of [`Self::recover_source_scattered()`].
350    ///
351    /// Columns of large shards are processed using up to 4 threads, which is beneficial when
352    /// erasure coding is the only thing running. Prefer [`Self::recover_source_scattered()`] when
353    /// calling from code that is already parallel to avoid contention.
354    ///
355    /// Small shards (like many record chunks) are not split into columns and are processed by a
356    /// single thread, the same way as [`Self::recover_source_scattered()`] does.
357    #[cfg(feature = "parallel")]
358    #[inline]
359    pub fn par_recover_source_scattered<
360        const NUM_SHARDS: usize,
361        const SHARD_BYTES: usize,
362        SourceShard,
363        ParityShard,
364    >(
365        &self,
366        source: [SourceShard; NUM_SHARDS],
367        source_present: &ShardsBitmap<NUM_SHARDS>,
368        parity: [Option<ParityShard>; NUM_SHARDS],
369    ) -> Result<(), ErasureCodingError>
370    where
371        SourceShard: AsMut<[u8; SHARD_BYTES]>,
372        ParityShard: AsRef<[u8; SHARD_BYTES]>,
373    {
374        let mut source = source;
375
376        Ok(recover_source::<NUM_SHARDS, SHARD_BYTES, true, _, _>(
377            &mut source.each_mut().map(|shard| shard.as_mut().as_mut_slice()),
378            source_present,
379            parity
380                .iter()
381                .map(|maybe_shard| maybe_shard.as_ref().map(|shard| shard.as_ref().as_slice())),
382        )?)
383    }
384
385    /// Recover missing source and parity shards in place, with everything stored contiguously.
386    ///
387    /// Use [`Self::recover_source()`] instead if parity shards are not needed.
388    #[inline]
389    pub fn recover_all<const NUM_SHARDS: usize, const SHARD_BYTES: usize>(
390        &self,
391        source: &mut [[u8; SHARD_BYTES]; NUM_SHARDS],
392        parity: &mut [[u8; SHARD_BYTES]; NUM_SHARDS],
393        present: &ShardsPresent<NUM_SHARDS>,
394    ) -> Result<(), ErasureCodingError> {
395        Ok(recover_all::<NUM_SHARDS, SHARD_BYTES, false, _, _>(
396            source, parity, present,
397        )?)
398    }
399
400    /// Parallel version of [`Self::recover_all()`].
401    ///
402    /// Columns of large shards are processed using up to 4 threads, which is beneficial when
403    /// erasure coding is the only thing running. Prefer [`Self::recover_all()`] when calling from
404    /// code that is already parallel to avoid contention.
405    ///
406    /// Small shards (like many record chunks) are not split into columns and are processed by a
407    /// single thread, the same way as [`Self::recover_all()`] does.
408    #[cfg(feature = "parallel")]
409    #[inline]
410    pub fn par_recover_all<const NUM_SHARDS: usize, const SHARD_BYTES: usize>(
411        &self,
412        source: &mut [[u8; SHARD_BYTES]; NUM_SHARDS],
413        parity: &mut [[u8; SHARD_BYTES]; NUM_SHARDS],
414        present: &ShardsPresent<NUM_SHARDS>,
415    ) -> Result<(), ErasureCodingError> {
416        Ok(recover_all::<NUM_SHARDS, SHARD_BYTES, true, _, _>(
417            source, parity, present,
418        )?)
419    }
420
421    /// Recover missing source and parity shards in place, where shards are not stored contiguously
422    #[inline]
423    pub fn recover_all_scattered<
424        const NUM_SHARDS: usize,
425        const SHARD_BYTES: usize,
426        SourceShard,
427        ParityShard,
428    >(
429        &self,
430        source: [SourceShard; NUM_SHARDS],
431        parity: [ParityShard; NUM_SHARDS],
432        present: &ShardsPresent<NUM_SHARDS>,
433    ) -> Result<(), ErasureCodingError>
434    where
435        SourceShard: AsMut<[u8; SHARD_BYTES]>,
436        ParityShard: AsMut<[u8; SHARD_BYTES]>,
437    {
438        let mut source = source;
439        let mut parity = parity;
440
441        Ok(recover_all::<NUM_SHARDS, SHARD_BYTES, false, _, _>(
442            &mut source.each_mut().map(|shard| shard.as_mut().as_mut_slice()),
443            &mut parity.each_mut().map(|shard| shard.as_mut().as_mut_slice()),
444            present,
445        )?)
446    }
447
448    /// Parallel version of [`Self::recover_all_scattered()`].
449    ///
450    /// Columns of large shards are processed using up to 4 threads, which is beneficial when
451    /// erasure coding is the only thing running. Prefer [`Self::recover_all_scattered()`] when
452    /// calling from code that is already parallel to avoid contention.
453    ///
454    /// Small shards (like many record chunks) are not split into columns and are processed by a
455    /// single thread, the same way as [`Self::recover_all_scattered()`] does.
456    #[cfg(feature = "parallel")]
457    #[inline]
458    pub fn par_recover_all_scattered<
459        const NUM_SHARDS: usize,
460        const SHARD_BYTES: usize,
461        SourceShard,
462        ParityShard,
463    >(
464        &self,
465        source: [SourceShard; NUM_SHARDS],
466        parity: [ParityShard; NUM_SHARDS],
467        present: &ShardsPresent<NUM_SHARDS>,
468    ) -> Result<(), ErasureCodingError>
469    where
470        SourceShard: AsMut<[u8; SHARD_BYTES]>,
471        ParityShard: AsMut<[u8; SHARD_BYTES]>,
472    {
473        let mut source = source;
474        let mut parity = parity;
475
476        Ok(recover_all::<NUM_SHARDS, SHARD_BYTES, true, _, _>(
477            &mut source.each_mut().map(|shard| shard.as_mut().as_mut_slice()),
478            &mut parity.each_mut().map(|shard| shard.as_mut().as_mut_slice()),
479            present,
480        )?)
481    }
482}
483
484/// Number of bytes of each shard that are processed at once.
485///
486/// Equals to `shard_bytes` when shards are not split into columns. Column boundaries are aligned to
487/// [`BLOCK_BYTES`]-byte blocks, which is what makes processing of columns independent. The last
488/// column may be narrower, but it ends where the shard ends, so it has the same layout as the shard
489/// itself.
490const fn column_bytes(num_shards: usize, shard_bytes: usize) -> usize {
491    let Some(column_bytes) = TARGET_COLUMNS_BYTES.checked_div(num_shards) else {
492        return shard_bytes;
493    };
494    let column_bytes = column_bytes / BLOCK_BYTES * BLOCK_BYTES;
495
496    if column_bytes < MIN_COLUMN_BYTES || column_bytes >= shard_bytes {
497        shard_bytes
498    } else {
499        column_bytes
500    }
501}
502
503fn extend<const NUM_SHARDS: usize, const SHARD_BYTES: usize, const PARALLEL: bool, Source, Parity>(
504    source: &[Source; NUM_SHARDS],
505    parity: &mut [Parity; NUM_SHARDS],
506) -> Result<(), Error>
507where
508    Source: AsRef<[u8]>,
509    Parity: AsMut<[u8]>,
510{
511    if const { column_bytes(NUM_SHARDS, SHARD_BYTES) == SHARD_BYTES } {
512        return encode(SHARD_BYTES, source, parity);
513    }
514
515    for_each_column::<NUM_SHARDS, PARALLEL, _, _, _>(
516        SHARD_BYTES,
517        const { column_bytes(NUM_SHARDS, SHARD_BYTES) },
518        source.each_ref().map(AsRef::as_ref),
519        parity.each_mut().map(AsMut::as_mut),
520        |shard_bytes, source, mut parity| encode(shard_bytes, &source, &mut parity),
521    )
522}
523
524/// `parity` must yield exactly `NUM_SHARDS` items
525fn recover_source<
526    'a,
527    const NUM_SHARDS: usize,
528    const SHARD_BYTES: usize,
529    const PARALLEL: bool,
530    Source,
531    Parity,
532>(
533    source: &mut [Source; NUM_SHARDS],
534    source_present: &ShardsBitmap<NUM_SHARDS>,
535    parity: Parity,
536) -> Result<(), Error>
537where
538    Source: AsRef<[u8]> + AsMut<[u8]>,
539    Parity: Iterator<Item = Option<&'a [u8]>>,
540{
541    if const { column_bytes(NUM_SHARDS, SHARD_BYTES) == SHARD_BYTES } {
542        return decode_source(SHARD_BYTES, source, source_present, parity);
543    }
544
545    let mut parity = parity;
546    let parity = array::from_fn::<_, NUM_SHARDS, _>(|_| parity.next().flatten());
547
548    for_each_column::<NUM_SHARDS, PARALLEL, _, _, _>(
549        SHARD_BYTES,
550        const { column_bytes(NUM_SHARDS, SHARD_BYTES) },
551        source.each_mut().map(AsMut::as_mut),
552        parity,
553        |shard_bytes, mut source, parity| {
554            decode_source(shard_bytes, &mut source, source_present, parity)
555        },
556    )
557}
558
559fn recover_all<
560    const NUM_SHARDS: usize,
561    const SHARD_BYTES: usize,
562    const PARALLEL: bool,
563    Source,
564    Parity,
565>(
566    source: &mut [Source; NUM_SHARDS],
567    parity: &mut [Parity; NUM_SHARDS],
568    present: &ShardsPresent<NUM_SHARDS>,
569) -> Result<(), Error>
570where
571    Source: AsRef<[u8]> + AsMut<[u8]>,
572    Parity: AsRef<[u8]> + AsMut<[u8]>,
573{
574    if const { column_bytes(NUM_SHARDS, SHARD_BYTES) == SHARD_BYTES } {
575        return decode_all(SHARD_BYTES, source, parity, present);
576    }
577
578    for_each_column::<NUM_SHARDS, PARALLEL, _, _, _>(
579        SHARD_BYTES,
580        const { column_bytes(NUM_SHARDS, SHARD_BYTES) },
581        source.each_mut().map(AsMut::as_mut),
582        parity.each_mut().map(AsMut::as_mut),
583        |shard_bytes, mut source, mut parity| {
584            decode_all(shard_bytes, &mut source, &mut parity, present)
585        },
586    )
587}
588
589/// Shard that can be split into columns
590trait IntoColumns: Sized {
591    /// Iterator over columns, yields at least as many columns as there are in the shard
592    type Columns: Iterator<Item = Self>;
593
594    fn into_columns(self, column_bytes: usize) -> Self::Columns;
595}
596
597impl<'a> IntoColumns for &'a [u8] {
598    type Columns = slice::Chunks<'a, u8>;
599
600    #[inline(always)]
601    fn into_columns(self, column_bytes: usize) -> Self::Columns {
602        self.chunks(column_bytes)
603    }
604}
605
606impl<'a> IntoColumns for &'a mut [u8] {
607    type Columns = slice::ChunksMut<'a, u8>;
608
609    #[inline(always)]
610    fn into_columns(self, column_bytes: usize) -> Self::Columns {
611        self.chunks_mut(column_bytes)
612    }
613}
614
615/// Columns of an optional shard, a missing shard has all columns missing
616struct OptionalColumns<'a>(Option<slice::Chunks<'a, u8>>);
617
618impl<'a> Iterator for OptionalColumns<'a> {
619    type Item = Option<&'a [u8]>;
620
621    #[inline(always)]
622    fn next(&mut self) -> Option<Self::Item> {
623        Some(self.0.as_mut().and_then(Iterator::next))
624    }
625}
626
627impl<'a> IntoColumns for Option<&'a [u8]> {
628    type Columns = OptionalColumns<'a>;
629
630    #[inline(always)]
631    fn into_columns(self, column_bytes: usize) -> Self::Columns {
632        OptionalColumns(self.map(|shard| shard.chunks(column_bytes)))
633    }
634}
635
636/// Split shards of `shard_bytes` bytes into columns, yielding the number of bytes in a column
637/// alongside the same column of every shard.
638///
639/// `column_bytes` must not be zero.
640// TODO: Once `ShardStorage` from https://github.com/AndersTrier/reed-solomon-simd/issues/64 lands,
641//  it will be possible to apply erasure coding to scattered shards directly without copying them
642//  into a separate work buffer. While natural, it is likely to be slower for columns: columns of
643//  large shards are far apart and may map to the same L2 cache sets, so a small contiguous work
644//  buffer reused across columns should still be preferred.
645fn columns<const NUM_SHARDS: usize, Shard>(
646    shard_bytes: usize,
647    column_bytes: usize,
648    shards: [Shard; NUM_SHARDS],
649) -> impl Iterator<Item = (usize, [Shard; NUM_SHARDS])>
650where
651    Shard: IntoColumns,
652{
653    let mut shards = shards.map(|shard| shard.into_columns(column_bytes));
654
655    (0..shard_bytes).step_by(column_bytes).map(move |offset| {
656        let column = shards.each_mut().map(|shard| {
657            shard
658                .next()
659                .expect("All shards have `shard_bytes` bytes and the same number of columns; qed")
660        });
661
662        (column_bytes.min(shard_bytes - offset), column)
663    })
664}
665
666/// Process shards of `shard_bytes` bytes in columns of `column_bytes` bytes.
667///
668/// Columns are processed one by one, or with `PARALLEL` shards are split into up to
669/// `MAX_PARALLELISM` regions of columns that are processed concurrently.
670fn for_each_column<const NUM_SHARDS: usize, const PARALLEL: bool, Source, Parity, Process>(
671    shard_bytes: usize,
672    column_bytes: usize,
673    source: [Source; NUM_SHARDS],
674    parity: [Parity; NUM_SHARDS],
675    process: Process,
676) -> Result<(), Error>
677where
678    Source: IntoColumns + Send,
679    Parity: IntoColumns + Send,
680    Process: Fn(usize, [Source; NUM_SHARDS], [Parity; NUM_SHARDS]) -> Result<(), Error> + Sync,
681{
682    #[cfg(feature = "parallel")]
683    if PARALLEL {
684        let num_columns = shard_bytes.div_ceil(column_bytes);
685        let region_bytes = num_columns.div_ceil(MAX_PARALLELISM) * column_bytes;
686        let process = &process;
687        let mut results = [const { Ok(()) }; MAX_PARALLELISM];
688
689        rayon::scope(|scope| {
690            let regions = columns(shard_bytes, region_bytes, source).zip(columns(
691                shard_bytes,
692                region_bytes,
693                parity,
694            ));
695
696            for (((region_bytes, source), (_, parity)), result) in regions.zip(&mut results) {
697                scope.spawn(move |_| {
698                    *result = for_each_column_sequential(
699                        region_bytes,
700                        column_bytes,
701                        source,
702                        parity,
703                        process,
704                    );
705                });
706            }
707        });
708
709        return results.into_iter().collect();
710    }
711
712    for_each_column_sequential(shard_bytes, column_bytes, source, parity, process)
713}
714
715fn for_each_column_sequential<const NUM_SHARDS: usize, Source, Parity, Process>(
716    shard_bytes: usize,
717    column_bytes: usize,
718    source: [Source; NUM_SHARDS],
719    parity: [Parity; NUM_SHARDS],
720    process: Process,
721) -> Result<(), Error>
722where
723    Source: IntoColumns,
724    Parity: IntoColumns,
725    Process: Fn(usize, [Source; NUM_SHARDS], [Parity; NUM_SHARDS]) -> Result<(), Error>,
726{
727    let columns =
728        columns(shard_bytes, column_bytes, source).zip(columns(shard_bytes, column_bytes, parity));
729
730    for ((shard_bytes, source), (_, parity)) in columns {
731        process(shard_bytes, source, parity)?;
732    }
733
734    Ok(())
735}
736
737/// Write parity shards for source shards
738fn encode<const NUM_SHARDS: usize, Source, Parity>(
739    shard_bytes: usize,
740    source: &[Source; NUM_SHARDS],
741    parity: &mut [Parity; NUM_SHARDS],
742) -> Result<(), Error>
743where
744    Source: AsRef<[u8]>,
745    Parity: AsMut<[u8]>,
746{
747    let mut encoder = new_encoder::<NUM_SHARDS>(shard_bytes)?;
748
749    for shard in source {
750        encoder.add_original_shard(shard)?;
751    }
752
753    let result = encoder.encode()?;
754
755    for (input, output) in result.recovery_iter().zip(parity) {
756        output.as_mut().copy_from_slice(input);
757    }
758
759    Ok(())
760}
761
762/// Recover missing source shards from present source and parity shards
763fn decode_source<'a, const NUM_SHARDS: usize, Source, Parity>(
764    shard_bytes: usize,
765    source: &mut [Source; NUM_SHARDS],
766    source_present: &ShardsBitmap<NUM_SHARDS>,
767    parity: Parity,
768) -> Result<(), Error>
769where
770    Source: AsRef<[u8]> + AsMut<[u8]>,
771    Parity: IntoIterator<Item = Option<&'a [u8]>>,
772{
773    let mut decoder = new_decoder::<NUM_SHARDS>(shard_bytes)?;
774
775    for (index, shard) in source.iter().enumerate() {
776        if source_present.get(index) {
777            decoder.add_original_shard(index, shard)?;
778        }
779    }
780    for (index, maybe_shard) in parity.into_iter().enumerate() {
781        if let Some(shard) = maybe_shard {
782            decoder.add_recovery_shard(index, shard)?;
783        }
784    }
785
786    let result = decoder.decode()?;
787
788    for (index, shard) in source.iter_mut().enumerate() {
789        if !source_present.get(index) {
790            shard
791                .as_mut()
792                .copy_from_slice(restored_original(&result, index));
793        }
794    }
795
796    Ok(())
797}
798
799/// Recover missing source and parity shards from present source and parity shards
800fn decode_all<const NUM_SHARDS: usize, Source, Parity>(
801    shard_bytes: usize,
802    source: &mut [Source; NUM_SHARDS],
803    parity: &mut [Parity; NUM_SHARDS],
804    present: &ShardsPresent<NUM_SHARDS>,
805) -> Result<(), Error>
806where
807    Source: AsRef<[u8]> + AsMut<[u8]>,
808    Parity: AsRef<[u8]> + AsMut<[u8]>,
809{
810    decode_source(
811        shard_bytes,
812        source,
813        &present.source,
814        parity
815            .iter()
816            .enumerate()
817            .map(|(index, shard)| present.parity.get(index).then_some(shard.as_ref())),
818    )?;
819
820    if present.parity.count() == NUM_SHARDS {
821        return Ok(());
822    }
823
824    // Source shards are complete at this point, so missing parity shards are simply encoded again
825    let mut encoder = new_encoder::<NUM_SHARDS>(shard_bytes)?;
826
827    for shard in source.iter() {
828        encoder.add_original_shard(shard)?;
829    }
830
831    let result = encoder.encode()?;
832
833    for (index, shard) in parity.iter_mut().enumerate() {
834        if !present.parity.get(index) {
835            shard.as_mut().copy_from_slice(recovery(&result, index));
836        }
837    }
838
839    Ok(())
840}
841
842// TODO: Encoder and decoder allocate a work buffer internally on every call, reuse it once in-place
843//  encoding from https://github.com/AndersTrier/reed-solomon-simd/issues/64 is available
844#[inline(always)]
845fn new_encoder<const NUM_SHARDS: usize>(
846    shard_bytes: usize,
847) -> Result<HighRateEncoder<DefaultEngine>, Error> {
848    HighRateEncoder::new(
849        NUM_SHARDS,
850        NUM_SHARDS,
851        shard_bytes,
852        DefaultEngine::new(),
853        None,
854    )
855}
856
857#[inline(always)]
858fn new_decoder<const NUM_SHARDS: usize>(
859    shard_bytes: usize,
860) -> Result<HighRateDecoder<DefaultEngine>, Error> {
861    HighRateDecoder::new(
862        NUM_SHARDS,
863        NUM_SHARDS,
864        shard_bytes,
865        DefaultEngine::new(),
866        None,
867    )
868}
869
870#[inline(always)]
871fn restored_original<'a>(
872    result: &'a reed_solomon_simd::DecoderResult<'_>,
873    index: usize,
874) -> &'a [u8] {
875    result
876        .restored_original(index)
877        .expect("Always corresponds to a missing source shard; qed")
878}
879
880#[inline(always)]
881fn recovery<'a>(result: &'a reed_solomon_simd::EncoderResult<'_>, index: usize) -> &'a [u8] {
882    result
883        .recovery(index)
884        .expect("Always corresponds to a missing parity shard; qed")
885}
ab_erasure_coding/lib.rs

ab_erasure_coding/
lib.rs
