Bring proof gen to fat clients

Cargo.toml

@@ -40,6 +40,7 @@ libp2p = { version = "0.56.0", features = ["kad", "identify", "ping", "autonat",
 libp2p-allow-block-list = "0.6.0"
 libp2p-webrtc-websys = "0.4.0"
 multihash = { version = "0.14.0", default-features = false, features = ["blake3", "sha3"] }
+rayon = "1.11.0"
 semver = "1.0.23"
 serde = { version = "1.0", features = ["derive"] }
 tokio = { version = "1.35", default-features = false, features = ["sync", "macros", "io-util", "rt", "time"] }

core/Cargo.toml

@@ -55,6 +55,8 @@ tar = "0.4"
 
 [target.'cfg(not(target_arch = "wasm32"))'.dependencies]
 avail-rust = { workspace = true, default-features = true }
+kate = { workspace = true, default-features = true }
+kate-recovery = { workspace = true, default-features = true }
 async-std = { workspace = true }
 color-eyre = { workspace = true, default-features = true }
 confy = { workspace = true }
@@ -81,7 +83,7 @@ tracing-error = "0.2"
 uuid = { workspace = true }
 void = { workspace = true }
 warp = { workspace = true }
-
+rayon = { workspace = true }
 # OpenTelemetry
 opentelemetry = { workspace = true }
 opentelemetry-otlp = { workspace = true }

core/src/fat_client.rs

@@ -10,11 +10,14 @@
 //! In case delay is configured, block processing is delayed for configured time.
 
 use async_trait::async_trait;
+use avail_core::AppExtrinsic;
 use avail_rust::{AvailHeader, H256};
 use codec::Encode;
 use color_eyre::{eyre::WrapErr, Result};
 use futures::future::join_all;
 #[cfg(feature = "multiproof")]
+use kate::Seed;
+#[cfg(feature = "multiproof")]
 use kate_recovery::data::MultiProofCell;
 #[cfg(not(feature = "multiproof"))]
 use kate_recovery::data::{self, SingleCell};
@@ -46,7 +49,12 @@ use crate::{
 pub trait Client {
 	async fn insert_cells_into_dht(&self, block: u32, cells: Vec<Cell>) -> Result<()>;
 	async fn insert_rows_into_dht(&self, block: u32, rows: Vec<(RowIndex, Vec<u8>)>) -> Result<()>;
-	async fn get_kate_proof(&self, hash: H256, positions: &[Position]) -> Result<Vec<Cell>>;
+	async fn get_kate_proof(
+		&self,
+		hash: H256,
+		positions: &[Position],
+		extrinsics: Vec<AppExtrinsic>,
+	) -> Result<Vec<Cell>>;
 }
 
 #[derive(Clone)]
@@ -102,13 +110,39 @@ impl<T: Database + Sync> Client for FatClient<T> {
 		self.p2p_client.insert_rows_into_dht(block, rows).await
 	}
 
-	async fn get_kate_proof(&self, hash: H256, positions: &[Position]) -> Result<Vec<Cell>> {
+	async fn get_kate_proof(
+		&self,
+		hash: H256,
+		positions: &[Position],
+		_extrinsics: Vec<AppExtrinsic>,
+	) -> Result<Vec<Cell>> {
 		#[cfg(feature = "multiproof")]
 		{
-			let cells: Vec<MultiProofCell> = self
-				.rpc_client
-				.request_kate_multi_proof(hash, positions)
-				.await?;
+			let block_length = self.rpc_client.query_block_length(Some(hash)).await?;
+			let seed = kate::Seed::default();
+
+			let cell_positions: Vec<(u32, u32)> =
+				positions.iter().map(|p| (p.row, p.col)).collect();
+
+			let multiproof_results =
+				crate::proof::multiproof(extrinsics, local_block_length, seed, cell_positions)?;
+
+			let cells: Vec<MultiProofCell> = multiproof_results
+				.into_iter()
+				.zip(positions.iter())
+				.map(|((proof, gcell_block), &position)| {
+					let (scalars, proof_bytes) = proof;
+					let raw_scalars: Vec<[u64; 4]> = scalars.into_iter().map(|s| s.0).collect();
+
+					MultiProofCell {
+						position,
+						scalars: raw_scalars,
+						proof: proof_bytes.0,
+						gcell_block,
+					}
+				})
+				.collect();
+
 			Ok(cells.into_iter().map(Cell::MultiProofCell).collect())
 		}
 
@@ -200,7 +234,7 @@ pub async fn process_block(
 	let positions = iter_partition_cells(partition, target_grid_dimensions);
 
 	let begin = Instant::now();
-	let get_kate_proof = |&n| client.get_kate_proof(header_hash, n);
+	let get_kate_proof = |&n| client.get_kate_proof(header_hash, n, vec![]);
 	let Partition { number, fraction } = cfg.block_matrix_partition;
 	info!(
 		block_number,
@@ -487,12 +521,14 @@ mod tests {
 		let mut mock_client = MockClient::new();
 		let cell_variants: Vec<Cell> = DEFAULT_CELLS.into_iter().collect();
 
-		mock_client.expect_get_kate_proof().returning(move |_, _| {
-			Box::pin({
-				let cells = cell_variants.clone();
-				async move { Ok(cells.to_vec()) }
-			})
-		});
+		mock_client
+			.expect_get_kate_proof()
+			.returning(move |_, _, _| {
+				Box::pin({
+					let cells = cell_variants.clone();
+					async move { Ok(cells.to_vec()) }
+				})
+			});
 		mock_client
 			.expect_insert_rows_into_dht()
 			.returning(|_, _| Box::pin(async move { Ok(()) }));

core/src/network/rpc/client.rs

@@ -1,4 +1,5 @@
 use avail_core::AppId;
+use avail_rust::avail::runtime_types::frame_system::limits::BlockLength;
 use avail_rust::{
 	avail::{self, runtime_types::sp_core::crypto::KeyTypeId},
 	primitives::kate::{Cells, GProof, GRawScalar, Rows},
@@ -758,6 +759,23 @@ impl<D: Database> Client<D> {
 		Ok(ver)
 	}
 
+	pub async fn query_block_length(&self, block_hash: Option<H256>) -> Result<BlockLength> {
+		self.with_retries(|client| async move {
+			let mut params = rpc_params![];
+			if let Some(hash) = block_hash {
+				params.push(hash)?;
+			}
+
+			client
+				.client
+				.rpc_client
+				.request("kate_blockLength", params)
+				.await
+				.map_err(Into::into)
+		})
+		.await
+	}
+
 	pub async fn get_validator_set_by_block_number(&self, block_num: u32) -> Result<Vec<Public>> {
 		let hash = self.get_block_hash(block_num).await?;
 		self.get_validator_set_by_hash(hash).await