Initial subspace-erasure-coding implementation

crates/subspace-core-primitives/src/lib.rs

@@ -198,20 +198,38 @@ impl From<&[u8; Self::SAFE_BYTES]> for Scalar {
     }
 }
 
+impl From<[u8; Self::SAFE_BYTES]> for Scalar {
+    fn from(value: [u8; Self::SAFE_BYTES]) -> Self {
+        Self::from(&value)
+    }
+}
+
 impl From<&[u8; Self::FULL_BYTES]> for Scalar {
     fn from(value: &[u8; Self::FULL_BYTES]) -> Self {
         Scalar(Fr::from_le_bytes_mod_order(value))
     }
 }
 
+impl From<[u8; Self::FULL_BYTES]> for Scalar {
+    fn from(value: [u8; Self::FULL_BYTES]) -> Self {
+        Self::from(&value)
+    }
+}
+
 impl From<&Scalar> for [u8; Scalar::FULL_BYTES] {
-    fn from(value: &Scalar) -> [u8; Scalar::FULL_BYTES] {
-        let mut bytes = [0u8; Scalar::FULL_BYTES];
+    fn from(value: &Scalar) -> Self {
+        let mut bytes = Self::default();
         value.write_to_bytes(&mut bytes);
         bytes
     }
 }
 
+impl From<Scalar> for [u8; Scalar::FULL_BYTES] {
+    fn from(value: Scalar) -> Self {
+        Self::from(&value)
+    }
+}
+
 impl Scalar {
     /// How many full bytes can be stored in BLS12-381 scalar (for instance before encoding). It is
     /// actually 254 bits, but bits are mut harder to work with and likely not worth it.

crates/subspace-erasure-coding/Cargo.toml

@@ -0,0 +1,27 @@
+[package]
+name = "subspace-erasure-coding"
+description = "Polynomial erasure coding implementation used in Subspace Network"
+license = "Apache-2.0"
+version = "0.1.0"
+authors = ["Nazar Mokrynskyi <nazar@mokrynskyi.com>"]
+edition = "2021"
+include = [
+    "/src",
+    "/Cargo.toml",
+]
+
+[dependencies]
+blst_from_scratch = { git = "https://github.com/sifraitech/rust-kzg", rev = "7eb52ca97576ea1eefe4dd2165f224c916f8c862", default-features = false }
+kzg = { git = "https://github.com/sifraitech/rust-kzg", rev = "7eb52ca97576ea1eefe4dd2165f224c916f8c862", default-features = false }
+subspace-core-primitives = { version = "0.1.0", path = "../subspace-core-primitives", default-features = false }
+
+[dev-dependencies]
+criterion = "0.4.0"
+rand = "0.8.5"
+
+[features]
+default = ["std"]
+std = [
+    "blst_from_scratch/std",
+    "subspace-core-primitives/std",
+]

crates/subspace-erasure-coding/src/lib.rs

@@ -0,0 +1,103 @@
+#![cfg_attr(not(feature = "std"), no_std)]
+
+extern crate alloc;
+
+#[cfg(all(test, features = "std"))]
+mod tests;
+
+use alloc::format;
+use alloc::string::{String, ToString};
+use alloc::vec::Vec;
+use blst_from_scratch::types::fft_settings::FsFFTSettings;
+use blst_from_scratch::types::fr::FsFr;
+use blst_from_scratch::types::poly::FsPoly;
+use core::num::NonZeroUsize;
+use kzg::{FFTSettings, PolyRecover, DAS};
+use subspace_core_primitives::Scalar;
+
+/// Erasure coding abstraction.
+///
+/// Supports creation of parity records and recovery of missing data.
+#[derive(Debug, Clone)]
+pub struct ErasureCoding {
+    fft_settings: 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 = FsFFTSettings::new(scale.get())?;
+
+        Ok(Self { fft_settings })
+    }
+
+    /// Extend sources using erasure coding.
+    ///
+    /// Returns parity data.
+    pub fn extend(&self, source: &[Scalar]) -> Result<Vec<Scalar>, String> {
+        // TODO: Once our scalars are based on `blst_from_scratch` we can use a bit of transmute to
+        //  avoid allocation here
+        // TODO: das_fft_extension modifies buffer internally, it needs to change to use
+        //  pre-allocated buffer instead of allocating a new one
+        let source = source
+            .iter()
+            .map(|scalar| {
+                FsFr::from_scalar(scalar.to_bytes())
+                    .map_err(|error| format!("Failed to convert scalar: {error}"))
+            })
+            .collect::<Result<Vec<_>, String>>()?;
+        let parity = self
+            .fft_settings
+            .das_fft_extension(&source)?
+            .into_iter()
+            .map(|scalar| {
+                // This is fine, scalar is guaranteed to be correct here
+                Scalar::from(scalar.to_scalar())
+            })
+            .collect();
+
+        Ok(parity)
+    }
+
+    /// 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> {
+        // TODO This is only necessary because upstream silently doesn't recover anything:
+        //  https://github.com/sifraitech/rust-kzg/issues/195
+        if shards.iter().filter(|scalar| scalar.is_some()).count() < self.fft_settings.max_width / 2
+        {
+            return Err("Impossible to recover, too many shards are missing".to_string());
+        }
+        // TODO: Once our scalars are based on `blst_from_scratch` we can use a bit of transmute to
+        //  avoid allocation here
+        let shards = shards
+            .iter()
+            .map(|maybe_scalar| {
+                maybe_scalar
+                    .map(|scalar| {
+                        FsFr::from_scalar(scalar.into())
+                            .map_err(|error| format!("Failed to convert scalar: {error}"))
+                    })
+                    .transpose()
+            })
+            .collect::<Result<Vec<_>, _>>()?;
+        let poly = <FsPoly as PolyRecover<FsFr, FsPoly, _>>::recover_poly_from_samples(
+            &shards,
+            &self.fft_settings,
+        )?;
+
+        Ok(poly
+            .coeffs
+            .iter()
+            .map(|scalar| {
+                // This is fine, scalar is guaranteed to be correct here
+                Scalar::from(scalar.to_scalar())
+            })
+            .collect())
+    }
+}

crates/subspace-erasure-coding/src/tests.rs

@@ -0,0 +1,109 @@
+use crate::ErasureCoding;
+use std::iter;
+use std::num::NonZeroUsize;
+use subspace_core_primitives::Scalar;
+
+// TODO: This could have been done in-place, once implemented can be exposed as a utility
+fn concatenated_to_interleaved<T>(input: Vec<T>) -> Vec<T>
+where
+    T: Clone,
+{
+    if input.len() <= 1 {
+        return input;
+    }
+
+    let (first_half, second_half) = input.split_at(input.len() / 2);
+
+    first_half
+        .iter()
+        .zip(second_half)
+        .flat_map(|(a, b)| [a, b])
+        .cloned()
+        .collect()
+}
+
+// TODO: This could have been done in-place, once implemented can be exposed as a utility
+fn interleaved_to_concatenated<T>(input: Vec<T>) -> Vec<T>
+where
+    T: Clone,
+{
+    let first_half = input.iter().step_by(2);
+    let second_half = input.iter().skip(1).step_by(2);
+
+    first_half.chain(second_half).cloned().collect()
+}
+
+#[test]
+fn basic() {
+    let scale = NonZeroUsize::new(8).unwrap();
+    let num_shards = 2usize.pow(scale.get() as u32);
+    let ec = ErasureCoding::new(scale).unwrap();
+
+    let source_shards = (0..num_shards / 2)
+        .map(|_| rand::random::<[u8; Scalar::SAFE_BYTES]>())
+        .map(Scalar::from)
+        .collect::<Vec<_>>();
+
+    let parity_shards = ec.extend(&source_shards).unwrap();
+
+    assert_ne!(source_shards, parity_shards);
+
+    let partial_shards = concatenated_to_interleaved(
+        iter::repeat(None)
+            .take(num_shards / 4)
+            .chain(source_shards.iter().skip(num_shards / 4).copied().map(Some))
+            .chain(parity_shards.iter().take(num_shards / 4).copied().map(Some))
+            .chain(iter::repeat(None).take(num_shards / 4))
+            .collect::<Vec<_>>(),
+    );
+
+    let recovered = interleaved_to_concatenated(ec.recover(&partial_shards).unwrap());
+
+    assert_eq!(
+        recovered,
+        source_shards
+            .iter()
+            .chain(&parity_shards)
+            .copied()
+            .collect::<Vec<_>>()
+    );
+}
+
+#[test]
+fn bad_shards_number() {
+    let scale = NonZeroUsize::new(8).unwrap();
+    let num_shards = 2usize.pow(scale.get() as u32);
+    let ec = ErasureCoding::new(scale).unwrap();
+
+    let source_shards = vec![Default::default(); num_shards - 1];
+
+    assert!(ec.extend(&source_shards).is_err());
+
+    let partial_shards = vec![Default::default(); num_shards - 1];
+    assert!(ec.recover(&partial_shards).is_err());
+}
+
+#[test]
+fn not_enough_partial() {
+    let scale = NonZeroUsize::new(8).unwrap();
+    let num_shards = 2usize.pow(scale.get() as u32);
+    let ec = ErasureCoding::new(scale).unwrap();
+
+    let mut partial_shards = vec![None; num_shards];
+
+    // Less than half is not sufficient
+    partial_shards
+        .iter_mut()
+        .take(num_shards / 2 - 1)
+        .for_each(|maybe_scalar| {
+            maybe_scalar.replace(Scalar::default());
+        });
+    assert!(ec.recover(&partial_shards).is_err());
+
+    // Any half is sufficient
+    partial_shards
+        .last_mut()
+        .unwrap()
+        .replace(Scalar::default());
+    assert!(ec.recover(&partial_shards).is_ok());
+}