Initial WASM Support (#254)

* add wasm32 target in nix * more try * wip fighting blst * wasm32 success * Replace Arc with Rc, no-default-features on ark-primitives * update CHANGELOG * add additional wasm32 target in CI * fix bench code * fix CI * add no-std check for non-wasm in CI & better std feature activation * wasm-aware MLE type for concurrency
EspressoSystems · May 13, 2023 · 45328b8 · 45328b8
1 parent 7323d81
commit 45328b8
Show file tree

Hide file tree

Showing 19 changed files with 163 additions and 80 deletions.
diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
@@ -32,6 +32,8 @@ jobs:
           override: true
           default: true
           components: rustfmt, clippy
+          target: |
+            wasm32-unknown-unknown
 
       - name: Install nightly toolchain
         uses: actions-rs/toolchain@v1
@@ -64,8 +66,12 @@ jobs:
       - name: Check all tests and binaries compilation
         run: cargo check --workspace --tests --lib --bins
 
-      - name: Check no_std compilation
-        run: cargo check --workspace --lib --no-default-features
+      - name: Check no_std support and WASM compilation
+        env:
+          RUSTFLAGS: -C target-cpu=generic
+        run: |
+          cargo check --no-default-features
+          cargo build --target wasm32-unknown-unknown --no-default-features
 
       - name: Test
         run: bash ./scripts/run_tests.sh

diff --git a/.gitignore b/.gitignore
@@ -5,6 +5,7 @@
 cargo-system-config.toml
 Cargo.lock
 *.org
+.pre-commit-config.yaml
 
 # Test coverage (grcov)
 default.profraw

diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -19,6 +19,7 @@ and follow [semantic versioning](https://semver.org/) for our releases.
 - [#233](https://github.com/EspressoSystems/jellyfish/pull/233) BLS aggregation APIs
 - [#234](https://github.com/EspressoSystems/jellyfish/pull/234) New `bytes_from_field_elements` util
 - [#231](https://github.com/EspressoSystems/jellyfish/pull/231) Implemented FK23 for fast amortized opening for univariate PCS
+- [#254](https://github.com/EspressoSystems/jellyfish/pull/254) Ensure `no_std` and target WASM support
 
 ### Changed
 

diff --git a/README.md b/README.md
@@ -32,12 +32,20 @@ cargo build
 Run an example:
 
 ```
-cargo run --release --example proof_of_exp
+cargo run --release --example proof_of_exp --features test-srs
 ```
 
 This is a simple example to prove and verify knowledge of exponent.
 It shows how one may compose a circuit, and then build a proof for the circuit.
 
+### WASM target
+
+Jellyfish is `no_std` compliant and compilable to WASM target environment, just run:
+
+```
+./scripts/build_wasm.sh
+```
+
 ### Tests
 
 ```

diff --git a/flake.lock b/flake.lock
diff --git a/flake.nix b/flake.nix
@@ -21,13 +21,32 @@
   outputs = { self, nixpkgs, flake-utils, flake-compat, rust-overlay, pre-commit-hooks, ... }:
     flake-utils.lib.eachDefaultSystem (system:
       let
-        overlays = [ (import rust-overlay) ];
+        overlays = [ 
+          (import rust-overlay)
+          (self: super: {
+            rustc = (super.rustc.override {
+                stdenv = self.stdenv.override {
+                    targetPlatform = super.stdenv.targetPlatform // {
+                        parsed = {
+                            cpu = { name = "wasm32"; };
+                            vendor = {name = "unknown";};
+                            kernel = {name = "unknown";};
+                            abi = {name = "unknown";};
+                        };
+                    };
+                };
+            }).overrideAttrs (attrs: {
+                configureFlags = attrs.configureFlags ++ ["--set=build.docs=false"];
+            });
+          })
+        ];
         pkgs = import nixpkgs { inherit system overlays; };
         nightlyToolchain = pkgs.rust-bin.selectLatestNightlyWith
           (toolchain: toolchain.minimal.override { extensions = [ "rustfmt" ]; });
 
         stableToolchain = pkgs.rust-bin.stable.latest.minimal.override {
           extensions = [ "clippy" "llvm-tools-preview" "rust-src" ];
+          targets = ["wasm32-unknown-unknown"];
         };
       in with pkgs;
       {
@@ -62,7 +81,8 @@
             };
           };
         };
-        devShell = mkShell {
+        devShell = clang15Stdenv.mkDerivation {
+          name = "clang15-nix-shell";
           buildInputs = [
             argbash
             openssl
@@ -72,7 +92,9 @@
             stableToolchain
             nightlyToolchain
             cargo-sort
-
+            clang-tools_15
+            clangStdenv
+            llvm_15
           ] ++ lib.optionals stdenv.isDarwin [ darwin.apple_sdk.frameworks.Security ];
 
           shellHook = ''
@@ -81,6 +103,11 @@
 
             # Ensure `cargo fmt` uses `rustfmt` from nightly.
             export RUSTFMT="${nightlyToolchain}/bin/rustfmt"
+
+            export C_INCLUDE_PATH="${llvmPackages_15.libclang.lib}/lib/clang/15.0.7/include"
+            export CC=$(which clang)
+            export AR=$(which llvm-ar)
+            export CFLAGS="-mcpu=generic"
           ''
           # install pre-commit hooks
           + self.check.${system}.pre-commit-check.shellHook;

diff --git a/plonk/Cargo.toml b/plonk/Cargo.toml
@@ -10,10 +10,6 @@ license = { workspace = true }
 rust-version = { workspace = true }
 
 [dependencies]
-ark-bls12-377 = "0.4.0"
-ark-bls12-381 = "0.4.0"
-ark-bn254 = "0.4.0"
-ark-bw6-761 = "0.4.0"
 ark-ec = "0.4.0"
 ark-ff = { version = "0.4.0", features = [ "asm" ] }
 ark-poly = "0.4.0"
@@ -31,13 +27,17 @@ jf-relation = { path = "../relation", default-features = false }
 jf-utils = { path = "../utilities" }
 merlin = { version = "3.0.0", default-features = false }
 num-bigint = { version = "0.4", default-features = false }
-rand_chacha = { version = "0.3.1" }
+rand_chacha = { version = "0.3.1", default-features = false }
 rayon = { version = "1.5.0", optional = true }
 serde = { version = "1.0", default-features = false, features = ["derive"] }
 sha3 = "^0.10"
 tagged-base64 = { git = "https://github.com/espressosystems/tagged-base64", tag = "0.3.0" }
 
 [dev-dependencies]
+ark-bls12-377 = "0.4.0"
+ark-bls12-381 = "0.4.0"
+ark-bn254 = "0.4.0"
+ark-bw6-761 = "0.4.0"
 ark-ed-on-bls12-377 = "0.4.0"
 ark-ed-on-bls12-381 = "0.4.0"
 ark-ed-on-bn254 = "0.4.0"
@@ -51,9 +51,13 @@ harness = false
 
 [features]
 default = ["parallel"]
-std = ["ark-std/std", "ark-serialize/std", "ark-ff/std", "ark-ec/std", "ark-poly/std"]
+std = [
+    "ark-std/std", "ark-serialize/std", "ark-ff/std", "ark-ec/std", "ark-poly/std",
+    "downcast-rs/std", "itertools/use_std", "jf-primitives/std", "jf-relation/std",
+    "jf-utils/std", "num-bigint/std", "rand_chacha/std", "sha3/std"
+]
 test_apis = [] # exposing apis for testing purpose
 parallel = ["ark-ff/parallel", "ark-ec/parallel", "ark-poly/parallel", 
             "jf-utils/parallel", "jf-relation/parallel", "jf-primitives/parallel",
-            "rayon" ]
+            "dep:rayon" ]
 test-srs = []
diff --git a/primitives/Cargo.toml b/primitives/Cargo.toml
@@ -12,7 +12,7 @@ ark-bls12-377 = "0.4.0"
 ark-bls12-381 = "0.4.0"
 ark-bn254 = "0.4.0"
 ark-bw6-761 = "0.4.0"
-ark-crypto-primitives = { version = "0.4.0", features = ["sponge"] }
+ark-crypto-primitives = { version = "0.4.0", default-features = false, features = ["sponge"] }
 ark-ec = "0.4.0"
 ark-ed-on-bls12-377 = "0.4.0"
 ark-ed-on-bls12-381 = "0.4.0"
@@ -21,8 +21,8 @@ ark-ff = "0.4.0"
 ark-poly = "0.4.0"
 ark-serialize = "0.4.0"
 ark-std = { version = "0.4.0", default-features = false }
-blst = "0.3.10"
-crypto_box = "0.8.1"
+blst = { git = "https://github.com/EspressoSystems/blst.git", branch = "no-std", default-features = false } # TODO: pin to a tag or commit
+crypto_box = { version = "0.8.1", default-features = false, features = ["alloc", "u32_backend"] }
 derivative = { version = "2", features = ["use_core"] }
 digest = { version = "0.10.1", default-features = false, features = ["alloc"] }
 displaydoc = { version = "0.2.3", default-features = false }
@@ -40,7 +40,7 @@ serde = { version = "1.0", default-features = false, features = ["derive"] }
 sha2 = { version = "0.10.1", default-features = false }
 sha3 = { version = "0.10.5", default-features = false }
 tagged-base64 = { git = "https://github.com/espressosystems/tagged-base64", tag = "0.3.0" }
-typenum = { version = "1.15.0", default-features = false }
+typenum = { version = "1.15.0", default-features = false, features = ["no_std"] }
 zeroize = { version = "1.5", default-features = false }
 
 [dev-dependencies]
@@ -54,7 +54,6 @@ ark-ed-on-bn254 = "0.4.0"
 bincode = "1.3"
 criterion = "0.4.0"
 hashbrown = "0.13.1"
-rand_core = { version = "^0.6.0", features = ["getrandom"] }
 
 [[bench]]
 name = "merkle_path"
@@ -66,7 +65,10 @@ harness = false
 
 [features]
 default = ["parallel"]
-std = []
+std = ["ark-std/std", "ark-serialize/std", "ark-ff/std", "ark-ec/std", 
+        "ark-poly/std", "ark-crypto-primitives/std", "num-bigint/std",
+        "num-traits/std", "sha2/std", "sha3/std", "zeroize/std", 
+        "itertools/use_std", "rand_chacha/std", "jf-utils/std", "jf-relation/std"]
 print-trace = ["ark-std/print-trace"] 
 parallel = ["ark-ff/parallel", "ark-ec/parallel", "jf-utils/parallel", 
             "jf-relation/parallel", "rayon" ]

diff --git a/primitives/benches/pcs.rs b/primitives/benches/pcs.rs
@@ -13,9 +13,8 @@ mod bench {
     use ark_bls12_381::{Bls12_381, Fr};
     use ark_ff::UniformRand;
     use ark_poly::{DenseMultilinearExtension, MultilinearExtension};
-    use ark_std::sync::Arc;
     use jf_primitives::pcs::{
-        prelude::{MultilinearKzgPCS, PCSError, PolynomialCommitmentScheme},
+        prelude::{MultilinearKzgPCS, PCSError, PolynomialCommitmentScheme, MLE},
         StructuredReferenceString,
     };
     use jf_utils::test_rng;
@@ -36,7 +35,7 @@ mod bench {
                 10
             };
 
-            let poly = Arc::new(DenseMultilinearExtension::rand(nv, &mut rng));
+            let poly = MLE::from(DenseMultilinearExtension::rand(nv, &mut rng));
             let (ml_ck, ml_vk) = uni_params.0.trim(nv)?;
             let (uni_ck, uni_vk) = uni_params.1.trim(nv)?;
             let ck = (ml_ck, uni_ck);

diff --git a/primitives/src/lib.rs b/primitives/src/lib.rs
@@ -16,6 +16,10 @@ extern crate std;
 #[macro_use]
 extern crate derivative;
 
+#[cfg(not(feature = "std"))]
+#[doc(hidden)]
+extern crate alloc;
+
 pub mod aead;
 pub mod circuit;
 pub mod commitment;

diff --git a/primitives/src/pcs/mod.rs b/primitives/src/pcs/mod.rs
@@ -19,7 +19,6 @@ use ark_std::{
     borrow::Borrow,
     fmt::Debug,
     hash::Hash,
-    path::Path,
     rand::{CryptoRng, RngCore},
     vec::Vec,
 };
@@ -32,13 +31,7 @@ pub trait PolynomialCommitmentScheme {
     /// Structured reference string
     type SRS: Clone + Debug + StructuredReferenceString;
     /// Polynomial and its associated types
-    type Polynomial: Clone
-        + Debug
-        + Hash
-        + PartialEq
-        + Eq
-        + CanonicalSerialize
-        + CanonicalDeserialize;
+    type Polynomial: Clone + Debug + Hash + PartialEq + Eq;
     /// Polynomial input domain
     type Point: Clone + Ord + Debug + Sync + Hash + PartialEq + Eq;
     /// Polynomial Evaluation
@@ -76,7 +69,7 @@ pub trait PolynomialCommitmentScheme {
     /// If `file=None`, we load the default choice of SRS.
     fn load_srs_from_file(
         supported_size: usize,
-        file: Option<&Path>,
+        file: Option<&str>,
     ) -> Result<Self::SRS, PCSError> {
         Self::SRS::load_srs_from_file(supported_size, file)
     }
@@ -218,7 +211,7 @@ pub trait StructuredReferenceString: Sized {
     /// implemented else where. We only load them into memory here.
     ///
     /// If `file=None`, we load the default choice of SRS.
-    fn load_srs_from_file(_supported_size: usize, _file: Option<&Path>) -> Result<Self, PCSError> {
+    fn load_srs_from_file(_supported_size: usize, _file: Option<&str>) -> Result<Self, PCSError> {
         unimplemented!("TODO: implement loading SRS from files");
     }
 }