1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116

#![cfg_attr(not(feature = "std"), no_std)]

#[cfg(test)]
mod tests;

extern crate alloc;

#[cfg(not(feature = "std"))]
use alloc::string::String;
use alloc::sync::Arc;
#[cfg(not(feature = "std"))]
use alloc::vec::Vec;
use core::num::NonZeroUsize;
use kzg::{FFTSettings, PolyRecover, DAS, FFTG1, G1};
use rust_kzg_blst::types::fft_settings::FsFFTSettings;
use rust_kzg_blst::types::g1::FsG1;
use rust_kzg_blst::types::poly::FsPoly;
use subspace_core_primitives::crypto::kzg::{Commitment, Polynomial};
use subspace_core_primitives::crypto::Scalar;

/// Erasure coding abstraction.
///
/// Supports creation of parity records and recovery of missing data.
#[derive(Debug, Clone)]
pub struct ErasureCoding {
    fft_settings: Arc<FsFFTSettings>,
}

impl ErasureCoding {
    /// Create new erasure coding instance.
    ///
    /// Number of shards supported is `2^scale`, half of shards are source data and the other half
    /// are parity.
    pub fn new(scale: NonZeroUsize) -> Result<Self, String> {
        let fft_settings = Arc::new(FsFFTSettings::new(scale.get())?);

        Ok(Self { fft_settings })
    }

    /// Max number of shards supported (both source and parity together)
    pub fn max_shards(&self) -> usize {
        self.fft_settings.max_width
    }

    /// Extend sources using erasure coding.
    ///
    /// Returns parity data.
    pub fn extend(&self, source: &[Scalar]) -> Result<Vec<Scalar>, String> {
        // TODO: das_fft_extension modifies buffer internally, it needs to change to use
        //  pre-allocated buffer instead of allocating a new one
        self.fft_settings
            .das_fft_extension(Scalar::slice_to_repr(source))
            .map(Scalar::vec_from_repr)
    }

    /// Recovery of missing shards from given shards (at least 1/2 should be `Some`).
    ///
    /// Both in input and output source shards are interleaved with parity shards:
    /// source, parity, source, parity, ...
    pub fn recover(&self, shards: &[Option<Scalar>]) -> Result<Vec<Scalar>, String> {
        let poly = FsPoly::recover_poly_from_samples(
            Scalar::slice_option_to_repr(shards),
            &self.fft_settings,
        )?;

        Ok(Scalar::vec_from_repr(poly.coeffs))
    }

    /// Recovery of missing shards from given shards (at least 1/2 should be `Some`) in form of
    /// normalized polynomial (allows to not do inverse FFT afterwards if polynomial is desired).
    ///
    /// Both in input and output source shards are interleaved with parity shards:
    /// source, parity, source, parity, ...
    pub fn recover_poly(&self, shards: &[Option<Scalar>]) -> Result<Polynomial, String> {
        let mut poly = Polynomial::from(FsPoly::recover_poly_coeffs_from_samples(
            Scalar::slice_option_to_repr(shards),
            &self.fft_settings,
        )?);

        poly.normalize();

        Ok(poly)
    }

    /// Recovery of source shards from given shards (at least 1/2 should be `Some`).
    ///
    /// The same as [`ErasureCoding::recover()`], but returns only source shards in form of an
    /// iterator.
    pub fn recover_source(
        &self,
        shards: &[Option<Scalar>],
    ) -> Result<impl ExactSizeIterator<Item = Scalar>, String> {
        Ok(self.recover(shards)?.into_iter().step_by(2))
    }

    /// Extend commitments using erasure coding.
    ///
    /// Returns both source and parity commitments interleaved.
    pub fn extend_commitments(
        &self,
        commitments: &[Commitment],
    ) -> Result<Vec<Commitment>, String> {
        // Inverse FFT to interpolate polynomial over source commitments
        let mut coeffs = self
            .fft_settings
            .fft_g1(Commitment::slice_to_repr(commitments), true)?;

        // Double the size
        coeffs.resize(coeffs.len() * 2, FsG1::identity());

        // FFT to get extended commitments
        self.fft_settings
            .fft_g1(&coeffs, false)
            .map(Commitment::vec_from_repr)
    }
}