Improve erasure coding testing

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

@@ -21,7 +21,7 @@ use rust_kzg_blst::types::fk20_multi_settings::FsFK20MultiSettings;
 use crate::erasure_coding::extend_poly;
 
 pub fn order_segments_row(segments: &Vec<Segment>, chunk_count: usize) -> Result<Vec<Option<SegmentData>>, String> {
-	if segments.len() > chunk_count * 2 {
+	if segments.len() > chunk_count * 2 || segments.len() == 0 {
 		return Err("segments x not equal".to_string());
 	}
 	let y = segments[0].position.y;
@@ -39,7 +39,7 @@ pub fn order_segments_col(
 	segments: &Vec<Segment>,
 	k: usize,
 ) -> Result<Vec<Option<SegmentData>>, String> {
-	if segments.len() > k * 2 {
+	if segments.len() > k * 2 || segments.len() == 0 {
 		return Err("segments x not equal".to_string());
 	}
 	let x = segments[0].position.x;
@@ -48,7 +48,7 @@ pub fn order_segments_col(
 		if segment.position.x != x {
 			return Err("segments x not equal".to_string());
 		}
-		ordered_segments[segment.position.x as usize] = Some(segment.content.clone());
+		ordered_segments[segment.position.y as usize] = Some(segment.content.clone());
 	}
 	Ok(ordered_segments)
 }

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

@@ -2,28 +2,23 @@ use crate::erasure_coding::*;
 use crate::extend_col::extend_segments_col;
 use crate::recovery::*;
 use crate::segment::*;
-// use blst_rust::types::g1::FsG1;
 use alloc::vec;
-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::Position;
 use melo_core_primitives::kzg::ReprConvert;
 use melo_core_primitives::kzg::{embedded_kzg_settings, KZGCommitment, KZG};
 use melo_core_primitives::segment::Segment;
+use melo_core_primitives::segment::SegmentData;
 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;
-use std::num::NonZeroUsize;
 
 fn reverse_bits_limited(length: usize, value: usize) -> usize {
 	let unused_bits = length.leading_zeros();
@@ -75,6 +70,170 @@ fn segment_datas_to_row_test() {
 	}
 }
 
+#[test]
+fn order_segments_col_test() {
+	// 构建随机 segment datas
+	let chunk_len: usize = 16;
+	let k: usize = 4;
+	let mut segment_datas: Vec<Option<SegmentData>> = Vec::new();
+	for _ in 0..k * 2 {
+		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;
+	// 构建segemnts
+	let segments_option = segment_datas
+		.iter()
+		.enumerate()
+		.map(|(i, segment_data)| {
+			let position = Position { x: 0, y: i as u32 };
+			match segment_data {
+				Some(segment_data) => Some(Segment { position, content: segment_data.clone() }),
+				None => None,
+			}
+		})
+		.collect::<Vec<_>>();
+	let segments = segments_option.iter().filter_map(|segment| segment.clone()).collect::<Vec<_>>();
+	let mut s_segments = segments.clone();
+	// 打乱顺序
+	s_segments.shuffle(&mut rand::thread_rng());
+	// 转换为 row
+	let col: Vec<Option<SegmentData>> = order_segments_col(&s_segments, k).unwrap();
+	// 验证是否正确
+	for i in 0..k * 2 {
+		// segments_option 如果是 some 则比较大小，如果是 None 则直接验证
+		if let Some(segment) = &segments_option[i] {
+			assert_eq!(col[i], Some(segment.content.clone()));
+		} else {
+			assert_eq!(col[i], None);
+		}
+	}
+	// 修改 s_segments 其中 1个 x
+	s_segments[1].position.x = 3;
+	// 转换为 row
+	let col: Result<Vec<Option<SegmentData>>, String> = order_segments_col(&s_segments, k);
+	// 验证是否失败
+	assert!(col.is_err());
+	// 向 s_segments 中 push 3个 随机 segment
+	for _ in 0..3 {
+		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 };
+		let position = Position { x: 0, y: 0 };
+		s_segments.push(Segment { position, content: segment_data });
+	}
+	// 转换为 row
+	let col: Result<Vec<Option<SegmentData>>, String> = order_segments_col(&s_segments, k);
+	// 验证是否失败
+	assert!(col.is_err());
+
+}
+
+#[test]
+fn poly_to_segment_vec_test() {
+	// 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);
+	// Get the commitment of poly
+	let kzg = KZG::new(embedded_kzg_settings());
+	let commitment = kzg.commit(&poly).unwrap();
+	// Convert to segments
+	let segments = poly_to_segment_vec(&poly, &kzg, 0, chunk_len).unwrap();
+
+	// Verify if it's correct
+	for i in 0..chunk_count {
+		let verify = segments[i].verify(&kzg, &commitment, chunk_count).unwrap();
+		assert!(verify);
+	}
+
+	// Convert segments to row
+	let mut row = segments
+		.into_iter()
+		.flat_map(|segment| segment.content.data)
+		.collect::<Vec<_>>();
+	// Reverse row
+	reverse_bit_order(&mut row);
+	// Convert row to coefficient form
+	let recovery_poly = kzg.get_fs().fft_fr(&BlsScalar::vec_to_repr(row), true).unwrap();
+	// Verify if it's correct
+	for i in 0..num_shards {
+		assert_eq!(recovery_poly[i], poly.0.coeffs[i]);
+	}
+}
+
+#[test]
+fn order_segments_row_test() {
+	// 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);
+	// Get the commitment of poly
+	let kzg = KZG::new(embedded_kzg_settings());
+	// Convert to segments
+	let segments = poly_to_segment_vec(&poly, &kzg, 0, chunk_len).unwrap();
+	let mut random_segments: Vec<Option<Segment>> = Vec::new();
+
+	let random_positions = random_vec(3 * chunk_count);
+	for i in 0..chunk_count * 2 {
+		let position = random_positions[i];
+		if position < 2 * chunk_count {
+			random_segments.push(Some(segments[position].clone()));
+		} else {
+			random_segments.push(None);
+		}
+	}
+	// 从 random_segments 获取有效的 segment
+	let mut s_segments = random_segments.iter().filter_map(|segment| segment.clone()).collect::<Vec<_>>();
+	// 随机洗牌 s_segments
+	s_segments.shuffle(&mut rand::thread_rng());
+	// Order segments
+	let ordered_segments = order_segments_row(&s_segments, chunk_count).unwrap();
+	for i in 0..chunk_count * 2 {
+		if let Some(segment_data) = &ordered_segments[i] {
+			assert_eq!(segment_data.data, segments[i].content.data);
+		}
+	}
+	// 获取 ordered_segments 中 some 的数量，是否和 s_segments 长度相等
+	let some_count = ordered_segments.iter().filter(|segment_data| segment_data.is_some()).count();
+	assert_eq!(some_count, s_segments.len());
+	// 修改 s_segments 其中 1个 y
+	s_segments[0].position.y = 3;
+	// Order segments
+	let ordered_segments = order_segments_row(&s_segments, chunk_count);
+	// Verify if it fails
+	assert!(ordered_segments.is_err());
+
+	s_segments[0].position.y = 0;
+	// 向 s_segments 中 push 2 * 个 随机 segment
+	for _ in 0..2 * 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 };
+		let position = Position { x: 0, y: 0 };
+		s_segments.push(Segment { position, content: segment_data });
+	}
+	// Order segments
+	let ordered_segments = order_segments_row(&s_segments, chunk_count);
+	// Verify if it fails
+	assert!(ordered_segments.is_err());
+
+}
+
 #[test]
 fn extend_poly_random_round_trip_test() {
 	let kzg = KZG::new(embedded_kzg_settings());
@@ -100,44 +259,43 @@ fn extend_poly_random_round_trip_test() {
 
 #[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());
+	// 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]