fix: Fix poly_to_segment_vec & recovery_row_from_segments

crates/melo-erasure-coding/src/erasure_coding.rs

@@ -35,22 +35,20 @@ pub fn extend(fs: &FsFFTSettings, source: &[BlsScalar]) -> Result<Vec<BlsScalar>
 		.map(BlsScalar::vec_from_repr)
 }
 
-pub fn extend_poly(fs: &FsFFTSettings, poly: &Polynomial) -> Result<Vec<FsFr>, String> {
+pub fn extend_poly(fs: &FsFFTSettings, poly: &Polynomial) -> Result<Vec<BlsScalar>, String> {
 	let mut coeffs = poly.0.coeffs.clone();
 	coeffs.resize(coeffs.len() * 2, FsFr::zero());
 	let mut extended_coeffs_fft = fs.fft_fr(&coeffs, false).unwrap();
 	reverse_bit_order(&mut extended_coeffs_fft);
-	Ok(extended_coeffs_fft)
+	Ok(BlsScalar::vec_from_repr(extended_coeffs_fft))
 }
 
 pub fn extend_fs_g1<T: ReprConvert<FsG1>>(
 	fs: &FsFFTSettings,
 	source: &[T],
 ) -> Result<Vec<T>, String> {
 	let mut coeffs = fs.fft_g1(T::slice_to_repr(source), true)?;
-
 	coeffs.resize(coeffs.len() * 2, FsG1::identity());
-
 	fs.fft_g1(&coeffs, false).map(T::vec_from_repr)
 }
 
@@ -64,6 +62,15 @@ pub fn recover_poly(
 	fs: &FsFFTSettings,
 	shards: &[Option<BlsScalar>],
 ) -> Result<Polynomial, String> {
-	let poly = FsPoly::recover_poly_from_samples(BlsScalar::slice_option_to_repr(shards), &fs)?;
-	Ok(Polynomial::from(poly))
+	// let poly = FsPoly::recover_poly_coeffs_from_samples(BlsScalar::slice_option_to_repr(shards), &fs)?;
+	// Ok(Polynomial::from(poly))
+
+	let mut poly = Polynomial::from(FsPoly::recover_poly_coeffs_from_samples(
+		BlsScalar::slice_option_to_repr(shards),
+		&fs,
+	)?);
+
+	poly.normalize();
+
+	Ok(poly)
 }

crates/melo-erasure-coding/src/recovery.rs

@@ -14,21 +14,27 @@
 
 extern crate alloc;
 
-use crate::erasure_coding::recover_poly;
+use crate::erasure_coding::{extend_poly, recover_poly};
 use crate::segment::{order_segments_row, segment_datas_to_row};
 use melo_core_primitives::kzg::{Position, KZG};
 use melo_core_primitives::segment::{Segment, SegmentData};
+use rust_kzg_blst::utils::reverse_bit_order;
 
 pub fn recovery_row_from_segments(
 	segments: &Vec<Segment>,
 	kzg: &KZG,
+	chunk_count: usize,
 ) -> Result<Vec<Segment>, String> {
 	let y = segments[0].position.y;
 	let segments_size = segments[0].size();
-	let order_segments = order_segments_row(&segments)?;
-	let row = segment_datas_to_row(&order_segments, segments_size);
+	let order_segments = order_segments_row(&segments, chunk_count)?;
+	let mut row = segment_datas_to_row(&order_segments, segments_size);
+	reverse_bit_order(&mut row);
 	let poly = recover_poly(kzg.get_fs(), &row)?;
-	let recovery_row = poly.to_bls_scalars();
+
+	let recovery_row = extend_poly(kzg.get_fs(), &poly)?;
+	// let mut recovery_row = recover(kzg.get_fs(), &row)?;
+	// reverse_bit_order(&mut recovery_row);
 	order_segments
 		.iter()
 		.enumerate()
@@ -39,14 +45,9 @@ pub fn recovery_row_from_segments(
 				None => {
 					let index = i * segments_size;
 					let data = recovery_row[index..(i + 1) * segments_size].to_vec();
-					let segment_data = SegmentData::from_data(
-						&position,
-						&data,
-						kzg,
-						&poly,
-						segments.len(),
-					)
-					.map_err(|e| e.to_string())?;
+					let segment_data =
+						SegmentData::from_data(&position, &data, kzg, &poly, segments.len())
+							.map_err(|e| e.to_string())?;
 					Ok(Segment { position, content: segment_data })
 				},
 			}

crates/melo-erasure-coding/src/segment.rs

@@ -14,19 +14,18 @@
 
 extern crate alloc;
 use kzg::FK20MultiSettings;
-use melo_core_primitives::config::FIELD_ELEMENTS_PER_BLOB;
-use melo_core_primitives::kzg::{BlsScalar, KZGProof, Polynomial, Position, ReprConvert, KZG};
+use melo_core_primitives::kzg::{BlsScalar, KZGProof, Polynomial, Position, KZG};
 use melo_core_primitives::segment::{Segment, SegmentData};
 use rust_kzg_blst::types::fk20_multi_settings::FsFK20MultiSettings;
 
 use crate::erasure_coding::extend_poly;
 
-pub fn order_segments_row(segments: &Vec<Segment>) -> Result<Vec<Option<SegmentData>>, String> {
-	if segments.len() > FIELD_ELEMENTS_PER_BLOB * 2 {
+pub fn order_segments_row(segments: &Vec<Segment>, chunk_count: usize) -> Result<Vec<Option<SegmentData>>, String> {
+	if segments.len() > chunk_count * 2 {
 		return Err("segments x not equal".to_string());
 	}
 	let y = segments[0].position.y;
-	let mut ordered_segments = vec![None; FIELD_ELEMENTS_PER_BLOB * 2];
+	let mut ordered_segments = vec![None; chunk_count * 2];
 	for segment in segments.iter() {
 		if segment.position.y != y {
 			return Err("segments y not equal".to_string());
@@ -82,15 +81,14 @@ pub fn poly_to_segment_vec(poly: &Polynomial, kzg: &KZG, y: usize, chunk_size: u
 	let poly_len = poly.0.coeffs.len();
 	let fk = FsFK20MultiSettings::new(&kzg.ks, 2 * poly_len, chunk_size).unwrap();
 	let all_proofs = fk.data_availability(&poly.0).unwrap();
-
-	let segments = extend_poly(&fk.kzg_settings.fs, &poly)?
+	let extended_poly = extend_poly(&fk.kzg_settings.fs, &poly)?;
+	let segments = extended_poly
 		.chunks(chunk_size)
 		.enumerate()
 		.map(|(i, chunk)| {
 			let position = Position { y: y as u32, x: i as u32 };
-			let content = chunk;
 			let proof = all_proofs[i];
-			Segment::new(position, BlsScalar::slice_from_repr(content), KZGProof(proof))
+			Segment::new(position, chunk, KZGProof(proof))
 		})
 		.collect::<Vec<_>>();
 

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

@@ -4,16 +4,22 @@ use crate::recovery::*;
 use crate::segment::*;
 // use blst_rust::types::g1::FsG1;
 use alloc::vec;
-use kzg::G1;
-use rust_kzg_blst::types::g1::FsG1;
+use melo_core_primitives::segment::SegmentData;
+use core::cell;
 use core::slice::Chunks;
 use kzg::FFTFr;
 use kzg::Fr;
+use kzg::G1;
 use melo_core_primitives::kzg::BlsScalar;
+use melo_core_primitives::kzg::Position;
+use melo_core_primitives::kzg::KZGProof;
 use melo_core_primitives::kzg::Polynomial;
 use melo_core_primitives::kzg::ReprConvert;
-use melo_core_primitives::kzg::{embedded_kzg_settings, KZG, KZGCommitment};
+use melo_core_primitives::kzg::{embedded_kzg_settings, KZGCommitment, KZG};
+use melo_core_primitives::segment::Segment;
+use rand::seq::SliceRandom;
 use rust_kzg_blst::types::fr::FsFr;
+use rust_kzg_blst::types::g1::FsG1;
 use rust_kzg_blst::types::poly::FsPoly;
 use rust_kzg_blst::utils::reverse_bit_order;
 use std::iter;
@@ -33,31 +39,106 @@ fn random_poly(s: usize) -> Polynomial {
 	Polynomial::from(poly)
 }
 
-// #[test]
-// fn recovery_row_from_segments_test() {
-// 	let scale = NonZeroUsize::new(4).unwrap();
-// 	let kzg = KZG::new(embedded_kzg_settings());
-// 	let num_shards = 2usize.pow(scale.get() as u32);
-// 	let source_shards = (0..num_shards / 2)
-// 		.map(|_| rand::random::<[u8; 31]>())
-// 		.map(BlsScalar::from)
-// 		.collect::<Vec<_>>();
+fn random_vec(s: usize) -> Vec<usize> {
+	let mut positions: Vec<usize> = (0..s).collect();
+	positions.shuffle(&mut rand::thread_rng());
+	positions
+}
 
-// 	let parity_shards = extend(kzg.get_fs(), &source_shards).unwrap();
+#[test]
+fn segment_datas_to_row_test() {
+	// 构建随机 segment datas
+	let chunk_len: usize = 16;
+	let chunk_count: usize = 4;
+	let num_shards = chunk_len * chunk_count;
+	let mut segment_datas: Vec<Option<SegmentData>> = Vec::new();
+	for _ in 0..chunk_count {
+		let data = (0..chunk_len)
+			.map(|_| rand::random::<[u8; 31]>())
+			.map(BlsScalar::from)
+			.collect::<Vec<_>>();
+		let proof = KZGProof(FsG1::rand());
+		let segment_data = SegmentData { data, proof };
+		segment_datas.push(Some(segment_data));
+	}
+	segment_datas[2] = None;
+	segment_datas[3] = None;
+	// 转换为 row
+	let row = segment_datas_to_row(&segment_datas, chunk_len);
+	// 验证是否正确
+	for i in 0..num_shards {
+		let data = match segment_datas[i / chunk_len] {
+			Some(ref segment_data) => Some(segment_data.data[i % chunk_len]),
+			None => None,
+		};
+		assert_eq!(row[i], data);
+	}
+}
 
-// 	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<_>>(),
-// 	);
+#[test]
+fn extend_poly_random_round_trip_test() {
+	let kzg = KZG::new(embedded_kzg_settings());
+	let num_shards = 16;
+	let poly = random_poly(num_shards);
 
-// 	let recovered = interleaved_to_concatenated(recover(kzg.get_fs(), &partial_shards).unwrap());
+	let extended_poly = extend_poly(kzg.get_fs(), &poly).unwrap();
+	assert_eq!(extended_poly.len(), 32);
 
-// 	assert_eq!(recovered, source_shards.iter().chain(&parity_shards).copied().collect::<Vec<_>>());
-// }
+	let random_positions = random_vec(num_shards * 2);
+	let mut cells = [None; 32];
+	for i in 0..num_shards {
+		let position = random_positions[i];
+		cells[position] = Some(extended_poly[position]);
+	}
+	reverse_bit_order(&mut cells);
+	let mut recovered_poly = recover(kzg.get_fs(), &cells.as_slice()).unwrap();
+	reverse_bit_order(&mut recovered_poly);
+	for i in 0..num_shards * 2 {
+		assert_eq!(recovered_poly[i].0, extended_poly[i].0);
+	}
+}
+
+#[test]
+fn recover_segments_random() {
+    // Build a random polynomial
+    let chunk_len: usize = 16;
+    let chunk_count: usize = 4;
+    let num_shards = chunk_len * chunk_count;
+
+    let poly = random_poly(num_shards);
+    // Convert the polynomial to segments
+    let kzg = KZG::new(embedded_kzg_settings());
+    let segments: Vec<Segment> = poly_to_segment_vec(&poly, &kzg, 0, chunk_len).unwrap();
+    assert_eq!(segments.len(), 8);
+
+    // Take most of them randomly
+    let mut random_segments: Vec<Segment> = Vec::new();
+    // Get a random Vec of length num_shards, where each number is unique and less than 2 * num_shards
+    let random_positions = random_vec(2 * chunk_count);
+
+    for i in 0..chunk_count {
+        random_segments.push(segments[random_positions[i]].clone());
+    }
+
+    // Recover segments
+    let recovered_segments =
+        recovery_row_from_segments(&random_segments, &kzg, chunk_count).unwrap();
+    assert_eq!(recovered_segments[0], segments[0]);
+
+    // Verify if the recovered segments are the same as the original segments
+    for i in 0..chunk_count {
+        assert_eq!(recovered_segments[i], segments[i]);
+    }
+
+    // Remove one segment from random_segments
+    random_segments.remove(0);
+    // Recover segments
+    let recovered_segments =
+        recovery_row_from_segments(&random_segments, &kzg, chunk_count);
+
+    // Verify if it fails
+    assert!(recovered_segments.is_err());
+}
 
 #[test]
 fn commit_multi_random() {
@@ -72,7 +153,7 @@ fn commit_multi_random() {
 	// Commit to the polynomial
 	let commitment = kzg.commit(&poly).unwrap();
 	// Compute the multi proofs
-	let proofs = kzg.all_proofs(&poly,chunk_len).unwrap();
+	let proofs = kzg.all_proofs(&poly, chunk_len).unwrap();
 
 	let mut extended_coeffs = poly.0.coeffs.clone();
 
@@ -167,33 +248,32 @@ fn extend_fs_g1_random() {
 
 #[test]
 fn extend_segments_col_random() {
-    // Build multiple polynomials with random coefficients
-    let chunk_len: usize = 16;
-    let chunk_count: usize = 4;
-    let num_shards = chunk_len * chunk_count;
-    let k: usize = 4;
-    let polys = (0..k).map(|_| random_poly(num_shards)).collect::<Vec<_>>();
-    // Commit to all polynomials
-    let kzg = KZG::new(embedded_kzg_settings());
-    let commitments = polys.iter().map(|poly| kzg.commit(poly).unwrap()).collect::<Vec<_>>();
-    // Extend polynomial commitments to twice the size
-    let extended_commitments = extend_fs_g1(kzg.get_fs(), &commitments).unwrap();
-    // Convert all polynomials to segments
-    let matrix = polys
-        .iter()
-        .enumerate()
-        .map(|(i, poly)| poly_to_segment_vec(&poly, &kzg, i, chunk_len).unwrap())
-        .collect::<Vec<_>>();
-    assert!(matrix[0][0].verify(&kzg, &commitments[0], chunk_count).unwrap());
-    // Pick a column from the segments
-    let pick_col_index: usize = 1;
-    let col = matrix.iter().map(|row| row[pick_col_index].clone()).collect::<Vec<_>>();
-    // Extend the column
-    let extended_col = extend_segments_col(kzg.get_fs(), &col).unwrap();
-
-    for i in 0..(chunk_count) {
-        let pick_s = extended_col[i].clone();
-        assert!(pick_s.verify(&kzg, &extended_commitments[i * 2 + 1], chunk_count).unwrap());
-    }
-
-}
+	// Build multiple polynomials with random coefficients
+	let chunk_len: usize = 16;
+	let chunk_count: usize = 4;
+	let num_shards = chunk_len * chunk_count;
+	let k: usize = 4;
+	let polys = (0..k).map(|_| random_poly(num_shards)).collect::<Vec<_>>();
+	// Commit to all polynomials
+	let kzg = KZG::new(embedded_kzg_settings());
+	let commitments = polys.iter().map(|poly| kzg.commit(poly).unwrap()).collect::<Vec<_>>();
+	// Extend polynomial commitments to twice the size
+	let extended_commitments = extend_fs_g1(kzg.get_fs(), &commitments).unwrap();
+	// Convert all polynomials to segments
+	let matrix = polys
+		.iter()
+		.enumerate()
+		.map(|(i, poly)| poly_to_segment_vec(&poly, &kzg, i, chunk_len).unwrap())
+		.collect::<Vec<_>>();
+	assert!(matrix[0][0].verify(&kzg, &commitments[0], chunk_count).unwrap());
+	// Pick a column from the segments
+	let pick_col_index: usize = 1;
+	let col = matrix.iter().map(|row| row[pick_col_index].clone()).collect::<Vec<_>>();
+	// Extend the column
+	let extended_col = extend_segments_col(kzg.get_fs(), &col).unwrap();
+
+	for i in 0..(chunk_count) {
+		let pick_s = extended_col[i].clone();
+		assert!(pick_s.verify(&kzg, &extended_commitments[i * 2 + 1], chunk_count).unwrap());
+	}
+}