chore: make WASM compatible

ICME-Lab · Jan 28, 2025 · d2cb6e4 · d2cb6e4
1 parent 5d04832
commit d2cb6e4
Show file tree

Hide file tree

Showing 7 changed files with 45 additions and 34 deletions.
diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml
@@ -28,3 +28,6 @@ jobs:
 
       - name: Run cargo clippy
         run: cargo clippy --verbose --all-targets --all-features -- -D warnings
+
+      - name: WASM build
+        run: cargo build --target wasm32-unknown-unknown
diff --git a/src/aggregation/tests.rs b/src/aggregation/tests.rs
@@ -4,7 +4,7 @@ use crate::{
     logging::init_logger,
     macros::{start_timer, stop_timer},
   },
-  wasm_ctx::{WASMArgsBuilder, WASMCtx, WasiWASMCtx, ZKWASMCtx},
+  wasm_ctx::{wasi::WasiWASMCtx, WASMArgsBuilder, WASMCtx, ZKWASMCtx},
   wasm_snark::{StepSize, WASMPublicParams, WasmSNARK, ZKWASMInstance},
 };
 use nova::{

diff --git a/src/sharding/tests.rs b/src/sharding/tests.rs
@@ -5,7 +5,7 @@ use crate::{
     macros::{start_timer, stop_timer},
     tracing::estimate_wasm,
   },
-  wasm_ctx::{TraceSliceValues, WASMArgsBuilder, WasiWASMCtx, ZKWASMCtx},
+  wasm_ctx::{wasi::WasiWASMCtx, TraceSliceValues, WASMArgsBuilder, ZKWASMCtx},
   wasm_snark::{StepSize, WASMPublicParams, WasmSNARK, ZKWASMInstance},
 };
 use nova::{

diff --git a/src/tests.rs b/src/tests.rs
@@ -1,6 +1,6 @@
 use super::{
   error::ZKWASMError,
-  wasm_ctx::{TraceSliceValues, WASMArgsBuilder, WASMCtx, WasiWASMCtx, ZKWASMCtx},
+  wasm_ctx::{wasi::WasiWASMCtx, TraceSliceValues, WASMArgsBuilder, WASMCtx, ZKWASMCtx},
   wasm_snark::{StepSize, WasmSNARK},
 };
 use crate::utils::{

diff --git a/src/utils/tracing.rs b/src/utils/tracing.rs
@@ -33,7 +33,7 @@ pub(crate) fn split_vector<T>(mut vec: Vec<T>, split_index: usize) -> (Vec<T>, V
 mod test {
   use wasmi::{BCGlobalIdx, WitnessVM};
 
-  use crate::wasm_ctx::{WASMArgsBuilder, WASMCtx, WasiWASMCtx, ZKWASMCtx};
+  use crate::wasm_ctx::{wasi::WasiWASMCtx, WASMArgsBuilder, WASMCtx, ZKWASMCtx};
   use std::{collections::HashMap, path::PathBuf};
 
   /// Count how many time an opcode gets used. Uses the J index of the opcode

diff --git a/src/wasm_ctx.rs b/src/wasm_ctx.rs
@@ -4,11 +4,9 @@ use crate::utils::{
   tracing::unwrap_rc_refcell,
   wasm::{decode_func_args, prepare_func_results, read_wasm_or_wat},
 };
-use rand::{rngs::StdRng, RngCore, SeedableRng};
 use serde::{Deserialize, Serialize};
 use std::{cell::RefCell, cmp, num::NonZeroUsize, path::PathBuf, rc::Rc};
 use wasmi::{Tracer, WitnessVM};
-use wasmi_wasi::{clocks_ctx, sched_ctx, Table, WasiCtx};
 
 /// Builder for [`WASMArgs`]. Constructs the arguments needed to construct a WASM execution context
 /// that will be used for proving.
@@ -290,41 +288,51 @@ impl ZKWASMCtx for WASMCtx {
   }
 }
 
-/// Wasm execution context
-#[derive(Debug, Clone)]
-pub struct WasiWASMCtx {
-  args: WASMArgs,
-}
+#[cfg(not(target_arch = "wasm32"))]
+/// Implementation of WASM execution context for zkVM using wasmi-wasi
+pub mod wasi {
+  use crate::error::ZKWASMError;
 
-impl WasiWASMCtx {
-  /// Create a new instance of [`WasiWASMCtx`]
-  pub fn new(args: WASMArgs) -> Self {
-    Self { args }
-  }
-}
-
-impl ZKWASMCtx for WasiWASMCtx {
-  type T = WasiCtx;
+  use super::{WASMArgs, ZKWASMCtx};
+  use rand::{rngs::StdRng, RngCore, SeedableRng};
+  use wasmi_wasi::{clocks_ctx, sched_ctx, Table, WasiCtx};
 
-  fn args(&self) -> &WASMArgs {
-    &self.args
+  /// Wasm execution context
+  #[derive(Debug, Clone)]
+  pub struct WasiWASMCtx {
+    args: WASMArgs,
   }
 
-  fn create_store(engine: &wasmi::Engine) -> wasmi::Store<Self::T> {
-    let wasi = WasiCtx::new(zkvm_random_ctx(), clocks_ctx(), sched_ctx(), Table::new());
-    wasmi::Store::new(engine, wasi)
+  impl WasiWASMCtx {
+    /// Create a new instance of [`WasiWASMCtx`]
+    pub fn new(args: WASMArgs) -> Self {
+      Self { args }
+    }
   }
 
-  fn create_linker(engine: &wasmi::Engine) -> Result<wasmi::Linker<Self::T>, ZKWASMError> {
-    let mut linker = <wasmi::Linker<WasiCtx>>::new(engine);
-    wasmi_wasi::add_to_linker(&mut linker, |ctx| ctx)?;
-    Ok(linker)
+  impl ZKWASMCtx for WasiWASMCtx {
+    type T = WasiCtx;
+
+    fn args(&self) -> &WASMArgs {
+      &self.args
+    }
+
+    fn create_store(engine: &wasmi::Engine) -> wasmi::Store<Self::T> {
+      let wasi = WasiCtx::new(zkvm_random_ctx(), clocks_ctx(), sched_ctx(), Table::new());
+      wasmi::Store::new(engine, wasi)
+    }
+
+    fn create_linker(engine: &wasmi::Engine) -> Result<wasmi::Linker<Self::T>, ZKWASMError> {
+      let mut linker = <wasmi::Linker<WasiCtx>>::new(engine);
+      wasmi_wasi::add_to_linker(&mut linker, |ctx| ctx)?;
+      Ok(linker)
+    }
   }
-}
 
-/// zkvm uses a seed to generate random numbers.
-pub fn zkvm_random_ctx() -> Box<dyn RngCore + Send + Sync> {
-  Box::new(StdRng::from_seed([0; 32]))
+  /// zkvm uses a seed to generate random numbers.
+  pub fn zkvm_random_ctx() -> Box<dyn RngCore + Send + Sync> {
+    Box::new(StdRng::from_seed([0; 32]))
+  }
 }
 
 /// # Initial Set (IS) Memory Sizes.

diff --git a/src/wasm_snark/mcc/tests.rs b/src/wasm_snark/mcc/tests.rs
@@ -8,7 +8,7 @@ use rand::{rngs::StdRng, RngCore, SeedableRng};
 use crate::{
   error::ZKWASMError,
   utils::logging::init_logger,
-  wasm_ctx::{WASMArgsBuilder, WASMCtx, WasiWASMCtx, ZKWASMCtx},
+  wasm_ctx::{wasi::WasiWASMCtx, WASMArgsBuilder, WASMCtx, ZKWASMCtx},
   wasm_snark::mcc::multiset_ops::step_RS_WS,
 };