From 6efd3fabcd07f4926fd9e517e5db01a7898aa781 Mon Sep 17 00:00:00 2001
From: Nazar Mokrynskyi <nazar@mokrynskyi.com>
Date: Fri, 12 May 2023 11:46:07 +0300
Subject: [PATCH 1/2] Update `rust-kzg` to latest revision with new APIs

---
 Cargo.lock                                           | 10 +++++-----
 Cargo.toml                                           |  2 +-
 crates/subspace-core-primitives/Cargo.toml           |  8 ++++----
 crates/subspace-core-primitives/src/crypto.rs        |  2 +-
 crates/subspace-core-primitives/src/crypto/kzg.rs    | 10 +++++-----
 .../subspace-core-primitives/src/crypto/kzg/tests.rs | 12 ++++++------
 crates/subspace-erasure-coding/Cargo.toml            | 10 +++++-----
 .../subspace-erasure-coding/benches/commitments.rs   |  2 +-
 crates/subspace-erasure-coding/src/lib.rs            |  8 ++++----
 crates/subspace-erasure-coding/src/tests.rs          |  2 +-
 10 files changed, 33 insertions(+), 33 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index bd3d8f9175..536454f93a 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -922,9 +922,9 @@ dependencies = [
 ]
 
 [[package]]
-name = "blst_from_scratch"
+name = "blst_rust"
 version = "0.1.0"
-source = "git+https://github.com/subspace/rust-kzg?rev=b71e7ace37420db889488cbd64d53eea38176111#b71e7ace37420db889488cbd64d53eea38176111"
+source = "git+https://github.com/subspace/rust-kzg?rev=1058cc8c8af8461b490dc212c41d7d506a746577#1058cc8c8af8461b490dc212c41d7d506a746577"
 dependencies = [
  "blst",
  "kzg",
@@ -4856,7 +4856,7 @@ dependencies = [
 [[package]]
 name = "kzg"
 version = "0.1.0"
-source = "git+https://github.com/subspace/rust-kzg?rev=b71e7ace37420db889488cbd64d53eea38176111#b71e7ace37420db889488cbd64d53eea38176111"
+source = "git+https://github.com/subspace/rust-kzg?rev=1058cc8c8af8461b490dc212c41d7d506a746577#1058cc8c8af8461b490dc212c41d7d506a746577"
 dependencies = [
  "blst",
  "sha2 0.10.6",
@@ -11266,7 +11266,7 @@ name = "subspace-core-primitives"
 version = "0.1.0"
 dependencies = [
  "blake2",
- "blst_from_scratch",
+ "blst_rust",
  "criterion",
  "derive_more",
  "hex",
@@ -11292,7 +11292,7 @@ dependencies = [
 name = "subspace-erasure-coding"
 version = "0.1.0"
 dependencies = [
- "blst_from_scratch",
+ "blst_rust",
  "criterion",
  "kzg",
  "rand 0.8.5",
diff --git a/Cargo.toml b/Cargo.toml
index 086aa99571..f55c0e20af 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -28,7 +28,7 @@ blake2 = { opt-level = 3 }
 blake3 = { opt-level = 3 }
 blake2b_simd = { opt-level = 3 }
 blst = { opt-level = 3 }
-blst_from_scratch = { opt-level = 3 }
+blst_rust = { opt-level = 3 }
 chacha20 = { opt-level = 3 }
 chacha20poly1305 = { opt-level = 3 }
 cranelift-codegen = { opt-level = 3 }
diff --git a/crates/subspace-core-primitives/Cargo.toml b/crates/subspace-core-primitives/Cargo.toml
index dad5860b1c..394d1ed5c1 100644
--- a/crates/subspace-core-primitives/Cargo.toml
+++ b/crates/subspace-core-primitives/Cargo.toml
@@ -18,11 +18,11 @@ bench = false
 [dependencies]
 blake2 = { version = "0.10.6", default-features = false }
 # TODO: Switch to upstream `main` once https://github.com/sifraitech/rust-kzg/pull/204 is merged and blst has upstream no_std support
-blst_from_scratch = { git = "https://github.com/subspace/rust-kzg", rev = "b71e7ace37420db889488cbd64d53eea38176111", default-features = false }
+blst_rust = { git = "https://github.com/subspace/rust-kzg", rev = "1058cc8c8af8461b490dc212c41d7d506a746577", default-features = false }
 derive_more = "0.99.17"
 hex = { version  = "0.4.3", default-features = false, features = ["alloc"] }
 # TODO: Switch to upstream `main` once https://github.com/sifraitech/rust-kzg/pull/204 is merged and blst has upstream no_std support
-kzg = { git = "https://github.com/subspace/rust-kzg", rev = "b71e7ace37420db889488cbd64d53eea38176111", default-features = false }
+kzg = { git = "https://github.com/subspace/rust-kzg", rev = "1058cc8c8af8461b490dc212c41d7d506a746577", default-features = false }
 num-traits = { version = "0.2.15", default-features = false }
 parity-scale-codec = { version = "3.4.0", default-features = false, features = ["derive", "max-encoded-len"] }
 parking_lot = { version = "0.12.1", optional = true }
@@ -53,7 +53,7 @@ default = [
 embedded-kzg-settings = []
 # Enables some APIs and internal parallelism for KZG
 parallel = [
-    "blst_from_scratch/parallel",
+    "blst_rust/parallel",
     "dep:rayon",
 ]
 serde = [
@@ -64,7 +64,7 @@ serde = [
 ]
 std = [
     "blake2/std",
-    "blst_from_scratch/std",
+    "blst_rust/std",
     "hex/std",
     "kzg/std",
     "num-traits/std",
diff --git a/crates/subspace-core-primitives/src/crypto.rs b/crates/subspace-core-primitives/src/crypto.rs
index 70d98ff2d4..5a279c2566 100644
--- a/crates/subspace-core-primitives/src/crypto.rs
+++ b/crates/subspace-core-primitives/src/crypto.rs
@@ -26,7 +26,7 @@ use alloc::vec::Vec;
 use blake2::digest::typenum::U32;
 use blake2::digest::{FixedOutput, Update};
 use blake2::{Blake2b, Blake2bMac, Digest};
-use blst_from_scratch::types::fr::FsFr;
+use blst_rust::types::fr::FsFr;
 use core::cmp::Ordering;
 use core::hash::{Hash, Hasher};
 use core::mem;
diff --git a/crates/subspace-core-primitives/src/crypto/kzg.rs b/crates/subspace-core-primitives/src/crypto/kzg.rs
index 774059bad9..ad9bca445d 100644
--- a/crates/subspace-core-primitives/src/crypto/kzg.rs
+++ b/crates/subspace-core-primitives/src/crypto/kzg.rs
@@ -13,11 +13,11 @@ use alloc::collections::BTreeMap;
 use alloc::string::{String, ToString};
 use alloc::sync::Arc;
 use alloc::vec::Vec;
-use blst_from_scratch::types::fft_settings::FsFFTSettings;
-use blst_from_scratch::types::g1::FsG1;
-use blst_from_scratch::types::g2::FsG2;
-use blst_from_scratch::types::kzg_settings::FsKZGSettings;
-use blst_from_scratch::types::poly::FsPoly;
+use blst_rust::types::fft_settings::FsFFTSettings;
+use blst_rust::types::g1::FsG1;
+use blst_rust::types::g2::FsG2;
+use blst_rust::types::kzg_settings::FsKZGSettings;
+use blst_rust::types::poly::FsPoly;
 use core::hash::{Hash, Hasher};
 use core::mem;
 use derive_more::{AsMut, AsRef, Deref, DerefMut, From, Into};
diff --git a/crates/subspace-core-primitives/src/crypto/kzg/tests.rs b/crates/subspace-core-primitives/src/crypto/kzg/tests.rs
index 76e25f10cb..1420234b7b 100644
--- a/crates/subspace-core-primitives/src/crypto/kzg/tests.rs
+++ b/crates/subspace-core-primitives/src/crypto/kzg/tests.rs
@@ -1,11 +1,11 @@
 use crate::crypto::kzg::{embedded_kzg_settings, Kzg};
 use crate::crypto::Scalar;
-use blst_from_scratch::consts::{G1_GENERATOR, G2_GENERATOR};
-use blst_from_scratch::types::fft_settings::FsFFTSettings;
-use blst_from_scratch::types::fr::FsFr;
-use blst_from_scratch::types::g1::FsG1;
-use blst_from_scratch::types::g2::FsG2;
-use blst_from_scratch::types::kzg_settings::FsKZGSettings;
+use blst_rust::consts::{G1_GENERATOR, G2_GENERATOR};
+use blst_rust::types::fft_settings::FsFFTSettings;
+use blst_rust::types::fr::FsFr;
+use blst_rust::types::g1::FsG1;
+use blst_rust::types::g2::FsG2;
+use blst_rust::types::kzg_settings::FsKZGSettings;
 use kzg::{FFTSettings, Fr, G1Mul, G2Mul};
 use rand::Rng;
 use rand_core::SeedableRng;
diff --git a/crates/subspace-erasure-coding/Cargo.toml b/crates/subspace-erasure-coding/Cargo.toml
index dd332fcfa3..0a0b370375 100644
--- a/crates/subspace-erasure-coding/Cargo.toml
+++ b/crates/subspace-erasure-coding/Cargo.toml
@@ -16,25 +16,25 @@ bench = false
 
 [dependencies]
 # TODO: Switch to upstream `main` once https://github.com/sifraitech/rust-kzg/pull/204 is merged and blst has upstream no_std support
-blst_from_scratch = { git = "https://github.com/subspace/rust-kzg", rev = "b71e7ace37420db889488cbd64d53eea38176111", default-features = false }
+blst_rust = { git = "https://github.com/subspace/rust-kzg", rev = "1058cc8c8af8461b490dc212c41d7d506a746577", default-features = false }
 # TODO: Switch to upstream `main` once https://github.com/sifraitech/rust-kzg/pull/204 is merged and blst has upstream no_std support
-kzg = { git = "https://github.com/subspace/rust-kzg", rev = "b71e7ace37420db889488cbd64d53eea38176111", default-features = false }
+kzg = { git = "https://github.com/subspace/rust-kzg", rev = "1058cc8c8af8461b490dc212c41d7d506a746577", default-features = false }
 subspace-core-primitives = { version = "0.1.0", path = "../subspace-core-primitives", default-features = false }
 
 [dev-dependencies]
 # TODO: Switch to upstream `main` once https://github.com/sifraitech/rust-kzg/pull/204 is merged and blst has upstream no_std support
-blst_from_scratch = { git = "https://github.com/subspace/rust-kzg", rev = "b71e7ace37420db889488cbd64d53eea38176111" }
+blst_rust = { git = "https://github.com/subspace/rust-kzg", rev = "1058cc8c8af8461b490dc212c41d7d506a746577" }
 criterion = "0.4.0"
 rand = "0.8.5"
 
 [features]
 default = ["std", "parallel"]
 std = [
-    "blst_from_scratch/std",
+    "blst_rust/std",
     "kzg/std",
     "subspace-core-primitives/std",
 ]
-parallel = ["blst_from_scratch/parallel"]
+parallel = ["blst_rust/parallel"]
 
 [[bench]]
 name = "commitments"
diff --git a/crates/subspace-erasure-coding/benches/commitments.rs b/crates/subspace-erasure-coding/benches/commitments.rs
index a6b7a1568d..9ef3e9781a 100644
--- a/crates/subspace-erasure-coding/benches/commitments.rs
+++ b/crates/subspace-erasure-coding/benches/commitments.rs
@@ -1,4 +1,4 @@
-use blst_from_scratch::types::g1::FsG1;
+use blst_rust::types::g1::FsG1;
 use criterion::{black_box, criterion_group, criterion_main, Criterion};
 use kzg::G1;
 use std::num::NonZeroUsize;
diff --git a/crates/subspace-erasure-coding/src/lib.rs b/crates/subspace-erasure-coding/src/lib.rs
index a70c08e698..4d6d439ee0 100644
--- a/crates/subspace-erasure-coding/src/lib.rs
+++ b/crates/subspace-erasure-coding/src/lib.rs
@@ -8,10 +8,10 @@ extern crate alloc;
 use alloc::string::{String, ToString};
 use alloc::sync::Arc;
 use alloc::vec::Vec;
-use blst_from_scratch::types::fft_settings::FsFFTSettings;
-use blst_from_scratch::types::fr::FsFr;
-use blst_from_scratch::types::g1::FsG1;
-use blst_from_scratch::types::poly::FsPoly;
+use blst_rust::types::fft_settings::FsFFTSettings;
+use blst_rust::types::fr::FsFr;
+use blst_rust::types::g1::FsG1;
+use blst_rust::types::poly::FsPoly;
 use core::num::NonZeroUsize;
 use kzg::{FFTSettings, PolyRecover, DAS, FFTG1, G1};
 use subspace_core_primitives::crypto::kzg::Commitment;
diff --git a/crates/subspace-erasure-coding/src/tests.rs b/crates/subspace-erasure-coding/src/tests.rs
index e2d09d0282..d15d26ebd8 100644
--- a/crates/subspace-erasure-coding/src/tests.rs
+++ b/crates/subspace-erasure-coding/src/tests.rs
@@ -1,5 +1,5 @@
 use crate::ErasureCoding;
-use blst_from_scratch::types::g1::FsG1;
+use blst_rust::types::g1::FsG1;
 use kzg::G1;
 use std::iter;
 use std::num::NonZeroUsize;

From 3793366f285f1d3599a0480addaa713a748dae66 Mon Sep 17 00:00:00 2001
From: Nazar Mokrynskyi <nazar@mokrynskyi.com>
Date: Fri, 12 May 2023 12:01:28 +0300
Subject: [PATCH 2/2] Remove extra FFTs using new APIs from rust-kzg

---
 .../src/crypto/kzg.rs                         |  2 +-
 crates/subspace-erasure-coding/src/lib.rs     | 28 +++++++----
 .../subspace-farmer-components/src/proving.rs | 46 +++----------------
 3 files changed, 27 insertions(+), 49 deletions(-)

diff --git a/crates/subspace-core-primitives/src/crypto/kzg.rs b/crates/subspace-core-primitives/src/crypto/kzg.rs
index ad9bca445d..e5443db70e 100644
--- a/crates/subspace-core-primitives/src/crypto/kzg.rs
+++ b/crates/subspace-core-primitives/src/crypto/kzg.rs
@@ -92,7 +92,7 @@ pub fn embedded_kzg_settings() -> FsKZGSettings {
 }
 
 /// Commitment to polynomial
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, From)]
 pub struct Polynomial(FsPoly);
 
 impl Polynomial {
diff --git a/crates/subspace-erasure-coding/src/lib.rs b/crates/subspace-erasure-coding/src/lib.rs
index 4d6d439ee0..a6d97ea39f 100644
--- a/crates/subspace-erasure-coding/src/lib.rs
+++ b/crates/subspace-erasure-coding/src/lib.rs
@@ -5,16 +5,15 @@ mod tests;
 
 extern crate alloc;
 
-use alloc::string::{String, ToString};
+use alloc::string::String;
 use alloc::sync::Arc;
 use alloc::vec::Vec;
 use blst_rust::types::fft_settings::FsFFTSettings;
-use blst_rust::types::fr::FsFr;
 use blst_rust::types::g1::FsG1;
 use blst_rust::types::poly::FsPoly;
 use core::num::NonZeroUsize;
 use kzg::{FFTSettings, PolyRecover, DAS, FFTG1, G1};
-use subspace_core_primitives::crypto::kzg::Commitment;
+use subspace_core_primitives::crypto::kzg::{Commitment, Polynomial};
 use subspace_core_primitives::crypto::Scalar;
 
 /// Erasure coding abstraction.
@@ -57,12 +56,7 @@ impl ErasureCoding {
     /// 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_none()).count() > shards.len() / 2 {
-            return Err("Impossible to recover, too many shards are missing".to_string());
-        }
-        let poly = <FsPoly as PolyRecover<FsFr, FsPoly, _>>::recover_poly_from_samples(
+        let poly = FsPoly::recover_poly_from_samples(
             Scalar::slice_option_to_repr(shards),
             &self.fft_settings,
         )?;
@@ -70,6 +64,22 @@ impl ErasureCoding {
         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
diff --git a/crates/subspace-farmer-components/src/proving.rs b/crates/subspace-farmer-components/src/proving.rs
index c244792d99..c845f6fddd 100644
--- a/crates/subspace-farmer-components/src/proving.rs
+++ b/crates/subspace-farmer-components/src/proving.rs
@@ -1,11 +1,8 @@
 use crate::auditing::ChunkCandidate;
-use crate::reading::{
-    read_record_metadata, read_sector_record_chunks, recover_extended_record_chunks, ReadingError,
-};
+use crate::reading::{read_record_metadata, read_sector_record_chunks, ReadingError};
 use crate::sector::{SectorContentsMap, SectorContentsMapFromBytesError, SectorMetadata};
 use std::collections::VecDeque;
 use std::marker::PhantomData;
-use std::mem::ManuallyDrop;
 use subspace_core_primitives::crypto::kzg::{Commitment, Kzg, Witness};
 use subspace_core_primitives::crypto::Scalar;
 use subspace_core_primitives::{
@@ -226,43 +223,14 @@ where
                     .expect("Within s-bucket range; qed")
                     .expect("Winning chunk was plotted; qed");
 
-                let extended_chunks = recover_extended_record_chunks(
-                    &sector_record_chunks,
-                    piece_offset,
-                    self.erasure_coding,
-                )?;
-
-                // A bit complicated way to avoid re-allocation in performance-sensitive place
-                let source_chunks = {
-                    let mut extended_chunks = ManuallyDrop::new(extended_chunks);
-
-                    // SAFETY: Original memory is not dropped, size of the data is statically known
-                    let mut extended_chunks = unsafe {
-                        Vec::from_raw_parts(
-                            extended_chunks.as_mut_ptr(),
-                            extended_chunks.len(),
-                            extended_chunks.len(),
-                        )
-                    };
-
-                    // Move source chunks into the first half of the vector
-                    for i in 0..Record::NUM_CHUNKS {
-                        extended_chunks[i] =
-                            extended_chunks[i * Record::NUM_S_BUCKETS / Record::NUM_CHUNKS];
-                    }
-
-                    // Shrink vector to just the source chunks without re-allocating
-                    extended_chunks.truncate(Record::NUM_CHUNKS);
-
-                    extended_chunks
-                };
-
-                let source_chunks_polynomial = self.kzg.poly(&source_chunks).map_err(|error| {
-                    ProvingError::FailedToCreatePolynomialForRecord {
+                let source_chunks_polynomial = self
+                    .erasure_coding
+                    .recover_poly(sector_record_chunks.as_slice())
+                    .map_err(|error| ReadingError::FailedToErasureDecodeRecord {
                         piece_offset,
                         error,
-                    }
-                })?;
+                    })?;
+                drop(sector_record_chunks);
 
                 let (record_commitment, record_witness) = read_record_metadata(
                     piece_offset,