Merge pull request #53 from availproject/marko/revert-to-parallel

Moved back to parallel for build proof

kate/src/com.rs

@@ -3,6 +3,7 @@ use std::{
 	convert::{TryFrom, TryInto},
 	mem::size_of,
 	num::TryFromIntError,
+	sync::Mutex,
 	time::Instant,
 };
 
@@ -75,6 +76,7 @@ pub enum Error {
 	/// The extended grid width does not fit cleanly into a domain for FFTs
 	ExtendedGridDomianSizeInvalid(usize),
 	IndexOutOfRange,
+	ConversionFailed,
 }
 
 impl From<TryFromIntError> for Error {
@@ -370,56 +372,80 @@ pub fn build_proof<M: Metrics>(
 	let prover_key = &prover_key;
 	let row_dom_x_pts = &row_dom_x_pts;
 
-	// generate proof only for requested cells
 	let total_start = Instant::now();
 
 	// attempt to parallelly compute proof for all requested cells
-	let cell_iter = cells.iter().zip(result_bytes.chunks_exact_mut(SPROOF_SIZE));
+	let cell_iter = cells
+		.into_par_iter()
+		.zip(result_bytes.par_chunks_exact_mut(SPROOF_SIZE));
+
+	let locked_errors = Mutex::new(Vec::<Error>::new());
 
-	for (cell, res) in cell_iter {
+	let get_cell_row = |cell: &Cell| -> Result<(Vec<BlsScalar>, usize, usize), Error> {
 		let r_index = usize::try_from(cell.row.0)?;
 		if r_index >= ext_rows || cell.col >= block_dims.cols {
 			return Err(Error::IndexOutOfRange);
-		} else {
-			let c_index = usize::try_from(cell.col.0)?;
-			let get_ext_data_matrix =
-				|j: usize| ext_data_matrix[r_index.saturating_add(j.saturating_mul(ext_rows))];
-
-			// construct polynomial per extended matrix row
-			#[cfg(feature = "parallel")]
-			let row = {
-				let mut row =
-					Vec::with_capacity(ext_cols.checked_add(1).ok_or(Error::BlockTooBig)?);
-				(0..ext_cols)
-					.into_par_iter()
-					.map(get_ext_data_matrix)
-					.collect_into_vec(&mut row);
-				row
-			};
-			#[cfg(not(feature = "parallel"))]
-			let row = (0..ext_cols)
+		}
+		let c_index = usize::try_from(cell.col.0)?;
+
+		let get_ext_data_matrix =
+			|j: usize| ext_data_matrix[r_index.saturating_add(j.saturating_mul(ext_rows))];
+
+		// construct polynomial per extended matrix row
+		#[cfg(feature = "parallel")]
+		let row: Vec<BlsScalar> = {
+			let mut row = Vec::with_capacity(ext_cols.checked_add(1).ok_or(Error::BlockTooBig)?);
+			(0..ext_cols)
+				.into_par_iter()
 				.map(get_ext_data_matrix)
-				.collect::<Vec<BlsScalar>>();
+				.collect_into_vec(&mut row);
+			row
+		};
+		#[cfg(not(feature = "parallel"))]
+		let row = (0..ext_cols)
+			.map(get_ext_data_matrix)
+			.collect::<Vec<BlsScalar>>();
+
+		Ok((row, r_index, c_index))
+	};
 
-			// row has to be a power of 2, otherwise interpolate() function panics TODO: cache evaluations
-			let poly = Evaluations::from_vec_and_domain(row, row_eval_domain).interpolate();
-			let witness = prover_key.compute_single_witness(&poly, &row_dom_x_pts[c_index]);
+	cell_iter.for_each(|(cell, res)| {
+		let result = get_cell_row(cell);
+		let Ok((row, r_index, c_index)) = result else {
+			if let Ok(mut errors) = locked_errors.lock() {
+				errors.push(result.err().expect("We checked before that this is OK. "))
+			}
+			return;
+		};
+
+		// row has to be a power of 2, otherwise interpolate() function panics TODO: cache evaluations
+		let poly = Evaluations::from_vec_and_domain(row, row_eval_domain).interpolate();
+		let witness = prover_key.compute_single_witness(&poly, &row_dom_x_pts[c_index]);
+
+		match prover_key.commit(&witness) {
+			Ok(commitment_to_witness) => {
+				let evaluated_point =
+					ext_data_matrix[r_index.saturating_add(c_index.saturating_mul(ext_rows))];
+
+				res[0..PROOF_SIZE].copy_from_slice(&commitment_to_witness.to_bytes());
+				res[PROOF_SIZE..].copy_from_slice(&evaluated_point.to_bytes());
+			},
+			Err(e) => {
+				if let Ok(mut errors) = locked_errors.lock() {
+					errors.push(Error::PlonkError(e))
+				}
+			},
+		};
+	});
 
-			match prover_key.commit(&witness) {
-				Ok(commitment_to_witness) => {
-					let evaluated_point =
-						ext_data_matrix[r_index.saturating_add(c_index.saturating_mul(ext_rows))];
+	metrics.proof_build_time(total_start.elapsed(), cells.len().saturated_into());
 
-					res[0..PROOF_SIZE].copy_from_slice(&commitment_to_witness.to_bytes());
-					res[PROOF_SIZE..].copy_from_slice(&evaluated_point.to_bytes());
-				},
-				Err(e) => return Err(Error::PlonkError(e)),
-			};
+	if let Ok(mut errors) = locked_errors.lock() {
+		if let Some(error) = errors.pop() {
+			return Err(error);
 		}
 	}
 
-	metrics.proof_build_time(total_start.elapsed(), cells.len().saturated_into());
-
 	Ok(result_bytes)
 }