From d6b78dde1bcd73959bfed0d110f640e2f67113cd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Aleksandar=20Terenti=C4=87?= <aleksandar@ethernal.tech>
Date: Mon, 25 Nov 2024 10:45:14 +0100
Subject: [PATCH] Enable WASM compilation on proof mod

---
 Cargo.lock        |  3 +++
 core/CHANGELOG.md |  1 +
 core/Cargo.toml   | 10 +++++-----
 core/src/lib.rs   |  1 -
 core/src/proof.rs | 45 +++++++++++++++++++++++++--------------------
 5 files changed, 34 insertions(+), 26 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index cdfb71f93..a468acb56 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1013,6 +1013,7 @@ dependencies = [
  "confy 0.5.1",
  "convert_case 0.6.0",
  "derive_more 1.0.0",
+ "dusk-bytes",
  "dusk-plonk",
  "futures",
  "getrandom",
@@ -1047,6 +1048,7 @@ dependencies = [
  "sysinfo",
  "test-case",
  "thiserror",
+ "thiserror-no-std",
  "threadpool",
  "tokio",
  "tokio-retry",
@@ -1058,6 +1060,7 @@ dependencies = [
  "uuid",
  "void",
  "warp",
+ "web-time",
 ]
 
 [[package]]
diff --git a/core/CHANGELOG.md b/core/CHANGELOG.md
index d24f35f55..67e880ce8 100644
--- a/core/CHANGELOG.md
+++ b/core/CHANGELOG.md
@@ -2,6 +2,7 @@
 
 ## 1.0.5
 
+- Enable WASM compilation on proof mod
 - Enable WASM compilation on utils and shutdown mods
 - Allocate new port on each new dial attempt
 - Set different dial conditions for bootstrap process and diagnostics API
diff --git a/core/Cargo.toml b/core/Cargo.toml
index db7f5a3c8..11eabe732 100644
--- a/core/Cargo.toml
+++ b/core/Cargo.toml
@@ -10,24 +10,23 @@ crate-type = ["cdylib", "rlib"]
 [dependencies]
 # Internal deps
 avail-rust = { workspace = true }
+dusk-plonk = { workspace = true }
 
 # 3rd-party
 better-panic = "0.3.0"
 codec = { package = "parity-scale-codec", version = "3", default-features = false, features = ["derive", "full", "bit-vec"] }
 color-eyre = { workspace = true }
 derive_more = { version = "1", features = ["from"] }
+dusk-bytes = "0.1.7"
 futures = { workspace = true }
 getrandom = { version = "0.2.15", features = ["js"] }
+itertools = "0.10.5"
 strip-ansi-escapes = "0.2.0"
 tokio = { workspace = true }
 tracing = { workspace = true }
 tracing-subscriber = { workspace = true }
 
 [target.'cfg(not(target_arch = "wasm32"))'.dependencies]
-# Internal deps
-dusk-plonk = { workspace = true }
-
-# 3rd-party
 async-std = { workspace = true }
 async-stream = "0.3.5"
 async-trait = { workspace = true }
@@ -39,7 +38,6 @@ derive_more = { version = "1", features = ["from"] }
 futures = { workspace = true }
 hex = { workspace = true }
 hyper = { version = "0.14.23", features = ["full", "http1"] }
-itertools = "0.10.5"
 jsonrpsee-core = { version = "0.21.0", features = ["client"] }
 libc = "0.2.150"
 libp2p = { workspace = true }
@@ -79,7 +77,9 @@ strum = { version = "0.26.3", features = ["derive"] }
 [target.'cfg(target_arch = "wasm32")'.dependencies]
 blake2b_simd = "1.0.2"
 rand = { workspace = true, features = ["std_rng"] }
+thiserror-no-std = "2.0.2"
 tokio_with_wasm = { version = "0.7.1", default-features = false, features = ["sync", "macros", "rt", "time"] }
+web-time = "1.1.0"
 
 [target.'cfg(not(target_arch = "wasm32"))'.dev-dependencies]
 hex-literal = "0.4.1"
diff --git a/core/src/lib.rs b/core/src/lib.rs
index 1ed4fdb79..36e891042 100644
--- a/core/src/lib.rs
+++ b/core/src/lib.rs
@@ -18,7 +18,6 @@ pub mod light_client;
 pub mod maintenance;
 #[cfg(not(target_arch = "wasm32"))]
 pub mod network;
-#[cfg(not(target_arch = "wasm32"))]
 pub mod proof;
 pub mod shutdown;
 #[cfg(not(target_arch = "wasm32"))]
diff --git a/core/src/proof.rs b/core/src/proof.rs
index 0fc78f5cb..b0727ed26 100644
--- a/core/src/proof.rs
+++ b/core/src/proof.rs
@@ -1,24 +1,29 @@
-//! Parallelized proof verification
-
 use avail_rust::kate_recovery::{
 	data::Cell,
 	matrix::{Dimensions, Position},
-	proof,
 };
 use color_eyre::eyre;
 use dusk_plonk::commitment_scheme::kzg10::PublicParameters;
+use futures::future::join_all;
 use itertools::{Either, Itertools};
 use std::sync::Arc;
-use tokio::{task::JoinSet, time::Instant};
-use tracing::{debug, Instrument};
+#[cfg(not(target_arch = "wasm32"))]
+use tokio::time::Instant;
+use tracing::debug;
+#[cfg(target_arch = "wasm32")]
+use web_time::Instant;
+
+mod core;
+
+use crate::utils::spawn_in_span;
 
 async fn verify_proof(
 	public_parameters: Arc<PublicParameters>,
 	dimensions: Dimensions,
 	commitment: [u8; 48],
 	cell: Cell,
-) -> Result<(Position, bool), proof::Error> {
-	proof::verify(&public_parameters, dimensions, &commitment, &cell)
+) -> Result<(Position, bool), core::Error> {
+	core::verify(&public_parameters, dimensions, &commitment, &cell)
 		.map(|verified| (cell.position, verified))
 }
 
@@ -36,24 +41,24 @@ pub async fn verify(
 
 	let start_time = Instant::now();
 
-	let mut tasks = JoinSet::new();
-
-	for cell in cells {
-		tasks.spawn(
-			verify_proof(
+	let tasks = cells
+		.iter()
+		.map(|cell| {
+			spawn_in_span(verify_proof(
 				public_parameters.clone(),
 				dimensions,
 				commitments[cell.position.row as usize],
 				cell.clone(),
-			)
-			.in_current_span(),
-		);
-	}
+			))
+		})
+		.collect::<Vec<_>>();
+
+	let join_results: Vec<_> = join_all(tasks)
+		.await
+		.into_iter()
+		.collect::<Result<_, _>>()?;
 
-	let mut results = Vec::with_capacity(cells.len());
-	while let Some(result) = tasks.join_next().await {
-		results.push(result??)
-	}
+	let results: Vec<(Position, bool)> = join_results.into_iter().collect::<Result<_, _>>()?;
 
 	debug!(block_num, duration = ?start_time.elapsed(), "Proof verification completed");