chore(rln-wasm): add browser benchmark with simple HTML file and default data, wokred with latest wasm-bindgen-rayon version

- Resolved overlap between stateless and merkletree feature flags.
- Updated every testcase related to stateless feature.
- Added compile-time feature check to avoid feature overlap.
- Added --no-default-features for all builds in nightly-release.yml
[(tested)](https://github.com/vacp2p/zerokit/actions/runs/16525062203).

---------

Co-authored-by: Ekaterina Broslavskaya <seemenkina@gmail.com>

.github/workflows/ci.yml

@@ -24,25 +24,21 @@ jobs:
   utils-test:
     strategy:
       matrix:
-        platform: [ ubuntu-latest, macos-latest ]
-        crate: [ utils ]
+        platform: [ubuntu-latest, macos-latest]
+        crate: [utils]
     runs-on: ${{ matrix.platform }}
     timeout-minutes: 60
 
-    name: test - ${{ matrix.crate }} - ${{ matrix.platform }}
+    name: Test - ${{ matrix.crate }} - ${{ matrix.platform }}
     steps:
       - name: Checkout sources
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
       - name: Install stable toolchain
-        uses: actions-rs/toolchain@v1
-        with:
-          profile: minimal
-          toolchain: stable
-          override: true
+        uses: dtolnay/rust-toolchain@stable
       - uses: Swatinem/rust-cache@v2
       - name: Install dependencies
         run: make installdeps
-      - name: cargo-make test
+      - name: Test utils
         run: |
           cargo make test --release
         working-directory: ${{ matrix.crate }}
@@ -50,26 +46,22 @@ jobs:
   rln-test:
     strategy:
       matrix:
-        platform: [ ubuntu-latest, macos-latest ]
-        crate: [ rln ]
-        feature: [ "default", "arkzkey", "stateless" ]
+        platform: [ubuntu-latest, macos-latest]
+        crate: [rln]
+        feature: ["default", "stateless"]
     runs-on: ${{ matrix.platform }}
     timeout-minutes: 60
 
-    name: test - ${{ matrix.crate }} - ${{ matrix.platform }} - ${{ matrix.feature }}
+    name: Test - ${{ matrix.crate }} - ${{ matrix.platform }} - ${{ matrix.feature }}
     steps:
       - name: Checkout sources
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
       - name: Install stable toolchain
-        uses: actions-rs/toolchain@v1
-        with:
-          profile: minimal
-          toolchain: stable
-          override: true
+        uses: dtolnay/rust-toolchain@stable
       - uses: Swatinem/rust-cache@v2
       - name: Install dependencies
         run: make installdeps
-      - name: cargo-make test
+      - name: Test rln
         run: |
           if [ ${{ matrix.feature }} == default ]; then
             cargo make test --release
@@ -78,74 +70,88 @@ jobs:
           fi
         working-directory: ${{ matrix.crate }}
 
-  rln-wasm:
+  rln-wasm-test:
     strategy:
       matrix:
-        platform: [ ubuntu-latest, macos-latest ]
-        feature: [ "default", "arkzkey" ]
+        platform: [ubuntu-latest, macos-latest]
+        crate: [rln-wasm]
+        feature: ["default"]
     runs-on: ${{ matrix.platform }}
     timeout-minutes: 60
 
-    name: test - rln-wasm - ${{ matrix.platform }} - ${{ matrix.feature }}
+    name: Test - ${{ matrix.crate }} - ${{ matrix.platform }} - ${{ matrix.feature }}
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
       - name: Install stable toolchain
-        uses: actions-rs/toolchain@v1
-        with:
-          profile: minimal
-          toolchain: stable
-          override: true
+        uses: dtolnay/rust-toolchain@stable
       - uses: Swatinem/rust-cache@v2
-      - name: Install Dependencies
+      - name: Install dependencies
         run: make installdeps
-      - name: cargo-make build
-        run: |
-          if [ ${{ matrix.feature }} == default ]; then
-            cargo make build
-          else
-            cargo make build_${{ matrix.feature }}
-          fi
-        working-directory: rln-wasm
-      - name: cargo-make test
-        run: |
-          if [ ${{ matrix.feature }} == default ]; then
-            cargo make test --release
-          else
-            cargo make test_${{ matrix.feature }} --release
-          fi
-        working-directory: rln-wasm
+      - name: Build rln-wasm
+        run: cargo make build
+        working-directory: ${{ matrix.crate }}
+      - name: Test rln-wasm on node
+        run: cargo make test --release
+        working-directory: ${{ matrix.crate }}
+      - name: Test rln-wasm on browser
+        run: cargo make test_browser --release
+        working-directory: ${{ matrix.crate }}
+
+  rln-wasm-parallel-test:
+    strategy:
+      matrix:
+        platform: [ubuntu-latest, macos-latest]
+        crate: [rln-wasm]
+        feature: ["parallel"]
+    runs-on: ${{ matrix.platform }}
+    timeout-minutes: 60
+
+    name: Test - ${{ matrix.crate }} - ${{ matrix.platform }} - ${{ matrix.feature }}
+    steps:
+      - uses: actions/checkout@v4
+      - name: Install nightly toolchain
+        uses: dtolnay/rust-toolchain@nightly
+        with:
+          components: rust-src
+          targets: wasm32-unknown-unknown
+      - uses: Swatinem/rust-cache@v2
+      - name: Install dependencies
+        run: make installdeps
+      - name: Build rln-wasm in parallel mode
+        run: cargo make build_parallel
+        working-directory: ${{ matrix.crate }}
+      - name: Test rln-wasm in parallel mode on browser
+        run: cargo make test_parallel --release
+        working-directory: ${{ matrix.crate }}
 
   lint:
     strategy:
       matrix:
         # we run lint tests only on ubuntu
-        platform: [ ubuntu-latest ]
-        crate: [ rln, rln-wasm, utils ]
+        platform: [ubuntu-latest]
+        crate: [rln, rln-wasm, utils]
     runs-on: ${{ matrix.platform }}
     timeout-minutes: 60
 
     name: lint - ${{ matrix.crate }} - ${{ matrix.platform }}
     steps:
       - name: Checkout sources
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
       - name: Install stable toolchain
-        uses: actions-rs/toolchain@v1
+        uses: dtolnay/rust-toolchain@stable
         with:
-          profile: minimal
-          toolchain: stable
-          override: true
           components: rustfmt, clippy
       - uses: Swatinem/rust-cache@v2
-      - name: Install Dependencies
+      - name: Install dependencies
         run: make installdeps
-      - name: cargo fmt
+      - name: Check formatting
         if: success() || failure()
         run: cargo fmt -- --check
         working-directory: ${{ matrix.crate }}
-      - name: cargo clippy
+      - name: Check clippy
         if: success() || failure()
         run: |
-          cargo clippy --release
+          cargo clippy --all-targets --release -- -D warnings
         working-directory: ${{ matrix.crate }}
 
   benchmark-utils:
@@ -154,15 +160,15 @@ jobs:
     strategy:
       matrix:
         # we run benchmark tests only on ubuntu
-        platform: [ ubuntu-latest ]
-        crate: [ utils ]
+        platform: [ubuntu-latest]
+        crate: [utils]
     runs-on: ${{ matrix.platform }}
     timeout-minutes: 60
 
-    name: benchmark - ${{ matrix.platform }} - ${{ matrix.crate }}
+    name: benchmark - ${{ matrix.crate }} - ${{ matrix.platform }}
     steps:
       - name: Checkout sources
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
       - uses: Swatinem/rust-cache@v2
       - uses: boa-dev/criterion-compare-action@v3
         with:
@@ -175,19 +181,19 @@ jobs:
     strategy:
       matrix:
         # we run benchmark tests only on ubuntu
-        platform: [ ubuntu-latest ]
-        crate: [ rln ]
-        feature: [ "default", "arkzkey" ]
+        platform: [ubuntu-latest]
+        crate: [rln]
+        feature: ["default"]
     runs-on: ${{ matrix.platform }}
     timeout-minutes: 60
 
-    name: benchmark - ${{ matrix.platform }} - ${{ matrix.crate }} - ${{ matrix.feature }}
+    name: benchmark - ${{ matrix.crate }} - ${{ matrix.platform }} - ${{ matrix.feature }}
     steps:
       - name: Checkout sources
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
       - uses: Swatinem/rust-cache@v2
       - uses: boa-dev/criterion-compare-action@v3
         with:
           branchName: ${{ github.base_ref }}
           cwd: ${{ matrix.crate }}
-          features: ${{ matrix.feature }}
+          features: ${{ matrix.feature }}

.github/workflows/nightly-release.yml

@@ -6,38 +6,38 @@ on:
 
 jobs:
   linux:
+    name: Linux build
+    runs-on: ubuntu-latest
     strategy:
       matrix:
-        feature: [ "default", "arkzkey", "stateless" ]
+        feature:
+          - "stateless"
+          - "stateless,parallel"
+          - "pmtree-ft"
+          - "pmtree-ft,parallel"
+          - "fullmerkletree"
+          - "fullmerkletree,parallel"
+          - "optimalmerkletree"
+          - "optimalmerkletree,parallel"
         target:
           - x86_64-unknown-linux-gnu
           - aarch64-unknown-linux-gnu
-          # - i686-unknown-linux-gnu
-        include:
-          - feature: stateless
-            cargo_args: --exclude rln-cli
-    name: Linux build
-    runs-on: ubuntu-latest
     steps:
       - name: Checkout sources
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
       - name: Install stable toolchain
-        uses: actions-rs/toolchain@v1
+        uses: dtolnay/rust-toolchain@stable
         with:
-          profile: minimal
-          toolchain: stable
-          override: true
           target: ${{ matrix.target }}
       - uses: Swatinem/rust-cache@v2
       - name: Install dependencies
         run: make installdeps
-      - name: cross build
+      - name: Cross build
         run: |
-          cross build --release --target ${{ matrix.target }} --features ${{ matrix.feature }} --workspace ${{ matrix.cargo_args }}
+          cross build --release --target ${{ matrix.target }} --no-default-features --features ${{ matrix.feature }} --workspace --exclude rln-cli
           mkdir release
           cp target/${{ matrix.target }}/release/librln* release/
           tar -czvf ${{ matrix.target }}-${{ matrix.feature }}-rln.tar.gz release/
-
       - name: Upload archive artifact
         uses: actions/upload-artifact@v4
         with:
@@ -50,33 +50,34 @@ jobs:
     runs-on: macos-latest
     strategy:
       matrix:
-        feature: [ "default", "arkzkey", "stateless" ]
+        feature:
+          - "stateless"
+          - "stateless,parallel"
+          - "pmtree-ft"
+          - "pmtree-ft,parallel"
+          - "fullmerkletree"
+          - "fullmerkletree,parallel"
+          - "optimalmerkletree"
+          - "optimalmerkletree,parallel"
         target:
           - x86_64-apple-darwin
           - aarch64-apple-darwin
-        include:
-          - feature: stateless
-            cargo_args: --exclude rln-cli
     steps:
       - name: Checkout sources
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
       - name: Install stable toolchain
-        uses: actions-rs/toolchain@v1
+        uses: dtolnay/rust-toolchain@stable
         with:
-          profile: minimal
-          toolchain: stable
-          override: true
           target: ${{ matrix.target }}
       - uses: Swatinem/rust-cache@v2
       - name: Install dependencies
         run: make installdeps
-      - name: cross build
+      - name: Cross build
         run: |
-          cross build --release --target ${{ matrix.target }} --features ${{ matrix.feature }} --workspace ${{ matrix.cargo_args }}
+          cross build --release --target ${{ matrix.target }} --no-default-features --features ${{ matrix.feature }} --workspace --exclude rln-cli
           mkdir release
           cp target/${{ matrix.target }}/release/librln* release/
           tar -czvf ${{ matrix.target }}-${{ matrix.feature }}-rln.tar.gz release/
-
       - name: Upload archive artifact
         uses: actions/upload-artifact@v4
         with:
@@ -85,47 +86,76 @@ jobs:
           retention-days: 2
 
   browser-rln-wasm:
-    name: Browser build (RLN WASM)
+    name: Browser build
     runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        feature:
+          - "default"
+          - "parallel"
     steps:
       - name: Checkout sources
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
       - name: Install stable toolchain
-        uses: actions-rs/toolchain@v1
+        uses: dtolnay/rust-toolchain@stable
         with:
-          profile: minimal
-          toolchain: stable
-          override: true
+          targets: wasm32-unknown-unknown
+      - name: Install nightly toolchain
+        uses: dtolnay/rust-toolchain@nightly
+        with:
+          targets: wasm32-unknown-unknown
+          components: rust-src
       - uses: Swatinem/rust-cache@v2
+        with:
+          key: wasm-${{ matrix.feature }}
       - name: Install dependencies
         run: make installdeps
-      - name: cross make build
+      - name: Install wasm-pack
+        run: curl https://rustwasm.github.io/wasm-pack/installer/init.sh -sSf | sh
+      - name: Install binaryen
         run: |
-          cross make build
-          mkdir release
-          cp pkg/** release/
-          tar -czvf browser-rln-wasm.tar.gz release/
-        working-directory: rln-wasm
+          sudo apt-get update
+          sudo apt-get install -y binaryen
+      - name: Build wasm package
+        run: |
+          if [[ ${{ matrix.feature }} == *parallel* ]]; then
+            env RUSTFLAGS="-C target-feature=+atomics,+bulk-memory,+mutable-globals" \
+            rustup run nightly wasm-pack build --release --target web --scope waku \
+            --features ${{ matrix.feature }} -Z build-std=panic_abort,std
 
+            wasm-bindgen --target web --split-linked-modules --out-dir ./pkg \
+            ./target/wasm32-unknown-unknown/release/rln_wasm.wasm
+          else
+            wasm-pack build --release --target web --scope waku --features ${{ matrix.feature }}
+          fi
+
+          sed -i.bak 's/rln-wasm/zerokit-rln-wasm/g' pkg/package.json && rm pkg/package.json.bak
+
+          wasm-opt pkg/rln_wasm_bg.wasm -Oz --strip-debug --strip-dwarf \
+          --remove-unused-module-elements --vacuum -o pkg/rln_wasm_bg.wasm
+
+          mkdir release
+          cp -r pkg/* release/
+          tar -czvf browser-rln-wasm-${{ matrix.feature }}.tar.gz release/
+        working-directory: rln-wasm
       - name: Upload archive artifact
         uses: actions/upload-artifact@v4
         with:
-          name: browser-rln-wasm-archive
-          path: rln-wasm/browser-rln-wasm.tar.gz
+          name: Browser-${{ matrix.feature }}-rln-wasm-archive
+          path: rln-wasm/browser-${{ matrix.feature }}-rln-wasm.tar.gz
           retention-days: 2
 
   prepare-prerelease:
     name: Prepare pre-release
-    needs: [ linux, macos, browser-rln-wasm ]
+    needs: [linux, macos, browser-rln-wasm]
     runs-on: ubuntu-latest
     steps:
       - name: Checkout code
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
         with:
           ref: master
       - name: Download artifacts
         uses: actions/download-artifact@v4
-
       - name: Delete tag
         uses: dev-drprasad/delete-tag-and-release@v0.2.1
         with:
@@ -133,7 +163,6 @@ jobs:
           tag_name: nightly
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-
       - name: Create prerelease
         run: |
           start_tag=$(gh release list -L 2 --exclude-drafts | grep -v nightly | cut -d$'\t' -f3 | sed -n '1p')
@@ -145,7 +174,6 @@ jobs:
             *-archive/*.tar.gz \
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-
       - name: Delete artifacts
         uses: geekyeggo/delete-artifact@v5
         with:

.github/workflows/sync-labels.yml

@@ -9,7 +9,7 @@ jobs:
   build:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
       - uses: micnncim/action-label-syncer@v1
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}

.gitignore

@@ -7,11 +7,11 @@ rln-cli/database
 
 # Generated by Cargo
 # will have compiled files and executables
-debug/
-target/
+/target
+/Cargo.lock
 
 # Generated by Nix
-result/
+result
 
 # These are backup files generated by rustfmt
 **/*.rs.bk

Cargo.toml

@@ -1,6 +1,6 @@
 [workspace]
-members = ["rln", "rln-cli", "rln-wasm", "utils"]
-default-members = ["rln", "rln-cli", "rln-wasm", "utils"]
+members = ["rln", "rln-cli", "utils"]
+exclude = ["rln-wasm"]
 resolver = "2"
 
 # Compilation profile for any non-workspace member.

Makefile

@@ -1,6 +1,6 @@
 .PHONY: all installdeps build test bench clean
 
-all: .pre-build build
+all: installdeps build
 
 .fetch-submodules:
 	@git submodule update --init --recursive
@@ -13,30 +13,27 @@ endif
 
 installdeps: .pre-build
 ifeq ($(shell uname),Darwin)
-	@brew update
-	@brew install cmake ninja
+	@brew install cmake ninja binaryen
 else ifeq ($(shell uname),Linux)
-	@sudo apt-get update
-	@sudo apt-get install -y cmake ninja-build
-endif
-	@if [ ! -d "$$HOME/.nvm" ]; then \
-		curl -o- https://raw.githubusercontent.com/nvm-sh/nvm/v0.40.2/install.sh | bash; \
+	@if [ -f /etc/os-release ] && grep -q "ID=nixos" /etc/os-release; then \
+		echo "Detected NixOS, skipping apt-get installation."; \
+	else \
+		sudo apt-get install -y cmake ninja-build binaryen; \
 	fi
-	@bash -c 'export NVM_DIR="$$HOME/.nvm" && \
-		[ -s "$$NVM_DIR/nvm.sh" ] && \. "$$NVM_DIR/nvm.sh" && \
-		nvm install 22.14.0 && \
-		nvm use 22.14.0'
-	@curl https://rustwasm.github.io/wasm-pack/installer/init.sh -sSf | sh
-	@echo "\033[1;32m>>> Now run this command to activate Node.js 22.14.0: \033[1;33msource $$HOME/.nvm/nvm.sh && nvm use 22.14.0\033[0m"
+endif
+	@which wasm-pack > /dev/null && wasm-pack --version | grep -q "0.13.1" || cargo install wasm-pack --version=0.13.1
+	@which wasm-bindgen > /dev/null && wasm-bindgen --version | grep -q "0.2.100" || cargo install wasm-bindgen-cli --version=0.2.100
+	@test -s "$$HOME/.nvm/nvm.sh" || curl -o- https://raw.githubusercontent.com/nvm-sh/nvm/v0.40.2/install.sh | bash
+	@bash -c '. "$$HOME/.nvm/nvm.sh"; [ "$$(node -v 2>/dev/null)" = "v22.14.0" ] || nvm install 22.14.0; nvm use 22.14.0; nvm alias default 22.14.0'
 
-build: .pre-build
+build: installdeps
 	@cargo make build
 
-test: .pre-build
+test: build
 	@cargo make test
 
-bench: .pre-build
+bench: build
 	@cargo make bench
 
 clean:
-	@cargo clean
+	@cargo clean

README.md

@@ -28,15 +28,18 @@ Zerokit currently focuses on RLN (Rate-Limiting Nullifier) implementation using
 
 ## Build and Test
 
-> [!IMPORTANT]
-> For WASM support or x32 architecture builds, use version `0.6.1`. The current version has dependency issues for these platforms. WASM support will return in a future release.
-
 ### Install Dependencies
 
 ```bash
 make installdeps
 ```
 
+#### Use Nix to install dependencies
+
+```bash
+nix develop
+```
+
 ### Build and Test All Crates
 
 ```bash
@@ -69,8 +72,8 @@ The execution graph file used by this code has been generated by means of the sa
 
 > [!IMPORTANT]
 > The circom-witnesscalc code fragments have been borrowed instead of depending on this crate,
-    because its types of input and output data were incompatible with the corresponding zerokit code fragments,
-    and circom-witnesscalc has some dependencies, which are redundant for our purpose.
+> because its types of input and output data were incompatible with the corresponding zerokit code fragments,
+> and circom-witnesscalc has some dependencies, which are redundant for our purpose.
 
 ## Documentation
 

flake.lock

@@ -18,7 +18,28 @@
     },
     "root": {
       "inputs": {
-        "nixpkgs": "nixpkgs"
+        "nixpkgs": "nixpkgs",
+        "rust-overlay": "rust-overlay"
+      }
+    },
+    "rust-overlay": {
+      "inputs": {
+        "nixpkgs": [
+          "nixpkgs"
+        ]
+      },
+      "locked": {
+        "lastModified": 1748399823,
+        "narHash": "sha256-kahD8D5hOXOsGbNdoLLnqCL887cjHkx98Izc37nDjlA=",
+        "owner": "oxalica",
+        "repo": "rust-overlay",
+        "rev": "d68a69dc71bc19beb3479800392112c2f6218159",
+        "type": "github"
+      },
+      "original": {
+        "owner": "oxalica",
+        "repo": "rust-overlay",
+        "type": "github"
       }
     }
   },

flake.nix

@@ -4,9 +4,13 @@
   inputs = {
     # Version 24.11
     nixpkgs.url = "github:NixOS/nixpkgs?rev=f44bd8ca21e026135061a0a57dcf3d0775b67a49";
+    rust-overlay = {
+      url = "github:oxalica/rust-overlay";
+      inputs.nixpkgs.follows = "nixpkgs";
+    };
   };
 
-  outputs = { self, nixpkgs }: 
+  outputs = { self, nixpkgs, rust-overlay }:
     let
       stableSystems = [
         "x86_64-linux" "aarch64-linux"
@@ -15,24 +19,51 @@
         "i686-windows"
       ];
       forAllSystems = nixpkgs.lib.genAttrs stableSystems;
-      pkgsFor = forAllSystems (system: import nixpkgs { inherit system; });
+      overlays = [
+        (import rust-overlay)
+        (f: p: { inherit rust-overlay; })
+      ];
+      pkgsFor = forAllSystems (system: import nixpkgs { inherit system overlays; });
     in rec
     {
       packages = forAllSystems (system: let
         pkgs = pkgsFor.${system};
+        buildPackage = pkgs.callPackage ./nix/default.nix;
+        buildRln = (buildPackage { src = self; project = "rln"; }).override;
       in rec {
-        zerokit-android-arm64 = pkgs.callPackage ./nix/default.nix { target-platform="aarch64-android-prebuilt"; rust-target= "aarch64-linux-android"; };
-        default = zerokit-android-arm64;
+        rln = buildRln
+
+        rln-linux-arm64 = buildRln {
+          target-platform = "aarch64-multiplatform";
+          rust-target = "aarch64-unknown-linux-gnu";
+        };
+
+        rln-android-arm64 = buildRln {
+          target-platform = "aarch64-android-prebuilt";
+          rust-target = "aarch64-linux-android";
+        };
+
+        rln-ios-arm64 = buildRln {
+          target-platform = "aarch64-darwin";
+          rust-target = "aarch64-apple-ios";
+        };
+
+        # TODO: Remove legacy name for RLN android library
+        zerokit-android-arm64 = rln-android-arm64;
+
+        default = rln;
       });
 
       devShells = forAllSystems (system: let
         pkgs = pkgsFor.${system};
       in {
         default = pkgs.mkShell {
-          inputsFrom = [
-            packages.${system}.default
+          buildInputs = with pkgs; [
+            git cmake cargo-make rustup
+            binaryen ninja gnuplot
+            rust-bin.stable.latest.default
           ];
         };
       });
     };
-}
+}

nix/default.nix

@@ -1,35 +1,61 @@
-{ 
+{
   pkgs,
-  target-platform ? "aarch64-android-prebuilt",
-  rust-target ? "aarch64-linux-android",
+  rust-overlay,
+  project,
+  src ? ../.,
+  release ? true,
+  target-platform ? null,
+  rust-target ? null,
 }:
 
-pkgs.pkgsCross.${target-platform}.rustPlatform.buildRustPackage {
-  pname = "zerokit";
-  version = "nightly";
+let
+  # Use cross-compilation if target-platform is specified.
+  targetPlatformPkgs = if target-platform != null
+    then pkgs.pkgsCross.${target-platform}
+    else pkgs;
 
-  src = ../.;
+  rust-bin = rust-overlay.lib.mkRustBin { } targetPlatformPkgs.buildPackages;
+
+  # Use Rust and Cargo versions from rust-overlay.
+  rustPlatform = targetPlatformPkgs.makeRustPlatform {
+    cargo = rust-bin.stable.latest.minimal;
+    rustc = rust-bin.stable.latest.minimal;
+  };
+in rustPlatform.buildRustPackage {
+  pname = "zerokit";
+  version = if src ? rev then src.rev else "nightly";
+
+  # Improve caching of sources
+  src = builtins.path { path = src; name = "zerokit"; };
 
   cargoLock = {
     lockFile = ../Cargo.lock;
     allowBuiltinFetchGit = true;
   };
 
+  doCheck = false;
+
   CARGO_HOME = "/tmp";
 
   buildPhase = ''
-    pushd rln
-    cargo rustc --crate-type=cdylib --release --lib --target=${rust-target}
-    popd
+    cargo build --lib \
+      ${if release             then "--release" else ""} \
+      ${if rust-target != null then "--target=${rust-target}" else ""} \
+      ${if features != null    then "--features=${features}" else ""} \
+      --manifest-path ${project}/Cargo.toml
   '';
 
   installPhase = ''
     mkdir -p $out/
-    cp ./target/${rust-target}/release/librln.so $out/
+    for file in $(find target -name 'librln.*' | grep -v deps/); do
+      mkdir -p $out/$(dirname $file)
+      cp -r $file $out/$file
+    done
   '';
 
+
   meta = with pkgs.lib; {
     description = "Zerokit";
     license = licenses.mit;
   };
-}
+}

rln-cli/Cargo.toml

@@ -13,15 +13,13 @@ path = "src/examples/stateless.rs"
 required-features = ["stateless"]
 
 [dependencies]
-rln = { path = "../rln", default-features = false }
-zerokit_utils = { path = "../utils" }
-clap = { version = "4.5.35", features = ["cargo", "derive", "env"] }
-clap_derive = { version = "4.5.32" }
-color-eyre = "0.6.3"
-serde_json = "1.0"
+rln = { path = "../rln", version = "0.8.0", default-features = false }
+zerokit_utils = { path = "../utils", version = "0.6.0", default-features = false }
+clap = { version = "4.5.41", features = ["cargo", "derive", "env"] }
+color-eyre = "0.6.5"
+serde_json = "1.0.141"
 serde = { version = "1.0", features = ["derive"] }
 
 [features]
-default = []
-arkzkey = ["rln/arkzkey"]
-stateless = ["rln/stateless"]
+default = ["rln/pmtree-ft", "rln/parallel"]
+stateless = ["rln/stateless", "rln/parallel"]

rln-cli/README.md

@@ -27,16 +27,6 @@ If the configuration file is empty, default settings will be used, but the tree
 
 We recommend using the example config, as all commands (except `new` and `create-with-params`) require an initialized RLN instance.
 
-## Feature Flags
-
-The CLI supports optional features. To enable the **arkzkey** feature, run:
-
-```bash
-cargo run --features arkzkey -- <SUBCOMMAND> [OPTIONS]
-```
-
-For more details, refer to the [Zerokit RLN Module](../rln/README.md) documentation.
-
 ## Relay Example
 
 The following [Example](src/examples/relay.rs) demonstrates how RLN enables spam prevention in anonymous environments for multple users.
@@ -47,12 +37,6 @@ You can run the example using the following command:
 cargo run --example relay
 ```
 
-or with the **arkzkey** feature flag:
-
-```bash
-cargo run --example relay --features arkzkey
-```
-
 You can also change **MESSAGE_LIMIT** and **TREEE_HEIGHT** in the [relay.rs](src/examples/relay.rs) file to see how the RLN instance behaves with different parameters.
 
 The customize **TREEE_HEIGHT** constant differs from the default value of `20` should follow [Custom Circuit Compilation](../rln/README.md#advanced-custom-circuit-compilation) instructions.
@@ -66,13 +50,7 @@ This example function similarly to the [Relay Example](#relay-example) but uses
 You can run the example using the following command:
 
 ```bash
-cargo run --example stateless --features stateless
-```
-
-or with the **arkzkey** feature flag:
-
-```bash
-cargo run --example stateless --features stateless,arkzkey
+cargo run --example stateless --no-default-features --features stateless
 ```
 
 ## CLI Commands

rln-cli/src/examples/relay.rs

@@ -6,13 +6,13 @@ use std::{
 };
 
 use clap::{Parser, Subcommand};
-use color_eyre::{eyre::eyre, Result};
+use color_eyre::{eyre::eyre, Report, Result};
 use rln::{
     circuit::Fr,
     hashers::{hash_to_field, poseidon_hash},
     protocol::{keygen, prepare_prove_input, prepare_verify_input},
     public::RLN,
-    utils::{bytes_le_to_fr, fr_to_bytes_le, generate_input_buffer},
+    utils::{fr_to_bytes_le, generate_input_buffer, IdSecret},
 };
 
 const MESSAGE_LIMIT: u32 = 1;
@@ -44,7 +44,7 @@ enum Commands {
 
 #[derive(Debug, Clone)]
 struct Identity {
-    identity_secret_hash: Fr,
+    identity_secret_hash: IdSecret,
     id_commitment: Fr,
 }
 
@@ -68,10 +68,7 @@ impl RLNSystem {
     fn new() -> Result<Self> {
         let mut resources: Vec<Vec<u8>> = Vec::new();
         let resources_path: PathBuf = format!("../rln/resources/tree_height_{TREEE_HEIGHT}").into();
-        #[cfg(feature = "arkzkey")]
         let filenames = ["rln_final.arkzkey", "graph.bin"];
-        #[cfg(not(feature = "arkzkey"))]
-        let filenames = ["rln_final.zkey", "graph.bin"];
         for filename in filenames {
             let fullpath = resources_path.join(Path::new(filename));
             let mut file = File::open(&fullpath)?;
@@ -103,7 +100,7 @@ impl RLNSystem {
         println!("Registered users:");
         for (index, identity) in &self.local_identities {
             println!("User Index: {index}");
-            println!("+ Identity Secret Hash: {}", identity.identity_secret_hash);
+            println!("+ Identity Secret Hash: {}", *identity.identity_secret_hash);
             println!("+ Identity Commitment: {}", identity.id_commitment);
             println!();
         }
@@ -118,7 +115,7 @@ impl RLNSystem {
         match self.rln.set_next_leaf(&mut buffer) {
             Ok(_) => {
                 println!("Registered User Index: {index}");
-                println!("+ Identity secret hash: {}", identity.identity_secret_hash);
+                println!("+ Identity secret hash: {}", *identity.identity_secret_hash);
                 println!("+ Identity commitment: {},", identity.id_commitment);
                 self.local_identities.insert(index, identity);
             }
@@ -143,7 +140,7 @@ impl RLNSystem {
         };
 
         let serialized = prepare_prove_input(
-            identity.identity_secret_hash,
+            identity.identity_secret_hash.clone(),
             user_index,
             Fr::from(MESSAGE_LIMIT),
             Fr::from(message_id),
@@ -182,7 +179,7 @@ impl RLNSystem {
             Ok(false) => {
                 println!("Verification failed: message_id must be unique within the epoch and satisfy 0 <= message_id < MESSAGE_LIMIT: {MESSAGE_LIMIT}");
             }
-            Err(err) => return Err(err),
+            Err(err) => return Err(Report::new(err)),
         }
         Ok(())
     }
@@ -211,7 +208,7 @@ impl RLNSystem {
         {
             Ok(_) => {
                 let output_data = output.into_inner();
-                let (leaked_identity_secret_hash, _) = bytes_le_to_fr(&output_data);
+                let (leaked_identity_secret_hash, _) = IdSecret::from_bytes_le(&output_data);
 
                 if let Some((user_index, identity)) = self
                     .local_identities
@@ -221,20 +218,21 @@ impl RLNSystem {
                     })
                     .map(|(index, identity)| (*index, identity))
                 {
-                    let real_identity_secret_hash = identity.identity_secret_hash;
+                    let real_identity_secret_hash = identity.identity_secret_hash.clone();
                     if leaked_identity_secret_hash != real_identity_secret_hash {
-                        Err(eyre!("identity secret hash mismatch {leaked_identity_secret_hash} != {real_identity_secret_hash}"))
+                        Err(eyre!("identity secret hash mismatch: leaked_identity_secret_hash != real_identity_secret_hash"))
                     } else {
-                        println!("DUPLICATE message ID detected! Reveal identity secret hash: {leaked_identity_secret_hash}");
+                        println!(
+                            "DUPLICATE message ID detected! Reveal identity secret hash: {}",
+                            *leaked_identity_secret_hash
+                        );
                         self.local_identities.remove(&user_index);
                         self.rln.delete_leaf(user_index)?;
                         println!("User index {user_index} has been SLASHED");
                         Ok(())
                     }
                 } else {
-                    Err(eyre!(
-                        "user identity secret hash {leaked_identity_secret_hash} not found"
-                    ))
+                    Err(eyre!("user identity secret hash ******** not found"))
                 }
             }
             Err(err) => Err(eyre!("Failed to recover identity secret: {err}")),

rln-cli/src/examples/stateless.rs

@@ -1,4 +1,5 @@
 #![cfg(feature = "stateless")]
+
 use std::{
     collections::HashMap,
     io::{stdin, stdout, Cursor, Write},
@@ -8,13 +9,12 @@ use clap::{Parser, Subcommand};
 use color_eyre::{eyre::eyre, Result};
 use rln::{
     circuit::{Fr, TEST_TREE_HEIGHT},
-    hashers::{hash_to_field, poseidon_hash},
-    poseidon_tree::PoseidonTree,
+    hashers::{hash_to_field, poseidon_hash, PoseidonHash},
     protocol::{keygen, prepare_verify_input, rln_witness_from_values, serialize_witness},
     public::RLN,
-    utils::{bytes_le_to_fr, fr_to_bytes_le},
+    utils::{fr_to_bytes_le, IdSecret},
 };
-use zerokit_utils::ZerokitMerkleTree;
+use zerokit_utils::{OptimalMerkleTree, ZerokitMerkleProof, ZerokitMerkleTree};
 
 const MESSAGE_LIMIT: u32 = 1;
 
@@ -45,7 +45,7 @@ enum Commands {
 
 #[derive(Debug, Clone)]
 struct Identity {
-    identity_secret_hash: Fr,
+    identity_secret_hash: IdSecret,
     id_commitment: Fr,
 }
 
@@ -61,7 +61,7 @@ impl Identity {
 
 struct RLNSystem {
     rln: RLN,
-    tree: PoseidonTree,
+    tree: OptimalMerkleTree<PoseidonHash>,
     used_nullifiers: HashMap<[u8; 32], Vec<u8>>,
     local_identities: HashMap<usize, Identity>,
 }
@@ -70,11 +70,12 @@ impl RLNSystem {
     fn new() -> Result<Self> {
         let rln = RLN::new()?;
         let default_leaf = Fr::from(0);
-        let tree = PoseidonTree::new(
+        let tree: OptimalMerkleTree<PoseidonHash> = OptimalMerkleTree::new(
             TEST_TREE_HEIGHT,
             default_leaf,
-            ConfigOf::<PoseidonTree>::default(),
-        )?;
+            ConfigOf::<OptimalMerkleTree<PoseidonHash>>::default(),
+        )
+        .unwrap();
 
         Ok(RLNSystem {
             rln,
@@ -93,7 +94,7 @@ impl RLNSystem {
         println!("Registered users:");
         for (index, identity) in &self.local_identities {
             println!("User Index: {index}");
-            println!("+ Identity Secret Hash: {}", identity.identity_secret_hash);
+            println!("+ Identity Secret Hash: {}", *identity.identity_secret_hash);
             println!("+ Identity Commitment: {}", identity.id_commitment);
             println!();
         }
@@ -107,7 +108,7 @@ impl RLNSystem {
         self.tree.update_next(rate_commitment)?;
 
         println!("Registered User Index: {index}");
-        println!("+ Identity secret hash: {}", identity.identity_secret_hash);
+        println!("+ Identity secret hash: {}", *identity.identity_secret_hash);
         println!("+ Identity commitment: {}", identity.id_commitment);
 
         self.local_identities.insert(index, identity);
@@ -130,8 +131,9 @@ impl RLNSystem {
         let x = hash_to_field(signal.as_bytes());
 
         let rln_witness = rln_witness_from_values(
-            identity.identity_secret_hash,
-            &merkle_proof,
+            identity.identity_secret_hash.clone(),
+            merkle_proof.get_path_elements(),
+            merkle_proof.get_path_index(),
             x,
             external_nullifier,
             Fr::from(MESSAGE_LIMIT),
@@ -208,7 +210,7 @@ impl RLNSystem {
         {
             Ok(_) => {
                 let output_data = output.into_inner();
-                let (leaked_identity_secret_hash, _) = bytes_le_to_fr(&output_data);
+                let (leaked_identity_secret_hash, _) = IdSecret::from_bytes_le(&output_data);
 
                 if let Some((user_index, identity)) = self
                     .local_identities
@@ -218,19 +220,19 @@ impl RLNSystem {
                     })
                     .map(|(index, identity)| (*index, identity))
                 {
-                    let real_identity_secret_hash = identity.identity_secret_hash;
+                    let real_identity_secret_hash = identity.identity_secret_hash.clone();
                     if leaked_identity_secret_hash != real_identity_secret_hash {
-                        Err(eyre!("identity secret hash mismatch {leaked_identity_secret_hash} != {real_identity_secret_hash}"))
+                        Err(eyre!("identity secret hash mismatch: leaked_identity_secret_hash != real_identity_secret_hash"))
                     } else {
-                        println!("DUPLICATE message ID detected! Reveal identity secret hash: {leaked_identity_secret_hash}");
+                        println!(
+                            "DUPLICATE message ID detected! Reveal identity secret hash: ********"
+                        );
                         self.local_identities.remove(&user_index);
                         println!("User index {user_index} has been SLASHED");
                         Ok(())
                     }
                 } else {
-                    Err(eyre!(
-                        "user identity secret hash {leaked_identity_secret_hash} not found"
-                    ))
+                    Err(eyre!("user identity secret hash ******** not found"))
                 }
             }
             Err(err) => Err(eyre!("Failed to recover identity secret: {err}")),

rln-cli/src/main.rs

@@ -51,10 +51,7 @@ fn main() -> Result<()> {
             resources_path,
         }) => {
             let mut resources: Vec<Vec<u8>> = Vec::new();
-            #[cfg(feature = "arkzkey")]
             let filenames = ["rln_final.arkzkey", "graph.bin"];
-            #[cfg(not(feature = "arkzkey"))]
-            let filenames = ["rln_final.zkey", "graph.bin"];
             for filename in filenames {
                 let fullpath = resources_path.join(Path::new(filename));
                 let mut file = File::open(&fullpath)?;
@@ -141,7 +138,7 @@ fn main() -> Result<()> {
                 .ok_or(Report::msg("no RLN instance initialized"))?
                 .prove(input_data, &mut output_buffer)?;
             let proof = output_buffer.into_inner();
-            println!("proof: {:?}", proof);
+            println!("proof: {proof:?}");
             Ok(())
         }
         Some(Commands::Verify { input }) => {
@@ -150,7 +147,7 @@ fn main() -> Result<()> {
                 .rln
                 .ok_or(Report::msg("no RLN instance initialized"))?
                 .verify(input_data)?;
-            println!("verified: {:?}", verified);
+            println!("verified: {verified:?}");
             Ok(())
         }
         Some(Commands::GenerateProof { input }) => {
@@ -161,7 +158,7 @@ fn main() -> Result<()> {
                 .ok_or(Report::msg("no RLN instance initialized"))?
                 .generate_rln_proof(input_data, &mut output_buffer)?;
             let proof = output_buffer.into_inner();
-            println!("proof: {:?}", proof);
+            println!("proof: {proof:?}");
             Ok(())
         }
         Some(Commands::VerifyWithRoots { input, roots }) => {
@@ -193,7 +190,7 @@ fn main() -> Result<()> {
             let output_buffer_inner = output_buffer.into_inner();
             let (path_elements, _) = bytes_le_to_vec_fr(&output_buffer_inner)?;
             for (index, element) in path_elements.iter().enumerate() {
-                println!("path element {}: {}", index, element);
+                println!("path element {index}: {element}");
             }
             Ok(())
         }

rln-wasm/Cargo.toml

@@ -1,39 +1,46 @@
 [package]
 name = "rln-wasm"
-version = "0.1.0"
+version = "0.2.0"
 edition = "2021"
 license = "MIT or Apache2"
 
 [lib]
 crate-type = ["cdylib", "rlib"]
-required-features = ["stateless"]
 
 [dependencies]
-rln = { path = "../rln", default-features = false }
-num-bigint = { version = "0.4.6", default-features = false, features = [
-    "rand",
-    "serde",
+rln = { path = "../rln", version = "0.8.0", default-features = false, features = [
+    "stateless",
 ] }
+zerokit_utils = { path = "../utils", version = "0.6.0", default-features = false }
+num-bigint = { version = "0.4.6", default-features = false }
+js-sys = "0.3.77"
 wasm-bindgen = "0.2.100"
 serde-wasm-bindgen = "0.6.5"
-js-sys = "0.3.77"
-serde_json = "1.0"
+wasm-bindgen-rayon = { version = "1.3.0", features = [
+    "no-bundler",
+], optional = true }
 
 # The `console_error_panic_xhook` crate provides better debugging of panics by
 # logging them with `console.error`. This is great for development, but requires
 # all the `std::fmt` and `std::panicking` infrastructure, so isn't great for
 # code size when deploying.
 console_error_panic_hook = { version = "0.1.7", optional = true }
-zerokit_utils = { path = "../utils" }
 
 [target.'cfg(target_arch = "wasm32")'.dependencies]
-getrandom = { version = "0.2.15", features = ["js"] }
+getrandom = { version = "0.2.16", features = ["js"] }
 
 [dev-dependencies]
+serde_json = "1.0.141"
 wasm-bindgen-test = "0.3.50"
 wasm-bindgen-futures = "0.4.50"
 
+[dev-dependencies.web-sys]
+version = "0.3.77"
+features = ["Window", "Navigator"]
+
 [features]
 default = ["console_error_panic_hook"]
-stateless = ["rln/stateless"]
-arkzkey = ["rln/arkzkey"]
+parallel = ["rln/parallel", "wasm-bindgen-rayon"]
+
+[package.metadata.docs.rs]
+all-features = true

rln-wasm/Makefile.toml

@@ -1,24 +1,61 @@
 [tasks.build]
 clear = true
-dependencies = ["pack_build", "pack_rename"]
+dependencies = ["pack_build", "pack_rename", "pack_resize"]
 
-[tasks.build_arkzkey]
+[tasks.build_parallel]
 clear = true
-dependencies = ["pack_build_arkzkey", "pack_rename"]
+dependencies = [
+    "pack_build_parallel",
+    "post_build_parallel",
+    "pack_rename",
+    "pack_resize",
+]
 
 [tasks.pack_build]
 command = "wasm-pack"
 args = ["build", "--release", "--target", "web", "--scope", "waku"]
-env = { "RUSTFLAGS" = "--cfg feature=\"stateless\"" }
 
-[tasks.pack_build_arkzkey]
-command = "wasm-pack"
-args = ["build", "--release", "--target", "web", "--scope", "waku"]
-env = { "RUSTFLAGS" = "--cfg feature=\"stateless\" --cfg feature=\"arkzkey\"" }
+[tasks.pack_build_parallel]
+command = "env"
+args = [
+    "RUSTFLAGS=-C target-feature=+atomics,+bulk-memory,+mutable-globals",
+    "rustup",
+    "run",
+    "nightly",
+    "wasm-pack",
+    "build",
+    "--release",
+    "--target",
+    "web",
+    "--scope",
+    "waku",
+    "--features",
+    "parallel",
+    "-Z",
+    "build-std=panic_abort,std",
+]
+
+[tasks.post_build_parallel]
+script = '''
+wasm-bindgen --target web --split-linked-modules --out-dir ./pkg ./target/wasm32-unknown-unknown/release/rln_wasm.wasm
+'''
 
 [tasks.pack_rename]
 script = "sed -i.bak 's/rln-wasm/zerokit-rln-wasm/g' pkg/package.json && rm pkg/package.json.bak"
 
+[tasks.pack_resize]
+command = "wasm-opt"
+args = [
+    "pkg/rln_wasm_bg.wasm",
+    "-Oz",
+    "--strip-debug",
+    "--strip-dwarf",
+    "--remove-unused-module-elements",
+    "--vacuum",
+    "-o",
+    "pkg/rln_wasm_bg.wasm",
+]
+
 [tasks.test]
 command = "wasm-pack"
 args = [
@@ -30,21 +67,44 @@ args = [
     "--",
     "--nocapture",
 ]
-env = { "RUSTFLAGS" = "--cfg feature=\"stateless\"" }
+dependencies = ["build"]
 
-[tasks.test_arkzkey]
+[tasks.test_browser]
 command = "wasm-pack"
 args = [
     "test",
     "--release",
-    "--node",
+    "--chrome",
+    "--headless",
     "--target",
     "wasm32-unknown-unknown",
     "--",
     "--nocapture",
 ]
-env = { "RUSTFLAGS" = "--cfg feature=\"stateless\" --cfg feature=\"arkzkey\"" }
-dependencies = ["build_arkzkey"]
+dependencies = ["build"]
+
+[tasks.test_parallel]
+command = "env"
+args = [
+    "RUSTFLAGS=-C target-feature=+atomics,+bulk-memory,+mutable-globals",
+    "rustup",
+    "run",
+    "nightly",
+    "wasm-pack",
+    "test",
+    "--release",
+    "--chrome",
+    "--headless",
+    "--target",
+    "wasm32-unknown-unknown",
+    "--features",
+    "parallel",
+    "-Z",
+    "build-std=panic_abort,std",
+    "--",
+    "--nocapture",
+]
+dependencies = ["build_parallel"]
 
 [tasks.bench]
 disabled = true

rln-wasm/README.md

@@ -1,22 +1,58 @@
 # RLN for WASM
 
-This library is used in [waku-org/js-rln](https://github.com/waku-org/js-rln/)
+[![npm version](https://badge.fury.io/js/@waku%2Fzerokit-rln-wasm.svg)](https://badge.fury.io/js/@waku%2Fzerokit-rln-wasm)
+[![License: MIT](https://img.shields.io/badge/License-MIT-blue.svg)](https://opensource.org/licenses/MIT)
+[![License: Apache 2.0](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](https://opensource.org/licenses/Apache-2.0)
 
-> **Note**: This project requires `wasm-pack` for compiling Rust to WebAssembly and `cargo-make` for running the build commands. Make sure both are installed before proceeding.
+The Zerokit RLN WASM Module provides WebAssembly bindings for working with
+Rate-Limiting Nullifier [RLN](https://rfc.vac.dev/spec/32/) zkSNARK proofs and primitives.
+This module is used by [waku-org/js-rln](https://github.com/waku-org/js-rln/) to enable
+RLN functionality in JavaScript/TypeScript applications.
 
-Install `wasm-pack`:
+## Install Dependencies
+
+> [!NOTE]
+> This project requires the following tools:
+>
+> - `wasm-pack` - for compiling Rust to WebAssembly
+> - `cargo-make` - for running build commands
+> - `nvm` - to install and manage Node.js
+>
+> Ensure all dependencies are installed before proceeding.
+
+### Manually
+
+#### Install `wasm-pack`
 
 ```bash
-curl https://rustwasm.github.io/wasm-pack/installer/init.sh -sSf | sh
+cargo install wasm-pack --version=0.13.1
 ```
 
-Install `cargo-make`
+#### Install `cargo-make`
 
 ```bash
 cargo install cargo-make
 ```
 
-Or install everything needed for `zerokit` at the root of the repository:
+#### Install `Node.js`
+
+If you don't have `nvm` (Node Version Manager), install it by following
+the [installation instructions](https://github.com/nvm-sh/nvm?tab=readme-ov-file#install--update-script).
+
+After installing `nvm`, install and use Node.js `v22.14.0`:
+
+```bash
+nvm install 22.14.0
+nvm use 22.14.0
+nvm alias default 22.14.0
+```
+
+If you already have Node.js installed,
+check your version with `node -v` command — the version must be strictly greater than 22.
+
+### Or install everything
+
+You can run the following command from the root of the repository to install all required dependencies for `zerokit`
 
 ```bash
 make installdeps
@@ -36,20 +72,81 @@ Compile zerokit for `wasm32-unknown-unknown`:
 cargo make build
 ```
 
-Or compile with the **arkzkey** feature enabled
-
-```bash
-cargo make build_arkzkey
-```
-
 ## Running tests and benchmarks
 
 ```bash
 cargo make test
 ```
 
-Or test with the **arkzkey** feature enabled
+If you want to run the tests in browser headless mode, you can use the following command:
 
 ```bash
-cargo make test_arkzkey
+cargo make test_browser
+```
+
+## Parallel computation
+
+The library supports parallel computation using the `wasm-bindgen-rayon` crate,
+enabling multi-threaded execution in the browser.
+
+> [!NOTE]
+> Parallel support is not enabled by default due to WebAssembly and browser limitations. \
+> Compiling this feature requires `nightly` Rust and the `wasm-bindgen-cli` tool.
+
+### Build Setup
+
+#### Install `nightly` Rust
+
+```bash
+rustup install nightly
+```
+
+#### Install `wasm-bindgen-cli`
+
+```bash
+cargo install wasm-bindgen-cli --version=0.2.100
+```
+
+### Build Commands
+
+To enable parallel computation for WebAssembly threads, you can use the following command:
+
+```bash
+cargo make build_parallel
+```
+
+### WebAssembly Threading Support
+
+Most modern browsers support WebAssembly threads,
+but they require the following headers to enable `SharedArrayBuffer`, which is necessary for multithreading:
+
+- Cross-Origin-Opener-Policy: same-origin
+- Cross-Origin-Embedder-Policy: require-corp
+
+Without these, the application will fall back to single-threaded mode.
+
+## Feature detection
+
+If you're targeting [older browser versions that didn't support WebAssembly threads yet](https://webassembly.org/roadmap/),
+you'll likely want to create two builds - one with thread support and one without -
+and use feature detection to choose the right one on the JavaScript side.
+
+You can use [wasm-feature-detect](https://github.com/GoogleChromeLabs/wasm-feature-detect) library for this purpose.
+For example, your code might look like this:
+
+```js
+import { threads } from 'wasm-feature-detect';
+
+let wasmPkg;
+
+if (await threads()) {
+  wasmPkg = await import('./pkg-with-threads/index.js');
+  await wasmPkg.default();
+  await wasmPkg.initThreadPool(navigator.hardwareConcurrency);
+} else {
+  wasmPkg = await import('./pkg-without-threads/index.js');
+  await wasmPkg.default();
+}
+
+wasmPkg.nowCallAnyExportedFuncs();
 ```

rln-wasm/benches/index.html

@@ -0,0 +1,561 @@
+<!DOCTYPE html>
+<html lang="en">
+  <head>
+    <meta charset="UTF-8" />
+    <meta name="viewport" content="width=device-width, initial-scale=1.0" />
+    <title>RLN WASM Benchmark</title>
+    <style>
+      body {
+        font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto,
+          sans-serif;
+        max-width: 700px;
+        margin: 0 auto;
+        padding: 20px;
+        color: #333;
+      }
+
+      h1 {
+        text-align: center;
+        margin-bottom: 30px;
+      }
+
+      .panel {
+        background: #f8f9fa;
+        border-radius: 8px;
+        padding: 20px;
+        margin-bottom: 20px;
+        box-shadow: 0 2px 4px rgba(0, 0, 0, 0.05);
+      }
+
+      .file-input {
+        margin-bottom: 15px;
+      }
+
+      label {
+        display: block;
+        margin-bottom: 5px;
+        font-weight: 500;
+      }
+
+      button {
+        background: #4361ee;
+        color: white;
+        border: none;
+        padding: 10px 20px;
+        border-radius: 4px;
+        font-size: 16px;
+        cursor: pointer;
+        margin-top: 10px;
+        width: 100%;
+      }
+
+      button:hover {
+        background: #3a56d4;
+      }
+
+      button:disabled {
+        background: #cccccc;
+        cursor: not-allowed;
+      }
+
+      .results-table {
+        width: 100%;
+        border-collapse: collapse;
+        margin-top: 10px;
+      }
+
+      .results-table th,
+      .results-table td {
+        padding: 12px;
+        text-align: left;
+        border-bottom: 1px solid #ddd;
+      }
+
+      .results-table th {
+        font-weight: 600;
+        background-color: #f1f3f5;
+      }
+
+      .operation {
+        font-weight: 500;
+      }
+
+      .time {
+        font-family: monospace;
+        text-align: right;
+      }
+
+      .status {
+        padding: 10px;
+        margin-top: 15px;
+        border-radius: 4px;
+        text-align: center;
+      }
+
+      .success {
+        background-color: #d4edda;
+        color: #155724;
+      }
+
+      .error {
+        background-color: #f8d7da;
+        color: #721c24;
+      }
+
+      .running {
+        background-color: #cce5ff;
+        color: #004085;
+      }
+
+      /* Download button style */
+      .download-btn {
+        background: #28a745;
+        margin-top: 15px;
+      }
+
+      .download-btn:hover {
+        background: #218838;
+      }
+
+      /* Summary section */
+      .summary {
+        margin-top: 15px;
+        padding-top: 15px;
+        border-top: 1px solid #ddd;
+      }
+
+      .summary h3 {
+        margin-top: 0;
+      }
+
+      /* Thread info */
+      .thread-info {
+        margin-bottom: 15px;
+        padding: 10px;
+        background-color: #f1f3f5;
+        border-radius: 4px;
+        font-weight: 500;
+      }
+    </style>
+    <script>
+      // Check if cross-origin isolation is available for SharedArrayBuffer
+      document.addEventListener("DOMContentLoaded", () => {
+        if (!crossOriginIsolated) {
+          const statusElement = document.getElementById("status");
+          if (statusElement) {
+            statusElement.innerHTML =
+              '<strong style="color: #721c24;">Error:</strong> ' +
+              "Cross-Origin Isolation is not enabled. Please run the server.js script with:<br>" +
+              "<code>node server.js</code><br>" +
+              'Then access this page via <a href="http://localhost:8000">http://localhost:8000</a>';
+            statusElement.className = "status error";
+          }
+        }
+      });
+    </script>
+  </head>
+
+  <body>
+    <h1>RLN WASM Benchmark</h1>
+
+    <div class="panel">
+      <div class="thread-info" id="threadInfo">Detecting CPU cores...</div>
+
+      <div
+        class="thread-mode-selector"
+        style="display: flex; gap: 10px; margin-bottom: 15px"
+      >
+        <button id="singleThreaded" class="thread-btn" style="flex: 1">
+          Single-Threaded Mode
+        </button>
+        <button
+          id="multiThreaded"
+          class="thread-btn"
+          style="flex: 1; background: #4361ee"
+        >
+          Multi-Threaded Mode
+        </button>
+      </div>
+
+      <button id="initThreads" class="init-btn">Initialize Thread Pool</button>
+
+      <div class="file-input">
+        <label for="zkeyFile">zKey File:</label>
+        <input type="file" id="zkeyFile" />
+      </div>
+      <div class="file-input">
+        <label for="rootFile">Root File:</label>
+        <input type="file" id="rootFile" />
+      </div>
+      <div class="file-input">
+        <label for="witnessFile">Witness File:</label>
+        <input type="file" id="witnessFile" />
+      </div>
+      <div class="file-input">
+        <label for="messageFile">Message File:</label>
+        <input type="file" id="messageFile" />
+      </div>
+      <div class="file-input">
+        <label for="proofFile">Proof File:</label>
+        <input type="file" id="proofFile" />
+      </div>
+
+      <button id="runBenchmark" disabled>Run Benchmark</button>
+
+      <div id="status" class="status">Please initialize thread pool first</div>
+    </div>
+
+    <div class="panel">
+      <h2>Results</h2>
+      <table class="results-table">
+        <thead>
+          <tr>
+            <th>Operation</th>
+            <th>Time (ms)</th>
+          </tr>
+        </thead>
+        <tbody id="results">
+          <!-- Results will be populated here -->
+        </tbody>
+      </table>
+
+      <div id="summarySection" class="summary" style="display: none">
+        <h3>Summary</h3>
+        <div id="summaryContent"></div>
+        <button id="downloadResults" class="download-btn">
+          Download Results
+        </button>
+      </div>
+    </div>
+
+    <script type="module">
+      import init, * as RLN from "../pkg/rln_wasm.js";
+
+      // Get DOM elements
+      const initThreadsBtn = document.getElementById("initThreads");
+      const runBtn = document.getElementById("runBenchmark");
+      const results = document.getElementById("results");
+      const status = document.getElementById("status");
+      const summarySection = document.getElementById("summarySection");
+      const summaryContent = document.getElementById("summaryContent");
+      const downloadBtn = document.getElementById("downloadResults");
+      const threadInfo = document.getElementById("threadInfo");
+
+      // Track benchmark operations
+      const benchmarks = [];
+      let threadPoolInitialized = false;
+      let cpuCount = navigator.hardwareConcurrency || 4; // Default to 4 if detection fails
+      let useMultiThreaded = true; // Default to multi-threaded mode
+
+      // Update the thread info display
+      updateThreadInfo();
+
+      // Function to update thread info display
+      function updateThreadInfo() {
+        if (useMultiThreaded) {
+          threadInfo.textContent = `Automatically detected ${cpuCount} CPU cores for optimal performance`;
+        } else {
+          threadInfo.textContent = `Using single-threaded mode (1 CPU core)`;
+        }
+      }
+
+      // Thread mode selection
+      document
+        .getElementById("singleThreaded")
+        .addEventListener("click", () => {
+          if (!threadPoolInitialized) {
+            useMultiThreaded = false;
+            document.getElementById("singleThreaded").style.background =
+              "#4361ee";
+            document.getElementById("multiThreaded").style.background =
+              "#6c757d";
+            updateThreadInfo();
+          } else {
+            updateStatus(
+              "Thread pool already initialized. Please refresh page to change mode.",
+              "error"
+            );
+          }
+        });
+
+      document.getElementById("multiThreaded").addEventListener("click", () => {
+        if (!threadPoolInitialized) {
+          useMultiThreaded = true;
+          document.getElementById("multiThreaded").style.background = "#4361ee";
+          document.getElementById("singleThreaded").style.background =
+            "#6c757d";
+          updateThreadInfo();
+        } else {
+          updateStatus(
+            "Thread pool already initialized. Please refresh page to change mode.",
+            "error"
+          );
+        }
+      });
+
+      // Initialize thread pool
+      initThreadsBtn.addEventListener("click", async () => {
+        if (!crossOriginIsolated) {
+          updateStatus(
+            "Cross-Origin Isolation is required. Please run with server.js",
+            "error"
+          );
+          return;
+        }
+
+        if (threadPoolInitialized) {
+          updateStatus("Thread pool already initialized.", "error");
+          return;
+        }
+
+        updateStatus(
+          `Initializing thread pool with ${cpuCount} cores...`,
+          "running"
+        );
+        initThreadsBtn.disabled = true;
+
+        try {
+          const start = performance.now();
+          await init();
+          await RLN.initThreadPool(useMultiThreaded ? cpuCount : 1);
+          const duration = performance.now() - start;
+
+          benchmarks.push({
+            name: "Initialize WASM Module",
+            duration: duration,
+            success: true,
+          });
+
+          updateResults();
+          threadPoolInitialized = true;
+          updateStatus(
+            `Thread pool initialized with ${
+              useMultiThreaded ? cpuCount : 1
+            } cores`,
+            "success"
+          );
+          runBtn.disabled = false;
+        } catch (error) {
+          console.error("Thread pool initialization error:", error);
+          updateStatus(
+            `Error initializing thread pool: ${error.message}`,
+            "error"
+          );
+        } finally {
+          initThreadsBtn.disabled = false;
+        }
+      });
+
+      // Measure operation time
+      async function benchmark(name, fn) {
+        updateStatus(`Running: ${name}...`, "running");
+
+        const start = performance.now();
+        try {
+          const result = await fn();
+          const duration = performance.now() - start;
+
+          // Record result
+          benchmarks.push({ name, duration, success: true });
+          updateResults();
+
+          return { result, duration };
+        } catch (error) {
+          const duration = performance.now() - start;
+          benchmarks.push({
+            name: `${name} (FAILED)`,
+            duration,
+            success: false,
+            error: error.message,
+          });
+          updateResults();
+          throw error;
+        }
+      }
+
+      // Update results table
+      function updateResults() {
+        results.innerHTML = "";
+        benchmarks.forEach((b) => {
+          const row = document.createElement("tr");
+
+          const nameCell = document.createElement("td");
+          nameCell.className = "operation";
+          nameCell.textContent = b.name;
+
+          const timeCell = document.createElement("td");
+          timeCell.className = "time";
+          timeCell.textContent = b.duration.toFixed(2);
+
+          row.appendChild(nameCell);
+          row.appendChild(timeCell);
+
+          if (!b.success) {
+            row.style.color = "#dc3545";
+          }
+
+          results.appendChild(row);
+        });
+      }
+
+      // Update status message
+      function updateStatus(message, type = "") {
+        status.textContent = message;
+        status.className = `status ${type}`;
+      }
+
+      // Show benchmark summary
+      function showSummary() {
+        if (benchmarks.length === 0) return;
+
+        const successfulOps = benchmarks.filter((b) => b.success).length;
+
+        let summaryHTML = `
+                <p><strong>Operations:</strong> ${successfulOps}/${
+          benchmarks.length
+        } successful</p>
+                <p><strong>CPU cores used:</strong> ${
+                  useMultiThreaded ? cpuCount : 1
+                }</p>
+                <p><strong>Mode:</strong> ${
+                  useMultiThreaded ? "Multi-threaded" : "Single-threaded"
+                }</p>
+            `;
+
+        summaryContent.innerHTML = summaryHTML;
+        summarySection.style.display = "block";
+      }
+
+      // Download results as JSON
+      downloadBtn.addEventListener("click", () => {
+        const dataStr = JSON.stringify(
+          {
+            timestamp: new Date().toISOString(),
+            mode: useMultiThreaded ? "multi-threaded" : "single-threaded",
+            cpuCount: useMultiThreaded ? cpuCount : 1,
+            operations: benchmarks,
+          },
+          null,
+          2
+        );
+
+        const blob = new Blob([dataStr], { type: "application/json" });
+        const url = URL.createObjectURL(blob);
+
+        const a = document.createElement("a");
+        a.href = url;
+        a.download = `rln-benchmark-${new Date()
+          .toISOString()
+          .slice(0, 19)}.json`;
+        document.body.appendChild(a);
+        a.click();
+        document.body.removeChild(a);
+        URL.revokeObjectURL(url);
+      });
+
+      // Read file as Uint8Array
+      async function readFile(file) {
+        const buffer = await file.arrayBuffer();
+        return new Uint8Array(buffer);
+      }
+
+      // Parse witness JSON
+      async function parseWitness(file) {
+        const text = await file.text();
+        try {
+          const data = JSON.parse(text);
+
+          if (!Array.isArray(data)) {
+            throw new Error("Witness JSON must be an array");
+          }
+
+          return data.map((value) => BigInt(value));
+        } catch (e) {
+          throw new Error(`Failed to parse witness JSON: ${e.message}`);
+        }
+      }
+
+      // Main benchmark runner
+      runBtn.addEventListener("click", async () => {
+        if (!threadPoolInitialized) {
+          updateStatus("Please initialize thread pool first", "error");
+          return;
+        }
+
+        const files = {
+          zkey: document.getElementById("zkeyFile").files[0],
+          root: document.getElementById("rootFile").files[0],
+          witness: document.getElementById("witnessFile").files[0],
+          message: document.getElementById("messageFile").files[0],
+          proof: document.getElementById("proofFile").files[0],
+        };
+
+        // Validation
+        if (!files.zkey || !files.root || !files.proof) {
+          updateStatus("Please select zKey, Root, and Proof files", "error");
+          return;
+        }
+
+        const canGenerate = files.witness && files.message;
+
+        try {
+          // Keep file references but clear previous benchmark results except initialization
+          const initBenchmark = benchmarks.find(
+            (b) => b.name === "Initialize WASM Module"
+          );
+          benchmarks.length = 0;
+          if (initBenchmark) {
+            benchmarks.push(initBenchmark);
+          }
+          updateResults();
+          summarySection.style.display = "none";
+          runBtn.disabled = true;
+          initThreadsBtn.disabled = true;
+
+          // Load files
+          const zkeyData = await readFile(files.zkey);
+          const rootData = await readFile(files.root);
+
+          // Create RLN instance
+          const { result: instance } = await benchmark(
+            "Create RLN Instance",
+            async () => {
+              return RLN.newRLN(zkeyData);
+            }
+          );
+
+          // Handle proof generation (if witness and message files provided)
+          if (canGenerate) {
+            const witnessData = await parseWitness(files.witness);
+            const messageData = await readFile(files.message);
+
+            await benchmark("Generate RLN Proof", async () => {
+              return RLN.generateRLNProofWithWitness(
+                instance,
+                witnessData,
+                messageData
+              );
+            });
+          }
+
+          // Verify uploaded proof (required)
+          const proofData = await readFile(files.proof);
+
+          await benchmark("Verify Proof", async () => {
+            return RLN.verifyWithRoots(instance, proofData, rootData);
+          });
+
+          updateStatus("Benchmark completed successfully!", "success");
+          showSummary();
+        } catch (error) {
+          console.error(error);
+          updateStatus(`Error: ${error.message}`, "error");
+          showSummary();
+        } finally {
+          runBtn.disabled = false;
+          initThreadsBtn.disabled = false;
+        }
+      });
+    </script>
+  </body>
+</html>

rln-wasm/benches/server.js

@@ -0,0 +1,83 @@
+const http = require("http");
+const fs = require("fs");
+const path = require("path");
+const url = require("url");
+
+const PORT = 8000;
+
+// MIME type mapping
+const MIME_TYPES = {
+  ".html": "text/html",
+  ".js": "text/javascript",
+  ".css": "text/css",
+  ".json": "application/json",
+  ".wasm": "application/wasm",
+  ".png": "image/png",
+  ".jpg": "image/jpeg",
+  ".gif": "image/gif",
+  ".svg": "image/svg+xml",
+};
+
+// Create HTTP server
+const server = http.createServer((req, res) => {
+  // Set COOP and COEP headers for SharedArrayBuffer support
+  res.setHeader("Cross-Origin-Opener-Policy", "same-origin");
+  res.setHeader("Cross-Origin-Embedder-Policy", "require-corp");
+
+  // Parse URL
+  const parsedUrl = url.parse(req.url);
+  let requestPath = parsedUrl.pathname;
+
+  // Ignore favicon
+  if (requestPath === "/favicon.ico") {
+    res.writeHead(204);
+    res.end();
+    return;
+  }
+
+  // Handle root path
+  let filePath = "." + requestPath;
+  if (filePath === "./") {
+    filePath = "./index.html";
+  }
+
+  // Handle pkg files (including snippets)
+  if (requestPath.startsWith("/pkg/")) {
+    filePath = ".." + requestPath;
+  }
+
+  // Determine content type based on file extension
+  const extname = path.extname(filePath);
+  const contentType = MIME_TYPES[extname] || "application/octet-stream";
+
+  // ❗ Block directory reads
+  if (fs.existsSync(filePath) && fs.statSync(filePath).isDirectory()) {
+    console.error(`Attempted directory read: ${filePath}`);
+    res.writeHead(403);
+    res.end("Forbidden: Cannot read directory directly");
+    return;
+  }
+
+  // Read and serve the file
+  fs.readFile(filePath, (error, content) => {
+    if (error) {
+      if (error.code === "ENOENT") {
+        console.error(`File not found: ${filePath}`);
+        res.writeHead(404);
+        res.end(`File not found: ${requestPath}`);
+      } else {
+        console.error(`Server error (${error.code}): ${filePath}`);
+        res.writeHead(500);
+        res.end(`Server Error: ${error.code}`);
+      }
+    } else {
+      res.writeHead(200, { "Content-Type": contentType });
+      res.end(content, "utf-8");
+    }
+  });
+});
+
+// Start the server
+server.listen(PORT, () => {
+  console.log(`Server running at http://localhost:${PORT}`);
+});

rln-wasm/resources/witness_calculator.js

@@ -1,4 +1,8 @@
-module.exports = async function builder(code, options) {
+// File generated with https://github.com/iden3/circom
+// following the instructions from:
+// https://github.com/vacp2p/zerokit/tree/master/rln#advanced-custom-circuit-compilation
+
+export async function builder(code, options) {
   options = options || {};
 
   let wasmModule;
@@ -101,7 +105,7 @@ module.exports = async function builder(code, options) {
     // Then append the value to the message we are creating
     msgStr += fromArray32(arr).toString();
   }
-};
+}
 
 class WitnessCalculator {
   constructor(instance, sanityCheck) {

rln-wasm/src/lib.rs

@@ -6,6 +6,9 @@ use rln::public::{hash, poseidon_hash, RLN};
 use std::vec::Vec;
 use wasm_bindgen::prelude::*;
 
+#[cfg(feature = "parallel")]
+pub use wasm_bindgen_rayon::init_thread_pool;
+
 #[wasm_bindgen(js_name = initPanicHook)]
 pub fn init_panic_hook() {
     console_error_panic_hook::set_once();

rln-wasm/src/utils.js

@@ -1,26 +0,0 @@
-const fs = require("fs");
-
-// Utils functions for loading circom witness calculator and reading files from test
-
-module.exports = {
-  read_file: function (path) {
-    return fs.readFileSync(path);
-  },
-
-  calculateWitness: async function (circom_path, inputs) {
-    const wc = require("../resources/witness_calculator.js");
-    const wasmFile = fs.readFileSync(circom_path);
-    const wasmFileBuffer = wasmFile.slice(
-      wasmFile.byteOffset,
-      wasmFile.byteOffset + wasmFile.byteLength
-    );
-    const witnessCalculator = await wc(wasmFileBuffer);
-    const calculatedWitness = await witnessCalculator.calculateWitness(
-      inputs,
-      false
-    );
-    return JSON.stringify(calculatedWitness, (key, value) =>
-      typeof value === "bigint" ? value.toString() : value
-    );
-  },
-};

rln-wasm/tests/browser.rs

@@ -0,0 +1,265 @@
+#![cfg(target_arch = "wasm32")]
+
+#[cfg(test)]
+mod tests {
+    use js_sys::{BigInt as JsBigInt, Date, Object, Uint8Array};
+    use rln::circuit::{Fr, TEST_TREE_HEIGHT};
+    use rln::hashers::{hash_to_field, poseidon_hash, PoseidonHash};
+    use rln::protocol::{prepare_verify_input, rln_witness_from_values, serialize_witness};
+    use rln::utils::{bytes_le_to_fr, fr_to_bytes_le, IdSecret};
+    use rln_wasm::{
+        wasm_generate_rln_proof_with_witness, wasm_key_gen, wasm_new, wasm_rln_witness_to_json,
+        wasm_verify_with_roots,
+    };
+    use wasm_bindgen::{prelude::wasm_bindgen, JsValue};
+    use wasm_bindgen_test::{console_log, wasm_bindgen_test, wasm_bindgen_test_configure};
+    use zerokit_utils::{
+        OptimalMerkleProof, OptimalMerkleTree, ZerokitMerkleProof, ZerokitMerkleTree,
+    };
+
+    #[cfg(feature = "parallel")]
+    use {rln_wasm::init_thread_pool, wasm_bindgen_futures::JsFuture, web_sys::window};
+
+    #[wasm_bindgen(inline_js = r#"
+    export function isThreadpoolSupported() {
+      return typeof SharedArrayBuffer !== 'undefined' &&
+             typeof Atomics !== 'undefined' &&
+             typeof crossOriginIsolated !== 'undefined' &&
+             crossOriginIsolated;
+    }
+
+    export function initWitnessCalculator(jsCode) {
+      const processedCode = jsCode
+        .replace(/export\s+async\s+function\s+builder/, 'async function builder')
+        .replace(/export\s*\{\s*builder\s*\};?/g, '');
+
+      const moduleFunc = new Function(processedCode + '\nreturn { builder };');
+      const witnessCalculatorModule = moduleFunc();
+
+      window.witnessCalculatorBuilder = witnessCalculatorModule.builder;
+
+      if (typeof window.witnessCalculatorBuilder !== 'function') {
+        return false;
+      }
+      return true;
+    }
+
+    export function readFile(data) {
+      return new Uint8Array(data);
+    }
+
+    export async function calculateWitness(circom_data, inputs) {
+      const wasmBuffer = circom_data instanceof Uint8Array ? circom_data : new Uint8Array(circom_data);
+      const witnessCalculator = await window.witnessCalculatorBuilder(wasmBuffer);
+      const calculatedWitness = await witnessCalculator.calculateWitness(inputs, false);
+      return JSON.stringify(calculatedWitness, (key, value) =>
+        typeof value === "bigint" ? value.toString() : value
+      );
+    }
+    "#)]
+    extern "C" {
+        #[wasm_bindgen(catch)]
+        fn isThreadpoolSupported() -> Result<bool, JsValue>;
+
+        #[wasm_bindgen(catch)]
+        fn initWitnessCalculator(js: &str) -> Result<bool, JsValue>;
+
+        #[wasm_bindgen(catch)]
+        fn readFile(data: &[u8]) -> Result<Uint8Array, JsValue>;
+
+        #[wasm_bindgen(catch)]
+        async fn calculateWitness(circom_data: &[u8], inputs: Object) -> Result<JsValue, JsValue>;
+    }
+
+    const WITNESS_CALCULATOR_JS: &str = include_str!("../resources/witness_calculator.js");
+
+    const ARKZKEY_BYTES: &[u8] =
+        include_bytes!("../../rln/resources/tree_height_20/rln_final.arkzkey");
+
+    const CIRCOM_BYTES: &[u8] = include_bytes!("../../rln/resources/tree_height_20/rln.wasm");
+
+    wasm_bindgen_test_configure!(run_in_browser);
+
+    #[wasm_bindgen_test]
+    pub async fn rln_wasm_benchmark() {
+        // Check if thread pool is supported
+        #[cfg(feature = "parallel")]
+        if !isThreadpoolSupported().expect("Failed to check thread pool support") {
+            panic!("Thread pool is NOT supported");
+        } else {
+            // Initialize thread pool
+            let cpu_count = window()
+                .expect("Failed to get window")
+                .navigator()
+                .hardware_concurrency() as usize;
+            JsFuture::from(init_thread_pool(cpu_count))
+                .await
+                .expect("Failed to initialize thread pool");
+        }
+
+        // Initialize witness calculator
+        initWitnessCalculator(WITNESS_CALCULATOR_JS)
+            .expect("Failed to initialize witness calculator");
+
+        let mut results = String::from("\nbenchmarks:\n");
+        let iterations = 10;
+
+        let zkey = readFile(&ARKZKEY_BYTES).expect("Failed to read zkey file");
+
+        // Benchmark wasm_new
+        let start_wasm_new = Date::now();
+        for _ in 0..iterations {
+            let _ = wasm_new(zkey.clone()).expect("Failed to create RLN instance");
+        }
+        let wasm_new_result = Date::now() - start_wasm_new;
+
+        // Create RLN instance for other benchmarks
+        let rln_instance = wasm_new(zkey).expect("Failed to create RLN instance");
+        let mut tree: OptimalMerkleTree<PoseidonHash> =
+            OptimalMerkleTree::default(TEST_TREE_HEIGHT).expect("Failed to create tree");
+
+        // Benchmark wasm_key_gen
+        let start_wasm_key_gen = Date::now();
+        for _ in 0..iterations {
+            let _ = wasm_key_gen(rln_instance).expect("Failed to generate keys");
+        }
+        let wasm_key_gen_result = Date::now() - start_wasm_key_gen;
+
+        // Generate identity pair for other benchmarks
+        let mem_keys = wasm_key_gen(rln_instance).expect("Failed to generate keys");
+        let id_key = mem_keys.subarray(0, 32);
+        let (identity_secret_hash, _) = IdSecret::from_bytes_le(&id_key.to_vec());
+        let (id_commitment, _) = bytes_le_to_fr(&mem_keys.subarray(32, 64).to_vec());
+
+        let epoch = hash_to_field(b"test-epoch");
+        let rln_identifier = hash_to_field(b"test-rln-identifier");
+        let external_nullifier = poseidon_hash(&[epoch, rln_identifier]);
+
+        let identity_index = tree.leaves_set();
+
+        let user_message_limit = Fr::from(100);
+
+        let rate_commitment = poseidon_hash(&[id_commitment, user_message_limit]);
+        tree.update_next(rate_commitment)
+            .expect("Failed to update tree");
+
+        let message_id = Fr::from(0);
+        let signal: [u8; 32] = [0; 32];
+        let x = hash_to_field(&signal);
+
+        let merkle_proof: OptimalMerkleProof<PoseidonHash> = tree
+            .proof(identity_index)
+            .expect("Failed to generate merkle proof");
+
+        let rln_witness = rln_witness_from_values(
+            identity_secret_hash,
+            merkle_proof.get_path_elements(),
+            merkle_proof.get_path_index(),
+            x,
+            external_nullifier,
+            user_message_limit,
+            message_id,
+        )
+        .expect("Failed to create RLN witness");
+
+        let serialized_witness =
+            serialize_witness(&rln_witness).expect("Failed to serialize witness");
+        let witness_buffer = Uint8Array::from(&serialized_witness[..]);
+
+        let json_inputs = wasm_rln_witness_to_json(rln_instance, witness_buffer.clone())
+            .expect("Failed to convert witness to JSON");
+
+        // Benchmark calculateWitness
+        let start_calculate_witness = Date::now();
+        for _ in 0..iterations {
+            let _ = calculateWitness(&CIRCOM_BYTES, json_inputs.clone())
+                .await
+                .expect("Failed to calculate witness");
+        }
+        let calculate_witness_result = Date::now() - start_calculate_witness;
+
+        // Calculate witness for other benchmarks
+        let calculated_witness_json = calculateWitness(&CIRCOM_BYTES, json_inputs)
+            .await
+            .expect("Failed to calculate witness")
+            .as_string()
+            .expect("Failed to convert calculated witness to string");
+        let calculated_witness_vec_str: Vec<String> =
+            serde_json::from_str(&calculated_witness_json).expect("Failed to parse JSON");
+        let calculated_witness: Vec<JsBigInt> = calculated_witness_vec_str
+            .iter()
+            .map(|x| JsBigInt::new(&x.into()).expect("Failed to create JsBigInt"))
+            .collect();
+
+        // Benchmark wasm_generate_rln_proof_with_witness
+        let start_wasm_generate_rln_proof_with_witness = Date::now();
+        for _ in 0..iterations {
+            let _ = wasm_generate_rln_proof_with_witness(
+                rln_instance,
+                calculated_witness.clone(),
+                witness_buffer.clone(),
+            )
+            .expect("Failed to generate proof");
+        }
+        let wasm_generate_rln_proof_with_witness_result =
+            Date::now() - start_wasm_generate_rln_proof_with_witness;
+
+        // Generate a proof for other benchmarks
+        let proof =
+            wasm_generate_rln_proof_with_witness(rln_instance, calculated_witness, witness_buffer)
+                .expect("Failed to generate proof");
+
+        let proof_data = proof.to_vec();
+        let verify_input = prepare_verify_input(proof_data, &signal);
+        let input_buffer = Uint8Array::from(&verify_input[..]);
+
+        let root = tree.root();
+        let roots_serialized = fr_to_bytes_le(&root);
+        let roots_buffer = Uint8Array::from(&roots_serialized[..]);
+
+        // Benchmark wasm_verify_with_roots
+        let start_wasm_verify_with_roots = Date::now();
+        for _ in 0..iterations {
+            let _ =
+                wasm_verify_with_roots(rln_instance, input_buffer.clone(), roots_buffer.clone())
+                    .expect("Failed to verify proof");
+        }
+        let wasm_verify_with_roots_result = Date::now() - start_wasm_verify_with_roots;
+
+        // Verify the proof with the root
+        let is_proof_valid = wasm_verify_with_roots(rln_instance, input_buffer, roots_buffer)
+            .expect("Failed to verify proof");
+        assert!(is_proof_valid, "verification failed");
+
+        // Format and display results
+        let format_duration = |duration_ms: f64| -> String {
+            let avg_ms = duration_ms / (iterations as f64);
+            if avg_ms >= 1000.0 {
+                format!("{:.3} s", avg_ms / 1000.0)
+            } else {
+                format!("{:.3} ms", avg_ms)
+            }
+        };
+
+        results.push_str(&format!("wasm_new: {}\n", format_duration(wasm_new_result)));
+        results.push_str(&format!(
+            "wasm_key_gen: {}\n",
+            format_duration(wasm_key_gen_result)
+        ));
+        results.push_str(&format!(
+            "calculateWitness: {}\n",
+            format_duration(calculate_witness_result)
+        ));
+        results.push_str(&format!(
+            "wasm_generate_rln_proof_with_witness: {}\n",
+            format_duration(wasm_generate_rln_proof_with_witness_result)
+        ));
+        results.push_str(&format!(
+            "wasm_verify_with_roots: {}\n",
+            format_duration(wasm_verify_with_roots_result)
+        ));
+
+        // Log the results
+        console_log!("{results}");
+    }
+}

rln-wasm/tests/node.rs

@@ -1,40 +1,91 @@
+#![cfg(not(feature = "parallel"))]
 #![cfg(target_arch = "wasm32")]
 
 #[cfg(test)]
 mod tests {
     use js_sys::{BigInt as JsBigInt, Date, Object, Uint8Array};
     use rln::circuit::{Fr, TEST_TREE_HEIGHT};
-    use rln::hashers::{hash_to_field, poseidon_hash};
-    use rln::poseidon_tree::PoseidonTree;
+    use rln::hashers::{hash_to_field, poseidon_hash, PoseidonHash};
     use rln::protocol::{prepare_verify_input, rln_witness_from_values, serialize_witness};
-    use rln::utils::{bytes_le_to_fr, fr_to_bytes_le};
-    use rln_wasm::*;
-    use wasm_bindgen::{prelude::*, JsValue};
-    use wasm_bindgen_test::wasm_bindgen_test;
-    use zerokit_utils::merkle_tree::merkle_tree::ZerokitMerkleTree;
+    use rln::utils::{bytes_le_to_fr, fr_to_bytes_le, IdSecret};
+    use rln_wasm::{
+        wasm_generate_rln_proof_with_witness, wasm_key_gen, wasm_new, wasm_rln_witness_to_json,
+        wasm_verify_with_roots,
+    };
+    use wasm_bindgen::{prelude::wasm_bindgen, JsValue};
+    use wasm_bindgen_test::{console_log, wasm_bindgen_test};
+    use zerokit_utils::{
+        OptimalMerkleProof, OptimalMerkleTree, ZerokitMerkleProof, ZerokitMerkleTree,
+    };
 
-    #[wasm_bindgen(module = "src/utils.js")]
+    const WITNESS_CALCULATOR_JS: &str = include_str!("../resources/witness_calculator.js");
+
+    #[wasm_bindgen(inline_js = r#"
+    const fs = require("fs");
+
+    let witnessCalculatorModule = null;
+
+    module.exports = {
+      initWitnessCalculator: function(code) {
+        const processedCode = code
+          .replace(/export\s+async\s+function\s+builder/, 'async function builder')
+          .replace(/export\s*\{\s*builder\s*\};?/g, '');
+
+        const moduleFunc = new Function(processedCode + '\nreturn { builder };');
+        witnessCalculatorModule = moduleFunc();
+
+        if (typeof witnessCalculatorModule.builder !== 'function') {
+          return false;
+        }
+        return true;
+      },
+
+      readFile: function (path) {
+        return fs.readFileSync(path);
+      },
+
+      calculateWitness: async function (circom_path, inputs) {
+        const wasmFile = fs.readFileSync(circom_path);
+        const wasmFileBuffer = wasmFile.slice(
+          wasmFile.byteOffset,
+          wasmFile.byteOffset + wasmFile.byteLength
+        );
+        const witnessCalculator = await witnessCalculatorModule.builder(wasmFileBuffer);
+        const calculatedWitness = await witnessCalculator.calculateWitness(
+          inputs,
+          false
+        );
+        return JSON.stringify(calculatedWitness, (key, value) =>
+          typeof value === "bigint" ? value.toString() : value
+        );
+      },
+    };
+    "#)]
     extern "C" {
         #[wasm_bindgen(catch)]
-        fn read_file(path: &str) -> Result<Uint8Array, JsValue>;
+        fn initWitnessCalculator(code: &str) -> Result<bool, JsValue>;
+
+        #[wasm_bindgen(catch)]
+        fn readFile(path: &str) -> Result<Uint8Array, JsValue>;
 
         #[wasm_bindgen(catch)]
         async fn calculateWitness(circom_path: &str, input: Object) -> Result<JsValue, JsValue>;
     }
 
-    #[cfg(feature = "arkzkey")]
-    const ZKEY_PATH: &str = "../rln/resources/tree_height_20/rln_final.arkzkey";
-    #[cfg(not(feature = "arkzkey"))]
-    const ZKEY_PATH: &str = "../rln/resources/tree_height_20/rln_final.zkey";
+    const ARKZKEY_PATH: &str = "../rln/resources/tree_height_20/rln_final.arkzkey";
 
     const CIRCOM_PATH: &str = "../rln/resources/tree_height_20/rln.wasm";
 
     #[wasm_bindgen_test]
     pub async fn rln_wasm_benchmark() {
+        // Initialize witness calculator
+        initWitnessCalculator(WITNESS_CALCULATOR_JS)
+            .expect("Failed to initialize witness calculator");
+
         let mut results = String::from("\nbenchmarks:\n");
         let iterations = 10;
 
-        let zkey = read_file(&ZKEY_PATH).expect("Failed to read zkey file");
+        let zkey = readFile(&ARKZKEY_PATH).expect("Failed to read zkey file");
 
         // Benchmark wasm_new
         let start_wasm_new = Date::now();
@@ -45,7 +96,8 @@ mod tests {
 
         // Create RLN instance for other benchmarks
         let rln_instance = wasm_new(zkey).expect("Failed to create RLN instance");
-        let mut tree = PoseidonTree::default(TEST_TREE_HEIGHT).expect("Failed to create tree");
+        let mut tree: OptimalMerkleTree<PoseidonHash> =
+            OptimalMerkleTree::default(TEST_TREE_HEIGHT).expect("Failed to create tree");
 
         // Benchmark wasm_key_gen
         let start_wasm_key_gen = Date::now();
@@ -57,7 +109,7 @@ mod tests {
         // Generate identity pair for other benchmarks
         let mem_keys = wasm_key_gen(rln_instance).expect("Failed to generate keys");
         let id_key = mem_keys.subarray(0, 32);
-        let (identity_secret_hash, _) = bytes_le_to_fr(&id_key.to_vec());
+        let (identity_secret_hash, _) = IdSecret::from_bytes_le(&id_key.to_vec());
         let (id_commitment, _) = bytes_le_to_fr(&mem_keys.subarray(32, 64).to_vec());
 
         let epoch = hash_to_field(b"test-epoch");
@@ -76,13 +128,14 @@ mod tests {
         let signal: [u8; 32] = [0; 32];
         let x = hash_to_field(&signal);
 
-        let merkle_proof = tree
+        let merkle_proof: OptimalMerkleProof<PoseidonHash> = tree
             .proof(identity_index)
             .expect("Failed to generate merkle proof");
 
         let rln_witness = rln_witness_from_values(
             identity_secret_hash,
-            &merkle_proof,
+            merkle_proof.get_path_elements(),
+            merkle_proof.get_path_index(),
             x,
             external_nullifier,
             user_message_limit,
@@ -175,7 +228,7 @@ mod tests {
             format_duration(wasm_key_gen_result)
         ));
         results.push_str(&format!(
-            "calculateWitness: {}\n",
+            "calculate_witness: {}\n",
             format_duration(calculate_witness_result)
         ));
         results.push_str(&format!(
@@ -188,102 +241,6 @@ mod tests {
         ));
 
         // Log the results
-        wasm_bindgen_test::console_log!("{results}");
-    }
-
-    #[wasm_bindgen_test]
-    pub async fn rln_wasm_test() {
-        // Read the zkey file
-        let zkey = read_file(&ZKEY_PATH).expect("Failed to read zkey file");
-
-        // Create RLN instance and separated tree
-        let rln_instance = wasm_new(zkey).expect("Failed to create RLN instance");
-        let mut tree = PoseidonTree::default(TEST_TREE_HEIGHT).expect("Failed to create tree");
-
-        // Setting up the epoch and rln_identifier
-        let epoch = hash_to_field(b"test-epoch");
-        let rln_identifier = hash_to_field(b"test-rln-identifier");
-        let external_nullifier = poseidon_hash(&[epoch, rln_identifier]);
-
-        // Generate identity pair
-        let mem_keys = wasm_key_gen(rln_instance).expect("Failed to generate keys");
-        let (identity_secret_hash, _) = bytes_le_to_fr(&mem_keys.subarray(0, 32).to_vec());
-        let (id_commitment, _) = bytes_le_to_fr(&mem_keys.subarray(32, 64).to_vec());
-
-        // Get index of the identity
-        let identity_index = tree.leaves_set();
-
-        // Setting up the user message limit
-        let user_message_limit = Fr::from(100);
-
-        // Updating the tree with the rate commitment
-        let rate_commitment = poseidon_hash(&[id_commitment, user_message_limit]);
-        tree.update_next(rate_commitment)
-            .expect("Failed to update tree");
-
-        // Generate merkle proof
-        let merkle_proof = tree
-            .proof(identity_index)
-            .expect("Failed to generate merkle proof");
-
-        // Create message id and signal
-        let message_id = Fr::from(0);
-        let signal: [u8; 32] = [0; 32];
-        let x = hash_to_field(&signal);
-
-        // Prepare input for witness calculation
-        let rln_witness = rln_witness_from_values(
-            identity_secret_hash,
-            &merkle_proof,
-            x,
-            external_nullifier,
-            user_message_limit,
-            message_id,
-        )
-        .expect("Failed to create RLN witness");
-
-        // Serialize the rln witness
-        let serialized_witness =
-            serialize_witness(&rln_witness).expect("Failed to serialize witness");
-        // Convert the serialized witness to a Uint8Array
-        let witness_buffer = Uint8Array::from(&serialized_witness[..]);
-
-        // Obtaining inputs that should be sent to circom witness calculator
-        let json_inputs = wasm_rln_witness_to_json(rln_instance, witness_buffer.clone())
-            .expect("Failed to convert witness to JSON");
-
-        // Calculating witness with JS
-        // (Using a JSON since wasm_bindgen does not like Result<Vec<JsBigInt>,JsValue>)
-        let calculated_witness_json = calculateWitness(&CIRCOM_PATH, json_inputs)
-            .await
-            .expect("Failed to calculate witness")
-            .as_string()
-            .expect("Failed to convert calculated witness to string");
-        let calculated_witness_vec_str: Vec<String> =
-            serde_json::from_str(&calculated_witness_json).expect("Failed to parse JSON");
-        let calculated_witness: Vec<JsBigInt> = calculated_witness_vec_str
-            .iter()
-            .map(|x| JsBigInt::new(&x.into()).expect("Failed to create JsBigInt"))
-            .collect();
-
-        // Generate a proof from the calculated witness
-        let proof =
-            wasm_generate_rln_proof_with_witness(rln_instance, calculated_witness, witness_buffer)
-                .expect("Failed to generate proof");
-
-        // Prepare the root for verification
-        let root = tree.root();
-        let roots_serialized = fr_to_bytes_le(&root);
-        let roots_buffer = Uint8Array::from(&roots_serialized[..]);
-
-        // Prepare input for proof verification
-        let proof_data = proof.to_vec();
-        let verify_input = prepare_verify_input(proof_data, &signal);
-        let input_buffer = Uint8Array::from(&verify_input[..]);
-
-        // Verify the proof with the root
-        let is_proof_valid = wasm_verify_with_roots(rln_instance, input_buffer, roots_buffer)
-            .expect("Failed to verify proof");
-        assert!(is_proof_valid, "verification failed");
+        console_log!("{results}");
     }
 }

rln/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "rln"
-version = "0.7.0"
+version = "0.8.0"
 edition = "2021"
 license = "MIT OR Apache-2.0"
 description = "APIs to manage, compute and verify zkSNARK proofs and RLN primitives"
@@ -9,7 +9,7 @@ homepage = "https://vac.dev"
 repository = "https://github.com/vacp2p/zerokit"
 
 [lib]
-crate-type = ["rlib", "staticlib"]
+crate-type = ["rlib", "staticlib", "cdylib"]
 bench = false
 
 # This flag disable cargo doctests, i.e. testing example code-snippets in documentation
@@ -19,78 +19,63 @@ doctest = false
 # ZKP Generation
 ark-bn254 = { version = "0.5.0", features = ["std"] }
 ark-relations = { version = "0.5.1", features = ["std"] }
-ark-ff = { version = "0.5.0", default-features = false, features = [
-    "parallel",
-] }
-ark-ec = { version = "0.5.0", default-features = false, features = [
-    "parallel",
-] }
-
-ark-std = { version = "0.5.0", default-features = false, features = [
-    "parallel",
-] }
-ark-poly = { version = "0.5.0", default-features = false, features = [
-    "parallel",
-] }
-ark-groth16 = { version = "0.5.0", default-features = false, features = [
-    "parallel",
-] }
-ark-serialize = { version = "0.5.0", default-features = false, features = [
-    "parallel",
-] }
+ark-ff = { version = "0.5.0", default-features = false }
+ark-ec = { version = "0.5.0", default-features = false }
+ark-std = { version = "0.5.0", default-features = false }
+ark-poly = { version = "0.5.0", default-features = false }
+ark-groth16 = { version = "0.5.0", default-features = false }
+ark-serialize = { version = "0.5.0", default-features = false }
 
 # error handling
-color-eyre = "0.6.3"
 thiserror = "2.0.12"
 
 # utilities
+rayon = { version = "1.10.0", optional = true }
 byteorder = "1.5.0"
 cfg-if = "1.0"
-num-bigint = { version = "0.4.6", default-features = false, features = [
-    "rand",
-    "std",
-] }
+num-bigint = { version = "0.4.6", default-features = false, features = ["std"] }
 num-traits = "0.2.19"
 once_cell = "1.21.3"
 lazy_static = "1.5.0"
 rand = "0.8.5"
 rand_chacha = "0.3.1"
-ruint = { version = "1.14.0", features = ["rand", "serde", "ark-ff-04"] }
+ruint = { version = "1.15.0", features = ["rand", "serde", "ark-ff-04"] }
 tiny-keccak = { version = "2.0.2", features = ["keccak"] }
-utils = { package = "zerokit_utils", version = "0.5.2", path = "../utils", default-features = false }
+zeroize = "1.8"
+utils = { package = "zerokit_utils", version = "0.6.0", path = "../utils", default-features = false }
 
 # serialization
-prost = "0.13.5"
-serde_json = "1.0"
+prost = "0.14.1"
+serde_json = "1.0.141"
 serde = { version = "1.0", features = ["derive"] }
 
 document-features = { version = "0.2.11", optional = true }
 
 [dev-dependencies]
-sled = "0.34.7"
-criterion = { version = "0.4.0", features = ["html_reports"] }
+criterion = { version = "0.7.0", features = ["html_reports"] }
 
 [features]
-default = ["pmtree-ft"]
-fullmerkletree = ["default"]
+default = ["parallel", "pmtree-ft"]
 stateless = []
-arkzkey = []
-
+parallel = [
+    "rayon",
+    "utils/parallel",
+    "ark-ff/parallel",
+    "ark-ec/parallel",
+    "ark-std/parallel",
+    "ark-poly/parallel",
+    "ark-groth16/parallel",
+    "ark-serialize/parallel",
+]
+fullmerkletree = []
+optimalmerkletree = []
 # Note: pmtree feature is still experimental
 pmtree-ft = ["utils/pmtree-ft"]
 
-[[bench]]
-name = "circuit_loading_arkzkey_benchmark"
-harness = false
-required-features = ["arkzkey"]
-
-[[bench]]
-name = "circuit_loading_benchmark"
-harness = false
-
 [[bench]]
 name = "pmtree_benchmark"
 harness = false
+required-features = ["pmtree-ft"]
 
 [[bench]]
 name = "poseidon_tree_benchmark"

rln/Makefile.toml

@@ -8,11 +8,7 @@ args = ["test", "--release", "--", "--nocapture"]
 
 [tasks.test_stateless]
 command = "cargo"
-args = ["test", "--release", "--features", "stateless"]
-
-[tasks.test_arkzkey]
-command = "cargo"
-args = ["test", "--release", "--features", "arkzkey"]
+args = ["test", "--release", "--no-default-features", "--features", "stateless"]
 
 [tasks.bench]
 command = "cargo"

rln/README.md

@@ -1,8 +1,12 @@
 # Zerokit RLN Module
 
 [![Crates.io](https://img.shields.io/crates/v/rln.svg)](https://crates.io/crates/rln)
+[![License: MIT](https://img.shields.io/badge/License-MIT-blue.svg)](https://opensource.org/licenses/MIT)
+[![License: Apache 2.0](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](https://opensource.org/licenses/Apache-2.0)
 
-The Zerokit RLN Module provides a Rust implementation for working with Rate-Limiting Nullifier [RLN](https://rfc.vac.dev/spec/32/) zkSNARK proofs and primitives. This module allows you to:
+The Zerokit RLN Module provides a Rust implementation for working with
+Rate-Limiting Nullifier [RLN](https://rfc.vac.dev/spec/32/) zkSNARK proofs and primitives.
+This module allows you to:
 
 - Generate and verify RLN proofs
 - Work with Merkle trees for commitment storage
@@ -11,7 +15,8 @@ The Zerokit RLN Module provides a Rust implementation for working with Rate-Limi
 ## Quick Start
 
 > [!IMPORTANT]
-> Version 0.6.1 is required for WASM support or x32 architecture. Current version doesn't support these platforms due to dependency issues. WASM support will return in a future release.
+> Version 0.7.0 is the only version that does not support WASM and x32 architecture.
+> WASM support is available in version 0.8.0 and above.
 
 ### Add RLN as dependency
 
@@ -24,9 +29,15 @@ rln = { git = "https://github.com/vacp2p/zerokit" }
 
 ## Basic Usage Example
 
-Note that we need to pass to RLN object constructor the path where the graph file (`graph.bin`, built for the input tree size), the corresponding proving key (`rln_final.zkey`) or (`rln_final_uncompr.arkzkey`) and verification key (`verification_key.arkvkey`, optional) are found.
+The RLN object constructor requires the following files:
 
-In the following we will use [cursors](https://doc.rust-lang.org/std/io/struct.Cursor.html) as readers/writers for interfacing with RLN public APIs.
+- `rln_final.arkzkey`: The proving key in arkzkey format.
+- `graph.bin`: The graph file built for the input tree size
+
+Additionally, `rln.wasm` is used for testing in the rln-wasm module.
+
+In the following we will use [cursors](https://doc.rust-lang.org/std/io/struct.Cursor.html)
+as readers/writers for interfacing with RLN public APIs.
 
 ```rust
 use std::io::Cursor;
@@ -74,7 +85,8 @@ fn main() {
     let signal = b"RLN is awesome";
 
     // 6. Prepare input for generate_rln_proof API
-    // input_data is [ identity_secret<32> | id_index<8> | external_nullifier<32> | user_message_limit<32> | message_id<32> | signal_len<8> | signal<var> ]
+    // input_data is [ identity_secret<32> | id_index<8> | external_nullifier<32>
+    //    | user_message_limit<32> | message_id<32> | signal_len<8> | signal<var> ]
     let prove_input = prepare_prove_input(
         identity_secret_hash,
         id_index,
@@ -92,11 +104,13 @@ fn main() {
         .unwrap();
 
     // We get the public outputs returned by the circuit evaluation
-    // The byte vector `proof_data` is serialized as `[ zk-proof | tree_root | external_nullifier | share_x | share_y | nullifier ]`.
+    // The byte vector `proof_data` is serialized as
+    //  `[ zk-proof | tree_root | external_nullifier | share_x | share_y | nullifier ]`.
     let proof_data = output_buffer.into_inner();
 
     // 8. Verify a RLN proof
-    // Input buffer is serialized as `[proof_data | signal_len | signal ]`, where `proof_data` is (computed as) the output obtained by `generate_rln_proof`.
+    // Input buffer is serialized as `[proof_data | signal_len | signal ]`,
+    //   where `proof_data` is (computed as) the output obtained by `generate_rln_proof`.
     let verify_data = prepare_verify_input(proof_data, signal);
 
     // We verify the zk-proof against the provided proof values
@@ -113,16 +127,20 @@ fn main() {
 The `external nullifier` includes two parameters.
 
 The first one is `epoch` and it's used to identify messages received in a certain time frame.
-It usually corresponds to the current UNIX time but can also be set to a random value or generated by a seed, provided that it corresponds to a field element.
+It usually corresponds to the current UNIX time but can also be set to a random value or generated by a seed,
+provided that it corresponds to a field element.
 
-The second one is `rln_identifier` and it's used to prevent a RLN ZK proof generated for one application to be re-used in another one.
+The second one is `rln_identifier` and it's used to prevent a RLN ZK proof generated
+for one application to be re-used in another one.
 
 ### Features
 
-- **Multiple Backend Support**: Choose between different zkey formats with feature flags
-  - `arkzkey`: Use the optimized Arkworks-compatible zkey format (faster loading)
-  - `stateless`: For stateless proof verification
-- **Pre-compiled Circuits**: Ready-to-use circuits with Merkle tree height of 20
+- **Stateless Mode**: Allows the use of RLN without maintaining state of the Merkle tree.
+- **Pre-compiled Circuits**: Ready-to-use circuits with Merkle tree depth of 20
+- **Wasm Support**: WebAssembly bindings via rln-wasm crate with features like:
+  - Browser and Node.js compatibility
+  - Optional parallel feature support using [wasm-bindgen-rayon](https://github.com/RReverser/wasm-bindgen-rayon)
+  - Headless browser testing capabilities
 
 ## Building and Testing
 
@@ -143,20 +161,22 @@ cargo make build
 # Test with default features
 cargo make test
 
-# Test with specific features
-cargo make test_arkzkey    # For arkzkey feature
-cargo make test_stateless  # For stateless feature
+# Test with stateless features
+cargo make test_stateless
 ```
 
 ## Advanced: Custom Circuit Compilation
 
-The `rln` (<https://github.com/rate-limiting-nullifier/circom-rln>) repository, which contains the RLN circuit implementation is using for pre-compiled RLN circuit for zerokit RLN.
+The `rln` (<https://github.com/rate-limiting-nullifier/circom-rln>) repository,
+which contains the RLN circuit implementation is using for pre-compiled RLN circuit for zerokit RLN.
 If you want to compile your own RLN circuit, you can follow the instructions below.
 
-### 1. Compile ZK Circuits for getting the zkey and verification key files
+### 1. Compile ZK Circuits for getting the zkey file
 
-This script actually generates not only the zkey and verification key files for the RLN circuit, but also the execution wasm file used for witness calculation.
-However, the wasm file is not needed for the `rln` module, because current implementation uses the iden3 graph file for witness calculation.
+This script actually generates not only the zkey file for the RLN circuit,
+but also the execution wasm file used for witness calculation.
+However, the wasm file is not needed for the `rln` module,
+because current implementation uses the iden3 graph file for witness calculation.
 This graph file is generated by the `circom-witnesscalc` tool in [step 2](#2-generate-witness-calculation-graph).
 
 To customize the circuit parameters, modify `circom-rln/circuits/rln.circom`:
@@ -169,19 +189,27 @@ component main { public [x, externalNullifier] } = RLN(N, M);
 
 Where:
 
-- `N`: Merkle tree height, determining the maximum membership capacity (2^N members).
+- `N`: Merkle tree depth, determining the maximum membership capacity (2^N members).
 
 - `M`: Bit size for range checks, setting an upper bound for the number of messages per epoch (2^M messages).
 
 > [!NOTE]
-> However, if `N` is too big, this might require a larger Powers of Tau ceremony than the one hardcoded in `./scripts/build-circuits.sh`, which is `2^14`. \
-> In such case, we refer to the official [Circom documentation](https://docs.circom.io/getting-started/proving-circuits/#powers-of-tau) for instructions on how to run an appropriate Powers of Tau ceremony and Phase 2 in order to compile the desired circuit. \
-> Additionally, while `M` sets an upper bound on the number of messages per epoch (`2^M`), you can configure lower message limit for your use case, as long as it satisfies `user_message_limit ≤ 2^M`. \
-> Currently, the `rln` module comes with a [pre-compiled](https://github.com/vacp2p/zerokit/tree/master/rln/resources) RLN circuit with a Merkle tree of height `20` and a bit size of `16`, allowing up to `2^20` registered members and a `2^16` message limit per epoch.
+> However, if `N` is too big, this might require a larger Powers of Tau ceremony
+> than the one hardcoded in `./scripts/build-circuits.sh`, which is `2^14`.
+> In such case, we refer to the official
+> [Circom documentation](https://docs.circom.io/getting-started/proving-circuits/#powers-of-tau)
+> for instructions on how to run an appropriate Powers of Tau ceremony and Phase 2 in order to compile the desired circuit. \
+> Additionally, while `M` sets an upper bound on the number of messages per epoch (`2^M`),
+> you can configure lower message limit for your use case, as long as it satisfies `user_message_limit ≤ 2^M`. \
+> Currently, the `rln` module comes with a [pre-compiled](https://github.com/vacp2p/zerokit/tree/master/rln/resources)
+> RLN circuit with a Merkle tree of depth `20` and a bit size of `16`,
+> allowing up to `2^20` registered members and a `2^16` message limit per epoch.
 
 #### Install circom compiler
 
-You can follow the instructions below or refer to the [installing Circom](https://docs.circom.io/getting-started/installation/#installing-circom) guide for more details, but make sure to use the specific version `v2.1.0`.
+You can follow the instructions below or refer to the
+[installing Circom](https://docs.circom.io/getting-started/installation/#installing-circom) guide for more details,
+but make sure to use the specific version `v2.1.0`.
 
 ```sh
 # Clone the circom repository
@@ -218,7 +246,8 @@ cp zkeyFiles/rln/final.zkey <path_to_rln_final.zkey>
 
 ### 2. Generate Witness Calculation Graph
 
-The execution graph file used for witness calculation can be compiled following instructions in the [circom-witnesscalc](https://github.com/iden3/circom-witnesscalc) repository.
+The execution graph file used for witness calculation can be compiled following instructions
+in the [circom-witnesscalc](https://github.com/iden3/circom-witnesscalc) repository.
 As mentioned in step 1, we should use `rln.circom` file from `circom-rln` repository.
 
 ```sh
@@ -235,11 +264,14 @@ cargo build
 cargo run --package circom_witnesscalc --bin build-circuit ../circom-rln/circuits/rln.circom <path_to_graph.bin>
 ```
 
-The `rln` module comes with [pre-compiled](https://github.com/vacp2p/zerokit/tree/master/rln/resources) execution graph files for the RLN circuit.
+The `rln` module comes with [pre-compiled](https://github.com/vacp2p/zerokit/tree/master/rln/resources)
+execution graph files for the RLN circuit.
 
-### 3. Generate Arkzkey Representation for zkey and verification key files
+### 3. Generate Arkzkey Representation for zkey file
 
-For faster loading, compile the zkey file into the arkzkey format using [ark-zkey](https://github.com/seemenkina/ark-zkey). This is fork of the [original](https://github.com/zkmopro/ark-zkey) repository with the uncompressed zkey support.
+For faster loading, compile the zkey file into the arkzkey format using
+[ark-zkey](https://github.com/seemenkina/ark-zkey).
+This is fork of the [original](https://github.com/zkmopro/ark-zkey) repository with the uncompressed arkzkey support.
 
 ```sh
 # Clone the ark-zkey repository
@@ -252,7 +284,21 @@ cd ark-zkey && cargo build
 cargo run --bin arkzkey-util <path_to_rln_final.zkey>
 ```
 
-Currently, the `rln` module comes with [pre-compiled](https://github.com/vacp2p/zerokit/tree/master/rln/resources) arkzkey keys for the RLN circuit.
+This will generate the `rln_final.arkzkey` file, which is used by the `rln` module.
+
+Currently, the `rln` module comes with
+[pre-compiled](https://github.com/vacp2p/zerokit/tree/master/rln/resources) arkzkey keys for the RLN circuit.
+
+> [!NOTE]
+> You can use this [convert_zkey.sh](./convert_zkey.sh) script
+> to automate the process of generating the arkzkey file from any zkey file
+
+Run the script as follows:
+
+```sh
+chmod +x ./convert_zkey.sh
+./convert_zkey.sh <path_to_rln_final.zkey>
+```
 
 ## Get involved
 

rln/benches/circuit_loading_arkzkey_benchmark.rs

@@ -1,25 +0,0 @@
-use criterion::{criterion_group, criterion_main, Criterion};
-use rln::circuit::{read_arkzkey_from_bytes_uncompressed, ARKZKEY_BYTES};
-
-pub fn uncompressed_bench(c: &mut Criterion) {
-    let arkzkey = ARKZKEY_BYTES.to_vec();
-    let size = arkzkey.len() as f32;
-    println!(
-        "Size of uncompressed arkzkey: {:.2?} MB",
-        size / 1024.0 / 1024.0
-    );
-
-    c.bench_function("arkzkey::arkzkey_from_raw_uncompressed", |b| {
-        b.iter(|| {
-            let r = read_arkzkey_from_bytes_uncompressed(&arkzkey);
-            assert_eq!(r.is_ok(), true);
-        })
-    });
-}
-
-criterion_group! {
-    name = benches;
-    config = Criterion::default().sample_size(10);
-    targets = uncompressed_bench
-}
-criterion_main!(benches);

rln/benches/circuit_loading_benchmark.rs

@@ -1,24 +0,0 @@
-use criterion::{criterion_group, criterion_main, Criterion};
-use rln::circuit::zkey::read_zkey;
-use std::io::Cursor;
-
-pub fn zkey_load_benchmark(c: &mut Criterion) {
-    let zkey = rln::circuit::ZKEY_BYTES.to_vec();
-    let size = zkey.len() as f32;
-    println!("Size of zkey: {:.2?} MB", size / 1024.0 / 1024.0);
-
-    c.bench_function("zkey::zkey_from_raw", |b| {
-        b.iter(|| {
-            let mut reader = Cursor::new(zkey.clone());
-            let r = read_zkey(&mut reader);
-            assert_eq!(r.is_ok(), true);
-        })
-    });
-}
-
-criterion_group! {
-    name = benches;
-    config = Criterion::default().sample_size(10);
-    targets = zkey_load_benchmark
-}
-criterion_main!(benches);

rln/benches/pmtree_benchmark.rs

@@ -5,7 +5,7 @@ use utils::ZerokitMerkleTree;
 pub fn pmtree_benchmark(c: &mut Criterion) {
     let mut tree = PmTree::default(2).unwrap();
 
-    let leaves: Vec<Fr> = (0..4).map(|s| Fr::from(s)).collect();
+    let leaves: Vec<Fr> = (0..4).map(Fr::from).collect();
 
     c.bench_function("Pmtree::set", |b| {
         b.iter(|| {
@@ -26,12 +26,6 @@ pub fn pmtree_benchmark(c: &mut Criterion) {
         })
     });
 
-    c.bench_function("Pmtree::compute_root", |b| {
-        b.iter(|| {
-            tree.compute_root().unwrap();
-        })
-    });
-
     c.bench_function("Pmtree::get", |b| {
         b.iter(|| {
             tree.get(0).unwrap();

rln/benches/poseidon_tree_benchmark.rs

@@ -6,7 +6,7 @@ use rln::{
 use utils::{FullMerkleTree, OptimalMerkleTree, ZerokitMerkleTree};
 
 pub fn get_leaves(n: u32) -> Vec<Fr> {
-    (0..n).map(|s| Fr::from(s)).collect()
+    (0..n).map(Fr::from).collect()
 }
 
 pub fn optimal_merkle_tree_poseidon_benchmark(c: &mut Criterion) {

rln/convert_zkey.sh

@@ -0,0 +1,53 @@
+#!/bin/bash
+
+# Convert zkey to arkzkey using /tmp directory
+# Usage: ./convert.sh <path_to_zkey_file>
+
+set -e
+
+# Check input
+if [ $# -eq 0 ]; then
+    echo "Usage: $0 <path_to_zkey_file>"
+    exit 1
+fi
+
+ZKEY_FILE="$1"
+
+if [ ! -f "$ZKEY_FILE" ]; then
+    echo "Error: File '$ZKEY_FILE' does not exist"
+    exit 1
+fi
+
+# Get absolute path before changing directories
+ZKEY_ABSOLUTE_PATH=$(realpath "$ZKEY_FILE")
+
+# Create temp directory in /tmp
+TEMP_DIR="/tmp/ark-zkey-$$"
+echo "Using temp directory: $TEMP_DIR"
+
+# Cleanup function
+cleanup() {
+    echo "Cleaning up temp directory: $TEMP_DIR"
+    rm -rf "$TEMP_DIR"
+}
+
+# Setup cleanup trap
+trap cleanup EXIT
+
+# Create temp directory and clone ark-zkey
+mkdir -p "$TEMP_DIR"
+cd "$TEMP_DIR"
+git clone https://github.com/seemenkina/ark-zkey.git
+cd ark-zkey
+cargo build
+
+# Convert
+cargo run --bin arkzkey-util "$ZKEY_ABSOLUTE_PATH"
+
+# Check if arkzkey file was created (tool creates it in same directory as input)
+ARKZKEY_FILE="${ZKEY_ABSOLUTE_PATH%.zkey}.arkzkey"
+
+if [ ! -f "$ARKZKEY_FILE" ]; then
+    echo "Could not find generated .arkzkey file at $ARKZKEY_FILE"
+    exit 1
+fi

rln/src/circuit/error.rs

@@ -0,0 +1,7 @@
+#[derive(Debug, thiserror::Error)]
+pub enum ZKeyReadError {
+    #[error("No proving key found!")]
+    EmptyBytes,
+    #[error("{0}")]
+    SerializationError(#[from] ark_serialize::SerializationError),
+}

rln/src/circuit/iden3calc.rs

@@ -8,19 +8,33 @@ pub mod storage;
 use ruint::aliases::U256;
 use std::collections::HashMap;
 use storage::deserialize_witnesscalc_graph;
+use zeroize::zeroize_flat_type;
 
+use crate::circuit::iden3calc::graph::fr_to_u256;
 use crate::circuit::Fr;
-use graph::{fr_to_u256, Node};
+use crate::utils::FrOrSecret;
+use graph::Node;
 
 pub type InputSignalsInfo = HashMap<String, (usize, usize)>;
 
-pub fn calc_witness<I: IntoIterator<Item = (String, Vec<Fr>)>>(
+pub fn calc_witness<I: IntoIterator<Item = (String, Vec<FrOrSecret>)>>(
     inputs: I,
     graph_data: &[u8],
 ) -> Vec<Fr> {
-    let inputs: HashMap<String, Vec<U256>> = inputs
+    let mut inputs: HashMap<String, Vec<U256>> = inputs
         .into_iter()
-        .map(|(key, value)| (key, value.iter().map(fr_to_u256).collect()))
+        .map(|(key, value)| {
+            (
+                key,
+                value
+                    .iter()
+                    .map(|f_| match f_ {
+                        FrOrSecret::IdSecret(s) => s.to_u256(),
+                        FrOrSecret::Fr(f) => fr_to_u256(f),
+                    })
+                    .collect(),
+            )
+        })
         .collect();
 
     let (nodes, signals, input_mapping): (Vec<Node>, Vec<usize>, InputSignalsInfo) =
@@ -28,8 +42,15 @@ pub fn calc_witness<I: IntoIterator<Item = (String, Vec<Fr>)>>(
 
     let mut inputs_buffer = get_inputs_buffer(get_inputs_size(&nodes));
     populate_inputs(&inputs, &input_mapping, &mut inputs_buffer);
-
-    graph::evaluate(&nodes, inputs_buffer.as_slice(), &signals)
+    if let Some(v) = inputs.get_mut("identitySecret") {
+        // ~== v[0] = U256::ZERO;
+        unsafe { zeroize_flat_type(v) };
+    }
+    let res = graph::evaluate(&nodes, inputs_buffer.as_slice(), &signals);
+    inputs_buffer.iter_mut().for_each(|i| {
+        unsafe { zeroize_flat_type(i) };
+    });
+    res
 }
 
 fn get_inputs_size(nodes: &[Node]) -> usize {
@@ -56,7 +77,7 @@ fn populate_inputs(
     for (key, value) in input_list {
         let (offset, len) = inputs_info[key];
         if len != value.len() {
-            panic!("Invalid input length for {}", key);
+            panic!("Invalid input length for {key}");
         }
 
         for (i, v) in value.iter().enumerate() {

rln/src/circuit/iden3calc/graph.rs

@@ -10,7 +10,7 @@ use std::{
     cmp::Ordering,
     collections::HashMap,
     error::Error,
-    ops::{BitAnd, BitOr, BitXor, Deref, Shl, Shr},
+    ops::{Deref, Shl, Shr},
 };
 
 use crate::circuit::iden3calc::proto;
@@ -944,14 +944,14 @@ mod tests {
         let x = M.div(uint!(2_U256));
 
         println!("x: {:?}", x.as_limbs());
-        println!("x: {}", M);
+        println!("x: {M}");
     }
 
     #[test]
     fn test_2() {
         let nodes: Vec<Node> = vec![];
         // let node = nodes[0];
-        let node = nodes.get(0);
-        println!("{:?}", node);
+        let node = nodes.first();
+        println!("{node:?}");
     }
 }

rln/src/circuit/iden3calc/storage.rs

@@ -419,13 +419,13 @@ mod tests {
         let mut r = WriteBackReader::new(std::io::Cursor::new(&data));
 
         let buf = &mut [0u8; 5];
-        r.read(buf).unwrap();
+        r.read_exact(buf).unwrap();
         assert_eq!(buf, &[1, 2, 3, 4, 5]);
 
         // return [4, 5] to reader
-        r.write(&buf[3..]).unwrap();
+        r.write_all(&buf[3..]).unwrap();
         // return [2, 3] to reader
-        r.write(&buf[1..3]).unwrap();
+        r.write_all(&buf[1..3]).unwrap();
 
         buf.fill(0);
 

rln/src/circuit/mod.rs

@@ -1,8 +1,8 @@
 // This crate provides interfaces for the zero-knowledge circuit and keys
 
+pub mod error;
 pub mod iden3calc;
 pub mod qap;
-pub mod zkey;
 
 use ::lazy_static::lazy_static;
 use ark_bn254::{
@@ -11,39 +11,22 @@ use ark_bn254::{
 };
 use ark_groth16::ProvingKey;
 use ark_relations::r1cs::ConstraintMatrices;
-use cfg_if::cfg_if;
-use color_eyre::{Report, Result};
 
+use crate::circuit::error::ZKeyReadError;
 use crate::circuit::iden3calc::calc_witness;
 
-#[cfg(feature = "arkzkey")]
-use {
-    ark_ff::Field, ark_serialize::CanonicalDeserialize, ark_serialize::CanonicalSerialize,
-    color_eyre::eyre::WrapErr,
-};
+use {ark_ff::Field, ark_serialize::CanonicalDeserialize, ark_serialize::CanonicalSerialize};
 
-#[cfg(not(feature = "arkzkey"))]
-use {crate::circuit::zkey::read_zkey, std::io::Cursor};
+use crate::utils::FrOrSecret;
 
-#[cfg(feature = "arkzkey")]
 pub const ARKZKEY_BYTES: &[u8] = include_bytes!("../../resources/tree_height_20/rln_final.arkzkey");
 
-pub const ZKEY_BYTES: &[u8] = include_bytes!("../../resources/tree_height_20/rln_final.zkey");
-
 #[cfg(not(target_arch = "wasm32"))]
 const GRAPH_BYTES: &[u8] = include_bytes!("../../resources/tree_height_20/graph.bin");
 
 lazy_static! {
-    static ref ZKEY: (ProvingKey<Curve>, ConstraintMatrices<Fr>) = {
-        cfg_if! {
-                if #[cfg(feature = "arkzkey")] {
-                    read_arkzkey_from_bytes_uncompressed(ARKZKEY_BYTES).expect("Failed to read arkzkey")
-                } else {
-                    let mut reader = Cursor::new(ZKEY_BYTES);
-                    read_zkey(&mut reader).expect("Failed to read zkey")
-                }
-        }
-    };
+    static ref ARKZKEY: (ProvingKey<Curve>, ConstraintMatrices<Fr>) =
+        read_arkzkey_from_bytes_uncompressed(ARKZKEY_BYTES).expect("Failed to read arkzkey");
 }
 
 pub const TEST_TREE_HEIGHT: usize = 20;
@@ -60,20 +43,14 @@ pub type G2Affine = ArkG2Affine;
 pub type G2Projective = ArkG2Projective;
 
 // Loads the proving key using a bytes vector
-pub fn zkey_from_raw(zkey_data: &[u8]) -> Result<(ProvingKey<Curve>, ConstraintMatrices<Fr>)> {
+pub fn zkey_from_raw(
+    zkey_data: &[u8],
+) -> Result<(ProvingKey<Curve>, ConstraintMatrices<Fr>), ZKeyReadError> {
     if zkey_data.is_empty() {
-        return Err(Report::msg("No proving key found!"));
+        return Err(ZKeyReadError::EmptyBytes);
     }
 
-    let proving_key_and_matrices = match () {
-        #[cfg(feature = "arkzkey")]
-        () => read_arkzkey_from_bytes_uncompressed(zkey_data)?,
-        #[cfg(not(feature = "arkzkey"))]
-        () => {
-            let mut reader = Cursor::new(zkey_data);
-            read_zkey(&mut reader)?
-        }
-    };
+    let proving_key_and_matrices = read_arkzkey_from_bytes_uncompressed(zkey_data)?;
 
     Ok(proving_key_and_matrices)
 }
@@ -81,10 +58,10 @@ pub fn zkey_from_raw(zkey_data: &[u8]) -> Result<(ProvingKey<Curve>, ConstraintM
 // Loads the proving key
 #[cfg(not(target_arch = "wasm32"))]
 pub fn zkey_from_folder() -> &'static (ProvingKey<Curve>, ConstraintMatrices<Fr>) {
-    &ZKEY
+    &ARKZKEY
 }
 
-pub fn calculate_rln_witness<I: IntoIterator<Item = (String, Vec<Fr>)>>(
+pub fn calculate_rln_witness<I: IntoIterator<Item = (String, Vec<FrOrSecret>)>>(
     inputs: I,
     graph_data: &[u8],
 ) -> Vec<Fr> {
@@ -101,11 +78,9 @@ pub fn graph_from_folder() -> &'static [u8] {
 // without print and allow to choose between compressed and uncompressed arkzkey
 ////////////////////////////////////////////////////////
 
-#[cfg(feature = "arkzkey")]
 #[derive(CanonicalSerialize, CanonicalDeserialize, Clone, Debug, PartialEq)]
 pub struct SerializableProvingKey(pub ProvingKey<Bn254>);
 
-#[cfg(feature = "arkzkey")]
 #[derive(CanonicalSerialize, CanonicalDeserialize, Clone, Debug, PartialEq)]
 pub struct SerializableConstraintMatrices<F: Field> {
     pub num_instance_variables: usize,
@@ -119,29 +94,25 @@ pub struct SerializableConstraintMatrices<F: Field> {
     pub c: SerializableMatrix<F>,
 }
 
-#[cfg(feature = "arkzkey")]
 #[derive(CanonicalSerialize, CanonicalDeserialize, Clone, Debug, PartialEq)]
 pub struct SerializableMatrix<F: Field> {
     pub data: Vec<Vec<(F, usize)>>,
 }
 
-#[cfg(feature = "arkzkey")]
 pub fn read_arkzkey_from_bytes_uncompressed(
     arkzkey_data: &[u8],
-) -> Result<(ProvingKey<Curve>, ConstraintMatrices<Fr>)> {
+) -> Result<(ProvingKey<Curve>, ConstraintMatrices<Fr>), ZKeyReadError> {
     if arkzkey_data.is_empty() {
-        return Err(Report::msg("No proving key found!"));
+        return Err(ZKeyReadError::EmptyBytes);
     }
 
     let mut cursor = std::io::Cursor::new(arkzkey_data);
 
     let serialized_proving_key =
-        SerializableProvingKey::deserialize_uncompressed_unchecked(&mut cursor)
-            .wrap_err("Failed to deserialize proving key")?;
+        SerializableProvingKey::deserialize_uncompressed_unchecked(&mut cursor)?;
 
     let serialized_constraint_matrices =
-        SerializableConstraintMatrices::deserialize_uncompressed_unchecked(&mut cursor)
-            .wrap_err("Failed to deserialize constraint matrices")?;
+        SerializableConstraintMatrices::deserialize_uncompressed_unchecked(&mut cursor)?;
 
     // Get on right form for API
     let proving_key: ProvingKey<Bn254> = serialized_proving_key.0;

rln/src/circuit/qap.rs

@@ -7,6 +7,12 @@ use ark_poly::EvaluationDomain;
 use ark_relations::r1cs::{ConstraintMatrices, ConstraintSystemRef, SynthesisError};
 use ark_std::{cfg_into_iter, cfg_iter, cfg_iter_mut, vec};
 
+#[cfg(feature = "parallel")]
+use rayon::iter::{
+    IndexedParallelIterator, IntoParallelIterator, IntoParallelRefIterator,
+    IntoParallelRefMutIterator, ParallelIterator,
+};
+
 /// Implements the witness map used by snarkjs. The arkworks witness map calculates the
 /// coefficients of H through computing (AB-C)/Z in the evaluation domain and going back to the
 /// coefficients domain. snarkjs instead precomputes the Lagrange form of the powers of tau bases

rln/src/circuit/zkey.rs

@@ -1,371 +0,0 @@
-// This file is based on the code by arkworks. Its preimage can be found here:
-// https://github.com/arkworks-rs/circom-compat/blob/3c95ed98e23a408b4d99a53e483a9bba39685a4e/src/zkey.rs
-
-//! ZKey Parsing
-//!
-//! Each ZKey file is broken into sections:
-//!  Header(1)
-//!       Prover Type 1 Groth
-//!  HeaderGroth(2)
-//!       n8q
-//!       q
-//!       n8r
-//!       r
-//!       NVars
-//!       NPub
-//!       DomainSize  (multiple of 2
-//!       alpha1
-//!       beta1
-//!       delta1
-//!       beta2
-//!       gamma2
-//!       delta2
-//!  IC(3)
-//!  Coefs(4)
-//!  PointsA(5)
-//!  PointsB1(6)
-//!  PointsB2(7)
-//!  PointsC(8)
-//!  PointsH(9)
-//!  Contributions(10)
-use ark_ff::{BigInteger256, PrimeField};
-use ark_relations::r1cs::ConstraintMatrices;
-use ark_serialize::{CanonicalDeserialize, SerializationError};
-use ark_std::log2;
-use byteorder::{LittleEndian, ReadBytesExt};
-
-use std::{
-    collections::HashMap,
-    io::{Read, Seek, SeekFrom},
-};
-
-use ark_bn254::{Bn254, Fq, Fq2, Fr, G1Affine, G2Affine};
-use ark_groth16::{ProvingKey, VerifyingKey};
-use num_traits::Zero;
-
-type IoResult<T> = Result<T, SerializationError>;
-
-#[derive(Clone, Debug)]
-struct Section {
-    position: u64,
-    #[allow(dead_code)]
-    size: usize,
-}
-
-/// Reads a SnarkJS ZKey file into an Arkworks ProvingKey.
-pub fn read_zkey<R: Read + Seek>(
-    reader: &mut R,
-) -> IoResult<(ProvingKey<Bn254>, ConstraintMatrices<Fr>)> {
-    let mut binfile = BinFile::new(reader)?;
-    let proving_key = binfile.proving_key()?;
-    let matrices = binfile.matrices()?;
-    Ok((proving_key, matrices))
-}
-
-#[derive(Debug)]
-struct BinFile<'a, R> {
-    #[allow(dead_code)]
-    ftype: String,
-    #[allow(dead_code)]
-    version: u32,
-    sections: HashMap<u32, Vec<Section>>,
-    reader: &'a mut R,
-}
-
-impl<'a, R: Read + Seek> BinFile<'a, R> {
-    fn new(reader: &'a mut R) -> IoResult<Self> {
-        let mut magic = [0u8; 4];
-        reader.read_exact(&mut magic)?;
-
-        let version = reader.read_u32::<LittleEndian>()?;
-
-        let num_sections = reader.read_u32::<LittleEndian>()?;
-
-        let mut sections = HashMap::new();
-        for _ in 0..num_sections {
-            let section_id = reader.read_u32::<LittleEndian>()?;
-            let section_length = reader.read_u64::<LittleEndian>()?;
-
-            let section = sections.entry(section_id).or_insert_with(Vec::new);
-            section.push(Section {
-                position: reader.stream_position()?,
-                size: section_length as usize,
-            });
-
-            reader.seek(SeekFrom::Current(section_length as i64))?;
-        }
-
-        Ok(Self {
-            ftype: std::str::from_utf8(&magic[..]).unwrap().to_string(),
-            version,
-            sections,
-            reader,
-        })
-    }
-
-    fn proving_key(&mut self) -> IoResult<ProvingKey<Bn254>> {
-        let header = self.groth_header()?;
-        let ic = self.ic(header.n_public)?;
-
-        let a_query = self.a_query(header.n_vars)?;
-        let b_g1_query = self.b_g1_query(header.n_vars)?;
-        let b_g2_query = self.b_g2_query(header.n_vars)?;
-        let l_query = self.l_query(header.n_vars - header.n_public - 1)?;
-        let h_query = self.h_query(header.domain_size as usize)?;
-
-        let vk = VerifyingKey::<Bn254> {
-            alpha_g1: header.verifying_key.alpha_g1,
-            beta_g2: header.verifying_key.beta_g2,
-            gamma_g2: header.verifying_key.gamma_g2,
-            delta_g2: header.verifying_key.delta_g2,
-            gamma_abc_g1: ic,
-        };
-
-        let pk = ProvingKey::<Bn254> {
-            vk,
-            beta_g1: header.verifying_key.beta_g1,
-            delta_g1: header.verifying_key.delta_g1,
-            a_query,
-            b_g1_query,
-            b_g2_query,
-            h_query,
-            l_query,
-        };
-
-        Ok(pk)
-    }
-
-    fn get_section(&self, id: u32) -> Section {
-        self.sections.get(&id).unwrap()[0].clone()
-    }
-
-    fn groth_header(&mut self) -> IoResult<HeaderGroth> {
-        let section = self.get_section(2);
-        let header = HeaderGroth::new(&mut self.reader, &section)?;
-        Ok(header)
-    }
-
-    fn ic(&mut self, n_public: usize) -> IoResult<Vec<G1Affine>> {
-        // the range is non-inclusive so we do +1 to get all inputs
-        self.g1_section(n_public + 1, 3)
-    }
-
-    /// Returns the [`ConstraintMatrices`] corresponding to the zkey
-    pub fn matrices(&mut self) -> IoResult<ConstraintMatrices<Fr>> {
-        let header = self.groth_header()?;
-
-        let section = self.get_section(4);
-        self.reader.seek(SeekFrom::Start(section.position))?;
-        let num_coeffs: u32 = self.reader.read_u32::<LittleEndian>()?;
-
-        // insantiate AB
-        let mut matrices = vec![vec![vec![]; header.domain_size as usize]; 2];
-        let mut max_constraint_index = 0;
-        for _ in 0..num_coeffs {
-            let matrix: u32 = self.reader.read_u32::<LittleEndian>()?;
-            let constraint: u32 = self.reader.read_u32::<LittleEndian>()?;
-            let signal: u32 = self.reader.read_u32::<LittleEndian>()?;
-
-            let value: Fr = deserialize_field_fr(&mut self.reader)?;
-            max_constraint_index = std::cmp::max(max_constraint_index, constraint);
-            matrices[matrix as usize][constraint as usize].push((value, signal as usize));
-        }
-
-        let num_constraints = max_constraint_index as usize - header.n_public;
-        // Remove the public input constraints, Arkworks adds them later
-        matrices.iter_mut().for_each(|m| {
-            m.truncate(num_constraints);
-        });
-        // This is taken from Arkworks' to_matrices() function
-        let a = matrices[0].clone();
-        let b = matrices[1].clone();
-        let a_num_non_zero: usize = a.iter().map(|lc| lc.len()).sum();
-        let b_num_non_zero: usize = b.iter().map(|lc| lc.len()).sum();
-        let matrices = ConstraintMatrices {
-            num_instance_variables: header.n_public + 1,
-            num_witness_variables: header.n_vars - header.n_public,
-            num_constraints,
-
-            a_num_non_zero,
-            b_num_non_zero,
-            c_num_non_zero: 0,
-
-            a,
-            b,
-            c: vec![],
-        };
-
-        Ok(matrices)
-    }
-
-    fn a_query(&mut self, n_vars: usize) -> IoResult<Vec<G1Affine>> {
-        self.g1_section(n_vars, 5)
-    }
-
-    fn b_g1_query(&mut self, n_vars: usize) -> IoResult<Vec<G1Affine>> {
-        self.g1_section(n_vars, 6)
-    }
-
-    fn b_g2_query(&mut self, n_vars: usize) -> IoResult<Vec<G2Affine>> {
-        self.g2_section(n_vars, 7)
-    }
-
-    fn l_query(&mut self, n_vars: usize) -> IoResult<Vec<G1Affine>> {
-        self.g1_section(n_vars, 8)
-    }
-
-    fn h_query(&mut self, n_vars: usize) -> IoResult<Vec<G1Affine>> {
-        self.g1_section(n_vars, 9)
-    }
-
-    fn g1_section(&mut self, num: usize, section_id: usize) -> IoResult<Vec<G1Affine>> {
-        let section = self.get_section(section_id as u32);
-        self.reader.seek(SeekFrom::Start(section.position))?;
-        deserialize_g1_vec(self.reader, num as u32)
-    }
-
-    fn g2_section(&mut self, num: usize, section_id: usize) -> IoResult<Vec<G2Affine>> {
-        let section = self.get_section(section_id as u32);
-        self.reader.seek(SeekFrom::Start(section.position))?;
-        deserialize_g2_vec(self.reader, num as u32)
-    }
-}
-
-#[derive(Default, Clone, Debug, CanonicalDeserialize)]
-pub struct ZVerifyingKey {
-    alpha_g1: G1Affine,
-    beta_g1: G1Affine,
-    beta_g2: G2Affine,
-    gamma_g2: G2Affine,
-    delta_g1: G1Affine,
-    delta_g2: G2Affine,
-}
-
-impl ZVerifyingKey {
-    fn new<R: Read>(reader: &mut R) -> IoResult<Self> {
-        let alpha_g1 = deserialize_g1(reader)?;
-        let beta_g1 = deserialize_g1(reader)?;
-        let beta_g2 = deserialize_g2(reader)?;
-        let gamma_g2 = deserialize_g2(reader)?;
-        let delta_g1 = deserialize_g1(reader)?;
-        let delta_g2 = deserialize_g2(reader)?;
-
-        Ok(Self {
-            alpha_g1,
-            beta_g1,
-            beta_g2,
-            gamma_g2,
-            delta_g1,
-            delta_g2,
-        })
-    }
-}
-
-#[derive(Clone, Debug)]
-struct HeaderGroth {
-    #[allow(dead_code)]
-    n8q: u32,
-    #[allow(dead_code)]
-    q: BigInteger256,
-    #[allow(dead_code)]
-    n8r: u32,
-    #[allow(dead_code)]
-    r: BigInteger256,
-
-    n_vars: usize,
-    n_public: usize,
-
-    domain_size: u32,
-    #[allow(dead_code)]
-    power: u32,
-
-    verifying_key: ZVerifyingKey,
-}
-
-impl HeaderGroth {
-    fn new<R: Read + Seek>(reader: &mut R, section: &Section) -> IoResult<Self> {
-        reader.seek(SeekFrom::Start(section.position))?;
-        Self::read(reader)
-    }
-
-    fn read<R: Read>(mut reader: &mut R) -> IoResult<Self> {
-        // TODO: Impl From<u32> in Arkworks
-        let n8q: u32 = u32::deserialize_uncompressed(&mut reader)?;
-        // group order r of Bn254
-        let q = BigInteger256::deserialize_uncompressed(&mut reader)?;
-
-        let n8r: u32 = u32::deserialize_uncompressed(&mut reader)?;
-        // Prime field modulus
-        let r = BigInteger256::deserialize_uncompressed(&mut reader)?;
-
-        let n_vars = u32::deserialize_uncompressed(&mut reader)? as usize;
-        let n_public = u32::deserialize_uncompressed(&mut reader)? as usize;
-
-        let domain_size: u32 = u32::deserialize_uncompressed(&mut reader)?;
-        let power = log2(domain_size as usize);
-
-        let verifying_key = ZVerifyingKey::new(&mut reader)?;
-
-        Ok(Self {
-            n8q,
-            q,
-            n8r,
-            r,
-            n_vars,
-            n_public,
-            domain_size,
-            power,
-            verifying_key,
-        })
-    }
-}
-
-// need to divide by R, since snarkjs outputs the zkey with coefficients
-// multiplieid by R^2
-fn deserialize_field_fr<R: Read>(reader: &mut R) -> IoResult<Fr> {
-    let bigint = BigInteger256::deserialize_uncompressed(reader)?;
-    Ok(Fr::new_unchecked(Fr::new_unchecked(bigint).into_bigint()))
-}
-
-// skips the multiplication by R because Circom points are already in Montgomery form
-fn deserialize_field<R: Read>(reader: &mut R) -> IoResult<Fq> {
-    let bigint = BigInteger256::deserialize_uncompressed(reader)?;
-    // if you use Fq::new it multiplies by R
-    Ok(Fq::new_unchecked(bigint))
-}
-
-pub fn deserialize_field2<R: Read>(reader: &mut R) -> IoResult<Fq2> {
-    let c0 = deserialize_field(reader)?;
-    let c1 = deserialize_field(reader)?;
-    Ok(Fq2::new(c0, c1))
-}
-
-fn deserialize_g1<R: Read>(reader: &mut R) -> IoResult<G1Affine> {
-    let x = deserialize_field(reader)?;
-    let y = deserialize_field(reader)?;
-    let infinity = x.is_zero() && y.is_zero();
-    if infinity {
-        Ok(G1Affine::identity())
-    } else {
-        Ok(G1Affine::new(x, y))
-    }
-}
-
-fn deserialize_g2<R: Read>(reader: &mut R) -> IoResult<G2Affine> {
-    let f1 = deserialize_field2(reader)?;
-    let f2 = deserialize_field2(reader)?;
-    let infinity = f1.is_zero() && f2.is_zero();
-    if infinity {
-        Ok(G2Affine::identity())
-    } else {
-        Ok(G2Affine::new(f1, f2))
-    }
-}
-
-fn deserialize_g1_vec<R: Read>(reader: &mut R, n_vars: u32) -> IoResult<Vec<G1Affine>> {
-    (0..n_vars).map(|_| deserialize_g1(reader)).collect()
-}
-
-fn deserialize_g2_vec<R: Read>(reader: &mut R, n_vars: u32) -> IoResult<Vec<G2Affine>> {
-    (0..n_vars).map(|_| deserialize_g2(reader)).collect()
-}

rln/src/error.rs

@@ -0,0 +1,77 @@
+use crate::circuit::error::ZKeyReadError;
+use ark_bn254::Fr;
+use ark_relations::r1cs::SynthesisError;
+use ark_serialize::SerializationError;
+use num_bigint::{BigInt, ParseBigIntError};
+use std::array::TryFromSliceError;
+use std::num::TryFromIntError;
+use std::string::FromUtf8Error;
+use thiserror::Error;
+use utils::error::{FromConfigError, ZerokitMerkleTreeError};
+
+#[derive(Debug, thiserror::Error)]
+pub enum ConversionError {
+    #[error("Expected radix 10 or 16")]
+    WrongRadix,
+    #[error("{0}")]
+    ParseBigInt(#[from] ParseBigIntError),
+    #[error("{0}")]
+    ToUsize(#[from] TryFromIntError),
+    #[error("{0}")]
+    FromSlice(#[from] TryFromSliceError),
+}
+
+#[derive(Error, Debug)]
+pub enum ProofError {
+    #[error("{0}")]
+    ProtocolError(#[from] ProtocolError),
+    #[error("Error producing proof: {0}")]
+    SynthesisError(#[from] SynthesisError),
+}
+
+#[derive(Debug, thiserror::Error)]
+pub enum ProtocolError {
+    #[error("{0}")]
+    Conversion(#[from] ConversionError),
+    #[error("Expected to read {0} bytes but read only {1} bytes")]
+    InvalidReadLen(usize, usize),
+    #[error("Cannot convert bigint {0:?} to biguint")]
+    BigUintConversion(BigInt),
+    #[error("{0}")]
+    JsonError(#[from] serde_json::Error),
+    #[error("Message id ({0}) is not within user_message_limit ({1})")]
+    InvalidMessageId(Fr, Fr),
+}
+
+#[derive(Debug, thiserror::Error)]
+pub enum ComputeIdSecretError {
+    /// Usually it means that the same signal is used to recover the user secret hash
+    #[error("Cannot recover secret: division by zero")]
+    DivisionByZero,
+}
+
+#[derive(Debug, thiserror::Error)]
+pub enum RLNError {
+    #[error("I/O error: {0}")]
+    IO(#[from] std::io::Error),
+    #[error("Utf8 error: {0}")]
+    Utf8(#[from] FromUtf8Error),
+    #[error("Serde json error: {0}")]
+    JSON(#[from] serde_json::Error),
+    #[error("Config error: {0}")]
+    Config(#[from] FromConfigError),
+    #[error("Serialization error: {0}")]
+    Serialization(#[from] SerializationError),
+    #[error("Merkle tree error: {0}")]
+    MerkleTree(#[from] ZerokitMerkleTreeError),
+    #[error("ZKey error: {0}")]
+    ZKey(#[from] ZKeyReadError),
+    #[error("Conversion error: {0}")]
+    Conversion(#[from] ConversionError),
+    #[error("Protocol error: {0}")]
+    Protocol(#[from] ProtocolError),
+    #[error("Proof error: {0}")]
+    Proof(#[from] ProofError),
+    #[error("Unable to extract secret")]
+    RecoverSecret(#[from] ComputeIdSecretError),
+}

rln/src/lib.rs

@@ -1,12 +1,35 @@
 pub mod circuit;
+pub mod error;
 #[cfg(not(target_arch = "wasm32"))]
 pub mod ffi;
 pub mod hashers;
 #[cfg(feature = "pmtree-ft")]
 pub mod pm_tree_adapter;
+#[cfg(not(feature = "stateless"))]
 pub mod poseidon_tree;
 pub mod protocol;
 pub mod public;
 #[cfg(test)]
 pub mod public_api_tests;
 pub mod utils;
+
+// Ensure that only one Merkle tree feature is enabled at a time
+#[cfg(any(
+    all(feature = "fullmerkletree", feature = "optimalmerkletree"),
+    all(feature = "fullmerkletree", feature = "pmtree-ft"),
+    all(feature = "optimalmerkletree", feature = "pmtree-ft"),
+))]
+compile_error!(
+    "Only one of `fullmerkletree`, `optimalmerkletree`, or `pmtree-ft` can be enabled at a time."
+);
+
+// Ensure that the `stateless` feature is not enabled with any Merkle tree features
+#[cfg(all(
+    feature = "stateless",
+    any(
+        feature = "fullmerkletree",
+        feature = "optimalmerkletree",
+        feature = "pmtree-ft"
+    )
+))]
+compile_error!("Cannot enable any Merkle tree features with stateless");

rln/src/pm_tree_adapter.rs

@@ -1,17 +1,15 @@
+use serde_json::Value;
 use std::fmt::Debug;
 use std::path::PathBuf;
 use std::str::FromStr;
 
-use color_eyre::{Report, Result};
-use serde_json::Value;
-
-use utils::pmtree::tree::Key;
-use utils::pmtree::{Database, Hasher};
-use utils::*;
-
 use crate::circuit::Fr;
 use crate::hashers::{poseidon_hash, PoseidonHash};
 use crate::utils::{bytes_le_to_fr, fr_to_bytes_le};
+use utils::error::{FromConfigError, ZerokitMerkleTreeError};
+use utils::pmtree::tree::Key;
+use utils::pmtree::{Database, Hasher, PmtreeErrorKind};
+use utils::{pmtree, Config, Mode, SledDB, ZerokitMerkleProof, ZerokitMerkleTree};
 
 const METADATA_KEY: [u8; 8] = *b"metadata";
 
@@ -63,9 +61,9 @@ fn get_tmp() -> bool {
 pub struct PmtreeConfig(Config);
 
 impl FromStr for PmtreeConfig {
-    type Err = Report;
+    type Err = FromConfigError;
 
-    fn from_str(s: &str) -> Result<Self> {
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
         let config: Value = serde_json::from_str(s)?;
 
         let path = config["path"].as_str();
@@ -85,10 +83,7 @@ impl FromStr for PmtreeConfig {
             && temporary.unwrap()
             && path.as_ref().unwrap().exists()
         {
-            return Err(Report::msg(format!(
-                "Path {:?} already exists, cannot use temporary",
-                path.unwrap()
-            )));
+            return Err(FromConfigError::PathExists);
         }
 
         let config = Config::new()
@@ -133,12 +128,16 @@ impl ZerokitMerkleTree for PmTree {
     type Hasher = PoseidonHash;
     type Config = PmtreeConfig;
 
-    fn default(depth: usize) -> Result<Self> {
+    fn default(depth: usize) -> Result<Self, ZerokitMerkleTreeError> {
         let default_config = PmtreeConfig::default();
         PmTree::new(depth, Self::Hasher::default_leaf(), default_config)
     }
 
-    fn new(depth: usize, _default_leaf: FrOf<Self::Hasher>, config: Self::Config) -> Result<Self> {
+    fn new(
+        depth: usize,
+        _default_leaf: FrOf<Self::Hasher>,
+        config: Self::Config,
+    ) -> Result<Self, ZerokitMerkleTreeError> {
         let tree_loaded = pmtree::MerkleTree::load(config.clone().0);
         let tree = match tree_loaded {
             Ok(tree) => tree,
@@ -168,14 +167,12 @@ impl ZerokitMerkleTree for PmTree {
         self.tree.root()
     }
 
-    fn compute_root(&mut self) -> Result<FrOf<Self::Hasher>> {
-        Ok(self.tree.root())
-    }
-
-    fn set(&mut self, index: usize, leaf: FrOf<Self::Hasher>) -> Result<()> {
-        self.tree
-            .set(index, leaf)
-            .map_err(|e| Report::msg(e.to_string()))?;
+    fn set(
+        &mut self,
+        index: usize,
+        leaf: FrOf<Self::Hasher>,
+    ) -> Result<(), ZerokitMerkleTreeError> {
+        self.tree.set(index, leaf)?;
         self.cached_leaves_indices[index] = 1;
         Ok(())
     }
@@ -184,27 +181,31 @@ impl ZerokitMerkleTree for PmTree {
         &mut self,
         start: usize,
         values: I,
-    ) -> Result<()> {
+    ) -> Result<(), ZerokitMerkleTreeError> {
         let v = values.into_iter().collect::<Vec<_>>();
-        self.tree
-            .set_range(start, v.clone().into_iter())
-            .map_err(|e| Report::msg(e.to_string()))?;
+        self.tree.set_range(start, v.clone().into_iter())?;
         for i in start..v.len() {
             self.cached_leaves_indices[i] = 1
         }
         Ok(())
     }
 
-    fn get(&self, index: usize) -> Result<FrOf<Self::Hasher>> {
-        self.tree.get(index).map_err(|e| Report::msg(e.to_string()))
+    fn get(&self, index: usize) -> Result<FrOf<Self::Hasher>, ZerokitMerkleTreeError> {
+        self.tree
+            .get(index)
+            .map_err(ZerokitMerkleTreeError::PmtreeErrorKind)
     }
 
-    fn get_subtree_root(&self, n: usize, index: usize) -> Result<FrOf<Self::Hasher>> {
+    fn get_subtree_root(
+        &self,
+        n: usize,
+        index: usize,
+    ) -> Result<FrOf<Self::Hasher>, ZerokitMerkleTreeError> {
         if n > self.depth() {
-            return Err(Report::msg("level exceeds depth size"));
+            return Err(ZerokitMerkleTreeError::InvalidLevel);
         }
         if index >= self.capacity() {
-            return Err(Report::msg("index exceeds set size"));
+            return Err(ZerokitMerkleTreeError::InvalidLeaf);
         }
         if n == 0 {
             Ok(self.root())
@@ -235,55 +236,71 @@ impl ZerokitMerkleTree for PmTree {
         start: usize,
         leaves: I,
         indices: J,
-    ) -> Result<()> {
+    ) -> Result<(), ZerokitMerkleTreeError> {
         let leaves = leaves.into_iter().collect::<Vec<_>>();
         let mut indices = indices.into_iter().collect::<Vec<_>>();
         indices.sort();
 
         match (leaves.len(), indices.len()) {
-            (0, 0) => Err(Report::msg("no leaves or indices to be removed")),
+            (0, 0) => Err(ZerokitMerkleTreeError::InvalidLeaf),
             (1, 0) => self.set(start, leaves[0]),
             (0, 1) => self.delete(indices[0]),
             (_, 0) => self.set_range(start, leaves.into_iter()),
-            (0, _) => self.remove_indices(&indices),
-            (_, _) => self.remove_indices_and_set_leaves(start, leaves, &indices),
+            (0, _) => self
+                .remove_indices(&indices)
+                .map_err(ZerokitMerkleTreeError::PmtreeErrorKind),
+            (_, _) => self
+                .remove_indices_and_set_leaves(start, leaves, &indices)
+                .map_err(ZerokitMerkleTreeError::PmtreeErrorKind),
         }
     }
 
-    fn update_next(&mut self, leaf: FrOf<Self::Hasher>) -> Result<()> {
+    fn update_next(&mut self, leaf: FrOf<Self::Hasher>) -> Result<(), ZerokitMerkleTreeError> {
         self.tree
             .update_next(leaf)
-            .map_err(|e| Report::msg(e.to_string()))
+            .map_err(ZerokitMerkleTreeError::PmtreeErrorKind)
     }
 
-    fn delete(&mut self, index: usize) -> Result<()> {
+    /// Delete a leaf in the merkle tree given its index
+    ///
+    /// Deleting a leaf is done by resetting it to its default value. Note that the next_index field
+    /// will not be changed (== previously used index cannot be reused - this to avoid replay
+    /// attacks or unexpected and very hard to tackle issues)
+    fn delete(&mut self, index: usize) -> Result<(), ZerokitMerkleTreeError> {
         self.tree
             .delete(index)
-            .map_err(|e| Report::msg(e.to_string()))?;
+            .map_err(ZerokitMerkleTreeError::PmtreeErrorKind)?;
         self.cached_leaves_indices[index] = 0;
         Ok(())
     }
 
-    fn proof(&self, index: usize) -> Result<Self::Proof> {
+    fn proof(&self, index: usize) -> Result<Self::Proof, ZerokitMerkleTreeError> {
         let proof = self.tree.proof(index)?;
         Ok(PmTreeProof { proof })
     }
 
-    fn verify(&self, leaf: &FrOf<Self::Hasher>, witness: &Self::Proof) -> Result<bool> {
+    fn verify(
+        &self,
+        leaf: &FrOf<Self::Hasher>,
+        witness: &Self::Proof,
+    ) -> Result<bool, ZerokitMerkleTreeError> {
         if self.tree.verify(leaf, &witness.proof) {
             Ok(true)
         } else {
-            Err(Report::msg("verify failed"))
+            Err(ZerokitMerkleTreeError::InvalidWitness)
         }
     }
 
-    fn set_metadata(&mut self, metadata: &[u8]) -> Result<()> {
-        self.tree.db.put(METADATA_KEY, metadata.to_vec())?;
+    fn set_metadata(&mut self, metadata: &[u8]) -> Result<(), ZerokitMerkleTreeError> {
+        self.tree
+            .db
+            .put(METADATA_KEY, metadata.to_vec())
+            .map_err(ZerokitMerkleTreeError::PmtreeErrorKind)?;
         self.metadata = metadata.to_vec();
         Ok(())
     }
 
-    fn metadata(&self) -> Result<Vec<u8>> {
+    fn metadata(&self) -> Result<Vec<u8>, ZerokitMerkleTreeError> {
         if !self.metadata.is_empty() {
             return Ok(self.metadata.clone());
         }
@@ -297,8 +314,11 @@ impl ZerokitMerkleTree for PmTree {
         Ok(data.unwrap())
     }
 
-    fn close_db_connection(&mut self) -> Result<()> {
-        self.tree.db.close().map_err(|e| Report::msg(e.to_string()))
+    fn close_db_connection(&mut self) -> Result<(), ZerokitMerkleTreeError> {
+        self.tree
+            .db
+            .close()
+            .map_err(ZerokitMerkleTreeError::PmtreeErrorKind)
     }
 }
 
@@ -306,15 +326,13 @@ type PmTreeHasher = <PmTree as ZerokitMerkleTree>::Hasher;
 type FrOfPmTreeHasher = FrOf<PmTreeHasher>;
 
 impl PmTree {
-    fn remove_indices(&mut self, indices: &[usize]) -> Result<()> {
+    fn remove_indices(&mut self, indices: &[usize]) -> Result<(), PmtreeErrorKind> {
         let start = indices[0];
         let end = indices.last().unwrap() + 1;
 
         let new_leaves = (start..end).map(|_| PmTreeHasher::default_leaf());
 
-        self.tree
-            .set_range(start, new_leaves)
-            .map_err(|e| Report::msg(e.to_string()))?;
+        self.tree.set_range(start, new_leaves)?;
 
         for i in start..end {
             self.cached_leaves_indices[i] = 0
@@ -327,7 +345,7 @@ impl PmTree {
         start: usize,
         leaves: Vec<FrOfPmTreeHasher>,
         indices: &[usize],
-    ) -> Result<()> {
+    ) -> Result<(), PmtreeErrorKind> {
         let min_index = *indices.first().unwrap();
         let max_index = start + leaves.len();
 
@@ -344,9 +362,7 @@ impl PmTree {
             set_values[start - min_index + i] = leaf;
         }
 
-        self.tree
-            .set_range(start, set_values)
-            .map_err(|e| Report::msg(e.to_string()))?;
+        self.tree.set_range(start, set_values)?;
 
         for i in indices {
             self.cached_leaves_indices[*i] = 0;

rln/src/poseidon_tree.rs

@@ -1,30 +1,32 @@
 // This crate defines the RLN module default Merkle tree implementation and its Hasher
+// Implementation inspired by https://github.com/worldcoin/semaphore-rs/blob/d462a4372f1fd9c27610f2acfe4841fab1d396aa/src/poseidon_tree.rs
 
-// Implementation inspired by https://github.com/worldcoin/semaphore-rs/blob/d462a4372f1fd9c27610f2acfe4841fab1d396aa/src/poseidon_tree.rs (no differences)
+#![cfg(not(feature = "stateless"))]
 
 use cfg_if::cfg_if;
 
-cfg_if! {
-    if #[cfg(feature = "pmtree-ft")] {
-        use crate::pm_tree_adapter::*;
-    } else {
-        use crate::hashers::{PoseidonHash};
-        use utils::merkle_tree::*;
-    }
-}
-
-// The zerokit RLN default Merkle tree implementation is the OptimalMerkleTree.
-// To switch to FullMerkleTree implementation, it is enough to enable the fullmerkletree feature
+// The zerokit RLN default Merkle tree implementation is the PMTree from the vacp2p_pmtree crate
+// To switch to FullMerkleTree or OptimalMerkleTree, enable the corresponding feature in the Cargo.toml file
 
 cfg_if! {
     if #[cfg(feature = "fullmerkletree")] {
+        use utils::{FullMerkleTree, FullMerkleProof};
+        use crate::hashers::PoseidonHash;
+
         pub type PoseidonTree = FullMerkleTree<PoseidonHash>;
         pub type MerkleProof = FullMerkleProof<PoseidonHash>;
-    } else if #[cfg(feature = "pmtree-ft")] {
-        pub type PoseidonTree =  PmTree;
-        pub type MerkleProof = PmTreeProof;
-    } else {
+    } else if #[cfg(feature = "optimalmerkletree")] {
+        use utils::{OptimalMerkleTree, OptimalMerkleProof};
+        use crate::hashers::PoseidonHash;
+
         pub type PoseidonTree = OptimalMerkleTree<PoseidonHash>;
         pub type MerkleProof = OptimalMerkleProof<PoseidonHash>;
+    } else if #[cfg(feature = "pmtree-ft")] {
+        use crate::pm_tree_adapter::{PmTree, PmTreeProof};
+
+        pub type PoseidonTree = PmTree;
+        pub type MerkleProof = PmTreeProof;
+    } else {
+        compile_error!("One of the features `fullmerkletree`, `optimalmerkletree`, or `pmtree-ft` must be enabled.");
     }
 }

rln/src/protocol.rs

@@ -1,27 +1,34 @@
 // This crate collects all the underlying primitives used to implement RLN
 
-use ark_bn254::Fr;
+#[cfg(not(feature = "stateless"))]
+use {
+    crate::error::ConversionError,
+    crate::poseidon_tree::PoseidonTree,
+    utils::{ZerokitMerkleProof, ZerokitMerkleTree},
+};
+
+use crate::circuit::{calculate_rln_witness, qap::CircomReduction, Curve};
+use crate::error::{ComputeIdSecretError, ProofError, ProtocolError};
+use crate::hashers::{hash_to_field, poseidon_hash};
+use crate::public::RLN_IDENTIFIER;
+use crate::utils::{
+    bytes_le_to_fr, bytes_le_to_vec_fr, bytes_le_to_vec_u8, fr_byte_size, fr_to_bytes_le,
+    normalize_usize, to_bigint, vec_fr_to_bytes_le, vec_u8_to_bytes_le, FrOrSecret, IdSecret,
+};
+use ark_bn254::{Fr, FrConfig};
+use ark_ff::{AdditiveGroup, Fp, MontBackend};
 use ark_groth16::{prepare_verifying_key, Groth16, Proof as ArkProof, ProvingKey, VerifyingKey};
-use ark_relations::r1cs::{ConstraintMatrices, SynthesisError};
+use ark_relations::r1cs::ConstraintMatrices;
 use ark_serialize::{CanonicalDeserialize, CanonicalSerialize};
 use ark_std::{rand::thread_rng, UniformRand};
-use color_eyre::{Report, Result};
 use num_bigint::BigInt;
 use rand::{Rng, SeedableRng};
 use rand_chacha::ChaCha20Rng;
 use serde::{Deserialize, Serialize};
 #[cfg(test)]
 use std::time::Instant;
-use thiserror::Error;
 use tiny_keccak::{Hasher as _, Keccak};
-
-use crate::circuit::{calculate_rln_witness, qap::CircomReduction, Curve};
-use crate::hashers::{hash_to_field, poseidon_hash};
-use crate::poseidon_tree::*;
-use crate::public::RLN_IDENTIFIER;
-use crate::utils::*;
-use utils::{ZerokitMerkleProof, ZerokitMerkleTree};
-
+use zeroize::Zeroize;
 ///////////////////////////////////////////////////////
 // RLN Witness data structure and utility functions
 ///////////////////////////////////////////////////////
@@ -29,7 +36,7 @@ use utils::{ZerokitMerkleProof, ZerokitMerkleTree};
 #[derive(Debug, PartialEq, Serialize, Deserialize)]
 pub struct RLNWitnessInput {
     #[serde(serialize_with = "ark_se", deserialize_with = "ark_de")]
-    identity_secret: Fr,
+    identity_secret: IdSecret,
     #[serde(serialize_with = "ark_se", deserialize_with = "ark_de")]
     user_message_limit: Fr,
     #[serde(serialize_with = "ark_se", deserialize_with = "ark_de")]
@@ -98,8 +105,8 @@ pub fn deserialize_identity_tuple(serialized: Vec<u8>) -> (Fr, Fr, Fr, Fr) {
 /// # Errors
 ///
 /// Returns an error if `rln_witness.message_id` is not within `rln_witness.user_message_limit`.
-/// input data is [ identity_secret<32> | user_message_limit<32> | message_id<32> | path_elements[<32>] | identity_path_index<8> | x<32> | external_nullifier<32> ]
-pub fn serialize_witness(rln_witness: &RLNWitnessInput) -> Result<Vec<u8>> {
+/// input data is [ identity_secret<32> | user_message_limit<32> | message_id<32> | path_elements<32> | identity_path_index<8> | x<32> | external_nullifier<32> ]
+pub fn serialize_witness(rln_witness: &RLNWitnessInput) -> Result<Vec<u8>, ProtocolError> {
     // Check if message_id is within user_message_limit
     message_id_range_check(&rln_witness.message_id, &rln_witness.user_message_limit)?;
 
@@ -111,11 +118,11 @@ pub fn serialize_witness(rln_witness: &RLNWitnessInput) -> Result<Vec<u8>> {
         fr_byte_size() * (5 + rln_witness.path_elements.len())
             + rln_witness.identity_path_index.len(),
     );
-    serialized.extend_from_slice(&fr_to_bytes_le(&rln_witness.identity_secret));
+    serialized.extend_from_slice(&rln_witness.identity_secret.to_bytes_le());
     serialized.extend_from_slice(&fr_to_bytes_le(&rln_witness.user_message_limit));
     serialized.extend_from_slice(&fr_to_bytes_le(&rln_witness.message_id));
-    serialized.extend_from_slice(&vec_fr_to_bytes_le(&rln_witness.path_elements)?);
-    serialized.extend_from_slice(&vec_u8_to_bytes_le(&rln_witness.identity_path_index)?);
+    serialized.extend_from_slice(&vec_fr_to_bytes_le(&rln_witness.path_elements));
+    serialized.extend_from_slice(&vec_u8_to_bytes_le(&rln_witness.identity_path_index));
     serialized.extend_from_slice(&fr_to_bytes_le(&rln_witness.x));
     serialized.extend_from_slice(&fr_to_bytes_le(&rln_witness.external_nullifier));
 
@@ -127,10 +134,10 @@ pub fn serialize_witness(rln_witness: &RLNWitnessInput) -> Result<Vec<u8>> {
 /// # Errors
 ///
 /// Returns an error if `message_id` is not within `user_message_limit`.
-pub fn deserialize_witness(serialized: &[u8]) -> Result<(RLNWitnessInput, usize)> {
+pub fn deserialize_witness(serialized: &[u8]) -> Result<(RLNWitnessInput, usize), ProtocolError> {
     let mut all_read: usize = 0;
 
-    let (identity_secret, read) = bytes_le_to_fr(&serialized[all_read..]);
+    let (identity_secret, read) = IdSecret::from_bytes_le(&serialized[all_read..]);
     all_read += read;
 
     let (user_message_limit, read) = bytes_le_to_fr(&serialized[all_read..]);
@@ -154,7 +161,7 @@ pub fn deserialize_witness(serialized: &[u8]) -> Result<(RLNWitnessInput, usize)
     all_read += read;
 
     if serialized.len() != all_read {
-        return Err(Report::msg("serialized length is not equal to all_read"));
+        return Err(ProtocolError::InvalidReadLen(serialized.len(), all_read));
     }
 
     Ok((
@@ -175,18 +182,22 @@ pub fn deserialize_witness(serialized: &[u8]) -> Result<(RLNWitnessInput, usize)
 // https://github.com/kilic/rln/blob/7ac74183f8b69b399e3bc96c1ae8ab61c026dc43/src/public.rs#L148
 // input_data is [ identity_secret<32> | id_index<8> | user_message_limit<32> | message_id<32> | external_nullifier<32> | signal_len<8> | signal<var> ]
 // return value is a rln witness populated according to this information
+#[cfg(not(feature = "stateless"))]
 pub fn proof_inputs_to_rln_witness(
     tree: &mut PoseidonTree,
     serialized: &[u8],
-) -> Result<(RLNWitnessInput, usize)> {
+) -> Result<(RLNWitnessInput, usize), ProtocolError> {
     let mut all_read: usize = 0;
 
-    let (identity_secret, read) = bytes_le_to_fr(&serialized[all_read..]);
+    let (identity_secret, read) = IdSecret::from_bytes_le(&serialized[all_read..]);
     all_read += read;
 
     let id_index = usize::try_from(u64::from_le_bytes(
-        serialized[all_read..all_read + 8].try_into()?,
-    ))?;
+        serialized[all_read..all_read + 8]
+            .try_into()
+            .map_err(ConversionError::FromSlice)?,
+    ))
+    .map_err(ConversionError::ToUsize)?;
     all_read += 8;
 
     let (user_message_limit, read) = bytes_le_to_fr(&serialized[all_read..]);
@@ -199,8 +210,11 @@ pub fn proof_inputs_to_rln_witness(
     all_read += read;
 
     let signal_len = usize::try_from(u64::from_le_bytes(
-        serialized[all_read..all_read + 8].try_into()?,
-    ))?;
+        serialized[all_read..all_read + 8]
+            .try_into()
+            .map_err(ConversionError::FromSlice)?,
+    ))
+    .map_err(ConversionError::ToUsize)?;
     all_read += 8;
 
     let signal: Vec<u8> = serialized[all_read..all_read + signal_len].to_vec();
@@ -231,18 +245,16 @@ pub fn proof_inputs_to_rln_witness(
 ///
 /// Returns an error if `message_id` is not within `user_message_limit`.
 pub fn rln_witness_from_values(
-    identity_secret: Fr,
-    merkle_proof: &MerkleProof,
+    identity_secret: IdSecret,
+    path_elements: Vec<Fp<MontBackend<FrConfig, 4>, 4>>,
+    identity_path_index: Vec<u8>,
     x: Fr,
     external_nullifier: Fr,
     user_message_limit: Fr,
     message_id: Fr,
-) -> Result<RLNWitnessInput> {
+) -> Result<RLNWitnessInput, ProtocolError> {
     message_id_range_check(&message_id, &user_message_limit)?;
 
-    let path_elements = merkle_proof.get_path_elements();
-    let identity_path_index = merkle_proof.get_path_index();
-
     Ok(RLNWitnessInput {
         identity_secret,
         path_elements,
@@ -257,7 +269,7 @@ pub fn rln_witness_from_values(
 pub fn random_rln_witness(tree_height: usize) -> RLNWitnessInput {
     let mut rng = thread_rng();
 
-    let identity_secret = hash_to_field(&rng.gen::<[u8; 32]>());
+    let identity_secret = IdSecret::rand(&mut rng);
     let x = hash_to_field(&rng.gen::<[u8; 32]>());
     let epoch = hash_to_field(&rng.gen::<[u8; 32]>());
     let rln_identifier = hash_to_field(RLN_IDENTIFIER); //hash_to_field(&rng.gen::<[u8; 32]>());
@@ -284,16 +296,24 @@ pub fn random_rln_witness(tree_height: usize) -> RLNWitnessInput {
     }
 }
 
-pub fn proof_values_from_witness(rln_witness: &RLNWitnessInput) -> Result<RLNProofValues> {
+pub fn proof_values_from_witness(
+    rln_witness: &RLNWitnessInput,
+) -> Result<RLNProofValues, ProtocolError> {
     message_id_range_check(&rln_witness.message_id, &rln_witness.user_message_limit)?;
 
     // y share
-    let a_0 = rln_witness.identity_secret;
-    let a_1 = poseidon_hash(&[a_0, rln_witness.external_nullifier, rln_witness.message_id]);
-    let y = a_0 + rln_witness.x * a_1;
+    let a_0 = &rln_witness.identity_secret;
+    let mut to_hash = [
+        *(a_0.clone()),
+        rln_witness.external_nullifier,
+        rln_witness.message_id,
+    ];
+    let a_1 = poseidon_hash(&to_hash);
+    let y = *(a_0.clone()) + rln_witness.x * a_1;
 
     // Nullifier
     let nullifier = poseidon_hash(&[a_1]);
+    to_hash[0].zeroize();
 
     // Merkle tree root computations
     let root = compute_tree_root(
@@ -361,7 +381,7 @@ pub fn deserialize_proof_values(serialized: &[u8]) -> (RLNProofValues, usize) {
 
 // input_data is [ identity_secret<32> | id_index<8> | user_message_limit<32> | message_id<32> | external_nullifier<32> | signal_len<8> | signal<var> ]
 pub fn prepare_prove_input(
-    identity_secret: Fr,
+    identity_secret: IdSecret,
     id_index: usize,
     user_message_limit: Fr,
     message_id: Fr,
@@ -374,7 +394,7 @@ pub fn prepare_prove_input(
     // - variable length signal data
     let mut serialized = Vec::with_capacity(fr_byte_size() * 4 + 16 + signal.len()); // length of 4 fr elements + 16 bytes (id_index + len) + signal length
 
-    serialized.extend_from_slice(&fr_to_bytes_le(&identity_secret));
+    serialized.extend_from_slice(&identity_secret.to_bytes_le());
     serialized.extend_from_slice(&normalize_usize(id_index));
     serialized.extend_from_slice(&fr_to_bytes_le(&user_message_limit));
     serialized.extend_from_slice(&fr_to_bytes_le(&message_id));
@@ -405,12 +425,15 @@ pub fn prepare_verify_input(proof_data: Vec<u8>, signal: &[u8]) -> Vec<u8> {
 ///////////////////////////////////////////////////////
 
 pub fn compute_tree_root(
-    identity_secret: &Fr,
+    identity_secret: &IdSecret,
     user_message_limit: &Fr,
     path_elements: &[Fr],
     identity_path_index: &[u8],
 ) -> Fr {
-    let id_commitment = poseidon_hash(&[*identity_secret]);
+    let mut to_hash = [*identity_secret.clone()];
+    let id_commitment = poseidon_hash(&to_hash);
+    to_hash[0].zeroize();
+
     let mut root = poseidon_hash(&[id_commitment, *user_message_limit]);
 
     for i in 0..identity_path_index.len() {
@@ -431,10 +454,12 @@ pub fn compute_tree_root(
 // Generates a tuple (identity_secret_hash, id_commitment) where
 // identity_secret_hash is random and id_commitment = PoseidonHash(identity_secret_hash)
 // RNG is instantiated using thread_rng()
-pub fn keygen() -> (Fr, Fr) {
+pub fn keygen() -> (IdSecret, Fr) {
     let mut rng = thread_rng();
-    let identity_secret_hash = Fr::rand(&mut rng);
-    let id_commitment = poseidon_hash(&[identity_secret_hash]);
+    let identity_secret_hash = IdSecret::rand(&mut rng);
+    let mut to_hash = [*identity_secret_hash.clone()];
+    let id_commitment = poseidon_hash(&to_hash);
+    to_hash[0].zeroize();
     (identity_secret_hash, id_commitment)
 }
 
@@ -502,7 +527,10 @@ pub fn extended_seeded_keygen(signal: &[u8]) -> (Fr, Fr, Fr, Fr) {
     )
 }
 
-pub fn compute_id_secret(share1: (Fr, Fr), share2: (Fr, Fr)) -> Result<Fr, String> {
+pub fn compute_id_secret(
+    share1: (Fr, Fr),
+    share2: (Fr, Fr),
+) -> Result<IdSecret, ComputeIdSecretError> {
     // Assuming a0 is the identity secret and a1 = poseidonHash([a0, external_nullifier]),
     // a (x,y) share satisfies the following relation
     // y = a_0 + x * a_1
@@ -512,30 +540,26 @@ pub fn compute_id_secret(share1: (Fr, Fr), share2: (Fr, Fr)) -> Result<Fr, Strin
     // If the two input shares were computed for the same external_nullifier and identity secret, we can recover the latter
     // y1 = a_0 + x1 * a_1
     // y2 = a_0 + x2 * a_1
-    let a_1 = (y1 - y2) / (x1 - x2);
-    let a_0 = y1 - x1 * a_1;
 
-    // If shares come from the same polynomial, a0 is correctly recovered and a1 = poseidonHash([a0, external_nullifier])
-    Ok(a_0)
+    if (x1 - x2) != Fr::ZERO {
+        let a_1 = (y1 - y2) / (x1 - x2);
+        let mut a_0 = y1 - x1 * a_1;
+
+        // If shares come from the same polynomial, a0 is correctly recovered and a1 = poseidonHash([a0, external_nullifier])
+        let id_secret = IdSecret::from(&mut a_0);
+        Ok(id_secret)
+    } else {
+        Err(ComputeIdSecretError::DivisionByZero)
+    }
 }
 
 ///////////////////////////////////////////////////////
 // zkSNARK utility functions
 ///////////////////////////////////////////////////////
 
-#[derive(Error, Debug)]
-pub enum ProofError {
-    #[error("Error reading circuit key: {0}")]
-    CircuitKeyError(#[from] Report),
-    #[error("Error producing witness: {0}")]
-    WitnessError(Report),
-    #[error("Error producing proof: {0}")]
-    SynthesisError(#[from] SynthesisError),
-}
-
 fn calculate_witness_element<E: ark_ec::pairing::Pairing>(
     witness: Vec<BigInt>,
-) -> Result<Vec<E::ScalarField>> {
+) -> Result<Vec<E::ScalarField>, ProtocolError> {
     use ark_ff::PrimeField;
     let modulus = <E::ScalarField as PrimeField>::MODULUS;
 
@@ -548,9 +572,9 @@ fn calculate_witness_element<E: ark_ec::pairing::Pairing>(
             modulus.into()
                 - w.abs()
                     .to_biguint()
-                    .ok_or(Report::msg("not a biguint value"))?
+                    .ok_or(ProtocolError::BigUintConversion(w))?
         } else {
-            w.to_biguint().ok_or(Report::msg("not a biguint value"))?
+            w.to_biguint().ok_or(ProtocolError::BigUintConversion(w))?
         };
         witness_vec.push(E::ScalarField::from(w))
     }
@@ -566,8 +590,7 @@ pub fn generate_proof_with_witness(
     #[cfg(test)]
     let now = Instant::now();
 
-    let full_assignment =
-        calculate_witness_element::<Curve>(witness).map_err(ProofError::WitnessError)?;
+    let full_assignment = calculate_witness_element::<Curve>(witness)?;
 
     #[cfg(test)]
     println!("witness generation took: {:.2?}", now.elapsed());
@@ -604,7 +627,7 @@ pub fn generate_proof_with_witness(
 /// Returns an error if `rln_witness.message_id` is not within `rln_witness.user_message_limit`.
 pub fn inputs_for_witness_calculation(
     rln_witness: &RLNWitnessInput,
-) -> Result<[(&str, Vec<Fr>); 7]> {
+) -> Result<[(&str, Vec<FrOrSecret>); 7], ProtocolError> {
     message_id_range_check(&rln_witness.message_id, &rln_witness.user_message_limit)?;
 
     let mut identity_path_index = Vec::with_capacity(rln_witness.identity_path_index.len());
@@ -614,13 +637,33 @@ pub fn inputs_for_witness_calculation(
         .for_each(|v| identity_path_index.push(Fr::from(*v)));
 
     Ok([
-        ("identitySecret", vec![rln_witness.identity_secret]),
-        ("userMessageLimit", vec![rln_witness.user_message_limit]),
-        ("messageId", vec![rln_witness.message_id]),
-        ("pathElements", rln_witness.path_elements.clone()),
-        ("identityPathIndex", identity_path_index),
-        ("x", vec![rln_witness.x]),
-        ("externalNullifier", vec![rln_witness.external_nullifier]),
+        (
+            "identitySecret",
+            vec![rln_witness.identity_secret.clone().into()],
+        ),
+        (
+            "userMessageLimit",
+            vec![rln_witness.user_message_limit.into()],
+        ),
+        ("messageId", vec![rln_witness.message_id.into()]),
+        (
+            "pathElements",
+            rln_witness
+                .path_elements
+                .iter()
+                .cloned()
+                .map(Into::into)
+                .collect(),
+        ),
+        (
+            "identityPathIndex",
+            identity_path_index.into_iter().map(Into::into).collect(),
+        ),
+        ("x", vec![rln_witness.x.into()]),
+        (
+            "externalNullifier",
+            vec![rln_witness.external_nullifier.into()],
+        ),
     ])
 }
 
@@ -732,7 +775,9 @@ where
 /// # Errors
 ///
 /// Returns an error if `rln_witness.message_id` is not within `rln_witness.user_message_limit`.
-pub fn rln_witness_from_json(input_json: serde_json::Value) -> Result<RLNWitnessInput> {
+pub fn rln_witness_from_json(
+    input_json: serde_json::Value,
+) -> Result<RLNWitnessInput, ProtocolError> {
     let rln_witness: RLNWitnessInput = serde_json::from_value(input_json).unwrap();
     message_id_range_check(&rln_witness.message_id, &rln_witness.user_message_limit)?;
 
@@ -744,7 +789,9 @@ pub fn rln_witness_from_json(input_json: serde_json::Value) -> Result<RLNWitness
 /// # Errors
 ///
 /// Returns an error if `message_id` is not within `user_message_limit`.
-pub fn rln_witness_to_json(rln_witness: &RLNWitnessInput) -> Result<serde_json::Value> {
+pub fn rln_witness_to_json(
+    rln_witness: &RLNWitnessInput,
+) -> Result<serde_json::Value, ProtocolError> {
     message_id_range_check(&rln_witness.message_id, &rln_witness.user_message_limit)?;
 
     let rln_witness_json = serde_json::to_value(rln_witness)?;
@@ -757,13 +804,15 @@ pub fn rln_witness_to_json(rln_witness: &RLNWitnessInput) -> Result<serde_json::
 /// # Errors
 ///
 /// Returns an error if `message_id` is not within `user_message_limit`.
-pub fn rln_witness_to_bigint_json(rln_witness: &RLNWitnessInput) -> Result<serde_json::Value> {
+pub fn rln_witness_to_bigint_json(
+    rln_witness: &RLNWitnessInput,
+) -> Result<serde_json::Value, ProtocolError> {
     message_id_range_check(&rln_witness.message_id, &rln_witness.user_message_limit)?;
 
     let mut path_elements = Vec::new();
 
     for v in rln_witness.path_elements.iter() {
-        path_elements.push(to_bigint(v)?.to_str_radix(10));
+        path_elements.push(to_bigint(v).to_str_radix(10));
     }
 
     let mut identity_path_index = Vec::new();
@@ -773,22 +822,26 @@ pub fn rln_witness_to_bigint_json(rln_witness: &RLNWitnessInput) -> Result<serde
         .for_each(|v| identity_path_index.push(BigInt::from(*v).to_str_radix(10)));
 
     let inputs = serde_json::json!({
-        "identitySecret": to_bigint(&rln_witness.identity_secret)?.to_str_radix(10),
-        "userMessageLimit": to_bigint(&rln_witness.user_message_limit)?.to_str_radix(10),
-        "messageId": to_bigint(&rln_witness.message_id)?.to_str_radix(10),
+        "identitySecret": to_bigint(&rln_witness.identity_secret).to_str_radix(10),
+        "userMessageLimit": to_bigint(&rln_witness.user_message_limit).to_str_radix(10),
+        "messageId": to_bigint(&rln_witness.message_id).to_str_radix(10),
         "pathElements": path_elements,
         "identityPathIndex": identity_path_index,
-        "x": to_bigint(&rln_witness.x)?.to_str_radix(10),
-        "externalNullifier":  to_bigint(&rln_witness.external_nullifier)?.to_str_radix(10),
+        "x": to_bigint(&rln_witness.x).to_str_radix(10),
+        "externalNullifier":  to_bigint(&rln_witness.external_nullifier).to_str_radix(10),
     });
 
     Ok(inputs)
 }
 
-pub fn message_id_range_check(message_id: &Fr, user_message_limit: &Fr) -> Result<()> {
+pub fn message_id_range_check(
+    message_id: &Fr,
+    user_message_limit: &Fr,
+) -> Result<(), ProtocolError> {
     if message_id > user_message_limit {
-        return Err(color_eyre::Report::msg(
-            "message_id is not within user_message_limit",
+        return Err(ProtocolError::InvalidMessageId(
+            *message_id,
+            *user_message_limit,
         ));
     }
     Ok(())

rln/src/public.rs

@@ -1,31 +1,43 @@
+use crate::circuit::{zkey_from_raw, Curve, Fr};
+use crate::hashers::{hash_to_field, poseidon_hash as utils_poseidon_hash};
+use crate::protocol::{
+    compute_id_secret, deserialize_proof_values, deserialize_witness, extended_keygen,
+    extended_seeded_keygen, keygen, proof_values_from_witness, rln_witness_to_bigint_json,
+    rln_witness_to_json, seeded_keygen, serialize_proof_values, verify_proof,
+};
+use crate::utils::{bytes_le_to_fr, bytes_le_to_vec_fr, fr_byte_size, fr_to_bytes_le};
+#[cfg(not(target_arch = "wasm32"))]
+use {
+    crate::{
+        circuit::{graph_from_folder, zkey_from_folder},
+        protocol::generate_proof,
+    },
+    std::default::Default,
+};
+
+#[cfg(target_arch = "wasm32")]
+use crate::protocol::generate_proof_with_witness;
+
 /// This is the main public API for RLN module. It is used by the FFI, and should be
 /// used by tests etc. as well
 #[cfg(not(feature = "stateless"))]
 use {
+    crate::protocol::{proof_inputs_to_rln_witness, serialize_witness},
+    crate::utils::{bytes_le_to_vec_u8, vec_fr_to_bytes_le, vec_u8_to_bytes_le},
     crate::{circuit::TEST_TREE_HEIGHT, poseidon_tree::PoseidonTree},
     serde_json::{json, Value},
     std::str::FromStr,
+    utils::error::ZerokitMerkleTreeError,
     utils::{Hasher, ZerokitMerkleProof, ZerokitMerkleTree},
 };
 
-use crate::circuit::{zkey_from_raw, Curve, Fr};
-use crate::hashers::{hash_to_field, poseidon_hash as utils_poseidon_hash};
-use crate::protocol::*;
-use crate::utils::*;
-#[cfg(not(target_arch = "wasm32"))]
-use {
-    crate::circuit::{graph_from_folder, zkey_from_folder},
-    std::default::Default,
-};
-
+use crate::error::{ConversionError, ProtocolError, RLNError};
 use ark_groth16::{Proof as ArkProof, ProvingKey, VerifyingKey};
 use ark_relations::r1cs::ConstraintMatrices;
 use ark_serialize::{CanonicalDeserialize, CanonicalSerialize, Read, Write};
-use color_eyre::{Report, Result};
-use std::io::Cursor;
-
 #[cfg(target_arch = "wasm32")]
 use num_bigint::BigInt;
+use std::io::Cursor;
 
 /// The application-specific RLN identifier.
 ///
@@ -64,7 +76,7 @@ impl RLN {
     /// let mut rln = RLN::new(tree_height, input);
     /// ```
     #[cfg(all(not(target_arch = "wasm32"), not(feature = "stateless")))]
-    pub fn new<R: Read>(tree_height: usize, mut input_data: R) -> Result<RLN> {
+    pub fn new<R: Read>(tree_height: usize, mut input_data: R) -> Result<RLN, RLNError> {
         // We read input
         let mut input: Vec<u8> = Vec::new();
         input_data.read_to_end(&mut input)?;
@@ -106,9 +118,8 @@ impl RLN {
     /// // We create a new RLN instance
     /// let mut rln = RLN::new();
     /// ```
-    #[cfg_attr(docsrs, doc(cfg(feature = "stateless")))]
     #[cfg(all(not(target_arch = "wasm32"), feature = "stateless"))]
-    pub fn new() -> Result<RLN> {
+    pub fn new() -> Result<RLN, RLNError> {
         let proving_key = zkey_from_folder().to_owned();
         let verification_key = proving_key.0.vk.to_owned();
         let graph_data = graph_from_folder().to_owned();
@@ -124,7 +135,7 @@ impl RLN {
     ///
     /// Input parameters are
     /// - `tree_height`: the height of the internal Merkle tree
-    /// - `zkey_vec`: a byte vector containing to the proving key (`rln_final.zkey`)  or (`rln_final.arkzkey`) as binary file
+    /// - `zkey_vec`: a byte vector containing to the proving key (`rln_final.arkzkey`) as binary file
     /// - `graph_data`: a byte vector containing the graph data (`graph.bin`) as binary file
     /// - `tree_config_input`: a reader for a string containing a json with the merkle tree configuration
     ///
@@ -137,7 +148,7 @@ impl RLN {
     /// let resources_folder = "./resources/tree_height_20/";
     ///
     /// let mut resources: Vec<Vec<u8>> = Vec::new();
-    /// for filename in ["rln_final.zkey", "graph.bin"] {
+    /// for filename in ["rln_final.arkzkey", "graph.bin"] {
     ///     let fullpath = format!("{resources_folder}{filename}");
     ///     let mut file = File::open(&fullpath).expect("no file found");
     ///     let metadata = std::fs::metadata(&fullpath).expect("unable to read metadata");
@@ -162,7 +173,7 @@ impl RLN {
         zkey_vec: Vec<u8>,
         graph_data: Vec<u8>,
         mut tree_config_input: R,
-    ) -> Result<RLN> {
+    ) -> Result<RLN, RLNError> {
         let proving_key = zkey_from_raw(&zkey_vec)?;
         let verification_key = proving_key.0.vk.to_owned();
 
@@ -195,7 +206,7 @@ impl RLN {
     /// Creates a new stateless RLN object by passing circuit resources as byte vectors.
     ///
     /// Input parameters are
-    /// - `zkey_vec`: a byte vector containing to the proving key (`rln_final.zkey`)  or (`rln_final.arkzkey`) as binary file
+    /// - `zkey_vec`: a byte vector containing to the proving key (`rln_final.arkzkey`) as binary file
     /// - `graph_data`: a byte vector containing the graph data (`graph.bin`) as binary file
     ///
     /// Example:
@@ -206,7 +217,7 @@ impl RLN {
     /// let resources_folder = "./resources/tree_height_20/";
     ///
     /// let mut resources: Vec<Vec<u8>> = Vec::new();
-    /// for filename in ["rln_final.zkey", "graph.bin"] {
+    /// for filename in ["rln_final.arkzkey", "graph.bin"] {
     ///     let fullpath = format!("{resources_folder}{filename}");
     ///     let mut file = File::open(&fullpath).expect("no file found");
     ///     let metadata = std::fs::metadata(&fullpath).expect("unable to read metadata");
@@ -221,7 +232,7 @@ impl RLN {
     /// );
     /// ```
     #[cfg(all(not(target_arch = "wasm32"), feature = "stateless"))]
-    pub fn new_with_params(zkey_vec: Vec<u8>, graph_data: Vec<u8>) -> Result<RLN> {
+    pub fn new_with_params(zkey_vec: Vec<u8>, graph_data: Vec<u8>) -> Result<RLN, RLNError> {
         let proving_key = zkey_from_raw(&zkey_vec)?;
         let verification_key = proving_key.0.vk.to_owned();
 
@@ -235,14 +246,14 @@ impl RLN {
     /// Creates a new stateless RLN object by passing circuit resources as a byte vector.
     ///
     /// Input parameters are
-    /// - `zkey_vec`: a byte vector containing the proving key (`rln_final.zkey`) or (`rln_final.arkzkey`) as binary file
+    /// - `zkey_vec`: a byte vector containing the proving key (`rln_final.arkzkey`) as binary file
     ///
     /// Example:
     /// ```
     /// use std::fs::File;
     /// use std::io::Read;
     ///
-    /// let zkey_path = "./resources/tree_height_20/rln_final.zkey";
+    /// let zkey_path = "./resources/tree_height_20/rln_final.arkzkey";
     ///
     /// let mut file = File::open(zkey_path).expect("Failed to open file");
     /// let metadata = std::fs::metadata(zkey_path).expect("Failed to read metadata");
@@ -252,7 +263,7 @@ impl RLN {
     /// let mut rln = RLN::new_with_params(zkey_vec)?;
     /// ```
     #[cfg(all(target_arch = "wasm32", feature = "stateless"))]
-    pub fn new_with_params(zkey_vec: Vec<u8>) -> Result<RLN> {
+    pub fn new_with_params(zkey_vec: Vec<u8>) -> Result<RLN, RLNError> {
         let proving_key = zkey_from_raw(&zkey_vec)?;
         let verification_key = proving_key.0.vk.to_owned();
 
@@ -272,7 +283,7 @@ impl RLN {
     /// Input values are:
     /// - `tree_height`: the height of the Merkle tree.
     #[cfg(not(feature = "stateless"))]
-    pub fn set_tree(&mut self, tree_height: usize) -> Result<()> {
+    pub fn set_tree(&mut self, tree_height: usize) -> Result<(), RLNError> {
         // We compute a default empty tree of desired height
         self.tree = PoseidonTree::default(tree_height)?;
 
@@ -303,7 +314,7 @@ impl RLN {
     /// rln.set_leaf(id_index, &mut buffer).unwrap();
     /// ```
     #[cfg(not(feature = "stateless"))]
-    pub fn set_leaf<R: Read>(&mut self, index: usize, mut input_data: R) -> Result<()> {
+    pub fn set_leaf<R: Read>(&mut self, index: usize, mut input_data: R) -> Result<(), RLNError> {
         // We read input
         let mut leaf_byte: Vec<u8> = Vec::new();
         input_data.read_to_end(&mut leaf_byte)?;
@@ -333,7 +344,7 @@ impl RLN {
     /// rln.get_leaf(id_index, &mut buffer).unwrap();
     /// let rate_commitment = deserialize_field_element(&buffer.into_inner()).unwrap();
     #[cfg(not(feature = "stateless"))]
-    pub fn get_leaf<W: Write>(&self, index: usize, mut output_data: W) -> Result<()> {
+    pub fn get_leaf<W: Write>(&self, index: usize, mut output_data: W) -> Result<(), RLNError> {
         // We get the leaf at input index
         let leaf = self.tree.get(index)?;
 
@@ -376,7 +387,11 @@ impl RLN {
     /// rln.set_leaves_from(index, &mut buffer).unwrap();
     /// ```
     #[cfg(not(feature = "stateless"))]
-    pub fn set_leaves_from<R: Read>(&mut self, index: usize, mut input_data: R) -> Result<()> {
+    pub fn set_leaves_from<R: Read>(
+        &mut self,
+        index: usize,
+        mut input_data: R,
+    ) -> Result<(), RLNError> {
         // We read input
         let mut leaves_byte: Vec<u8> = Vec::new();
         input_data.read_to_end(&mut leaves_byte)?;
@@ -385,8 +400,7 @@ impl RLN {
 
         // We set the leaves
         self.tree
-            .override_range(index, leaves.into_iter(), [].into_iter())
-            .map_err(|_| Report::msg("Could not set leaves"))?;
+            .override_range(index, leaves.into_iter(), [].into_iter())?;
         Ok(())
     }
 
@@ -397,7 +411,7 @@ impl RLN {
     /// Input values are:
     /// - `input_data`: a reader for the serialization of multiple leaf values (serialization done with [`rln::utils::vec_fr_to_bytes_le`](crate::utils::vec_fr_to_bytes_le))
     #[cfg(not(feature = "stateless"))]
-    pub fn init_tree_with_leaves<R: Read>(&mut self, input_data: R) -> Result<()> {
+    pub fn init_tree_with_leaves<R: Read>(&mut self, input_data: R) -> Result<(), RLNError> {
         // reset the tree
         // NOTE: this requires the tree to be initialized with the correct height initially
         // TODO: accept tree_height as a parameter and initialize the tree with that height
@@ -453,7 +467,7 @@ impl RLN {
         index: usize,
         mut input_leaves: R,
         mut input_indices: R,
-    ) -> Result<()> {
+    ) -> Result<(), RLNError> {
         // We read input
         let mut leaves_byte: Vec<u8> = Vec::new();
         input_leaves.read_to_end(&mut leaves_byte)?;
@@ -468,8 +482,7 @@ impl RLN {
 
         // We set the leaves
         self.tree
-            .override_range(index, leaves.into_iter(), indices.into_iter())
-            .map_err(|e| Report::msg(format!("Could not perform the batch operation: {e}")))?;
+            .override_range(index, leaves.into_iter(), indices.into_iter())?;
         Ok(())
     }
 
@@ -522,7 +535,7 @@ impl RLN {
     /// rln.set_next_leaf(&mut buffer).unwrap();
     /// ```
     #[cfg(not(feature = "stateless"))]
-    pub fn set_next_leaf<R: Read>(&mut self, mut input_data: R) -> Result<()> {
+    pub fn set_next_leaf<R: Read>(&mut self, mut input_data: R) -> Result<(), RLNError> {
         // We read input
         let mut leaf_byte: Vec<u8> = Vec::new();
         input_data.read_to_end(&mut leaf_byte)?;
@@ -548,7 +561,7 @@ impl RLN {
     /// rln.delete_leaf(index).unwrap();
     /// ```
     #[cfg(not(feature = "stateless"))]
-    pub fn delete_leaf(&mut self, index: usize) -> Result<()> {
+    pub fn delete_leaf(&mut self, index: usize) -> Result<(), RLNError> {
         self.tree.delete(index)?;
         Ok(())
     }
@@ -566,7 +579,7 @@ impl RLN {
     /// rln.set_metadata(metadata).unwrap();
     /// ```
     #[cfg(not(feature = "stateless"))]
-    pub fn set_metadata(&mut self, metadata: &[u8]) -> Result<()> {
+    pub fn set_metadata(&mut self, metadata: &[u8]) -> Result<(), RLNError> {
         self.tree.set_metadata(metadata)?;
         Ok(())
     }
@@ -586,7 +599,7 @@ impl RLN {
     /// let metadata = buffer.into_inner();
     /// ```
     #[cfg(not(feature = "stateless"))]
-    pub fn get_metadata<W: Write>(&self, mut output_data: W) -> Result<()> {
+    pub fn get_metadata<W: Write>(&self, mut output_data: W) -> Result<(), RLNError> {
         let metadata = self.tree.metadata()?;
         output_data.write_all(&metadata)?;
         Ok(())
@@ -606,10 +619,9 @@ impl RLN {
     /// let (root, _) = bytes_le_to_fr(&buffer.into_inner());
     /// ```
     #[cfg(not(feature = "stateless"))]
-    pub fn get_root<W: Write>(&self, mut output_data: W) -> Result<()> {
+    pub fn get_root<W: Write>(&self, mut output_data: W) -> Result<(), RLNError> {
         let root = self.tree.root();
         output_data.write_all(&fr_to_bytes_le(&root))?;
-
         Ok(())
     }
 
@@ -634,7 +646,7 @@ impl RLN {
         level: usize,
         index: usize,
         mut output_data: W,
-    ) -> Result<()> {
+    ) -> Result<(), RLNError> {
         let subroot = self.tree.get_subtree_root(level, index)?;
         output_data.write_all(&fr_to_bytes_le(&subroot))?;
 
@@ -663,13 +675,14 @@ impl RLN {
     /// let (identity_path_index, _) = bytes_le_to_vec_u8(&buffer_inner[read..].to_vec());
     /// ```
     #[cfg(not(feature = "stateless"))]
-    pub fn get_proof<W: Write>(&self, index: usize, mut output_data: W) -> Result<()> {
+    pub fn get_proof<W: Write>(&self, index: usize, mut output_data: W) -> Result<(), RLNError> {
         let merkle_proof = self.tree.proof(index).expect("proof should exist");
         let path_elements = merkle_proof.get_path_elements();
         let identity_path_index = merkle_proof.get_path_index();
 
-        output_data.write_all(&vec_fr_to_bytes_le(&path_elements)?)?;
-        output_data.write_all(&vec_u8_to_bytes_le(&identity_path_index)?)?;
+        // Note: unwrap safe - vec_fr_to_bytes_le & vec_u8_to_bytes_le are infallible
+        output_data.write_all(&vec_fr_to_bytes_le(&path_elements))?;
+        output_data.write_all(&vec_u8_to_bytes_le(&identity_path_index))?;
 
         Ok(())
     }
@@ -707,7 +720,7 @@ impl RLN {
     /// assert_eq!(idxs, [0, 1, 2, 3, 4]);
     /// ```
     #[cfg(not(feature = "stateless"))]
-    pub fn get_empty_leaves_indices<W: Write>(&self, mut output_data: W) -> Result<()> {
+    pub fn get_empty_leaves_indices<W: Write>(&self, mut output_data: W) -> Result<(), RLNError> {
         let idxs = self.tree.get_empty_leaves_indices();
         idxs.serialize_compressed(&mut output_data)?;
         Ok(())
@@ -742,7 +755,7 @@ impl RLN {
         &mut self,
         mut input_data: R,
         mut output_data: W,
-    ) -> Result<()> {
+    ) -> Result<(), RLNError> {
         // We read input RLN witness and we serialize_compressed it
         let mut serialized_witness: Vec<u8> = Vec::new();
         input_data.read_to_end(&mut serialized_witness)?;
@@ -792,7 +805,7 @@ impl RLN {
     ///
     /// assert!(verified);
     /// ```
-    pub fn verify<R: Read>(&self, mut input_data: R) -> Result<bool> {
+    pub fn verify<R: Read>(&self, mut input_data: R) -> Result<bool, RLNError> {
         // Input data is serialized for Curve as:
         // serialized_proof (compressed, 4*32 bytes) || serialized_proof_values (6*32 bytes), i.e.
         // [ proof<128> | root<32> | external_nullifier<32> | x<32> | y<32> | nullifier<32> ]
@@ -863,7 +876,7 @@ impl RLN {
         &mut self,
         mut input_data: R,
         mut output_data: W,
-    ) -> Result<()> {
+    ) -> Result<(), RLNError> {
         // We read input RLN witness and we serialize_compressed it
         let mut witness_byte: Vec<u8> = Vec::new();
         input_data.read_to_end(&mut witness_byte)?;
@@ -888,7 +901,7 @@ impl RLN {
         &mut self,
         mut input_data: R,
         mut output_data: W,
-    ) -> Result<()> {
+    ) -> Result<(), RLNError> {
         let mut serialized_witness: Vec<u8> = Vec::new();
         input_data.read_to_end(&mut serialized_witness)?;
         let (rln_witness, _) = deserialize_witness(&serialized_witness)?;
@@ -912,7 +925,7 @@ impl RLN {
         calculated_witness: Vec<BigInt>,
         serialized_witness: Vec<u8>,
         mut output_data: W,
-    ) -> Result<()> {
+    ) -> Result<(), RLNError> {
         let (rln_witness, _) = deserialize_witness(&serialized_witness[..])?;
         let proof_values = proof_values_from_witness(&rln_witness)?;
 
@@ -953,7 +966,7 @@ impl RLN {
     /// assert!(verified);
     /// ```
     #[cfg(not(feature = "stateless"))]
-    pub fn verify_rln_proof<R: Read>(&self, mut input_data: R) -> Result<bool> {
+    pub fn verify_rln_proof<R: Read>(&self, mut input_data: R) -> Result<bool, RLNError> {
         let mut serialized: Vec<u8> = Vec::new();
         input_data.read_to_end(&mut serialized)?;
         let mut all_read = 0;
@@ -964,8 +977,11 @@ impl RLN {
         all_read += read;
 
         let signal_len = usize::try_from(u64::from_le_bytes(
-            serialized[all_read..all_read + 8].try_into()?,
-        ))?;
+            serialized[all_read..all_read + 8]
+                .try_into()
+                .map_err(ConversionError::FromSlice)?,
+        ))
+        .map_err(ConversionError::from)?;
         all_read += 8;
 
         let signal: Vec<u8> = serialized[all_read..all_read + signal_len].to_vec();
@@ -1028,7 +1044,11 @@ impl RLN {
     ///
     /// assert!(verified);
     /// ```
-    pub fn verify_with_roots<R: Read>(&self, mut input_data: R, mut roots_data: R) -> Result<bool> {
+    pub fn verify_with_roots<R: Read>(
+        &self,
+        mut input_data: R,
+        mut roots_data: R,
+    ) -> Result<bool, RLNError> {
         let mut serialized: Vec<u8> = Vec::new();
         input_data.read_to_end(&mut serialized)?;
         let mut all_read = 0;
@@ -1039,8 +1059,11 @@ impl RLN {
         all_read += read;
 
         let signal_len = usize::try_from(u64::from_le_bytes(
-            serialized[all_read..all_read + 8].try_into()?,
-        ))?;
+            serialized[all_read..all_read + 8]
+                .try_into()
+                .map_err(ConversionError::FromSlice)?,
+        ))
+        .map_err(ConversionError::ToUsize)?;
         all_read += 8;
 
         let signal: Vec<u8> = serialized[all_read..all_read + signal_len].to_vec();
@@ -1109,9 +1132,9 @@ impl RLN {
     /// // We serialize_compressed the keygen output
     /// let (identity_secret_hash, id_commitment) = deserialize_identity_pair(buffer.into_inner());
     /// ```
-    pub fn key_gen<W: Write>(&self, mut output_data: W) -> Result<()> {
+    pub fn key_gen<W: Write>(&self, mut output_data: W) -> Result<(), RLNError> {
         let (identity_secret_hash, id_commitment) = keygen();
-        output_data.write_all(&fr_to_bytes_le(&identity_secret_hash))?;
+        output_data.write_all(&identity_secret_hash.to_bytes_le())?;
         output_data.write_all(&fr_to_bytes_le(&id_commitment))?;
 
         Ok(())
@@ -1139,7 +1162,7 @@ impl RLN {
     /// // We serialize_compressed the keygen output
     /// let (identity_trapdoor, identity_nullifier, identity_secret_hash, id_commitment) = deserialize_identity_tuple(buffer.into_inner());
     /// ```
-    pub fn extended_key_gen<W: Write>(&self, mut output_data: W) -> Result<()> {
+    pub fn extended_key_gen<W: Write>(&self, mut output_data: W) -> Result<(), RLNError> {
         let (identity_trapdoor, identity_nullifier, identity_secret_hash, id_commitment) =
             extended_keygen();
         output_data.write_all(&fr_to_bytes_le(&identity_trapdoor))?;
@@ -1178,7 +1201,7 @@ impl RLN {
         &self,
         mut input_data: R,
         mut output_data: W,
-    ) -> Result<()> {
+    ) -> Result<(), RLNError> {
         let mut serialized: Vec<u8> = Vec::new();
         input_data.read_to_end(&mut serialized)?;
 
@@ -1221,7 +1244,7 @@ impl RLN {
         &self,
         mut input_data: R,
         mut output_data: W,
-    ) -> Result<()> {
+    ) -> Result<(), RLNError> {
         let mut serialized: Vec<u8> = Vec::new();
         input_data.read_to_end(&mut serialized)?;
 
@@ -1274,7 +1297,7 @@ impl RLN {
         mut input_proof_data_1: R,
         mut input_proof_data_2: R,
         mut output_data: W,
-    ) -> Result<()> {
+    ) -> Result<(), RLNError> {
         // We serialize_compressed the two proofs, and we get the corresponding RLNProofValues objects
         let mut serialized: Vec<u8> = Vec::new();
         input_proof_data_1.read_to_end(&mut serialized)?;
@@ -1298,14 +1321,11 @@ impl RLN {
             let share2 = (proof_values_2.x, proof_values_2.y);
 
             // We recover the secret
-            let recovered_identity_secret_hash = compute_id_secret(share1, share2);
+            let recovered_identity_secret_hash =
+                compute_id_secret(share1, share2).map_err(RLNError::RecoverSecret)?;
 
             // If an identity secret hash is recovered, we write it to output_data, otherwise nothing will be written.
-            if let Ok(identity_secret_hash) = recovered_identity_secret_hash {
-                output_data.write_all(&fr_to_bytes_le(&identity_secret_hash))?;
-            } else {
-                return Err(Report::msg("could not extract secret"));
-            }
+            output_data.write_all(&recovered_identity_secret_hash.to_bytes_le())?;
         }
 
         Ok(())
@@ -1318,13 +1338,16 @@ impl RLN {
     ///
     /// The function returns the corresponding [`RLNWitnessInput`] object serialized using [`rln::protocol::serialize_witness`](crate::protocol::serialize_witness).
     #[cfg(not(feature = "stateless"))]
-    pub fn get_serialized_rln_witness<R: Read>(&mut self, mut input_data: R) -> Result<Vec<u8>> {
+    pub fn get_serialized_rln_witness<R: Read>(
+        &mut self,
+        mut input_data: R,
+    ) -> Result<Vec<u8>, RLNError> {
         // We read input RLN witness and we serialize_compressed it
         let mut witness_byte: Vec<u8> = Vec::new();
         input_data.read_to_end(&mut witness_byte)?;
         let (rln_witness, _) = proof_inputs_to_rln_witness(&mut self.tree, &witness_byte)?;
 
-        serialize_witness(&rln_witness)
+        serialize_witness(&rln_witness).map_err(RLNError::Protocol)
     }
 
     /// Converts a byte serialization of a [`RLNWitnessInput`] object to the corresponding JSON serialization.
@@ -1333,7 +1356,10 @@ impl RLN {
     /// - `serialized_witness`: the byte serialization of a [`RLNWitnessInput`] object (serialization done with  [`rln::protocol::serialize_witness`](crate::protocol::serialize_witness)).
     ///
     /// The function returns the corresponding JSON encoding of the input [`RLNWitnessInput`] object.
-    pub fn get_rln_witness_json(&mut self, serialized_witness: &[u8]) -> Result<serde_json::Value> {
+    pub fn get_rln_witness_json(
+        &mut self,
+        serialized_witness: &[u8],
+    ) -> Result<serde_json::Value, ProtocolError> {
         let (rln_witness, _) = deserialize_witness(serialized_witness)?;
         rln_witness_to_json(&rln_witness)
     }
@@ -1348,7 +1374,7 @@ impl RLN {
     pub fn get_rln_witness_bigint_json(
         &mut self,
         serialized_witness: &[u8],
-    ) -> Result<serde_json::Value> {
+    ) -> Result<serde_json::Value, ProtocolError> {
         let (rln_witness, _) = deserialize_witness(serialized_witness)?;
         rln_witness_to_bigint_json(&rln_witness)
     }
@@ -1358,7 +1384,7 @@ impl RLN {
     /// If not called, the connection will be closed when the RLN object is dropped.
     /// This improves robustness of the tree.
     #[cfg(not(feature = "stateless"))]
-    pub fn flush(&mut self) -> Result<()> {
+    pub fn flush(&mut self) -> Result<(), ZerokitMerkleTreeError> {
         self.tree.close_db_connection()
     }
 }
@@ -1399,7 +1425,10 @@ impl Default for RLN {
 /// // We serialize_compressed the keygen output
 /// let field_element = deserialize_field_element(output_buffer.into_inner());
 /// ```
-pub fn hash<R: Read, W: Write>(mut input_data: R, mut output_data: W) -> Result<()> {
+pub fn hash<R: Read, W: Write>(
+    mut input_data: R,
+    mut output_data: W,
+) -> Result<(), std::io::Error> {
     let mut serialized: Vec<u8> = Vec::new();
     input_data.read_to_end(&mut serialized)?;
 
@@ -1432,7 +1461,10 @@ pub fn hash<R: Read, W: Write>(mut input_data: R, mut output_data: W) -> Result<
 /// // We serialize_compressed the hash output
 /// let hash_result = deserialize_field_element(output_buffer.into_inner());
 /// ```
-pub fn poseidon_hash<R: Read, W: Write>(mut input_data: R, mut output_data: W) -> Result<()> {
+pub fn poseidon_hash<R: Read, W: Write>(
+    mut input_data: R,
+    mut output_data: W,
+) -> Result<(), RLNError> {
     let mut serialized: Vec<u8> = Vec::new();
     input_data.read_to_end(&mut serialized)?;
 

rln/src/public_api_tests.rs

@@ -1,16 +1,21 @@
 use crate::circuit::TEST_TREE_HEIGHT;
-use crate::protocol::*;
+use crate::protocol::{
+    proof_values_from_witness, random_rln_witness, serialize_proof_values, serialize_witness,
+    verify_proof, RLNProofValues,
+};
 use crate::public::RLN;
-use crate::utils::*;
+use crate::utils::str_to_fr;
 use ark_groth16::Proof as ArkProof;
 use ark_serialize::CanonicalDeserialize;
+use serde_json::{json, Value};
 use std::io::Cursor;
 use std::str::FromStr;
 
-use serde_json::{json, Value};
+#[cfg(not(feature = "stateless"))]
+use crate::utils::generate_input_buffer;
 
 fn fq_from_str(s: &str) -> ark_bn254::Fq {
-    ark_bn254::Fq::from_str(&s).unwrap()
+    ark_bn254::Fq::from_str(s).unwrap()
 }
 
 fn g1_from_str(g1: &[String]) -> ark_bn254::G1Affine {
@@ -40,7 +45,7 @@ fn value_to_string_vec(value: &Value) -> Vec<String> {
     value
         .as_array()
         .unwrap()
-        .into_iter()
+        .iter()
         .map(|val| val.as_str().unwrap().to_string())
         .collect()
 }
@@ -87,7 +92,7 @@ fn test_groth16_proof_hardcoded() {
                 .as_array()
                 .unwrap()
                 .iter()
-                .map(|item| value_to_string_vec(item))
+                .map(value_to_string_vec)
                 .collect::<Vec<Vec<String>>>(),
         ),
         c: g1_from_str(&value_to_string_vec(&valid_snarkjs_proof["pi_c"])),
@@ -199,7 +204,7 @@ mod tree_test {
             // We check if the number of leaves set is consistent
             assert_eq!(rln.tree.leaves_set(), i);
 
-            let mut buffer = Cursor::new(fr_to_bytes_le(&leaf));
+            let mut buffer = Cursor::new(fr_to_bytes_le(leaf));
             rln.set_leaf(i, &mut buffer).unwrap();
         }
 
@@ -213,7 +218,7 @@ mod tree_test {
 
         // We add leaves one by one using the internal index (new leaves goes in next available position)
         for leaf in &leaves {
-            let mut buffer = Cursor::new(fr_to_bytes_le(&leaf));
+            let mut buffer = Cursor::new(fr_to_bytes_le(leaf));
             rln.set_next_leaf(&mut buffer).unwrap();
         }
 
@@ -231,7 +236,7 @@ mod tree_test {
         rln.set_tree(tree_height).unwrap();
 
         // We add leaves in a batch into the tree
-        let mut buffer = Cursor::new(vec_fr_to_bytes_le(&leaves).unwrap());
+        let mut buffer = Cursor::new(vec_fr_to_bytes_le(&leaves));
         rln.init_tree_with_leaves(&mut buffer).unwrap();
 
         // We check if number of leaves set is consistent
@@ -289,7 +294,7 @@ mod tree_test {
         let mut rln = RLN::new(tree_height, generate_input_buffer()).unwrap();
 
         // We add leaves in a batch into the tree
-        let mut buffer = Cursor::new(vec_fr_to_bytes_le(&leaves).unwrap());
+        let mut buffer = Cursor::new(vec_fr_to_bytes_le(&leaves));
         rln.init_tree_with_leaves(&mut buffer).unwrap();
 
         // We check if number of leaves set is consistent
@@ -303,11 +308,11 @@ mod tree_test {
         // `init_tree_with_leaves` resets the tree to the height it was initialized with, using `set_tree`
 
         // We add leaves in a batch starting from index 0..set_index
-        let mut buffer = Cursor::new(vec_fr_to_bytes_le(&leaves[0..set_index]).unwrap());
+        let mut buffer = Cursor::new(vec_fr_to_bytes_le(&leaves[0..set_index]));
         rln.init_tree_with_leaves(&mut buffer).unwrap();
 
         // We add the remaining n leaves in a batch starting from index m
-        let mut buffer = Cursor::new(vec_fr_to_bytes_le(&leaves[set_index..]).unwrap());
+        let mut buffer = Cursor::new(vec_fr_to_bytes_le(&leaves[set_index..]));
         rln.set_leaves_from(set_index, &mut buffer).unwrap();
 
         // We check if number of leaves set is consistent
@@ -325,7 +330,7 @@ mod tree_test {
 
         // We add leaves one by one using the internal index (new leaves goes in next available position)
         for leaf in &leaves {
-            let mut buffer = Cursor::new(fr_to_bytes_le(&leaf));
+            let mut buffer = Cursor::new(fr_to_bytes_le(leaf));
             rln.set_next_leaf(&mut buffer).unwrap();
         }
 
@@ -359,7 +364,7 @@ mod tree_test {
         let mut rln = RLN::new(tree_height, generate_input_buffer()).unwrap();
 
         // We add leaves in a batch into the tree
-        let mut buffer = Cursor::new(vec_fr_to_bytes_le(&leaves).unwrap());
+        let mut buffer = Cursor::new(vec_fr_to_bytes_le(&leaves));
         rln.init_tree_with_leaves(&mut buffer).unwrap();
 
         // We check if number of leaves set is consistent
@@ -377,8 +382,8 @@ mod tree_test {
         let last_leaf_index = no_of_leaves - 1;
         let indices = vec![last_leaf_index as u8];
         let last_leaf = vec![*last_leaf];
-        let indices_buffer = Cursor::new(vec_u8_to_bytes_le(&indices).unwrap());
-        let leaves_buffer = Cursor::new(vec_fr_to_bytes_le(&last_leaf).unwrap());
+        let indices_buffer = Cursor::new(vec_u8_to_bytes_le(&indices));
+        let leaves_buffer = Cursor::new(vec_fr_to_bytes_le(&last_leaf));
 
         rln.atomic_operation(last_leaf_index, leaves_buffer, indices_buffer)
             .unwrap();
@@ -408,7 +413,7 @@ mod tree_test {
         let mut rln = RLN::new(tree_height, generate_input_buffer()).unwrap();
 
         // We add leaves in a batch into the tree
-        let mut buffer = Cursor::new(vec_fr_to_bytes_le(&leaves).unwrap());
+        let mut buffer = Cursor::new(vec_fr_to_bytes_le(&leaves));
         rln.init_tree_with_leaves(&mut buffer).unwrap();
 
         // We check if number of leaves set is consistent
@@ -422,8 +427,8 @@ mod tree_test {
         let zero_index = 0;
         let indices = vec![zero_index as u8];
         let zero_leaf: Vec<Fr> = vec![];
-        let indices_buffer = Cursor::new(vec_u8_to_bytes_le(&indices).unwrap());
-        let leaves_buffer = Cursor::new(vec_fr_to_bytes_le(&zero_leaf).unwrap());
+        let indices_buffer = Cursor::new(vec_u8_to_bytes_le(&indices));
+        let leaves_buffer = Cursor::new(vec_fr_to_bytes_le(&zero_leaf));
         rln.atomic_operation(0, leaves_buffer, indices_buffer)
             .unwrap();
 
@@ -452,7 +457,7 @@ mod tree_test {
         let mut rln = RLN::new(tree_height, generate_input_buffer()).unwrap();
 
         // We add leaves in a batch into the tree
-        let mut buffer = Cursor::new(vec_fr_to_bytes_le(&leaves).unwrap());
+        let mut buffer = Cursor::new(vec_fr_to_bytes_le(&leaves));
         rln.init_tree_with_leaves(&mut buffer).unwrap();
 
         // We check if number of leaves set is consistent
@@ -466,8 +471,8 @@ mod tree_test {
         let set_index = rng.gen_range(0..no_of_leaves) as usize;
         let indices = vec![set_index as u8];
         let zero_leaf: Vec<Fr> = vec![];
-        let indices_buffer = Cursor::new(vec_u8_to_bytes_le(&indices).unwrap());
-        let leaves_buffer = Cursor::new(vec_fr_to_bytes_le(&zero_leaf).unwrap());
+        let indices_buffer = Cursor::new(vec_u8_to_bytes_le(&indices));
+        let leaves_buffer = Cursor::new(vec_fr_to_bytes_le(&zero_leaf));
         rln.atomic_operation(0, leaves_buffer, indices_buffer)
             .unwrap();
 
@@ -509,7 +514,7 @@ mod tree_test {
         let (root_empty, _) = bytes_le_to_fr(&buffer.into_inner());
 
         // We add leaves in a batch into the tree
-        let mut buffer = Cursor::new(vec_fr_to_bytes_le(&leaves).unwrap());
+        let mut buffer = Cursor::new(vec_fr_to_bytes_le(&leaves));
 
         #[allow(unused_must_use)]
         rln.set_leaves_from(bad_index, &mut buffer)
@@ -599,7 +604,7 @@ mod tree_test {
         let mut rln = RLN::new(tree_height, generate_input_buffer()).unwrap();
 
         // We add leaves in a batch into the tree
-        let mut buffer = Cursor::new(vec_fr_to_bytes_le(&leaves).unwrap());
+        let mut buffer = Cursor::new(vec_fr_to_bytes_le(&leaves));
         rln.init_tree_with_leaves(&mut buffer).unwrap();
 
         // Generate identity pair
@@ -671,7 +676,7 @@ mod tree_test {
         let mut rln = RLN::new(tree_height, generate_input_buffer()).unwrap();
 
         // We add leaves in a batch into the tree
-        let mut buffer = Cursor::new(vec_fr_to_bytes_le(&leaves).unwrap());
+        let mut buffer = Cursor::new(vec_fr_to_bytes_le(&leaves));
         rln.init_tree_with_leaves(&mut buffer).unwrap();
 
         // Generate identity pair
@@ -754,7 +759,7 @@ mod tree_test {
         let mut rln = RLN::new(tree_height, generate_input_buffer()).unwrap();
 
         // We add leaves in a batch into the tree
-        let mut buffer = Cursor::new(vec_fr_to_bytes_le(&leaves).unwrap());
+        let mut buffer = Cursor::new(vec_fr_to_bytes_le(&leaves));
         rln.init_tree_with_leaves(&mut buffer).unwrap();
 
         // Generate identity pair
@@ -823,7 +828,7 @@ mod tree_test {
             .verify_with_roots(&mut input_buffer.clone(), &mut roots_buffer)
             .unwrap();
 
-        assert_eq!(verified, false);
+        assert!(!verified);
 
         // We get the root of the tree obtained adding one leaf per time
         let mut buffer = Cursor::new(Vec::<u8>::new());
@@ -877,7 +882,7 @@ mod tree_test {
         // We prepare input for generate_rln_proof API
         // input_data is [ identity_secret<32> | id_index<8> | user_message_limit<32> | message_id<32> | external_nullifier<32> | signal_len<8> | signal<var> ]
         let prove_input1 = prepare_prove_input(
-            identity_secret_hash,
+            identity_secret_hash.clone(),
             identity_index,
             user_message_limit,
             message_id,
@@ -886,7 +891,7 @@ mod tree_test {
         );
 
         let prove_input2 = prepare_prove_input(
-            identity_secret_hash,
+            identity_secret_hash.clone(),
             identity_index,
             user_message_limit,
             message_id,
@@ -929,7 +934,7 @@ mod tree_test {
 
         // We check if the recovered identity secret hash corresponds to the original one
         let (recovered_identity_secret_hash, _) = bytes_le_to_fr(&serialized_identity_secret_hash);
-        assert_eq!(recovered_identity_secret_hash, identity_secret_hash);
+        assert_eq!(recovered_identity_secret_hash, *identity_secret_hash);
 
         // We now test that computing identity_secret_hash is unsuccessful if shares computed from two different identity secret hashes but within same epoch are passed
 
@@ -979,20 +984,22 @@ mod tree_test {
 
         // ensure that the recovered secret does not match with either of the
         // used secrets in proof generation
-        assert_ne!(recovered_identity_secret_hash_new, identity_secret_hash_new);
+        assert_ne!(
+            recovered_identity_secret_hash_new,
+            *identity_secret_hash_new
+        );
     }
 }
 
 #[cfg(feature = "stateless")]
 mod stateless_test {
     use crate::circuit::{Fr, TEST_TREE_HEIGHT};
-    use crate::hashers::{hash_to_field, poseidon_hash as utils_poseidon_hash};
-    use crate::poseidon_tree::PoseidonTree;
+    use crate::hashers::{hash_to_field, poseidon_hash as utils_poseidon_hash, PoseidonHash};
     use crate::protocol::*;
     use crate::public::RLN;
     use crate::utils::*;
     use std::io::Cursor;
-    use utils::ZerokitMerkleTree;
+    use utils::{OptimalMerkleTree, ZerokitMerkleProof, ZerokitMerkleTree};
 
     use ark_std::{rand::thread_rng, UniformRand};
     use rand::Rng;
@@ -1005,10 +1012,10 @@ mod stateless_test {
         let mut rln = RLN::new().unwrap();
 
         let default_leaf = Fr::from(0);
-        let mut tree = PoseidonTree::new(
+        let mut tree: OptimalMerkleTree<PoseidonHash> = OptimalMerkleTree::new(
             TEST_TREE_HEIGHT,
             default_leaf,
-            ConfigOf::<PoseidonTree>::default(),
+            ConfigOf::<OptimalMerkleTree<PoseidonHash>>::default(),
         )
         .unwrap();
 
@@ -1039,7 +1046,8 @@ mod stateless_test {
 
         let rln_witness = rln_witness_from_values(
             identity_secret_hash,
-            &merkle_proof,
+            merkle_proof.get_path_elements(),
+            merkle_proof.get_path_index(),
             x,
             external_nullifier,
             user_message_limit,
@@ -1081,7 +1089,7 @@ mod stateless_test {
             .verify_with_roots(&mut input_buffer.clone(), &mut roots_buffer)
             .unwrap();
 
-        assert_eq!(verified, false);
+        assert!(!verified);
 
         // We get the root of the tree obtained adding one leaf per time
         let root = tree.root();
@@ -1102,10 +1110,10 @@ mod stateless_test {
         let mut rln = RLN::new().unwrap();
 
         let default_leaf = Fr::from(0);
-        let mut tree = PoseidonTree::new(
+        let mut tree: OptimalMerkleTree<PoseidonHash> = OptimalMerkleTree::new(
             TEST_TREE_HEIGHT,
             default_leaf,
-            ConfigOf::<PoseidonTree>::default(),
+            ConfigOf::<OptimalMerkleTree<PoseidonHash>>::default(),
         )
         .unwrap();
 
@@ -1133,8 +1141,9 @@ mod stateless_test {
         let merkle_proof = tree.proof(identity_index).expect("proof should exist");
 
         let rln_witness1 = rln_witness_from_values(
-            identity_secret_hash,
-            &merkle_proof,
+            identity_secret_hash.clone(),
+            merkle_proof.get_path_elements(),
+            merkle_proof.get_path_index(),
             x1,
             external_nullifier,
             user_message_limit,
@@ -1143,8 +1152,9 @@ mod stateless_test {
         .unwrap();
 
         let rln_witness2 = rln_witness_from_values(
-            identity_secret_hash,
-            &merkle_proof,
+            identity_secret_hash.clone(),
+            merkle_proof.get_path_elements(),
+            merkle_proof.get_path_index(),
             x2,
             external_nullifier,
             user_message_limit,
@@ -1183,7 +1193,7 @@ mod stateless_test {
 
         // We check if the recovered identity secret hash corresponds to the original one
         let (recovered_identity_secret_hash, _) = bytes_le_to_fr(&serialized_identity_secret_hash);
-        assert_eq!(recovered_identity_secret_hash, identity_secret_hash);
+        assert_eq!(recovered_identity_secret_hash, *identity_secret_hash);
 
         // We now test that computing identity_secret_hash is unsuccessful if shares computed from two different identity secret hashes but within same epoch are passed
 
@@ -1199,8 +1209,9 @@ mod stateless_test {
         let merkle_proof_new = tree.proof(identity_index_new).expect("proof should exist");
 
         let rln_witness3 = rln_witness_from_values(
-            identity_secret_hash_new,
-            &merkle_proof_new,
+            identity_secret_hash_new.clone(),
+            merkle_proof_new.get_path_elements(),
+            merkle_proof_new.get_path_index(),
             x3,
             external_nullifier,
             user_message_limit,
@@ -1227,7 +1238,7 @@ mod stateless_test {
 
         let serialized_identity_secret_hash = output_buffer.into_inner();
         let (recovered_identity_secret_hash_new, _) =
-            bytes_le_to_fr(&serialized_identity_secret_hash);
+            IdSecret::from_bytes_le(&serialized_identity_secret_hash);
 
         // ensure that the recovered secret does not match with either of the
         // used secrets in proof generation

rln/src/utils.rs

@@ -1,29 +1,34 @@
 // This crate provides cross-module useful utilities (mainly type conversions) not necessarily specific to RLN
 
+use crate::circuit::Fr;
+use crate::error::ConversionError;
 use ark_ff::PrimeField;
-use color_eyre::{Report, Result};
+use ark_serialize::{CanonicalDeserialize, CanonicalSerialize};
+use ark_std::UniformRand;
 use num_bigint::{BigInt, BigUint};
 use num_traits::Num;
+use rand::Rng;
+use ruint::aliases::U256;
 use serde_json::json;
 use std::io::Cursor;
-
-use crate::circuit::Fr;
+use std::ops::Deref;
+use zeroize::{Zeroize, ZeroizeOnDrop, Zeroizing};
 
 #[inline(always)]
-pub fn to_bigint(el: &Fr) -> Result<BigInt> {
-    Ok(BigUint::from(*el).into())
+pub fn to_bigint(el: &Fr) -> BigInt {
+    BigUint::from(*el).into()
 }
 
 #[inline(always)]
-pub fn fr_byte_size() -> usize {
+pub const fn fr_byte_size() -> usize {
     let mbs = <Fr as PrimeField>::MODULUS_BIT_SIZE;
     ((mbs + 64 - (mbs % 64)) / 8) as usize
 }
 
 #[inline(always)]
-pub fn str_to_fr(input: &str, radix: u32) -> Result<Fr> {
+pub fn str_to_fr(input: &str, radix: u32) -> Result<Fr, ConversionError> {
     if !(radix == 10 || radix == 16) {
-        return Err(Report::msg("wrong radix"));
+        return Err(ConversionError::WrongRadix);
     }
 
     // We remove any quote present and we trim
@@ -58,7 +63,7 @@ pub fn fr_to_bytes_le(input: &Fr) -> Vec<u8> {
 }
 
 #[inline(always)]
-pub fn vec_fr_to_bytes_le(input: &[Fr]) -> Result<Vec<u8>> {
+pub fn vec_fr_to_bytes_le(input: &[Fr]) -> Vec<u8> {
     // Calculate capacity for Vec:
     // - 8 bytes for normalized vector length (usize)
     // - each Fr element requires fr_byte_size() bytes (typically 32 bytes)
@@ -72,11 +77,11 @@ pub fn vec_fr_to_bytes_le(input: &[Fr]) -> Result<Vec<u8>> {
         bytes.extend_from_slice(&fr_to_bytes_le(el));
     }
 
-    Ok(bytes)
+    bytes
 }
 
 #[inline(always)]
-pub fn vec_u8_to_bytes_le(input: &[u8]) -> Result<Vec<u8>> {
+pub fn vec_u8_to_bytes_le(input: &[u8]) -> Vec<u8> {
     // Calculate capacity for Vec:
     // - 8 bytes for normalized vector length (usize)
     // - variable length input data
@@ -88,11 +93,11 @@ pub fn vec_u8_to_bytes_le(input: &[u8]) -> Result<Vec<u8>> {
     // We store the input
     bytes.extend_from_slice(input);
 
-    Ok(bytes)
+    bytes
 }
 
 #[inline(always)]
-pub fn bytes_le_to_vec_u8(input: &[u8]) -> Result<(Vec<u8>, usize)> {
+pub fn bytes_le_to_vec_u8(input: &[u8]) -> Result<(Vec<u8>, usize), ConversionError> {
     let mut read: usize = 0;
 
     let len = usize::try_from(u64::from_le_bytes(input[0..8].try_into()?))?;
@@ -105,12 +110,11 @@ pub fn bytes_le_to_vec_u8(input: &[u8]) -> Result<(Vec<u8>, usize)> {
 }
 
 #[inline(always)]
-pub fn bytes_le_to_vec_fr(input: &[u8]) -> Result<(Vec<Fr>, usize)> {
+pub fn bytes_le_to_vec_fr(input: &[u8]) -> Result<(Vec<Fr>, usize), ConversionError> {
     let mut read: usize = 0;
-    let mut res: Vec<Fr> = Vec::new();
-
     let len = usize::try_from(u64::from_le_bytes(input[0..8].try_into()?))?;
     read += 8;
+    let mut res: Vec<Fr> = Vec::with_capacity(len);
 
     let el_size = fr_byte_size();
     for i in 0..len {
@@ -123,7 +127,7 @@ pub fn bytes_le_to_vec_fr(input: &[u8]) -> Result<(Vec<Fr>, usize)> {
 }
 
 #[inline(always)]
-pub fn bytes_le_to_vec_usize(input: &[u8]) -> Result<Vec<usize>> {
+pub fn bytes_le_to_vec_usize(input: &[u8]) -> Result<Vec<usize>, ConversionError> {
     let nof_elem = usize::try_from(u64::from_le_bytes(input[0..8].try_into()?))?;
     if nof_elem == 0 {
         Ok(vec![])
@@ -151,3 +155,81 @@ pub fn normalize_usize(input: usize) -> [u8; 8] {
 pub fn generate_input_buffer() -> Cursor<String> {
     Cursor::new(json!({}).to_string())
 }
+
+#[derive(
+    Debug, Zeroize, ZeroizeOnDrop, Clone, PartialEq, CanonicalSerialize, CanonicalDeserialize,
+)]
+pub struct IdSecret(ark_bn254::Fr);
+
+impl IdSecret {
+    pub fn rand<R: Rng + ?Sized>(rng: &mut R) -> Self {
+        let mut fr = Fr::rand(rng);
+        let res = Self::from(&mut fr);
+        // No need to zeroize fr (already zeroiz'ed in from implementation)
+        #[allow(clippy::let_and_return)]
+        res
+    }
+
+    pub fn from_bytes_le(input: &[u8]) -> (Self, usize) {
+        let el_size = fr_byte_size();
+        let b_uint = BigUint::from_bytes_le(&input[0..el_size]);
+        let mut fr = Fr::from(b_uint);
+        let res = IdSecret::from(&mut fr);
+        // Note: no zeroize on b_uint as it has been moved
+        (res, el_size)
+    }
+
+    pub(crate) fn to_bytes_le(&self) -> Zeroizing<Vec<u8>> {
+        let input_biguint: BigUint = self.0.into();
+        let mut res = input_biguint.to_bytes_le();
+        res.resize(fr_byte_size(), 0);
+        Zeroizing::new(res)
+    }
+
+    /// Warning: this can leak the secret value
+    /// Warning: Leaked value is of type 'U256' which implement Copy (every copy will not be zeroized)
+    pub(crate) fn to_u256(&self) -> U256 {
+        let mut big_int = self.0.into_bigint();
+        let res = U256::from_limbs(big_int.0);
+        big_int.zeroize();
+        res
+    }
+}
+
+impl From<&mut Fr> for IdSecret {
+    fn from(value: &mut Fr) -> Self {
+        let id_secret = Self(*value);
+        value.zeroize();
+        id_secret
+    }
+}
+
+impl Deref for IdSecret {
+    type Target = Fr;
+
+    /// Deref to &Fr
+    ///
+    /// Warning: this can leak the secret value
+    /// Warning: Leaked value is of type 'Fr' which implement Copy (every copy will not be zeroized)
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+#[derive(Debug, Zeroize, ZeroizeOnDrop)]
+pub enum FrOrSecret {
+    IdSecret(IdSecret),
+    Fr(Fr),
+}
+
+impl From<Fr> for FrOrSecret {
+    fn from(value: Fr) -> Self {
+        FrOrSecret::Fr(value)
+    }
+}
+
+impl From<IdSecret> for FrOrSecret {
+    fn from(value: IdSecret) -> Self {
+        FrOrSecret::IdSecret(value)
+    }
+}

rln/tests/ffi.rs

@@ -3,7 +3,7 @@
 mod test {
     use ark_std::{rand::thread_rng, UniformRand};
     use rand::Rng;
-    use rln::circuit::*;
+    use rln::circuit::{Fr, TEST_TREE_HEIGHT};
     use rln::ffi::{hash as ffi_hash, poseidon_hash as ffi_poseidon_hash, *};
     use rln::hashers::{hash_to_field, poseidon_hash as utils_poseidon_hash, ROUND_PARAMS};
     use rln::protocol::*;
@@ -14,6 +14,7 @@ mod test {
     use std::io::Read;
     use std::mem::MaybeUninit;
     use std::time::{Duration, Instant};
+    use zeroize::Zeroize;
 
     const NO_OF_LEAVES: usize = 256;
 
@@ -27,7 +28,7 @@ mod test {
     }
 
     fn set_leaves_init(rln_pointer: &mut RLN, leaves: &[Fr]) {
-        let leaves_ser = vec_fr_to_bytes_le(&leaves).unwrap();
+        let leaves_ser = vec_fr_to_bytes_le(leaves);
         let input_buffer = &Buffer::from(leaves_ser.as_ref());
         let success = init_tree_with_leaves(rln_pointer, input_buffer);
         assert!(success, "init tree with leaves call failed");
@@ -49,14 +50,14 @@ mod test {
         root
     }
 
-    fn identity_pair_gen(rln_pointer: &mut RLN) -> (Fr, Fr) {
+    fn identity_pair_gen(rln_pointer: &mut RLN) -> (IdSecret, Fr) {
         let mut output_buffer = MaybeUninit::<Buffer>::uninit();
         let success = key_gen(rln_pointer, output_buffer.as_mut_ptr());
         assert!(success, "key gen call failed");
         let output_buffer = unsafe { output_buffer.assume_init() };
         let result_data = <&[u8]>::from(&output_buffer).to_vec();
-        let (identity_secret_hash, read) = bytes_le_to_fr(&result_data);
-        let (id_commitment, _) = bytes_le_to_fr(&result_data[read..].to_vec());
+        let (identity_secret_hash, read) = IdSecret::from_bytes_le(&result_data);
+        let (id_commitment, _) = bytes_le_to_fr(&result_data[read..]);
         (identity_secret_hash, id_commitment)
     }
 
@@ -80,7 +81,7 @@ mod test {
         // We first add leaves one by one specifying the index
         for (i, leaf) in leaves.iter().enumerate() {
             // We prepare the rate_commitment and we set the leaf at provided index
-            let leaf_ser = fr_to_bytes_le(&leaf);
+            let leaf_ser = fr_to_bytes_le(leaf);
             let input_buffer = &Buffer::from(leaf_ser.as_ref());
             let success = set_leaf(rln_pointer, i, input_buffer);
             assert!(success, "set leaf call failed");
@@ -95,7 +96,7 @@ mod test {
 
         // We add leaves one by one using the internal index (new leaves goes in next available position)
         for leaf in &leaves {
-            let leaf_ser = fr_to_bytes_le(&leaf);
+            let leaf_ser = fr_to_bytes_le(leaf);
             let input_buffer = &Buffer::from(leaf_ser.as_ref());
             let success = set_next_leaf(rln_pointer, input_buffer);
             assert!(success, "set next leaf call failed");
@@ -156,7 +157,7 @@ mod test {
         // random number between 0..no_of_leaves
         let mut rng = thread_rng();
         let set_index = rng.gen_range(0..NO_OF_LEAVES) as usize;
-        println!("set_index: {}", set_index);
+        println!("set_index: {set_index}");
 
         // We add leaves in a batch into the tree
         set_leaves_init(rln_pointer, &leaves);
@@ -170,7 +171,7 @@ mod test {
         set_leaves_init(rln_pointer, &leaves[0..set_index]);
 
         // We add the remaining n leaves in a batch starting from index set_index
-        let leaves_n = vec_fr_to_bytes_le(&leaves[set_index..]).unwrap();
+        let leaves_n = vec_fr_to_bytes_le(&leaves[set_index..]);
         let buffer = &Buffer::from(leaves_n.as_ref());
         let success = set_leaves_from(rln_pointer, set_index, buffer);
         assert!(success, "set leaves from call failed");
@@ -188,7 +189,7 @@ mod test {
 
         // We add leaves one by one using the internal index (new leaves goes in next available position)
         for leaf in &leaves {
-            let leaf_ser = fr_to_bytes_le(&leaf);
+            let leaf_ser = fr_to_bytes_le(leaf);
             let input_buffer = &Buffer::from(leaf_ser.as_ref());
             let success = set_next_leaf(rln_pointer, input_buffer);
             assert!(success, "set next leaf call failed");
@@ -220,17 +221,12 @@ mod test {
         let last_leaf_index = NO_OF_LEAVES - 1;
         let indices = vec![last_leaf_index as u8];
         let last_leaf = vec![*last_leaf];
-        let indices = vec_u8_to_bytes_le(&indices).unwrap();
+        let indices = vec_u8_to_bytes_le(&indices);
         let indices_buffer = &Buffer::from(indices.as_ref());
-        let leaves = vec_fr_to_bytes_le(&last_leaf).unwrap();
+        let leaves = vec_fr_to_bytes_le(&last_leaf);
         let leaves_buffer = &Buffer::from(leaves.as_ref());
 
-        let success = atomic_operation(
-            rln_pointer,
-            last_leaf_index as usize,
-            leaves_buffer,
-            indices_buffer,
-        );
+        let success = atomic_operation(rln_pointer, last_leaf_index, leaves_buffer, indices_buffer);
         assert!(success, "atomic operation call failed");
 
         // We get the root of the tree obtained after a no-op
@@ -253,7 +249,7 @@ mod test {
         let root_empty = get_tree_root(rln_pointer);
 
         // We add leaves in a batch into the tree
-        let leaves = vec_fr_to_bytes_le(&leaves).unwrap();
+        let leaves = vec_fr_to_bytes_le(&leaves);
         let buffer = &Buffer::from(leaves.as_ref());
         let success = set_leaves_from(rln_pointer, bad_index, buffer);
         assert!(!success, "set leaves from call succeeded");
@@ -271,8 +267,11 @@ mod test {
         let rln_pointer = create_rln_instance();
 
         // generate identity
-        let identity_secret_hash = hash_to_field(b"test-merkle-proof");
-        let id_commitment = utils_poseidon_hash(&[identity_secret_hash]);
+        let mut identity_secret_hash_ = hash_to_field(b"test-merkle-proof");
+        let identity_secret_hash = IdSecret::from(&mut identity_secret_hash_);
+        let mut to_hash = [*identity_secret_hash.clone()];
+        let id_commitment = utils_poseidon_hash(&to_hash);
+        to_hash[0].zeroize();
         let user_message_limit = Fr::from(100);
         let rate_commitment = utils_poseidon_hash(&[id_commitment, user_message_limit]);
 
@@ -305,7 +304,7 @@ mod test {
         let result_data = <&[u8]>::from(&output_buffer).to_vec();
 
         let (path_elements, read) = bytes_le_to_vec_fr(&result_data).unwrap();
-        let (identity_path_index, _) = bytes_le_to_vec_u8(&result_data[read..].to_vec()).unwrap();
+        let (identity_path_index, _) = bytes_le_to_vec_u8(&result_data[read..]).unwrap();
 
         // We check correct computation of the path and indexes
         let expected_path_elements: Vec<Fr> = [
@@ -393,7 +392,7 @@ mod test {
             let success = verify(rln_pointer, input_buffer, proof_is_valid_ptr);
             verify_time += now.elapsed().as_nanos();
             assert!(success, "verify call failed");
-            assert_eq!(proof_is_valid, true);
+            assert!(proof_is_valid);
         }
 
         println!(
@@ -415,12 +414,9 @@ mod test {
         // We obtain the root from the RLN instance
         let root_rln_folder = get_tree_root(rln_pointer);
 
-        #[cfg(feature = "arkzkey")]
         let zkey_path = "./resources/tree_height_20/rln_final.arkzkey";
-        #[cfg(not(feature = "arkzkey"))]
-        let zkey_path = "./resources/tree_height_20/rln_final.zkey";
-        let mut zkey_file = File::open(&zkey_path).expect("no file found");
-        let metadata = std::fs::metadata(&zkey_path).expect("unable to read metadata");
+        let mut zkey_file = File::open(zkey_path).expect("no file found");
+        let metadata = std::fs::metadata(zkey_path).expect("unable to read metadata");
         let mut zkey_buffer = vec![0; metadata.len() as usize];
         zkey_file
             .read_exact(&mut zkey_buffer)
@@ -429,8 +425,8 @@ mod test {
         let zkey_data = &Buffer::from(&zkey_buffer[..]);
 
         let graph_data = "./resources/tree_height_20/graph.bin";
-        let mut graph_file = File::open(&graph_data).expect("no file found");
-        let metadata = std::fs::metadata(&graph_data).expect("unable to read metadata");
+        let mut graph_file = File::open(graph_data).expect("no file found");
+        let metadata = std::fs::metadata(graph_data).expect("unable to read metadata");
         let mut graph_buffer = vec![0; metadata.len() as usize];
         graph_file
             .read_exact(&mut graph_buffer)
@@ -525,7 +521,7 @@ mod test {
         let proof_is_valid_ptr = &mut proof_is_valid as *mut bool;
         let success = verify_rln_proof(rln_pointer, input_buffer, proof_is_valid_ptr);
         assert!(success, "verify call failed");
-        assert_eq!(proof_is_valid, true);
+        assert!(proof_is_valid);
     }
 
     #[test]
@@ -600,7 +596,7 @@ mod test {
             verify_with_roots(rln_pointer, input_buffer, roots_buffer, proof_is_valid_ptr);
         assert!(success, "verify call failed");
         // Proof should be valid
-        assert_eq!(proof_is_valid, true);
+        assert!(proof_is_valid);
 
         // We then try to verify against some random values not containing the correct one.
         for _ in 0..5 {
@@ -614,7 +610,7 @@ mod test {
             verify_with_roots(rln_pointer, input_buffer, roots_buffer, proof_is_valid_ptr);
         assert!(success, "verify call failed");
         // Proof should be invalid.
-        assert_eq!(proof_is_valid, false);
+        assert!(!proof_is_valid);
 
         // We finally include the correct root
         // We get the root of the tree obtained adding one leaf per time
@@ -630,7 +626,7 @@ mod test {
             verify_with_roots(rln_pointer, input_buffer, roots_buffer, proof_is_valid_ptr);
         assert!(success, "verify call failed");
         // Proof should be valid.
-        assert_eq!(proof_is_valid, true);
+        assert!(proof_is_valid);
     }
 
     #[test]
@@ -674,7 +670,7 @@ mod test {
         // We prepare input for generate_rln_proof API
         // input_data is [ identity_secret<32> | id_index<8> | user_message_limit<32> | message_id<32> | external_nullifier<32> | signal_len<8> | signal<var> ]
         let prove_input1 = prepare_prove_input(
-            identity_secret_hash,
+            identity_secret_hash.clone(),
             identity_index,
             user_message_limit,
             message_id,
@@ -683,7 +679,7 @@ mod test {
         );
 
         let prove_input2 = prepare_prove_input(
-            identity_secret_hash,
+            identity_secret_hash.clone(),
             identity_index,
             user_message_limit,
             message_id,
@@ -718,7 +714,7 @@ mod test {
 
         // We check if the recovered identity secret hash corresponds to the original one
         let (recovered_identity_secret_hash, _) = bytes_le_to_fr(&serialized_identity_secret_hash);
-        assert_eq!(recovered_identity_secret_hash, identity_secret_hash);
+        assert_eq!(recovered_identity_secret_hash, *identity_secret_hash);
 
         // We now test that computing identity_secret_hash is unsuccessful if shares computed from two different identity secret hashes but within same epoch are passed
 
@@ -772,7 +768,10 @@ mod test {
 
         // ensure that the recovered secret does not match with either of the
         // used secrets in proof generation
-        assert_ne!(recovered_identity_secret_hash_new, identity_secret_hash_new);
+        assert_ne!(
+            recovered_identity_secret_hash_new,
+            *identity_secret_hash_new
+        );
     }
 
     #[test]
@@ -790,7 +789,7 @@ mod test {
         let output_buffer = unsafe { output_buffer.assume_init() };
         let result_data = <&[u8]>::from(&output_buffer).to_vec();
         let (identity_secret_hash, read) = bytes_le_to_fr(&result_data);
-        let (id_commitment, _) = bytes_le_to_fr(&result_data[read..].to_vec());
+        let (id_commitment, _) = bytes_le_to_fr(&result_data[read..]);
 
         // We check against expected values
         let expected_identity_secret_hash_seed_bytes = str_to_fr(
@@ -892,7 +891,7 @@ mod test {
         for _ in 0..number_of_inputs {
             inputs.push(Fr::rand(&mut rng));
         }
-        let inputs_ser = vec_fr_to_bytes_le(&inputs).unwrap();
+        let inputs_ser = vec_fr_to_bytes_le(&inputs);
         let input_buffer = &Buffer::from(inputs_ser.as_ref());
 
         let expected_hash = utils_poseidon_hash(inputs.as_ref());
@@ -991,14 +990,15 @@ mod stateless_test {
     use rln::circuit::*;
     use rln::ffi::generate_rln_proof_with_witness;
     use rln::ffi::{hash as ffi_hash, poseidon_hash as ffi_poseidon_hash, *};
-    use rln::hashers::{hash_to_field, poseidon_hash as utils_poseidon_hash, ROUND_PARAMS};
-    use rln::poseidon_tree::PoseidonTree;
+    use rln::hashers::{
+        hash_to_field, poseidon_hash as utils_poseidon_hash, PoseidonHash, ROUND_PARAMS,
+    };
     use rln::protocol::*;
     use rln::public::RLN;
     use rln::utils::*;
     use std::mem::MaybeUninit;
     use std::time::{Duration, Instant};
-    use utils::ZerokitMerkleTree;
+    use utils::{OptimalMerkleTree, ZerokitMerkleProof, ZerokitMerkleTree};
 
     type ConfigOf<T> = <T as ZerokitMerkleTree>::Config;
 
@@ -1009,14 +1009,14 @@ mod stateless_test {
         unsafe { &mut *rln_pointer.assume_init() }
     }
 
-    fn identity_pair_gen(rln_pointer: &mut RLN) -> (Fr, Fr) {
+    fn identity_pair_gen(rln_pointer: &mut RLN) -> (IdSecret, Fr) {
         let mut output_buffer = MaybeUninit::<Buffer>::uninit();
         let success = key_gen(rln_pointer, output_buffer.as_mut_ptr());
         assert!(success, "key gen call failed");
         let output_buffer = unsafe { output_buffer.assume_init() };
         let result_data = <&[u8]>::from(&output_buffer).to_vec();
-        let (identity_secret_hash, read) = bytes_le_to_fr(&result_data);
-        let (id_commitment, _) = bytes_le_to_fr(&result_data[read..].to_vec());
+        let (identity_secret_hash, read) = IdSecret::from_bytes_le(&result_data);
+        let (id_commitment, _) = bytes_le_to_fr(&result_data[read..]);
         (identity_secret_hash, id_commitment)
     }
 
@@ -1033,10 +1033,10 @@ mod stateless_test {
     #[test]
     fn test_recover_id_secret_stateless_ffi() {
         let default_leaf = Fr::from(0);
-        let mut tree = PoseidonTree::new(
+        let mut tree: OptimalMerkleTree<PoseidonHash> = OptimalMerkleTree::new(
             TEST_TREE_HEIGHT,
             default_leaf,
-            ConfigOf::<PoseidonTree>::default(),
+            ConfigOf::<OptimalMerkleTree<PoseidonHash>>::default(),
         )
         .unwrap();
 
@@ -1068,8 +1068,9 @@ mod stateless_test {
 
         // We prepare input for generate_rln_proof API
         let rln_witness1 = rln_witness_from_values(
-            identity_secret_hash,
-            &merkle_proof,
+            identity_secret_hash.clone(),
+            merkle_proof.get_path_elements(),
+            merkle_proof.get_path_index(),
             x1,
             external_nullifier,
             user_message_limit,
@@ -1079,8 +1080,9 @@ mod stateless_test {
         let serialized1 = serialize_witness(&rln_witness1).unwrap();
 
         let rln_witness2 = rln_witness_from_values(
-            identity_secret_hash,
-            &merkle_proof,
+            identity_secret_hash.clone(),
+            merkle_proof.get_path_elements(),
+            merkle_proof.get_path_index(),
             x2,
             external_nullifier,
             user_message_limit,
@@ -1115,7 +1117,8 @@ mod stateless_test {
         assert!(!serialized_identity_secret_hash.is_empty());
 
         // We check if the recovered identity secret hash corresponds to the original one
-        let (recovered_identity_secret_hash, _) = bytes_le_to_fr(&serialized_identity_secret_hash);
+        let (recovered_identity_secret_hash, _) =
+            IdSecret::from_bytes_le(&serialized_identity_secret_hash);
         assert_eq!(recovered_identity_secret_hash, identity_secret_hash);
 
         // We now test that computing identity_secret_hash is unsuccessful if shares computed from two different identity secret hashes but within same epoch are passed
@@ -1133,8 +1136,9 @@ mod stateless_test {
         let merkle_proof_new = tree.proof(identity_index_new).expect("proof should exist");
 
         let rln_witness3 = rln_witness_from_values(
-            identity_secret_hash_new,
-            &merkle_proof_new,
+            identity_secret_hash_new.clone(),
+            merkle_proof_new.get_path_elements(),
+            merkle_proof_new.get_path_index(),
             x3,
             external_nullifier,
             user_message_limit,
@@ -1166,16 +1170,19 @@ mod stateless_test {
 
         // ensure that the recovered secret does not match with either of the
         // used secrets in proof generation
-        assert_ne!(recovered_identity_secret_hash_new, identity_secret_hash_new);
+        assert_ne!(
+            recovered_identity_secret_hash_new,
+            *identity_secret_hash_new
+        );
     }
 
     #[test]
     fn test_verify_with_roots_stateless_ffi() {
         let default_leaf = Fr::from(0);
-        let mut tree = PoseidonTree::new(
+        let mut tree: OptimalMerkleTree<PoseidonHash> = OptimalMerkleTree::new(
             TEST_TREE_HEIGHT,
             default_leaf,
-            ConfigOf::<PoseidonTree>::default(),
+            ConfigOf::<OptimalMerkleTree<PoseidonHash>>::default(),
         )
         .unwrap();
 
@@ -1205,7 +1212,8 @@ mod stateless_test {
         // We prepare input for generate_rln_proof API
         let rln_witness = rln_witness_from_values(
             identity_secret_hash,
-            &merkle_proof,
+            merkle_proof.get_path_elements(),
+            merkle_proof.get_path_index(),
             x,
             external_nullifier,
             user_message_limit,
@@ -1229,7 +1237,7 @@ mod stateless_test {
             verify_with_roots(rln_pointer, input_buffer, roots_buffer, proof_is_valid_ptr);
         assert!(success, "verify call failed");
         // Proof should be valid
-        assert_eq!(proof_is_valid, true);
+        assert!(proof_is_valid);
 
         // We serialize in the roots buffer some random values and we check that the proof is not verified since doesn't contain the correct root the proof refers to
         for _ in 0..5 {
@@ -1243,7 +1251,7 @@ mod stateless_test {
             verify_with_roots(rln_pointer, input_buffer, roots_buffer, proof_is_valid_ptr);
         assert!(success, "verify call failed");
         // Proof should be invalid.
-        assert_eq!(proof_is_valid, false);
+        assert!(!proof_is_valid);
 
         // We get the root of the tree obtained adding one leaf per time
         let root = tree.root();
@@ -1258,7 +1266,7 @@ mod stateless_test {
             verify_with_roots(rln_pointer, input_buffer, roots_buffer, proof_is_valid_ptr);
         assert!(success, "verify call failed");
         // Proof should be valid.
-        assert_eq!(proof_is_valid, true);
+        assert!(proof_is_valid);
     }
 
     #[test]
@@ -1303,7 +1311,7 @@ mod stateless_test {
             let success = verify(rln_pointer, input_buffer, proof_is_valid_ptr);
             verify_time += now.elapsed().as_nanos();
             assert!(success, "verify call failed");
-            assert_eq!(proof_is_valid, true);
+            assert!(proof_is_valid);
         }
 
         println!(
@@ -1331,7 +1339,7 @@ mod stateless_test {
         let output_buffer = unsafe { output_buffer.assume_init() };
         let result_data = <&[u8]>::from(&output_buffer).to_vec();
         let (identity_secret_hash, read) = bytes_le_to_fr(&result_data);
-        let (id_commitment, _) = bytes_le_to_fr(&result_data[read..].to_vec());
+        let (id_commitment, _) = bytes_le_to_fr(&result_data[read..]);
 
         // We check against expected values
         let expected_identity_secret_hash_seed_bytes = str_to_fr(
@@ -1433,7 +1441,7 @@ mod stateless_test {
         for _ in 0..number_of_inputs {
             inputs.push(Fr::rand(&mut rng));
         }
-        let inputs_ser = vec_fr_to_bytes_le(&inputs).unwrap();
+        let inputs_ser = vec_fr_to_bytes_le(&inputs);
         let input_buffer = &Buffer::from(inputs_ser.as_ref());
 
         let expected_hash = utils_poseidon_hash(inputs.as_ref());

rln/tests/poseidon_tree.rs

@@ -1,28 +1,37 @@
 ////////////////////////////////////////////////////////////
-/// Tests
+// Tests
 ////////////////////////////////////////////////////////////
 
+#![cfg(not(feature = "stateless"))]
+
 #[cfg(test)]
 mod test {
     use rln::hashers::{poseidon_hash, PoseidonHash};
-    use rln::{circuit::*, poseidon_tree::PoseidonTree};
+    use rln::{
+        circuit::{Fr, TEST_TREE_HEIGHT},
+        poseidon_tree::PoseidonTree,
+    };
     use utils::{FullMerkleTree, OptimalMerkleTree, ZerokitMerkleProof, ZerokitMerkleTree};
 
     #[test]
-    // The test is checked correctness for `FullMerkleTree` and `OptimalMerkleTree` with Poseidon hash
+    // The test checked correctness for `FullMerkleTree` and `OptimalMerkleTree` with Poseidon hash
     fn test_zerokit_merkle_implementations() {
         let sample_size = 100;
-        let leaves: Vec<Fr> = (0..sample_size).map(|s| Fr::from(s)).collect();
+        let leaves: Vec<Fr> = (0..sample_size).map(Fr::from).collect();
 
         let mut tree_full = FullMerkleTree::<PoseidonHash>::default(TEST_TREE_HEIGHT).unwrap();
         let mut tree_opt = OptimalMerkleTree::<PoseidonHash>::default(TEST_TREE_HEIGHT).unwrap();
 
-        for i in 0..sample_size.try_into().unwrap() {
-            tree_full.set(i, leaves[i]).unwrap();
+        for (i, leave) in leaves
+            .into_iter()
+            .enumerate()
+            .take(sample_size.try_into().unwrap())
+        {
+            tree_full.set(i, leave).unwrap();
             let proof = tree_full.proof(i).expect("index should be set");
             assert_eq!(proof.leaf_index(), i);
 
-            tree_opt.set(i, leaves[i]).unwrap();
+            tree_opt.set(i, leave).unwrap();
             assert_eq!(tree_opt.root(), tree_full.root());
             let proof = tree_opt.proof(i).expect("index should be set");
             assert_eq!(proof.leaf_index(), i);
@@ -98,7 +107,7 @@ mod test {
 
         // check remove_indices_and_set_leaves inside override_range function
         assert!(tree.get_empty_leaves_indices().is_empty());
-        let leaves_2: Vec<Fr> = (0..2).map(|s| Fr::from(s as i32)).collect();
+        let leaves_2: Vec<Fr> = (0..2).map(Fr::from).collect();
         tree.override_range(0, leaves_2.clone().into_iter(), [0, 1, 2, 3].into_iter())
             .unwrap();
         assert_eq!(tree.get_empty_leaves_indices(), vec![2, 3]);
@@ -113,7 +122,7 @@ mod test {
             .unwrap();
         assert_eq!(tree.get_empty_leaves_indices(), vec![2, 3]);
 
-        let leaves_4: Vec<Fr> = (0..4).map(|s| Fr::from(s as i32)).collect();
+        let leaves_4: Vec<Fr> = (0..4).map(Fr::from).collect();
         // check if the indexes for write and delete are the same
         tree.override_range(0, leaves_4.clone().into_iter(), [0, 1, 2, 3].into_iter())
             .unwrap();

rln/tests/protocol.rs

@@ -1,3 +1,5 @@
+#![cfg(not(feature = "stateless"))]
+
 #[cfg(test)]
 mod test {
     use ark_ff::BigInt;
@@ -5,7 +7,12 @@ mod test {
     use rln::circuit::{Fr, TEST_TREE_HEIGHT};
     use rln::hashers::{hash_to_field, poseidon_hash};
     use rln::poseidon_tree::PoseidonTree;
-    use rln::protocol::*;
+    use rln::protocol::{
+        deserialize_proof_values, deserialize_witness, generate_proof, keygen,
+        proof_values_from_witness, rln_witness_from_json, rln_witness_from_values,
+        rln_witness_to_json, seeded_keygen, serialize_proof_values, serialize_witness,
+        verify_proof, RLNWitnessInput,
+    };
     use rln::utils::str_to_fr;
     use utils::{ZerokitMerkleProof, ZerokitMerkleTree};
 
@@ -29,7 +36,7 @@ mod test {
             ConfigOf::<PoseidonTree>::default(),
         )
         .unwrap();
-        tree.set(leaf_index, rate_commitment.into()).unwrap();
+        tree.set(leaf_index, rate_commitment).unwrap();
 
         // We check correct computation of the root
         let root = tree.root();
@@ -100,7 +107,7 @@ mod test {
             ConfigOf::<PoseidonTree>::default(),
         )
         .unwrap();
-        tree.set(leaf_index, rate_commitment.into()).unwrap();
+        tree.set(leaf_index, rate_commitment).unwrap();
 
         let merkle_proof = tree.proof(leaf_index).expect("proof should exist");
 
@@ -114,7 +121,8 @@ mod test {
 
         rln_witness_from_values(
             identity_secret_hash,
-            &merkle_proof,
+            merkle_proof.get_path_elements(),
+            merkle_proof.get_path_index(),
             x,
             external_nullifier,
             user_message_limit,
@@ -137,11 +145,11 @@ mod test {
         assert_eq!(rln_witness_deser, rln_witness);
 
         // Let's generate a zkSNARK proof
-        let proof = generate_proof(&proving_key, &rln_witness_deser, &graph_data).unwrap();
+        let proof = generate_proof(proving_key, &rln_witness_deser, graph_data).unwrap();
         let proof_values = proof_values_from_witness(&rln_witness_deser).unwrap();
 
         // Let's verify the proof
-        let verified = verify_proof(&verification_key, &proof, &proof_values);
+        let verified = verify_proof(verification_key, &proof, &proof_values);
 
         assert!(verified.unwrap());
     }
@@ -160,12 +168,12 @@ mod test {
         let graph_data = graph_from_folder();
 
         // Let's generate a zkSNARK proof
-        let proof = generate_proof(&proving_key, &rln_witness_deser, &graph_data).unwrap();
+        let proof = generate_proof(proving_key, &rln_witness_deser, graph_data).unwrap();
 
         let proof_values = proof_values_from_witness(&rln_witness_deser).unwrap();
 
         // Let's verify the proof
-        let success = verify_proof(&verification_key, &proof, &proof_values).unwrap();
+        let success = verify_proof(verification_key, &proof, &proof_values).unwrap();
 
         assert!(success);
     }

rln/tests/public.rs

@@ -3,7 +3,15 @@ mod test {
     #[cfg(not(feature = "stateless"))]
     use {
         ark_ff::BigInt,
-        rln::{circuit::TEST_TREE_HEIGHT, protocol::compute_tree_root},
+        rln::{
+            circuit::TEST_TREE_HEIGHT,
+            protocol::compute_tree_root,
+            utils::{
+                bytes_le_to_vec_fr, bytes_le_to_vec_u8, bytes_le_to_vec_usize, fr_to_bytes_le,
+                generate_input_buffer, IdSecret,
+            },
+        },
+        zeroize::Zeroize,
     };
 
     use ark_std::{rand::thread_rng, UniformRand};
@@ -12,7 +20,7 @@ mod test {
     use rln::hashers::{hash_to_field, poseidon_hash as utils_poseidon_hash, ROUND_PARAMS};
     use rln::protocol::deserialize_identity_tuple;
     use rln::public::{hash as public_hash, poseidon_hash as public_poseidon_hash, RLN};
-    use rln::utils::*;
+    use rln::utils::{bytes_le_to_fr, str_to_fr, vec_fr_to_bytes_le};
     use std::io::Cursor;
 
     #[test]
@@ -25,8 +33,13 @@ mod test {
         let mut rln = RLN::new(TEST_TREE_HEIGHT, generate_input_buffer()).unwrap();
 
         // generate identity
-        let identity_secret_hash = hash_to_field(b"test-merkle-proof");
-        let id_commitment = utils_poseidon_hash(&vec![identity_secret_hash]);
+        let mut identity_secret_hash_ = hash_to_field(b"test-merkle-proof");
+        let identity_secret_hash = IdSecret::from(&mut identity_secret_hash_);
+
+        let mut to_hash = [*identity_secret_hash.clone()];
+        let id_commitment = utils_poseidon_hash(&to_hash);
+        to_hash[0].zeroize();
+
         let rate_commitment = utils_poseidon_hash(&[id_commitment, user_message_limit.into()]);
 
         // check that leaves indices is empty
@@ -66,7 +79,7 @@ mod test {
 
         let buffer_inner = buffer.into_inner();
         let (path_elements, read) = bytes_le_to_vec_fr(&buffer_inner).unwrap();
-        let (identity_path_index, _) = bytes_le_to_vec_u8(&buffer_inner[read..].to_vec()).unwrap();
+        let (identity_path_index, _) = bytes_le_to_vec_u8(&buffer_inner[read..]).unwrap();
 
         // We check correct computation of the path and indexes
         let expected_path_elements: Vec<Fr> = [
@@ -148,7 +161,7 @@ mod test {
         let serialized_output = output_buffer.into_inner();
 
         let (identity_secret_hash, read) = bytes_le_to_fr(&serialized_output);
-        let (id_commitment, _) = bytes_le_to_fr(&serialized_output[read..].to_vec());
+        let (id_commitment, _) = bytes_le_to_fr(&serialized_output[read..]);
 
         // We check against expected values
         let expected_identity_secret_hash_seed_bytes = str_to_fr(
@@ -243,7 +256,7 @@ mod test {
         }
         let expected_hash = utils_poseidon_hash(&inputs);
 
-        let mut input_buffer = Cursor::new(vec_fr_to_bytes_le(&inputs).unwrap());
+        let mut input_buffer = Cursor::new(vec_fr_to_bytes_le(&inputs));
         let mut output_buffer = Cursor::new(Vec::<u8>::new());
 
         public_poseidon_hash(&mut input_buffer, &mut output_buffer).unwrap();

utils/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "zerokit_utils"
-version = "0.5.2"
+version = "0.6.0"
 edition = "2021"
 license = "MIT OR Apache-2.0"
 description = "Various utilities for Zerokit"
@@ -12,28 +12,27 @@ repository = "https://github.com/vacp2p/zerokit"
 bench = false
 
 [dependencies]
-ark-ff = { version = "0.5.0", default-features = false, features = [
-    "parallel",
-] }
-num-bigint = { version = "0.4.6", default-features = false, features = [
-    "rand",
-] }
-color-eyre = "0.6.3"
-pmtree = { package = "vacp2p_pmtree", version = "2.0.2", optional = true }
+ark-ff = { version = "0.5.0", default-features = false }
+num-bigint = { version = "0.4.6", default-features = false }
+# pmtree = { package = "vacp2p_pmtree", version = "2.0.2", optional = true }
+pmtree = { git = "https://github.com/vacp2p/pmtree", branch = "upgrade-rayon-version", package = "vacp2p_pmtree", optional = true }
 sled = "0.34.7"
-serde = "1.0"
+serde_json = "1.0.141"
 lazy_static = "1.5.0"
-hex = "0.4"
+hex = "0.4.3"
+rayon = "1.10.0"
+thiserror = "2.0"
 
 [dev-dependencies]
 ark-bn254 = { version = "0.5.0", features = ["std"] }
 num-traits = "0.2.19"
-hex-literal = "1.0.0"
+hex-literal = "0.4.1"
 tiny-keccak = { version = "2.0.2", features = ["keccak"] }
-criterion = { version = "0.4.0", features = ["html_reports"] }
+criterion = { version = "0.7.0", features = ["html_reports"] }
 
 [features]
 default = []
+parallel = ["ark-ff/parallel"]
 pmtree-ft = ["pmtree"]
 
 [[bench]]
@@ -43,3 +42,6 @@ harness = false
 [[bench]]
 name = "poseidon_benchmark"
 harness = false
+
+[package.metadata.docs.rs]
+all-features = true

utils/README.md

@@ -1,39 +1,45 @@
 # Zerokit Utils Crate
 
 [![Crates.io](https://img.shields.io/crates/v/zerokit_utils.svg)](https://crates.io/crates/zerokit_utils)
+[![License: MIT](https://img.shields.io/badge/License-MIT-blue.svg)](https://opensource.org/licenses/MIT)
+[![License: Apache 2.0](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](https://opensource.org/licenses/Apache-2.0)
 
-Cryptographic primitives for zero-knowledge applications, featuring efficient Merkle tree implementations and a Poseidon hash function.
+**Zerokit Utils** provides essential cryptographic primitives optimized for zero-knowledge applications.
+This crate features efficient Merkle tree implementations and a Poseidon hash function,
+designed to be robust and performant.
 
 ## Overview
 
 This crate provides core cryptographic components optimized for zero-knowledge proof systems:
 
-1. Multiple Merkle tree implementations with different space/time tradeoffs
-2. A Poseidon hash implementation
+- **Multiple Merkle Trees**: Various implementations optimised for the trade-off between space and time.
+- **Poseidon Hash Function**: An efficient hashing algorithm suitable for ZK contexts, with customizable parameters.
+- **Parallel Performance**: Leverages Rayon for significant speed-ups in Merkle tree computations.
+- **Arkworks Compatibility**: Poseidon hash implementation is designed to work seamlessly
+  with Arkworks field traits and data structures.
 
 ## Merkle Tree Implementations
 
-The crate supports two interchangeable Merkle tree implementations:
+Merkle trees are fundamental data structures for verifying data integrity and set membership.
+Zerokit Utils offers two interchangeable implementations:
 
-- **FullMerkleTree**
-  - Stores each tree node in memory
-- **OptimalMerkleTree**
-  - Only stores nodes used to prove accumulation of set leaves
+### Understanding Merkle Tree Terminology
 
-### Implementation notes
+To better understand the structure and parameters of our Merkle trees, here's a quick glossary:
 
-Glossary:
+- **Depth (`depth`)**: level of leaves if we count from root.
+  If the root is at level 0, leaves are at level `depth`.
+- **Number of Levels**: `depth + 1`.
+- **Capacity (Number of Leaves)**: $2^{\text{depth}}$. This is the maximum number of leaves the tree can hold.
+- **Total Number of Nodes**: $2^{(\text{depth} + 1)} - 1$ for a full binary tree.
 
-* depth: level of leaves if we count from levels from 0
-* number of levels: depth + 1
-* capacity (== number of leaves) -- 1 << depth
-* total number of nodes: 1 << (depth + 1)) - 1
+**Example for a tree with `depth: 3`**:
 
-So for instance:
-* depth: 3
-* number of levels: 4
-* capacity (number of leaves): 8
-* total number of nodes: 15
+- Number of Levels: 4 (levels 0, 1, 2, 3)
+- Capacity (Number of Leaves): $2^3 = 8$
+- Total Number of Nodes: $2^{(3+1)} - 1 = 15$
+
+Visual representation of a Merkle tree with `depth: 3`:
 
 ```mermaid
 flowchart TD
@@ -53,39 +59,55 @@ flowchart TD
     N6 -->|Leaf| L8
 ```
 
+### Available Implementations
+
+- **FullMerkleTree**
+  - Stores all tree nodes in memory.
+  - Use Case: Use when memory is abundant and operation speed is critical.
+
+- **OptimalMerkleTree**
+  - Stores only the nodes required to prove the accumulation of set leaves (i.e., authentication paths).
+  - Use Case: Suited for environments where memory efficiency is a higher priority than raw speed.
+
+#### Parallel Processing with Rayon
+
+Both `OptimalMerkleTree` and `FullMerkleTree` internally utilize the Rayon crate
+to accelerate computations through data parallelism.
+This can lead to significant performance improvements, particularly during updates to large Merkle trees.
 
 ## Poseidon Hash Implementation
 
-This crate provides an implementation to compute the Poseidon hash round constants and MDS matrices:
+This crate provides an implementation for computing Poseidon hash round constants and MDS matrices.
+Key characteristics include:
 
-- **Customizable parameters**: Supports different security levels and input sizes
-- **Arkworks-friendly**: Adapted to work over arkworks field traits and custom data structures
+- **Customizable parameters**: Supports various security levels and input sizes,
+  allowing you to tailor the hash function to your specific needs.
+- **Arkworks-friendly**: Adapted to integrate smoothly with Arkworks field traits and custom data structures.
 
-### Security Note
+### ⚠️ Security Note
 
-The MDS matrices are generated iteratively using the Grain LFSR until certain criteria are met.
-According to the paper, such matrices must respect specific conditions which are checked by 3 different algorithms in the reference implementation.
+The MDS matrices used in the Poseidon hash function are generated iteratively
+using the Grain LFSR (Linear Feedback Shift Register) algorithm until specific cryptographic criteria are met.
 
-These validation algorithms are not currently implemented in this crate.
-For the hardcoded parameters, the first random matrix generated satisfies these conditions.
-If using different parameters, you should check against the reference implementation how many matrices are generated before outputting the correct one,
-and pass this number to the `skip_matrices` parameter of the `find_poseidon_ark_and_mds` function.
+- The reference Poseidon implementation includes validation algorithms to ensure these criteria are satisfied.
+  These validation algorithms are not currently implemented in this crate.
+- For the hardcoded parameters provided within this crate,
+  the initially generated random matrix has been verified to meet these conditions.
+- If you intend to use custom parameters, it is crucial to verify your generated MDS matrix.
+  You should consult the Poseidon reference implementation to determine
+  how many matrices are typically skipped before a valid one is found.
+  This count should then be passed as the `skip_matrices parameter` to the `find_poseidon_ark_and_mds`
+  function in this crate.
 
 ## Installation
 
-Add Zerokit Utils to your Rust project:
+Add zerokit-utils as a dependency to your Cargo.toml file:
 
 ```toml
 [dependencies]
-zerokit-utils = "0.5.1"
+zerokit-utils = "0.6.0"
 ```
 
-## Performance Considerations
-
-- **FullMerkleTree**: Use when memory is abundant and operation speed is critical
-- **OptimalMerkleTree**: Use when memory efficiency is more important than raw speed
-- **Poseidon**: Offers a good balance between security and performance for ZK applications
-
 ## Building and Testing
 
 ```bash

utils/benches/merkle_tree_benchmark.rs

@@ -1,7 +1,7 @@
-use criterion::{criterion_group, criterion_main, Criterion};
-use hex_literal::hex;
-use lazy_static::lazy_static;
 use std::{fmt::Display, str::FromStr};
+
+use criterion::{criterion_group, criterion_main, Criterion};
+use lazy_static::lazy_static;
 use tiny_keccak::{Hasher as _, Keccak};
 use zerokit_utils::{
     FullMerkleConfig, FullMerkleTree, Hasher, OptimalMerkleConfig, OptimalMerkleTree,
@@ -47,55 +47,78 @@ impl FromStr for TestFr {
 }
 
 lazy_static! {
-    static ref LEAVES: [TestFr; 4] = [
-        hex!("0000000000000000000000000000000000000000000000000000000000000001"),
-        hex!("0000000000000000000000000000000000000000000000000000000000000002"),
-        hex!("0000000000000000000000000000000000000000000000000000000000000003"),
-        hex!("0000000000000000000000000000000000000000000000000000000000000004"),
-    ]
-    .map(TestFr);
+    static ref LEAVES: Vec<TestFr> = {
+        let mut leaves = Vec::with_capacity(1 << 20);
+        for i in 0..(1 << 20) {
+            let mut bytes = [0u8; 32];
+            bytes[28..].copy_from_slice(&(i as u32).to_be_bytes());
+            leaves.push(TestFr(bytes));
+        }
+        leaves
+    };
+    static ref INDICES: Vec<usize> = (0..(1 << 20)).collect();
 }
 
+const NOF_LEAVES: usize = 8192;
+
 pub fn optimal_merkle_tree_benchmark(c: &mut Criterion) {
     let mut tree =
-        OptimalMerkleTree::<Keccak256>::new(2, TestFr([0; 32]), OptimalMerkleConfig::default())
+        OptimalMerkleTree::<Keccak256>::new(20, TestFr([0; 32]), OptimalMerkleConfig::default())
             .unwrap();
 
+    for i in 0..NOF_LEAVES {
+        tree.set(i, LEAVES[i % LEAVES.len()]).unwrap();
+    }
+
     c.bench_function("OptimalMerkleTree::set", |b| {
+        let mut index = NOF_LEAVES;
         b.iter(|| {
-            tree.set(0, LEAVES[0]).unwrap();
+            tree.set(index % (1 << 20), LEAVES[index % LEAVES.len()])
+                .unwrap();
+            index = (index + 1) % (1 << 20);
         })
     });
 
     c.bench_function("OptimalMerkleTree::delete", |b| {
+        let mut index = 0;
         b.iter(|| {
-            tree.delete(0).unwrap();
+            tree.delete(index % NOF_LEAVES).unwrap();
+            tree.set(index % NOF_LEAVES, LEAVES[index % LEAVES.len()])
+                .unwrap();
+            index = (index + 1) % NOF_LEAVES;
         })
     });
 
     c.bench_function("OptimalMerkleTree::override_range", |b| {
+        let mut offset = 0;
         b.iter(|| {
-            tree.override_range(0, LEAVES.into_iter(), [0, 1, 2, 3].into_iter())
-                .unwrap();
-        })
-    });
-
-    c.bench_function("OptimalMerkleTree::compute_root", |b| {
-        b.iter(|| {
-            tree.compute_root().unwrap();
+            let range = offset..offset + NOF_LEAVES;
+            tree.override_range(
+                offset,
+                LEAVES[range.clone()].iter().cloned(),
+                INDICES[range.clone()].iter().cloned(),
+            )
+            .unwrap();
+            offset = (offset + NOF_LEAVES) % (1 << 20);
         })
     });
 
     c.bench_function("OptimalMerkleTree::get", |b| {
+        let mut index = 0;
         b.iter(|| {
-            tree.get(0).unwrap();
+            tree.get(index % NOF_LEAVES).unwrap();
+            index = (index + 1) % NOF_LEAVES;
         })
     });
 
-    // check intermediate node getter which required additional computation of sub root index
     c.bench_function("OptimalMerkleTree::get_subtree_root", |b| {
+        let mut level = 1;
+        let mut index = 0;
         b.iter(|| {
-            tree.get_subtree_root(1, 0).unwrap();
+            tree.get_subtree_root(level % 20, index % (1 << (20 - (level % 20))))
+                .unwrap();
+            index = (index + 1) % (1 << (20 - (level % 20)));
+            level = 1 + (level % 20);
         })
     });
 
@@ -108,43 +131,61 @@ pub fn optimal_merkle_tree_benchmark(c: &mut Criterion) {
 
 pub fn full_merkle_tree_benchmark(c: &mut Criterion) {
     let mut tree =
-        FullMerkleTree::<Keccak256>::new(2, TestFr([0; 32]), FullMerkleConfig::default()).unwrap();
+        FullMerkleTree::<Keccak256>::new(20, TestFr([0; 32]), FullMerkleConfig::default()).unwrap();
+
+    for i in 0..NOF_LEAVES {
+        tree.set(i, LEAVES[i % LEAVES.len()]).unwrap();
+    }
 
     c.bench_function("FullMerkleTree::set", |b| {
+        let mut index = NOF_LEAVES;
         b.iter(|| {
-            tree.set(0, LEAVES[0]).unwrap();
+            tree.set(index % (1 << 20), LEAVES[index % LEAVES.len()])
+                .unwrap();
+            index = (index + 1) % (1 << 20);
         })
     });
 
     c.bench_function("FullMerkleTree::delete", |b| {
+        let mut index = 0;
         b.iter(|| {
-            tree.delete(0).unwrap();
+            tree.delete(index % NOF_LEAVES).unwrap();
+            tree.set(index % NOF_LEAVES, LEAVES[index % LEAVES.len()])
+                .unwrap();
+            index = (index + 1) % NOF_LEAVES;
         })
     });
 
     c.bench_function("FullMerkleTree::override_range", |b| {
+        let mut offset = 0;
         b.iter(|| {
-            tree.override_range(0, LEAVES.into_iter(), [0, 1, 2, 3].into_iter())
-                .unwrap();
-        })
-    });
-
-    c.bench_function("FullMerkleTree::compute_root", |b| {
-        b.iter(|| {
-            tree.compute_root().unwrap();
+            let range = offset..offset + NOF_LEAVES;
+            tree.override_range(
+                offset,
+                LEAVES[range.clone()].iter().cloned(),
+                INDICES[range.clone()].iter().cloned(),
+            )
+            .unwrap();
+            offset = (offset + NOF_LEAVES) % (1 << 20);
         })
     });
 
     c.bench_function("FullMerkleTree::get", |b| {
+        let mut index = 0;
         b.iter(|| {
-            tree.get(0).unwrap();
+            tree.get(index % NOF_LEAVES).unwrap();
+            index = (index + 1) % NOF_LEAVES;
         })
     });
 
-    // check intermediate node getter which required additional computation of sub root index
     c.bench_function("FullMerkleTree::get_subtree_root", |b| {
+        let mut level = 1;
+        let mut index = 0;
         b.iter(|| {
-            tree.get_subtree_root(1, 0).unwrap();
+            tree.get_subtree_root(level % 20, index % (1 << (20 - (level % 20))))
+                .unwrap();
+            index = (index + 1) % (1 << (20 - (level % 20)));
+            level = 1 + (level % 20);
         })
     });
 

utils/benches/poseidon_benchmark.rs

@@ -1,7 +1,7 @@
+use std::hint::black_box;
+
 use ark_bn254::Fr;
-use criterion::{
-    black_box, criterion_group, criterion_main, BatchSize, BenchmarkId, Criterion, Throughput,
-};
+use criterion::{criterion_group, criterion_main, BatchSize, BenchmarkId, Criterion, Throughput};
 use zerokit_utils::Poseidon;
 
 const ROUND_PARAMS: [(usize, usize, usize, usize); 8] = [

utils/src/merkle_tree/error.rs

@@ -0,0 +1,31 @@
+#[derive(thiserror::Error, Debug)]
+pub enum ZerokitMerkleTreeError {
+    #[error("Invalid index")]
+    InvalidIndex,
+    // InvalidProof,
+    #[error("Leaf index out of bounds")]
+    InvalidLeaf,
+    #[error("Level exceeds tree depth")]
+    InvalidLevel,
+    #[error("Subtree index out of bounds")]
+    InvalidSubTreeIndex,
+    #[error("Start level is != from end level")]
+    InvalidStartAndEndLevel,
+    #[error("set_range got too many leaves")]
+    TooManySet,
+    #[error("Unknown error while computing merkle proof")]
+    ComputingProofError,
+    #[error("Invalid witness length (!= tree depth)")]
+    InvalidWitness,
+    #[cfg(feature = "pmtree-ft")]
+    #[error("Pmtree error: {0}")]
+    PmtreeErrorKind(#[from] pmtree::PmtreeErrorKind),
+}
+
+#[derive(Debug, thiserror::Error)]
+pub enum FromConfigError {
+    #[error("Error while reading pmtree config: {0}")]
+    JsonError(#[from] serde_json::Error),
+    #[error("Error while creating pmtree config: path already exists")]
+    PathExists,
+}

utils/src/merkle_tree/full_merkle_tree.rs

@@ -1,28 +1,29 @@
-use crate::merkle_tree::{FrOf, Hasher, ZerokitMerkleProof, ZerokitMerkleTree};
-use color_eyre::{Report, Result};
 use std::{
     cmp::max,
     fmt::Debug,
-    iter::{once, repeat_n, successors},
+    iter::{once, repeat_n},
     str::FromStr,
 };
 
+use rayon::iter::{IntoParallelIterator, ParallelIterator};
+
+use crate::merkle_tree::{
+    error::{FromConfigError, ZerokitMerkleTreeError},
+    FrOf, Hasher, ZerokitMerkleProof, ZerokitMerkleTree, MIN_PARALLEL_NODES,
+};
 ////////////////////////////////////////////////////////////
 ///// Full Merkle Tree Implementation
 ////////////////////////////////////////////////////////////
 
 /// Merkle tree with all leaf and intermediate hashes stored
 #[derive(Clone, PartialEq, Eq, Debug)]
-pub struct FullMerkleTree<H: Hasher> {
+pub struct FullMerkleTree<H>
+where
+    H: Hasher,
+{
     /// The depth of the tree, i.e. the number of levels from leaf to root
     depth: usize,
 
-    /// The nodes cached from the empty part of the tree (where leaves are set to default).
-    /// Since the rightmost part of the tree is usually changed much later than its creation,
-    /// we can prove accumulation of elements in the leftmost part, with no need to initialize the full tree
-    /// and by caching few intermediate nodes to the root computed from default leaves
-    cached_nodes: Vec<H::Fr>,
-
     /// The tree nodes
     nodes: Vec<H::Fr>,
 
@@ -30,11 +31,11 @@ pub struct FullMerkleTree<H: Hasher> {
     /// Set to 0 if the leaf is empty and set to 1 in otherwise.
     cached_leaves_indices: Vec<u8>,
 
-    // The next available (i.e., never used) tree index. Equivalently, the number of leaves added to the tree
-    // (deletions leave next_index unchanged)
+    /// The next available (i.e., never used) tree index. Equivalently, the number of leaves added to the tree
+    /// (deletions leave next_index unchanged)
     next_index: usize,
 
-    // metadata that an application may use to store additional information
+    /// metadata that an application may use to store additional information
     metadata: Vec<u8>,
 }
 
@@ -56,9 +57,9 @@ pub struct FullMerkleProof<H: Hasher>(pub Vec<FullMerkleBranch<H>>);
 pub struct FullMerkleConfig(());
 
 impl FromStr for FullMerkleConfig {
-    type Err = Report;
+    type Err = FromConfigError;
 
-    fn from_str(_s: &str) -> Result<Self> {
+    fn from_str(_s: &str) -> Result<Self, Self::Err> {
         Ok(FullMerkleConfig::default())
     }
 }
@@ -72,85 +73,89 @@ where
     type Hasher = H;
     type Config = FullMerkleConfig;
 
-    fn default(depth: usize) -> Result<Self> {
+    fn default(depth: usize) -> Result<Self, ZerokitMerkleTreeError> {
         FullMerkleTree::<H>::new(depth, Self::Hasher::default_leaf(), Self::Config::default())
     }
 
     /// Creates a new `MerkleTree`
     /// depth - the height of the tree made only of hash nodes. 2^depth is the maximum number of leaves hash nodes
-    fn new(depth: usize, initial_leaf: FrOf<Self::Hasher>, _config: Self::Config) -> Result<Self> {
+    fn new(
+        depth: usize,
+        default_leaf: FrOf<Self::Hasher>,
+        _config: Self::Config,
+    ) -> Result<Self, ZerokitMerkleTreeError> {
         // Compute cache node values, leaf to root
-        let cached_nodes = successors(Some(initial_leaf), |prev| Some(H::hash(&[*prev, *prev])))
-            .take(depth + 1)
-            .collect::<Vec<_>>();
+        let mut cached_nodes: Vec<H::Fr> = Vec::with_capacity(depth + 1);
+        cached_nodes.push(default_leaf);
+        for i in 0..depth {
+            cached_nodes.push(H::hash(&[cached_nodes[i]; 2]));
+        }
+        cached_nodes.reverse();
 
         // Compute node values
         let nodes = cached_nodes
             .iter()
-            .rev()
             .enumerate()
             .flat_map(|(levels, hash)| repeat_n(hash, 1 << levels))
             .cloned()
             .collect::<Vec<_>>();
         debug_assert!(nodes.len() == (1 << (depth + 1)) - 1);
 
-        let next_index = 0;
-
         Ok(Self {
             depth,
-            cached_nodes,
             nodes,
             cached_leaves_indices: vec![0; 1 << depth],
-            next_index,
+            next_index: 0,
             metadata: Vec::new(),
         })
     }
 
-    fn close_db_connection(&mut self) -> Result<()> {
+    fn close_db_connection(&mut self) -> Result<(), ZerokitMerkleTreeError> {
         Ok(())
     }
 
-    // Returns the depth of the tree
+    /// Returns the depth of the tree
     fn depth(&self) -> usize {
         self.depth
     }
 
-    // Returns the capacity of the tree, i.e. the maximum number of accumulatable leaves
+    /// Returns the capacity of the tree, i.e. the maximum number of accumulatable leaves
     fn capacity(&self) -> usize {
         1 << self.depth
     }
 
-    // Returns the total number of leaves set
+    /// Returns the total number of leaves set
     fn leaves_set(&self) -> usize {
         self.next_index
     }
 
-    // Returns the root of the tree
+    /// Returns the root of the tree
     fn root(&self) -> FrOf<Self::Hasher> {
         self.nodes[0]
     }
 
-    // Sets a leaf at the specified tree index
-    fn set(&mut self, leaf: usize, hash: FrOf<Self::Hasher>) -> Result<()> {
+    /// Sets a leaf at the specified tree index
+    fn set(&mut self, leaf: usize, hash: FrOf<Self::Hasher>) -> Result<(), ZerokitMerkleTreeError> {
         self.set_range(leaf, once(hash))?;
         self.next_index = max(self.next_index, leaf + 1);
         Ok(())
     }
 
-    // Get a leaf from the specified tree index
-    fn get(&self, leaf: usize) -> Result<FrOf<Self::Hasher>> {
+    /// Get a leaf from the specified tree index
+    fn get(&self, leaf: usize) -> Result<FrOf<Self::Hasher>, ZerokitMerkleTreeError> {
         if leaf >= self.capacity() {
-            return Err(Report::msg("leaf index out of bounds"));
+            return Err(ZerokitMerkleTreeError::InvalidLeaf);
         }
         Ok(self.nodes[self.capacity() + leaf - 1])
     }
 
-    fn get_subtree_root(&self, n: usize, index: usize) -> Result<H::Fr> {
+    /// Returns the root of the subtree at level n and index
+    fn get_subtree_root(&self, n: usize, index: usize) -> Result<H::Fr, ZerokitMerkleTreeError> {
         if n > self.depth() {
-            return Err(Report::msg("level exceeds depth size"));
+            return Err(ZerokitMerkleTreeError::InvalidIndex);
         }
         if index >= self.capacity() {
-            return Err(Report::msg("index exceeds set size"));
+            return Err(ZerokitMerkleTreeError::InvalidLeaf);
         }
         if n == 0 {
             Ok(self.root())
@@ -160,7 +165,7 @@ where
             let mut idx = self.capacity() + index - 1;
             let mut nd = self.depth;
             loop {
-                let parent = self.parent(idx).unwrap();
+                let parent = self.parent(idx).expect("parent should exist");
                 nd -= 1;
                 if nd == n {
                     return Ok(self.nodes[parent]);
@@ -171,6 +176,8 @@ where
             }
         }
     }
+
+    /// Returns the indices of the leaves that are empty
     fn get_empty_leaves_indices(&self) -> Vec<usize> {
         self.cached_leaves_indices
             .iter()
@@ -181,40 +188,45 @@ where
             .collect()
     }
 
-    // Sets tree nodes, starting from start index
-    // Function proper of FullMerkleTree implementation
-    fn set_range<I: IntoIterator<Item = FrOf<Self::Hasher>>>(
+    /// Sets multiple leaves from the specified tree index
+    fn set_range<I: ExactSizeIterator<Item = FrOf<Self::Hasher>>>(
         &mut self,
         start: usize,
-        hashes: I,
-    ) -> Result<()> {
+        leaves: I,
+    ) -> Result<(), ZerokitMerkleTreeError> {
         let index = self.capacity() + start - 1;
         let mut count = 0;
-        // first count number of hashes, and check that they fit in the tree
+        // first count number of leaves, and check that they fit in the tree
         // then insert into the tree
-        let hashes = hashes.into_iter().collect::<Vec<_>>();
-        if hashes.len() + start > self.capacity() {
-            return Err(Report::msg("provided hashes do not fit in the tree"));
+        let leaves = leaves.into_iter().collect::<Vec<_>>();
+        if leaves.len() + start > self.capacity() {
+            return Err(ZerokitMerkleTreeError::TooManySet);
         }
-        hashes.into_iter().for_each(|hash| {
+        leaves.into_iter().for_each(|hash| {
             self.nodes[index + count] = hash;
             self.cached_leaves_indices[start + count] = 1;
             count += 1;
         });
         if count != 0 {
-            self.update_nodes(index, index + (count - 1))?;
+            self.update_hashes(index, index + (count - 1))?;
             self.next_index = max(self.next_index, start + count);
         }
         Ok(())
     }
 
-    fn override_range<I, J>(&mut self, start: usize, leaves: I, indices: J) -> Result<()>
+    /// Overrides a range of leaves while resetting specified indices to default and preserving unaffected values.
+    fn override_range<I, J>(
+        &mut self,
+        start: usize,
+        leaves: I,
+        indices: J,
+    ) -> Result<(), ZerokitMerkleTreeError>
     where
-        I: IntoIterator<Item = FrOf<Self::Hasher>>,
-        J: IntoIterator<Item = usize>,
+        I: ExactSizeIterator<Item = FrOf<Self::Hasher>>,
+        J: ExactSizeIterator<Item = usize>,
     {
         let indices = indices.into_iter().collect::<Vec<_>>();
-        let min_index = *indices.first().unwrap();
+        let min_index = *indices.first().expect("indices should not be empty");
         let leaves_vec = leaves.into_iter().collect::<Vec<_>>();
 
         let max_index = start + leaves_vec.len();
@@ -237,17 +249,16 @@ where
         }
 
         self.set_range(start, set_values.into_iter())
-            .map_err(|e| Report::msg(e.to_string()))
     }
 
-    // Sets a leaf at the next available index
-    fn update_next(&mut self, leaf: FrOf<Self::Hasher>) -> Result<()> {
+    /// Sets a leaf at the next available index
+    fn update_next(&mut self, leaf: FrOf<Self::Hasher>) -> Result<(), ZerokitMerkleTreeError> {
         self.set(self.next_index, leaf)?;
         Ok(())
     }
 
-    // Deletes a leaf at a certain index by setting it to its default value (next_index is not updated)
-    fn delete(&mut self, index: usize) -> Result<()> {
+    /// Deletes a leaf at a certain index by setting it to its default value (next_index is not updated)
+    fn delete(&mut self, index: usize) -> Result<(), ZerokitMerkleTreeError> {
         // We reset the leaf only if we previously set a leaf at that index
         if index < self.next_index {
             self.set(index, H::default_leaf())?;
@@ -257,9 +268,9 @@ where
     }
 
     // Computes a merkle proof the leaf at the specified index
-    fn proof(&self, leaf: usize) -> Result<FullMerkleProof<H>> {
+    fn proof(&self, leaf: usize) -> Result<FullMerkleProof<H>, ZerokitMerkleTreeError> {
         if leaf >= self.capacity() {
-            return Err(Report::msg("index exceeds set size"));
+            return Err(ZerokitMerkleTreeError::InvalidLeaf);
         }
         let mut index = self.capacity() + leaf - 1;
         let mut path = Vec::with_capacity(self.depth + 1);
@@ -276,30 +287,29 @@ where
     }
 
     // Verifies a Merkle proof with respect to the input leaf and the tree root
-    fn verify(&self, hash: &FrOf<Self::Hasher>, proof: &FullMerkleProof<H>) -> Result<bool> {
+    fn verify(
+        &self,
+        hash: &FrOf<Self::Hasher>,
+        proof: &FullMerkleProof<H>,
+    ) -> Result<bool, ZerokitMerkleTreeError> {
         Ok(proof.compute_root_from(hash) == self.root())
     }
 
-    fn compute_root(&mut self) -> Result<FrOf<Self::Hasher>> {
-        Ok(self.root())
-    }
-
-    fn set_metadata(&mut self, metadata: &[u8]) -> Result<()> {
+    fn set_metadata(&mut self, metadata: &[u8]) -> Result<(), ZerokitMerkleTreeError> {
         self.metadata = metadata.to_vec();
         Ok(())
     }
 
-    fn metadata(&self) -> Result<Vec<u8>> {
+    fn metadata(&self) -> Result<Vec<u8>, ZerokitMerkleTreeError> {
         Ok(self.metadata.to_vec())
     }
 }
 
+// Utilities for updating the tree nodes
 impl<H: Hasher> FullMerkleTree<H>
 where
     H: Hasher,
 {
-    // Utilities for updating the tree nodes
-
     /// For a given node index, return the parent node index
     /// Returns None if there is no parent (root node)
     fn parent(&self, index: usize) -> Option<usize> {
@@ -315,23 +325,60 @@ where
         (index << 1) + 1
     }
 
+    /// Returns the depth level of a node based on its index in the flattened tree.
     fn levels(&self, index: usize) -> usize {
         // `n.next_power_of_two()` will return `n` iff `n` is a power of two.
         // The extra offset corrects this.
         (index + 2).next_power_of_two().trailing_zeros() as usize - 1
     }
 
-    fn update_nodes(&mut self, start: usize, end: usize) -> Result<()> {
-        if self.levels(start) != self.levels(end) {
-            return Err(Report::msg("self.levels(start) != self.levels(end)"));
+    /// Updates parent hashes after modifying a range of nodes at the same level.
+    ///
+    /// - `start_index`: The first index at the current level that was updated.
+    /// - `end_index`: The last index (inclusive) at the same level that was updated.
+    fn update_hashes(
+        &mut self,
+        start_index: usize,
+        end_index: usize,
+    ) -> Result<(), ZerokitMerkleTreeError> {
+        // Ensure the range is within the same tree level
+        if self.levels(start_index) != self.levels(end_index) {
+            return Err(ZerokitMerkleTreeError::InvalidStartAndEndLevel);
         }
-        if let (Some(start), Some(end)) = (self.parent(start), self.parent(end)) {
-            for parent in start..=end {
-                let child = self.first_child(parent);
-                self.nodes[parent] = H::hash(&[self.nodes[child], self.nodes[child + 1]]);
+
+        // Compute parent indices for the range
+        if let (Some(start_parent), Some(end_parent)) =
+            (self.parent(start_index), self.parent(end_index))
+        {
+            // Use parallel processing when the number of pairs exceeds the threshold
+            if end_parent - start_parent + 1 >= MIN_PARALLEL_NODES {
+                let updates: Vec<(usize, H::Fr)> = (start_parent..=end_parent)
+                    .into_par_iter()
+                    .map(|parent| {
+                        let left_child = self.first_child(parent);
+                        let right_child = left_child + 1;
+                        let hash = H::hash(&[self.nodes[left_child], self.nodes[right_child]]);
+                        (parent, hash)
+                    })
+                    .collect();
+
+                for (parent, hash) in updates {
+                    self.nodes[parent] = hash;
+                }
+            } else {
+                // Otherwise, fallback to sequential update for small ranges
+                for parent in start_parent..=end_parent {
+                    let left_child = self.first_child(parent);
+                    let right_child = left_child + 1;
+                    self.nodes[parent] =
+                        H::hash(&[self.nodes[left_child], self.nodes[right_child]]);
+                }
             }
-            self.update_nodes(start, end)?;
+
+            // Recurse to update upper levels
+            self.update_hashes(start_parent, end_parent)?;
         }
+
         Ok(())
     }
 }

utils/src/merkle_tree/merkle_tree.rs

@@ -8,20 +8,25 @@
 // and https://github.com/worldcoin/semaphore-rs/blob/d462a4372f1fd9c27610f2acfe4841fab1d396aa/src/merkle_tree.rs
 
 //!
-//! # To do
+//! # TODO
 //!
 //! * Disk based storage backend (using mmaped files should be easy)
 //! * Implement serialization for tree and Merkle proof
 
-use std::str::FromStr;
+use crate::merkle_tree::error::ZerokitMerkleTreeError;
+use std::{
+    fmt::{Debug, Display},
+    str::FromStr,
+};
 
-use color_eyre::Result;
+/// Enables parallel hashing when there are at least 8 nodes (4 pairs to hash), justifying the overhead.
+pub const MIN_PARALLEL_NODES: usize = 8;
 
 /// In the Hasher trait we define the node type, the default leaf
 /// and the hash function used to initialize a Merkle Tree implementation
 pub trait Hasher {
     /// Type of the leaf and tree node
-    type Fr: Clone + Copy + Eq + Default + std::fmt::Debug + std::fmt::Display + FromStr;
+    type Fr: Clone + Copy + Eq + Default + Debug + Display + FromStr + Send + Sync;
 
     /// Returns the default tree leaf
     fn default_leaf() -> Self::Fr;
@@ -39,35 +44,52 @@ pub trait ZerokitMerkleTree {
     type Hasher: Hasher;
     type Config: Default + FromStr;
 
-    fn default(depth: usize) -> Result<Self>
+    fn default(depth: usize) -> Result<Self, ZerokitMerkleTreeError>
     where
         Self: Sized;
-    fn new(depth: usize, default_leaf: FrOf<Self::Hasher>, config: Self::Config) -> Result<Self>
+    fn new(
+        depth: usize,
+        default_leaf: FrOf<Self::Hasher>,
+        config: Self::Config,
+    ) -> Result<Self, ZerokitMerkleTreeError>
     where
         Self: Sized;
     fn depth(&self) -> usize;
     fn capacity(&self) -> usize;
     fn leaves_set(&self) -> usize;
     fn root(&self) -> FrOf<Self::Hasher>;
-    fn compute_root(&mut self) -> Result<FrOf<Self::Hasher>>;
-    fn get_subtree_root(&self, n: usize, index: usize) -> Result<FrOf<Self::Hasher>>;
-    fn set(&mut self, index: usize, leaf: FrOf<Self::Hasher>) -> Result<()>;
-    fn set_range<I>(&mut self, start: usize, leaves: I) -> Result<()>
+    fn get_subtree_root(
+        &self,
+        n: usize,
+        index: usize,
+    ) -> Result<FrOf<Self::Hasher>, ZerokitMerkleTreeError>;
+    fn set(&mut self, index: usize, leaf: FrOf<Self::Hasher>)
+        -> Result<(), ZerokitMerkleTreeError>;
+    fn set_range<I>(&mut self, start: usize, leaves: I) -> Result<(), ZerokitMerkleTreeError>
     where
         I: ExactSizeIterator<Item = FrOf<Self::Hasher>>;
-    fn get(&self, index: usize) -> Result<FrOf<Self::Hasher>>;
+    fn get(&self, index: usize) -> Result<FrOf<Self::Hasher>, ZerokitMerkleTreeError>;
     fn get_empty_leaves_indices(&self) -> Vec<usize>;
-    fn override_range<I, J>(&mut self, start: usize, leaves: I, to_remove_indices: J) -> Result<()>
+    fn override_range<I, J>(
+        &mut self,
+        start: usize,
+        leaves: I,
+        to_remove_indices: J,
+    ) -> Result<(), ZerokitMerkleTreeError>
     where
         I: ExactSizeIterator<Item = FrOf<Self::Hasher>>,
         J: ExactSizeIterator<Item = usize>;
-    fn update_next(&mut self, leaf: FrOf<Self::Hasher>) -> Result<()>;
-    fn delete(&mut self, index: usize) -> Result<()>;
-    fn proof(&self, index: usize) -> Result<Self::Proof>;
-    fn verify(&self, leaf: &FrOf<Self::Hasher>, witness: &Self::Proof) -> Result<bool>;
-    fn set_metadata(&mut self, metadata: &[u8]) -> Result<()>;
-    fn metadata(&self) -> Result<Vec<u8>>;
-    fn close_db_connection(&mut self) -> Result<()>;
+    fn update_next(&mut self, leaf: FrOf<Self::Hasher>) -> Result<(), ZerokitMerkleTreeError>;
+    fn delete(&mut self, index: usize) -> Result<(), ZerokitMerkleTreeError>;
+    fn proof(&self, index: usize) -> Result<Self::Proof, ZerokitMerkleTreeError>;
+    fn verify(
+        &self,
+        leaf: &FrOf<Self::Hasher>,
+        witness: &Self::Proof,
+    ) -> Result<bool, ZerokitMerkleTreeError>;
+    fn set_metadata(&mut self, metadata: &[u8]) -> Result<(), ZerokitMerkleTreeError>;
+    fn metadata(&self) -> Result<Vec<u8>, ZerokitMerkleTreeError>;
+    fn close_db_connection(&mut self) -> Result<(), ZerokitMerkleTreeError>;
 }
 
 pub trait ZerokitMerkleProof {

utils/src/merkle_tree/mod.rs

@@ -1,7 +1,11 @@
+pub mod error;
 pub mod full_merkle_tree;
 #[allow(clippy::module_inception)]
 pub mod merkle_tree;
 pub mod optimal_merkle_tree;
-pub use self::full_merkle_tree::*;
-pub use self::merkle_tree::*;
-pub use self::optimal_merkle_tree::*;
+
+pub use self::full_merkle_tree::{FullMerkleConfig, FullMerkleProof, FullMerkleTree};
+pub use self::merkle_tree::{
+    FrOf, Hasher, ZerokitMerkleProof, ZerokitMerkleTree, MIN_PARALLEL_NODES,
+};
+pub use self::optimal_merkle_tree::{OptimalMerkleConfig, OptimalMerkleProof, OptimalMerkleTree};

utils/src/merkle_tree/optimal_merkle_tree.rs

@@ -1,10 +1,11 @@
-use crate::merkle_tree::{Hasher, ZerokitMerkleProof, ZerokitMerkleTree};
-use crate::FrOf;
-use color_eyre::{Report, Result};
-use std::cmp::min;
-use std::collections::HashMap;
-use std::str::FromStr;
-use std::{cmp::max, fmt::Debug};
+use std::{cmp::max, collections::HashMap, fmt::Debug, str::FromStr};
+
+use rayon::iter::{IntoParallelIterator, ParallelIterator};
+
+use crate::merkle_tree::{
+    error::{FromConfigError, ZerokitMerkleTreeError},
+    FrOf, Hasher, ZerokitMerkleProof, ZerokitMerkleTree, MIN_PARALLEL_NODES,
+};
 ////////////////////////////////////////////////////////////
 ///// Optimal Merkle Tree Implementation
 ////////////////////////////////////////////////////////////
@@ -31,11 +32,11 @@ where
     /// Set to 0 if the leaf is empty and set to 1 in otherwise.
     cached_leaves_indices: Vec<u8>,
 
-    // The next available (i.e., never used) tree index. Equivalently, the number of leaves added to the tree
-    // (deletions leave next_index unchanged)
+    /// The next available (i.e., never used) tree index. Equivalently, the number of leaves added to the tree
+    /// (deletions leave next_index unchanged)
     next_index: usize,
 
-    // metadata that an application may use to store additional information
+    /// metadata that an application may use to store additional information
     metadata: Vec<u8>,
 }
 
@@ -48,17 +49,14 @@ pub struct OptimalMerkleProof<H: Hasher>(pub Vec<(H::Fr, u8)>);
 pub struct OptimalMerkleConfig(());
 
 impl FromStr for OptimalMerkleConfig {
-    type Err = Report;
+    type Err = FromConfigError;
 
-    fn from_str(_s: &str) -> Result<Self> {
+    fn from_str(_s: &str) -> Result<Self, Self::Err> {
         Ok(OptimalMerkleConfig::default())
     }
 }
 
-////////////////////////////////////////////////////////////
-///// Implementations
-////////////////////////////////////////////////////////////
-
+/// Implementations
 impl<H: Hasher> ZerokitMerkleTree for OptimalMerkleTree<H>
 where
     H: Hasher,
@@ -67,22 +65,28 @@ where
     type Hasher = H;
     type Config = OptimalMerkleConfig;
 
-    fn default(depth: usize) -> Result<Self> {
+    fn default(depth: usize) -> Result<Self, ZerokitMerkleTreeError> {
         OptimalMerkleTree::<H>::new(depth, H::default_leaf(), Self::Config::default())
     }
 
     /// Creates a new `MerkleTree`
     /// depth - the height of the tree made only of hash nodes. 2^depth is the maximum number of leaves hash nodes
-    fn new(depth: usize, default_leaf: H::Fr, _config: Self::Config) -> Result<Self> {
+    fn new(
+        depth: usize,
+        default_leaf: H::Fr,
+        _config: Self::Config,
+    ) -> Result<Self, ZerokitMerkleTreeError> {
+        // Compute cache node values, leaf to root
         let mut cached_nodes: Vec<H::Fr> = Vec::with_capacity(depth + 1);
         cached_nodes.push(default_leaf);
         for i in 0..depth {
             cached_nodes.push(H::hash(&[cached_nodes[i]; 2]));
         }
         cached_nodes.reverse();
+
         Ok(OptimalMerkleTree {
-            cached_nodes,
             depth,
+            cached_nodes,
             nodes: HashMap::with_capacity(1 << depth),
             cached_leaves_indices: vec![0; 1 << depth],
             next_index: 0,
@@ -90,36 +94,57 @@ where
         })
     }
 
-    fn close_db_connection(&mut self) -> Result<()> {
+    fn close_db_connection(&mut self) -> Result<(), ZerokitMerkleTreeError> {
         Ok(())
     }
 
-    // Returns the depth of the tree
+    /// Returns the depth of the tree
     fn depth(&self) -> usize {
         self.depth
     }
 
-    // Returns the capacity of the tree, i.e. the maximum number of accumulatable leaves
+    /// Returns the capacity of the tree, i.e. the maximum number of accumulatable leaves
     fn capacity(&self) -> usize {
         1 << self.depth
     }
 
-    // Returns the total number of leaves set
+    /// Returns the total number of leaves set
     fn leaves_set(&self) -> usize {
         self.next_index
     }
 
-    // Returns the root of the tree
+    /// Returns the root of the tree
     fn root(&self) -> H::Fr {
         self.get_node(0, 0)
     }
 
-    fn get_subtree_root(&self, n: usize, index: usize) -> Result<H::Fr> {
+    /// Sets a leaf at the specified tree index
+    fn set(&mut self, index: usize, leaf: H::Fr) -> Result<(), ZerokitMerkleTreeError> {
+        if index >= self.capacity() {
+            return Err(ZerokitMerkleTreeError::InvalidLeaf);
+        }
+        self.nodes.insert((self.depth, index), leaf);
+        self.update_hashes(index, 1)?;
+        self.next_index = max(self.next_index, index + 1);
+        self.cached_leaves_indices[index] = 1;
+        Ok(())
+    }
+
+    /// Get a leaf from the specified tree index
+    fn get(&self, index: usize) -> Result<H::Fr, ZerokitMerkleTreeError> {
+        if index >= self.capacity() {
+            return Err(ZerokitMerkleTreeError::InvalidLeaf);
+        }
+        Ok(self.get_node(self.depth, index))
+    }
+
+    /// Returns the root of the subtree at level n and index
+    fn get_subtree_root(&self, n: usize, index: usize) -> Result<H::Fr, ZerokitMerkleTreeError> {
         if n > self.depth() {
-            return Err(Report::msg("level exceeds depth size"));
+            return Err(ZerokitMerkleTreeError::InvalidLevel);
         }
         if index >= self.capacity() {
-            return Err(Report::msg("index exceeds set size"));
+            return Err(ZerokitMerkleTreeError::InvalidLeaf);
         }
         if n == 0 {
             Ok(self.root())
@@ -130,26 +155,7 @@ where
         }
     }
 
-    // Sets a leaf at the specified tree index
-    fn set(&mut self, index: usize, leaf: H::Fr) -> Result<()> {
-        if index >= self.capacity() {
-            return Err(Report::msg("index exceeds set size"));
-        }
-        self.nodes.insert((self.depth, index), leaf);
-        self.update_hashes(index, 1)?;
-        self.next_index = max(self.next_index, index + 1);
-        self.cached_leaves_indices[index] = 1;
-        Ok(())
-    }
-
-    // Get a leaf from the specified tree index
-    fn get(&self, index: usize) -> Result<H::Fr> {
-        if index >= self.capacity() {
-            return Err(Report::msg("index exceeds set size"));
-        }
-        Ok(self.get_node(self.depth, index))
-    }
-
+    /// Returns the indices of the leaves that are empty
     fn get_empty_leaves_indices(&self) -> Vec<usize> {
         self.cached_leaves_indices
             .iter()
@@ -160,16 +166,16 @@ where
             .collect()
     }
 
-    // Sets multiple leaves from the specified tree index
+    /// Sets multiple leaves from the specified tree index
     fn set_range<I: ExactSizeIterator<Item = H::Fr>>(
         &mut self,
         start: usize,
         leaves: I,
-    ) -> Result<()> {
+    ) -> Result<(), ZerokitMerkleTreeError> {
         // check if the range is valid
         let leaves_len = leaves.len();
         if start + leaves_len > self.capacity() {
-            return Err(Report::msg("provided range exceeds set size"));
+            return Err(ZerokitMerkleTreeError::TooManySet);
         }
         for (i, leaf) in leaves.enumerate() {
             self.nodes.insert((self.depth, start + i), leaf);
@@ -180,13 +186,19 @@ where
         Ok(())
     }
 
-    fn override_range<I, J>(&mut self, start: usize, leaves: I, indices: J) -> Result<()>
+    /// Overrides a range of leaves while resetting specified indices to default and preserving unaffected values.
+    fn override_range<I, J>(
+        &mut self,
+        start: usize,
+        leaves: I,
+        indices: J,
+    ) -> Result<(), ZerokitMerkleTreeError>
     where
         I: ExactSizeIterator<Item = FrOf<Self::Hasher>>,
         J: ExactSizeIterator<Item = usize>,
     {
         let indices = indices.into_iter().collect::<Vec<_>>();
-        let min_index = *indices.first().unwrap();
+        let min_index = *indices.first().expect("indices should not be empty");
         let leaves_vec = leaves.into_iter().collect::<Vec<_>>();
 
         let max_index = start + leaves_vec.len();
@@ -195,7 +207,7 @@ where
 
         for i in min_index..start {
             if !indices.contains(&i) {
-                let value = self.get_leaf(i);
+                let value = self.get(i)?;
                 set_values[i - min_index] = value;
             }
         }
@@ -209,17 +221,16 @@ where
         }
 
         self.set_range(start, set_values.into_iter())
-            .map_err(|e| Report::msg(e.to_string()))
     }
 
-    // Sets a leaf at the next available index
-    fn update_next(&mut self, leaf: H::Fr) -> Result<()> {
+    /// Sets a leaf at the next available index
+    fn update_next(&mut self, leaf: H::Fr) -> Result<(), ZerokitMerkleTreeError> {
         self.set(self.next_index, leaf)?;
         Ok(())
     }
 
-    // Deletes a leaf at a certain index by setting it to its default value (next_index is not updated)
-    fn delete(&mut self, index: usize) -> Result<()> {
+    /// Deletes a leaf at a certain index by setting it to its default value (next_index is not updated)
+    fn delete(&mut self, index: usize) -> Result<(), ZerokitMerkleTreeError> {
         // We reset the leaf only if we previously set a leaf at that index
         if index < self.next_index {
             self.set(index, H::default_leaf())?;
@@ -228,17 +239,20 @@ where
         Ok(())
     }
 
-    // Computes a merkle proof the leaf at the specified index
-    fn proof(&self, index: usize) -> Result<Self::Proof> {
+    /// Computes a merkle proof the leaf at the specified index
+    fn proof(&self, index: usize) -> Result<Self::Proof, ZerokitMerkleTreeError> {
         if index >= self.capacity() {
-            return Err(Report::msg("index exceeds set size"));
+            return Err(ZerokitMerkleTreeError::InvalidLeaf);
         }
         let mut witness = Vec::<(H::Fr, u8)>::with_capacity(self.depth);
         let mut i = index;
         let mut depth = self.depth;
         loop {
             i ^= 1;
-            witness.push((self.get_node(depth, i), (1 - (i & 1)).try_into().unwrap()));
+            witness.push((
+                self.get_node(depth, i),
+                (1 - (i & 1)).try_into().expect("0 or 1 expected"),
+            ));
             i >>= 1;
             depth -= 1;
             if depth == 0 {
@@ -246,130 +260,100 @@ where
             }
         }
         if i != 0 {
-            Err(Report::msg("i != 0"))
+            Err(ZerokitMerkleTreeError::ComputingProofError)
         } else {
             Ok(OptimalMerkleProof(witness))
         }
     }
 
-    // Verifies a Merkle proof with respect to the input leaf and the tree root
-    fn verify(&self, leaf: &H::Fr, witness: &Self::Proof) -> Result<bool> {
+    /// Verifies a Merkle proof with respect to the input leaf and the tree root
+    fn verify(&self, leaf: &H::Fr, witness: &Self::Proof) -> Result<bool, ZerokitMerkleTreeError> {
         if witness.length() != self.depth {
-            return Err(Report::msg("witness length doesn't match tree depth"));
+            return Err(ZerokitMerkleTreeError::InvalidWitness);
         }
         let expected_root = witness.compute_root_from(leaf);
         Ok(expected_root.eq(&self.root()))
     }
 
-    fn compute_root(&mut self) -> Result<FrOf<Self::Hasher>> {
-        self.recalculate_from(0)?;
-        Ok(self.root())
-    }
-
-    fn set_metadata(&mut self, metadata: &[u8]) -> Result<()> {
+    fn set_metadata(&mut self, metadata: &[u8]) -> Result<(), ZerokitMerkleTreeError> {
         self.metadata = metadata.to_vec();
         Ok(())
     }
 
-    fn metadata(&self) -> Result<Vec<u8>> {
+    fn metadata(&self) -> Result<Vec<u8>, ZerokitMerkleTreeError> {
         Ok(self.metadata.to_vec())
     }
 }
 
+// Utilities for updating the tree nodes
 impl<H: Hasher> OptimalMerkleTree<H>
 where
     H: Hasher,
 {
-    // Utilities for updating the tree nodes
-
+    /// Returns the value of a node at a specific (depth, index).
+    /// Falls back to a cached default if the node hasn't been set.
     fn get_node(&self, depth: usize, index: usize) -> H::Fr {
-        let node = *self
+        *self
             .nodes
             .get(&(depth, index))
-            .unwrap_or_else(|| &self.cached_nodes[depth]);
-        node
+            .unwrap_or_else(|| &self.cached_nodes[depth])
     }
 
-    pub fn get_leaf(&self, index: usize) -> H::Fr {
-        self.get_node(self.depth, index)
-    }
-
-    fn hash_couple(&mut self, depth: usize, index: usize) -> H::Fr {
+    /// Computes the hash of a node’s two children at the given depth.
+    /// If the index is odd, it is rounded down to the nearest even index.
+    fn hash_couple(&self, depth: usize, index: usize) -> H::Fr {
         let b = index & !1;
         H::hash(&[self.get_node(depth, b), self.get_node(depth, b + 1)])
     }
 
-    fn recalculate_from(&mut self, index: usize) -> Result<()> {
-        let mut i = index;
-        let mut depth = self.depth;
-        loop {
-            let h = self.hash_couple(depth, i);
-            i >>= 1;
-            depth -= 1;
-            self.nodes.insert((depth, i), h);
-            self.cached_leaves_indices[index] = 1;
-            if depth == 0 {
-                break;
-            }
-        }
-        if depth != 0 {
-            return Err(Report::msg("did not reach the depth"));
-        }
-        if i != 0 {
-            return Err(Report::msg("did not go through all indexes"));
-        }
-        Ok(())
-    }
-
-    /// Update hashes after some leaves have been set or updated
-    /// index - first leaf index (which has been set or updated)
-    /// length - number of elements set or updated
-    fn update_hashes(&mut self, index: usize, length: usize) -> Result<()> {
-        // parent depth & index (used to store in the tree)
-        let mut parent_depth = self.depth - 1; // tree depth (or leaves depth) - 1
-        let mut parent_index = index >> 1;
-        let mut parent_index_bak = parent_index;
-        // maximum index at this depth
-        let parent_max_index_0 = (1 << parent_depth) / 2;
-        // Based on given length (number of elements we will update)
-        // we could restrict the parent_max_index
-        let current_index_max = if (index + length) % 2 == 0 {
-            index + length + 2
-        } else {
-            index + length + 1
-        };
-        let mut parent_max_index = min(current_index_max >> 1, parent_max_index_0);
-
-        // current depth & index (used to compute the hash)
-        // current depth initially == tree depth (or leaves depth)
+    /// Updates parent hashes after modifying a range of leaf nodes.
+    ///
+    /// - `start`: Starting leaf index that was updated.
+    /// - `length`: Number of consecutive leaves that were updated.
+    fn update_hashes(&mut self, start: usize, length: usize) -> Result<(), ZerokitMerkleTreeError> {
+        // Start at the leaf level
         let mut current_depth = self.depth;
-        let mut current_index = if index % 2 == 0 { index } else { index - 1 };
-        let mut current_index_bak = current_index;
 
-        loop {
-            // Hash 2 values at (current depth, current_index) & (current_depth, current_index + 1)
-            let n_hash = self.hash_couple(current_depth, current_index);
-            // Insert this hash at (parent_depth, parent_index)
-            self.nodes.insert((parent_depth, parent_index), n_hash);
+        // Round down to include the left sibling in the pair (if start is odd)
+        let mut current_index = start & !1;
 
-            if parent_depth == 0 {
-                // We just set the root hash of the tree - nothing to do anymore
-                break;
-            }
-            // Incr parent index
-            parent_index += 1;
-            // Incr current index (+2 because we've just hashed current index & current_index + 1)
-            current_index += 2;
-            if parent_index >= parent_max_index {
-                // reset (aka decr depth & reset indexes)
-                parent_depth -= 1;
-                parent_index = parent_index_bak >> 1;
-                parent_index_bak = parent_index;
-                parent_max_index >>= 1;
-                current_depth -= 1;
-                current_index = current_index_bak >> 1;
-                current_index_bak = current_index;
+        // Compute the max index at this level, round up to include the last updated leaf’s right sibling (if start + length is odd)
+        let mut current_index_max = (start + length + 1) & !1;
+
+        // Traverse from the leaf level up to the root
+        while current_depth > 0 {
+            // Compute the parent level (one level above the current)
+            let parent_depth = current_depth - 1;
+
+            // Use parallel processing when the number of pairs exceeds the threshold
+            if current_index_max - current_index >= MIN_PARALLEL_NODES {
+                let updates: Vec<((usize, usize), H::Fr)> = (current_index..current_index_max)
+                    .step_by(2)
+                    .collect::<Vec<_>>()
+                    .into_par_iter()
+                    .map(|index| {
+                        // Hash two child nodes at positions (current_depth, index) and (current_depth, index + 1)
+                        let hash = self.hash_couple(current_depth, index);
+                        // Return the computed parent hash and its position at
+                        ((parent_depth, index >> 1), hash)
+                    })
+                    .collect();
+
+                for (parent, hash) in updates {
+                    self.nodes.insert(parent, hash);
+                }
+            } else {
+                // Otherwise, fallback to sequential update for small ranges
+                for index in (current_index..current_index_max).step_by(2) {
+                    let hash = self.hash_couple(current_depth, index);
+                    self.nodes.insert((parent_depth, index >> 1), hash);
+                }
             }
+
+            // Move up one level in the tree
+            current_index >>= 1;
+            current_index_max = (current_index_max + 1) >> 1;
+            current_depth -= 1;
         }
 
         Ok(())
@@ -383,7 +367,7 @@ where
     type Index = u8;
     type Hasher = H;
 
-    // Returns the length of a Merkle proof
+    /// Returns the length of a Merkle proof
     fn length(&self) -> usize {
         self.0.len()
     }

utils/src/pm_tree/mod.rs

@@ -1,4 +1,4 @@
 pub mod sled_adapter;
-pub use self::sled_adapter::*;
+pub use self::sled_adapter::SledDB;
 pub use pmtree;
-pub use sled::*;
+pub use sled::{Config, Mode};

utils/src/poseidon/mod.rs

@@ -1,5 +1,4 @@
 pub mod poseidon_hash;
-pub use self::poseidon_hash::*;
+pub use poseidon_hash::Poseidon;
 
 pub mod poseidon_constants;
-pub use self::poseidon_constants::*;

utils/tests/merkle_tree.rs

@@ -6,7 +6,7 @@ pub mod test {
     use tiny_keccak::{Hasher as _, Keccak};
     use zerokit_utils::{
         FullMerkleConfig, FullMerkleTree, Hasher, OptimalMerkleConfig, OptimalMerkleTree,
-        ZerokitMerkleProof, ZerokitMerkleTree,
+        ZerokitMerkleProof, ZerokitMerkleTree, MIN_PARALLEL_NODES,
     };
     #[derive(Clone, Copy, Eq, PartialEq)]
     struct Keccak256;
@@ -42,7 +42,7 @@ pub mod test {
         type Err = std::string::FromUtf8Error;
 
         fn from_str(s: &str) -> Result<Self, Self::Err> {
-            Ok(TestFr(s.as_bytes().try_into().unwrap()))
+            Ok(TestFr(s.as_bytes().try_into().expect("Invalid length")))
         }
     }
 
@@ -50,7 +50,7 @@ pub mod test {
         fn from(value: u32) -> Self {
             let mut bytes: Vec<u8> = vec![0; 28];
             bytes.extend_from_slice(&value.to_be_bytes());
-            TestFr(bytes.as_slice().try_into().unwrap())
+            TestFr(bytes.as_slice().try_into().expect("Invalid length"))
         }
     }
 
@@ -58,12 +58,12 @@ pub mod test {
 
     fn default_full_merkle_tree(depth: usize) -> FullMerkleTree<Keccak256> {
         FullMerkleTree::<Keccak256>::new(depth, TestFr([0; 32]), FullMerkleConfig::default())
-            .unwrap()
+            .expect("Failed to create FullMerkleTree")
     }
 
     fn default_optimal_merkle_tree(depth: usize) -> OptimalMerkleTree<Keccak256> {
         OptimalMerkleTree::<Keccak256>::new(depth, TestFr([0; 32]), OptimalMerkleConfig::default())
-            .unwrap()
+            .expect("Failed to create OptimalMerkleTree")
     }
 
     #[test]
@@ -83,21 +83,101 @@ pub mod test {
         let nof_leaves = 4;
         let leaves: Vec<TestFr> = (1..=nof_leaves as u32).map(TestFr::from).collect();
 
-        let mut tree = default_full_merkle_tree(DEFAULT_DEPTH);
-        assert_eq!(tree.root(), default_tree_root);
+        let mut tree_full = default_full_merkle_tree(DEFAULT_DEPTH);
+        assert_eq!(tree_full.root(), default_tree_root);
         for i in 0..nof_leaves {
-            tree.set(i, leaves[i]).unwrap();
-            assert_eq!(tree.root(), roots[i]);
+            tree_full.set(i, leaves[i]).expect("Failed to set leaf");
+            assert_eq!(tree_full.root(), roots[i]);
         }
 
-        let mut tree = default_optimal_merkle_tree(DEFAULT_DEPTH);
-        assert_eq!(tree.root(), default_tree_root);
+        let mut tree_opt = default_optimal_merkle_tree(DEFAULT_DEPTH);
+        assert_eq!(tree_opt.root(), default_tree_root);
         for i in 0..nof_leaves {
-            tree.set(i, leaves[i]).unwrap();
-            assert_eq!(tree.root(), roots[i]);
+            tree_opt.set(i, leaves[i]).expect("Failed to set leaf");
+            assert_eq!(tree_opt.root(), roots[i]);
         }
     }
 
+    #[test]
+    fn test_set_range() {
+        let depth = 4;
+        let leaves: Vec<TestFr> = (0..(1 << depth) as u32).map(TestFr::from).collect();
+
+        let mut tree_full = default_full_merkle_tree(depth);
+        let root_before = tree_full.root();
+        tree_full
+            .set_range(0, leaves.iter().cloned())
+            .expect("Failed to set leaves");
+        let root_after = tree_full.root();
+        assert_ne!(root_before, root_after);
+
+        let mut tree_opt = default_optimal_merkle_tree(depth);
+        let root_before = tree_opt.root();
+        tree_opt
+            .set_range(0, leaves.iter().cloned())
+            .expect("Failed to set leaves");
+        let root_after = tree_opt.root();
+        assert_ne!(root_before, root_after);
+    }
+
+    #[test]
+    fn test_update_next() {
+        let mut tree_full = default_full_merkle_tree(DEFAULT_DEPTH);
+        let mut tree_opt = default_optimal_merkle_tree(DEFAULT_DEPTH);
+
+        for i in 0..4 {
+            let leaf = TestFr::from(i as u32);
+            tree_full.update_next(leaf).expect("Failed to update leaf");
+            tree_opt.update_next(leaf).expect("Failed to update leaf");
+            assert_eq!(tree_full.get(i).expect("Failed to get leaf"), leaf);
+            assert_eq!(tree_opt.get(i).expect("Failed to get leaf"), leaf);
+        }
+
+        assert_eq!(tree_full.leaves_set(), 4);
+        assert_eq!(tree_opt.leaves_set(), 4);
+    }
+
+    #[test]
+    fn test_delete_and_reset() {
+        let index = 1;
+        let original_leaf = TestFr::from(42);
+        let new_leaf = TestFr::from(99);
+
+        let mut tree_full = default_full_merkle_tree(DEFAULT_DEPTH);
+        tree_full
+            .set(index, original_leaf)
+            .expect("Failed to set leaf");
+        let root_with_original = tree_full.root();
+
+        tree_full.delete(index).expect("Failed to delete leaf");
+        let root_after_delete = tree_full.root();
+        assert_ne!(root_with_original, root_after_delete);
+
+        tree_full.set(index, new_leaf).expect("Failed to set leaf");
+        let root_after_reset = tree_full.root();
+
+        assert_ne!(root_after_delete, root_after_reset);
+        assert_ne!(root_with_original, root_after_reset);
+        assert_eq!(tree_full.get(index).expect("Failed to get leaf"), new_leaf);
+
+        let mut tree_opt = default_optimal_merkle_tree(DEFAULT_DEPTH);
+        tree_opt
+            .set(index, original_leaf)
+            .expect("Failed to set leaf");
+        let root_with_original = tree_opt.root();
+
+        tree_opt.delete(index).expect("Failed to delete leaf");
+        let root_after_delete = tree_opt.root();
+        assert_ne!(root_with_original, root_after_delete);
+
+        tree_opt.set(index, new_leaf).expect("Failed to set leaf");
+        let root_after_reset = tree_opt.root();
+
+        assert_ne!(root_after_delete, root_after_reset);
+        assert_ne!(root_with_original, root_after_reset);
+        assert_eq!(tree_opt.get(index).expect("Failed to get leaf"), new_leaf);
+    }
+
     #[test]
     fn test_get_empty_leaves_indices() {
         let depth = 4;
@@ -123,31 +203,29 @@ pub mod test {
             assert_eq!(tree_full.get_empty_leaves_indices(), vec_idxs);
         }
 
-        // Check situation when the number of items to insert is less than the number of items to delete
+        // check situation when the number of items to insert is less than the number of items to delete
         tree_full
             .override_range(0, leaves_2.clone().into_iter(), [0, 1, 2, 3].into_iter())
-            .unwrap();
+            .expect("Failed to override range");
 
         // check if the indexes for write and delete are the same
         tree_full
             .override_range(0, leaves_4.clone().into_iter(), [0, 1, 2, 3].into_iter())
-            .unwrap();
-        assert_eq!(tree_full.get_empty_leaves_indices(), vec![]);
+            .expect("Failed to override range");
+        assert_eq!(tree_full.get_empty_leaves_indices(), Vec::<usize>::new());
 
         // check if indexes for deletion are before indexes for overwriting
         tree_full
             .override_range(4, leaves_4.clone().into_iter(), [0, 1, 2, 3].into_iter())
-            .unwrap();
+            .expect("Failed to override range");
         assert_eq!(tree_full.get_empty_leaves_indices(), vec![0, 1, 2, 3]);
 
         // check if the indices for write and delete do not overlap completely
         tree_full
             .override_range(2, leaves_4.clone().into_iter(), [0, 1, 2, 3].into_iter())
-            .unwrap();
+            .expect("Failed to override range");
         assert_eq!(tree_full.get_empty_leaves_indices(), vec![0, 1]);
 
-        //// Optimal Merkle Tree Trest
-
         let mut tree_opt = default_optimal_merkle_tree(depth);
         let _ = tree_opt.set_range(0, leaves.clone().into_iter());
         assert!(tree_opt.get_empty_leaves_indices().is_empty());
@@ -164,27 +242,27 @@ pub mod test {
             assert_eq!(tree_opt.get_empty_leaves_indices(), vec_idxs);
         }
 
-        // Check situation when the number of items to insert is less than the number of items to delete
+        // check situation when the number of items to insert is less than the number of items to delete
         tree_opt
             .override_range(0, leaves_2.clone().into_iter(), [0, 1, 2, 3].into_iter())
-            .unwrap();
+            .expect("Failed to override range");
 
         // check if the indexes for write and delete are the same
         tree_opt
             .override_range(0, leaves_4.clone().into_iter(), [0, 1, 2, 3].into_iter())
-            .unwrap();
-        assert_eq!(tree_opt.get_empty_leaves_indices(), vec![]);
+            .expect("Failed to override range");
+        assert_eq!(tree_opt.get_empty_leaves_indices(), Vec::<usize>::new());
 
         // check if indexes for deletion are before indexes for overwriting
         tree_opt
             .override_range(4, leaves_4.clone().into_iter(), [0, 1, 2, 3].into_iter())
-            .unwrap();
+            .expect("Failed to override range");
         assert_eq!(tree_opt.get_empty_leaves_indices(), vec![0, 1, 2, 3]);
 
         // check if the indices for write and delete do not overlap completely
         tree_opt
             .override_range(2, leaves_4.clone().into_iter(), [0, 1, 2, 3].into_iter())
-            .unwrap();
+            .expect("Failed to override range");
         assert_eq!(tree_opt.get_empty_leaves_indices(), vec![0, 1]);
     }
 
@@ -194,18 +272,25 @@ pub mod test {
         let nof_leaves: usize = 4;
         let leaves: Vec<TestFr> = (0..nof_leaves as u32).map(TestFr::from).collect();
 
-        let mut tree_full = default_optimal_merkle_tree(depth);
+        let mut tree_full = default_full_merkle_tree(depth);
         let _ = tree_full.set_range(0, leaves.iter().cloned());
 
         for i in 0..nof_leaves {
             // check leaves
             assert_eq!(
-                tree_full.get(i).unwrap(),
-                tree_full.get_subtree_root(depth, i).unwrap()
+                tree_full.get(i).expect("Failed to get leaf"),
+                tree_full
+                    .get_subtree_root(depth, i)
+                    .expect("Failed to get subtree root")
             );
 
             // check root
-            assert_eq!(tree_full.root(), tree_full.get_subtree_root(0, i).unwrap());
+            assert_eq!(
+                tree_full.root(),
+                tree_full
+                    .get_subtree_root(0, i)
+                    .expect("Failed to get subtree root")
+            );
         }
 
         // check intermediate nodes
@@ -215,26 +300,39 @@ pub mod test {
                 let idx_r = (i + 1) * (1 << (depth - n));
                 let idx_sr = idx_l;
 
-                let prev_l = tree_full.get_subtree_root(n, idx_l).unwrap();
-                let prev_r = tree_full.get_subtree_root(n, idx_r).unwrap();
-                let subroot = tree_full.get_subtree_root(n - 1, idx_sr).unwrap();
+                let prev_l = tree_full
+                    .get_subtree_root(n, idx_l)
+                    .expect("Failed to get subtree root");
+                let prev_r = tree_full
+                    .get_subtree_root(n, idx_r)
+                    .expect("Failed to get subtree root");
+                let subroot = tree_full
+                    .get_subtree_root(n - 1, idx_sr)
+                    .expect("Failed to get subtree root");
 
                 // check intermediate nodes
                 assert_eq!(Keccak256::hash(&[prev_l, prev_r]), subroot);
             }
         }
 
-        let mut tree_opt = default_full_merkle_tree(depth);
+        let mut tree_opt = default_optimal_merkle_tree(depth);
         let _ = tree_opt.set_range(0, leaves.iter().cloned());
 
         for i in 0..nof_leaves {
             // check leaves
             assert_eq!(
-                tree_opt.get(i).unwrap(),
-                tree_opt.get_subtree_root(depth, i).unwrap()
+                tree_opt.get(i).expect("Failed to get leaf"),
+                tree_opt
+                    .get_subtree_root(depth, i)
+                    .expect("Failed to get subtree root")
             );
             // check root
-            assert_eq!(tree_opt.root(), tree_opt.get_subtree_root(0, i).unwrap());
+            assert_eq!(
+                tree_opt.root(),
+                tree_opt
+                    .get_subtree_root(0, i)
+                    .expect("Failed to get subtree root")
+            );
         }
 
         // check intermediate nodes
@@ -244,9 +342,15 @@ pub mod test {
                 let idx_r = (i + 1) * (1 << (depth - n));
                 let idx_sr = idx_l;
 
-                let prev_l = tree_opt.get_subtree_root(n, idx_l).unwrap();
-                let prev_r = tree_opt.get_subtree_root(n, idx_r).unwrap();
-                let subroot = tree_opt.get_subtree_root(n - 1, idx_sr).unwrap();
+                let prev_l = tree_opt
+                    .get_subtree_root(n, idx_l)
+                    .expect("Failed to get subtree root");
+                let prev_r = tree_opt
+                    .get_subtree_root(n, idx_r)
+                    .expect("Failed to get subtree root");
+                let subroot = tree_opt
+                    .get_subtree_root(n - 1, idx_sr)
+                    .expect("Failed to get subtree root");
 
                 // check intermediate nodes
                 assert_eq!(Keccak256::hash(&[prev_l, prev_r]), subroot);
@@ -259,61 +363,83 @@ pub mod test {
         let nof_leaves = 4;
         let leaves: Vec<TestFr> = (0..nof_leaves as u32).map(TestFr::from).collect();
 
-        // We thest the FullMerkleTree implementation
-        let mut tree = default_full_merkle_tree(DEFAULT_DEPTH);
+        // We test the FullMerkleTree implementation
+        let mut tree_full = default_full_merkle_tree(DEFAULT_DEPTH);
         for i in 0..nof_leaves {
             // We set the leaves
-            tree.set(i, leaves[i]).unwrap();
+            tree_full.set(i, leaves[i]).expect("Failed to set leaf");
 
             // We compute a merkle proof
-            let proof = tree.proof(i).expect("index should be set");
+            let proof = tree_full.proof(i).expect("Failed to compute proof");
 
             // We verify if the merkle proof corresponds to the right leaf index
             assert_eq!(proof.leaf_index(), i);
 
             // We verify the proof
-            assert!(tree.verify(&leaves[i], &proof).unwrap());
+            assert!(tree_full
+                .verify(&leaves[i], &proof)
+                .expect("Failed to verify proof"));
 
             // We ensure that the Merkle proof and the leaf generate the same root as the tree
-            assert_eq!(proof.compute_root_from(&leaves[i]), tree.root());
+            assert_eq!(proof.compute_root_from(&leaves[i]), tree_full.root());
 
             // We check that the proof is not valid for another leaf
-            assert!(!tree.verify(&leaves[(i + 1) % nof_leaves], &proof).unwrap());
+            assert!(!tree_full
+                .verify(&leaves[(i + 1) % nof_leaves], &proof)
+                .expect("Failed to verify proof"));
         }
 
         // We test the OptimalMerkleTree implementation
-        let mut tree = default_optimal_merkle_tree(DEFAULT_DEPTH);
+        let mut tree_opt = default_optimal_merkle_tree(DEFAULT_DEPTH);
         for i in 0..nof_leaves {
             // We set the leaves
-            tree.set(i, leaves[i]).unwrap();
+            tree_opt.set(i, leaves[i]).expect("Failed to set leaf");
 
             // We compute a merkle proof
-            let proof = tree.proof(i).expect("index should be set");
+            let proof = tree_opt.proof(i).expect("Failed to compute proof");
 
             // We verify if the merkle proof corresponds to the right leaf index
             assert_eq!(proof.leaf_index(), i);
 
             // We verify the proof
-            assert!(tree.verify(&leaves[i], &proof).unwrap());
+            assert!(tree_opt
+                .verify(&leaves[i], &proof)
+                .expect("Failed to verify proof"));
 
             // We ensure that the Merkle proof and the leaf generate the same root as the tree
-            assert_eq!(proof.compute_root_from(&leaves[i]), tree.root());
+            assert_eq!(proof.compute_root_from(&leaves[i]), tree_opt.root());
 
             // We check that the proof is not valid for another leaf
-            assert!(!tree.verify(&leaves[(i + 1) % nof_leaves], &proof).unwrap());
+            assert!(!tree_opt
+                .verify(&leaves[(i + 1) % nof_leaves], &proof)
+                .expect("Failed to verify proof"));
         }
     }
 
+    #[test]
+    fn test_proof_fail() {
+        let tree_full = default_full_merkle_tree(DEFAULT_DEPTH);
+        let tree_opt = default_optimal_merkle_tree(DEFAULT_DEPTH);
+
+        let invalid_leaf = TestFr::from(12345);
+
+        let proof_full = tree_full.proof(0).expect("Failed to compute proof");
+        let proof_opt = tree_opt.proof(0).expect("Failed to compute proof");
+
+        // Should fail because no leaf was set
+        assert!(!tree_full
+            .verify(&invalid_leaf, &proof_full)
+            .expect("Failed to verify proof"));
+        assert!(!tree_opt
+            .verify(&invalid_leaf, &proof_opt)
+            .expect("Failed to verify proof"));
+    }
+
     #[test]
     fn test_override_range() {
         let nof_leaves = 4;
         let leaves: Vec<TestFr> = (0..nof_leaves as u32).map(TestFr::from).collect();
 
-        let mut tree = default_optimal_merkle_tree(DEFAULT_DEPTH);
-
-        // We set the leaves
-        tree.set_range(0, leaves.iter().cloned()).unwrap();
-
         let new_leaves = [
             hex!("0000000000000000000000000000000000000000000000000000000000000005"),
             hex!("0000000000000000000000000000000000000000000000000000000000000006"),
@@ -322,17 +448,70 @@ pub mod test {
 
         let to_delete_indices: [usize; 2] = [0, 1];
 
-        // We override the leaves
-        tree.override_range(
-            0, // start from the end of the initial leaves
-            new_leaves.iter().cloned(),
-            to_delete_indices.iter().cloned(),
-        )
-        .unwrap();
+        let mut tree_full = default_full_merkle_tree(DEFAULT_DEPTH);
+        tree_full
+            .set_range(0, leaves.iter().cloned())
+            .expect("Failed to set leaves");
+
+        tree_full
+            .override_range(
+                0,
+                new_leaves.iter().cloned(),
+                to_delete_indices.iter().cloned(),
+            )
+            .expect("Failed to override range");
 
-        // ensure that the leaves are set correctly
         for (i, &new_leaf) in new_leaves.iter().enumerate() {
-            assert_eq!(tree.get_leaf(i), new_leaf);
+            assert_eq!(tree_full.get(i).expect("Failed to get leaf"), new_leaf);
+        }
+
+        let mut tree_opt = default_optimal_merkle_tree(DEFAULT_DEPTH);
+        tree_opt
+            .set_range(0, leaves.iter().cloned())
+            .expect("Failed to set leaves");
+
+        tree_opt
+            .override_range(
+                0,
+                new_leaves.iter().cloned(),
+                to_delete_indices.iter().cloned(),
+            )
+            .expect("Failed to override range");
+
+        for (i, &new_leaf) in new_leaves.iter().enumerate() {
+            assert_eq!(tree_opt.get(i).expect("Failed to get leaf"), new_leaf);
+        }
+    }
+
+    #[test]
+    fn test_override_range_parallel_triggered() {
+        let depth = 13;
+        let nof_leaves = 8192;
+
+        // number of leaves larger than MIN_PARALLEL_NODES to trigger parallel hashing
+        assert!(MIN_PARALLEL_NODES < nof_leaves);
+
+        let leaves: Vec<TestFr> = (0..nof_leaves as u32).map(TestFr::from).collect();
+        let indices: Vec<usize> = (0..nof_leaves).collect();
+
+        let mut tree_full = default_full_merkle_tree(depth);
+
+        tree_full
+            .override_range(0, leaves.iter().cloned(), indices.iter().cloned())
+            .expect("Failed to override range");
+
+        for (i, &leaf) in leaves.iter().enumerate() {
+            assert_eq!(tree_full.get(i).expect("Failed to get leaf"), leaf);
+        }
+
+        let mut tree_opt = default_optimal_merkle_tree(depth);
+
+        tree_opt
+            .override_range(0, leaves.iter().cloned(), indices.iter().cloned())
+            .expect("Failed to override range");
+
+        for (i, &leaf) in leaves.iter().enumerate() {
+            assert_eq!(tree_opt.get(i).expect("Failed to get leaf"), leaf);
         }
     }
 }