chore: Release

fix: include tries & exp backoff to catch WouldBlock errors

.github/workflows/ci.yml

@@ -28,18 +28,18 @@ on:
 name: Tests
 
 jobs:
-
-  multiplier:
+  test:
     strategy:
       matrix:
         platform: [ubuntu-latest, macos-latest]
+        crate: [multiplier, semaphore, rln, utils]
     runs-on: ${{ matrix.platform }}
     timeout-minutes: 60
   
-    name: multiplier - ${{ matrix.platform }}
+    name: test - ${{ matrix.crate }} - ${{ matrix.platform }}
     steps:
       - name: Checkout sources
-        uses: actions/checkout@v2
+        uses: actions/checkout@v3
       - name: Install stable toolchain
         uses: actions-rs/toolchain@v1
         with:
@@ -52,33 +52,8 @@ jobs:
       - name: cargo-make test
         run: |
             cargo make test --release
-        working-directory: multiplier
+        working-directory: ${{ matrix.crate }}
   
-  semaphore:
-    strategy:
-      matrix:
-        platform: [ubuntu-latest, macos-latest]
-    runs-on: ${{ matrix.platform }}
-    timeout-minutes: 60
-  
-    name: semaphore - ${{ matrix.platform }}
-    steps:
-      - name: Checkout sources
-        uses: actions/checkout@v2
-      - name: Install stable toolchain
-        uses: actions-rs/toolchain@v1
-        with:
-          profile: minimal
-          toolchain: stable
-          override: true
-      - uses: Swatinem/rust-cache@v2
-      - name: Install Dependencies
-        run: make installdeps
-      - name: cargo-make test
-        run: |
-            cargo make test --release
-        working-directory: semaphore
-
   rln-wasm:
     strategy:
       matrix:
@@ -86,9 +61,9 @@ jobs:
     runs-on: ${{ matrix.platform }}
     timeout-minutes: 60
   
-    name: rln-wasm - ${{ matrix.platform }}
+    name: test - rln-wasm - ${{ matrix.platform }}
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v3
       - name: Install stable toolchain
         uses: actions-rs/toolchain@v1
         with:
@@ -102,71 +77,22 @@ jobs:
         uses: jetli/wasm-pack-action@v0.3.0
       - run: cargo make build
         working-directory: rln-wasm
-      - run: cargo-make test
+      - run: cargo make test --release
         working-directory: rln-wasm
 
-  rln:
-    strategy:
-      matrix:
-        platform: [ubuntu-latest, macos-latest]
-    runs-on: ${{ matrix.platform }}
-    timeout-minutes: 60
-  
-    name: rln - ${{ matrix.platform }}
-    steps:
-      - name: Checkout sources
-        uses: actions/checkout@v2
-      - name: Install stable toolchain
-        uses: actions-rs/toolchain@v1
-        with:
-          profile: minimal
-          toolchain: stable
-          override: true
-      - uses: Swatinem/rust-cache@v2
-      - name: Install Dependencies
-        run: make installdeps
-      - name: cargo-make test
-        run: |
-            cargo make test --release
-        working-directory: rln
-
-  utils:
-    strategy:
-      matrix:
-        platform: [ubuntu-latest, macos-latest]
-    runs-on: ${{ matrix.platform }}
-    timeout-minutes: 60
-  
-    name: utils - ${{ matrix.platform }}
-    steps:
-      - name: Checkout sources
-        uses: actions/checkout@v2
-      - name: Install stable toolchain
-        uses: actions-rs/toolchain@v1
-        with:
-          profile: minimal
-          toolchain: stable
-          override: true
-      - uses: Swatinem/rust-cache@v2
-      - name: Install Dependencies
-        run: make installdeps
-      - name: cargo-make test
-        run: |
-            cargo make test --release
-        working-directory: utils
-
   lint:
     strategy:
       matrix:
         # we run lint tests only on ubuntu
         platform: [ubuntu-latest]
+        crate: [multiplier, semaphore, rln, utils]
     runs-on: ${{ matrix.platform }}
     timeout-minutes: 60
 
-    name: lint - ${{ matrix.platform }}  
+    name: lint - ${{ matrix.crate }} - ${{ matrix.platform }}  
     steps:
       - name: Checkout sources
-        uses: actions/checkout@v2
+        uses: actions/checkout@v3
       - name: Install stable toolchain
         uses: actions-rs/toolchain@v1
         with:
@@ -179,27 +105,33 @@ jobs:
         run: make installdeps
       - name: cargo fmt
         if: success() || failure()
-        run: cargo fmt --all -- --check
-      - name: multiplier - cargo clippy
+        run: cargo fmt -- --check
+        working-directory: ${{ matrix.crate }}
+      - name: cargo clippy
         if: success() || failure()
         run: |
             cargo clippy --release -- -D warnings
-        working-directory: multiplier
-      - name: semaphore - cargo clippy
-        if: success() || failure()
-        run: |
-            cargo clippy --release -- -D warnings
-        working-directory: semaphore
-      - name: rln - cargo clippy
-        if: success() || failure()
-        run: |
-            cargo clippy --release -- -D warnings
-        working-directory: rln
-      - name: utils - cargo clippy
-        if: success() || failure()
-        run: |
-            cargo clippy --release -- -D warnings
-        working-directory: utils
+        working-directory: ${{ matrix.crate }}
         # We skip clippy on rln-wasm, since wasm target is managed by cargo make
         # Currently not treating warnings as error, too noisy
-        # -- -D warnings
+        # -- -D warnings
+  benchmark:
+    # run only in pull requests
+    if: github.event_name == 'pull_request'
+    strategy:
+      matrix:
+        # we run benchmark tests only on ubuntu
+        platform: [ubuntu-latest]
+        crate: [rln, utils]
+    runs-on: ${{ matrix.platform }}
+    timeout-minutes: 60
+
+    name: benchmark - ${{ matrix.platform }} - ${{ matrix.crate }}
+    steps:
+      - name: Checkout sources
+        uses: actions/checkout@v3
+      - uses: Swatinem/rust-cache@v2
+      - uses: boa-dev/criterion-compare-action@v3
+        with:
+          branchName: ${{ github.base_ref }}
+          cwd: ${{ matrix.crate }}

.github/workflows/nightly-release.yml

@@ -16,7 +16,7 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Checkout sources
-        uses: actions/checkout@v2
+        uses: actions/checkout@v3
       - name: Install stable toolchain
         uses: actions-rs/toolchain@v1
         with:
@@ -51,7 +51,7 @@ jobs:
           - aarch64-apple-darwin
     steps:
       - name: Checkout sources
-        uses: actions/checkout@v2
+        uses: actions/checkout@v3
       - name: Install stable toolchain
         uses: actions-rs/toolchain@v1
         with:
@@ -81,7 +81,7 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Checkout sources
-        uses: actions/checkout@v2
+        uses: actions/checkout@v3
       - name: Install stable toolchain
         uses: actions-rs/toolchain@v1
         with:
@@ -115,14 +115,14 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Checkout code
-        uses: actions/checkout@v2
+        uses: actions/checkout@v3
         with:
           ref: master
       - name: Download artifacts
         uses: actions/download-artifact@v2
           
       - name: Delete tag
-        uses: dev-drprasad/delete-tag-and-release@v0.2.0
+        uses: dev-drprasad/delete-tag-and-release@v0.2.1
         with:
           delete_release: true
           tag_name: nightly
@@ -135,6 +135,7 @@ jobs:
           gh release create nightly --prerelease --target master \
             --title 'Nightly build ("master" branch)' \
             --generate-notes \
+            --draft=false \
             --notes-start-tag $start_tag \
             *-archive/*.tar.gz \
         env:

.github/workflows/sync-labels.yml

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

.gitignore

@@ -8,10 +8,7 @@ rln/pmtree_db
 # will have compiled files and executables
 debug/
 target/
-
-# Remove Cargo.lock from gitignore if creating an executable, leave it for libraries
-# More information here https://doc.rust-lang.org/cargo/guide/cargo-toml-vs-cargo-lock.html
-Cargo.lock
+wabt/
 
 # These are backup files generated by rustfmt
 **/*.rs.bk

Cargo.toml

@@ -4,9 +4,11 @@ members = [
   "private-settlement",
   "semaphore",
   "rln",
+  "rln-cli",
   "rln-wasm",
   "utils",
 ]
+resolver = "2"
 
 # Compilation profile for any non-workspace member.
 # Dependencies are optimized, even in a dev build. This improves dev performance
@@ -19,4 +21,4 @@ opt-level = 3
 opt-level = "s"
 
 [profile.release.package."semaphore"]
-codegen-units = 1
+codegen-units = 1

Makefile

@@ -8,10 +8,20 @@ all: .pre-build build
 .pre-build: .fetch-submodules
 	@cargo install cargo-make
 ifdef CI
-	@cargo install cross --git https://github.com/cross-rs/cross --branch main
+	@cargo install cross --git https://github.com/cross-rs/cross.git --rev 1511a28
 endif
 
 installdeps: .pre-build
+ifeq ($(shell uname),Darwin)
+# commented due to https://github.com/orgs/Homebrew/discussions/4612	
+# @brew update 
+	@brew install cmake ninja
+else ifeq ($(shell uname),Linux)
+	@sudo apt-get update
+	@sudo apt-get install -y cmake ninja-build
+endif
+	@git clone --recursive https://github.com/WebAssembly/wabt.git
+	@cd wabt && mkdir build && cd build && cmake .. -GNinja && ninja && sudo ninja install
 
 build: .pre-build
 	@cargo make build

multiplier/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "multiplier"
-version = "0.1.0"
+version = "0.3.0"
 edition = "2018"
 license = "MIT OR Apache-2.0"
 

private-settlement/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "private-settlement"
-version = "0.1.0"
+version = "0.3.0"
 edition = "2021"
 license = "MIT OR Apache-2.0"
 

rln-cli/Cargo.toml

@@ -0,0 +1,13 @@
+[package]
+name = "rln-cli"
+version = "0.3.0"
+edition = "2021"
+
+[dependencies]
+rln = { path = "../rln" }
+clap = { version = "4.2.7", features = ["cargo", "derive", "env"]}
+clap_derive = { version = "=4.2.0" }
+color-eyre = "=0.6.2"
+# serialization
+serde_json = "1.0.48"
+serde = { version = "1.0.130", features = ["derive"] }

rln-cli/src/commands.rs

@@ -0,0 +1,67 @@
+use std::path::PathBuf;
+
+use clap::Subcommand;
+
+#[derive(Subcommand)]
+pub(crate) enum Commands {
+    New {
+        tree_height: usize,
+        /// Sets a custom config file
+        #[arg(short, long)]
+        config: PathBuf,
+    },
+    NewWithParams {
+        tree_height: usize,
+        /// Sets a custom config file
+        #[arg(short, long)]
+        config: PathBuf,
+        #[arg(short, long)]
+        tree_config_input: PathBuf,
+    },
+    SetTree {
+        tree_height: usize,
+    },
+    SetLeaf {
+        index: usize,
+        #[arg(short, long)]
+        file: PathBuf,
+    },
+    SetMultipleLeaves {
+        index: usize,
+        #[arg(short, long)]
+        file: PathBuf,
+    },
+    ResetMultipleLeaves {
+        #[arg(short, long)]
+        file: PathBuf,
+    },
+    SetNextLeaf {
+        #[arg(short, long)]
+        file: PathBuf,
+    },
+    DeleteLeaf {
+        index: usize,
+    },
+    GetRoot,
+    GetProof {
+        index: usize,
+    },
+    Prove {
+        #[arg(short, long)]
+        input: PathBuf,
+    },
+    Verify {
+        #[arg(short, long)]
+        file: PathBuf,
+    },
+    GenerateProof {
+        #[arg(short, long)]
+        input: PathBuf,
+    },
+    VerifyWithRoots {
+        #[arg(short, long)]
+        input: PathBuf,
+        #[arg(short, long)]
+        roots: PathBuf,
+    },
+}

rln-cli/src/config.rs

@@ -0,0 +1,29 @@
+use color_eyre::Result;
+use serde::{Deserialize, Serialize};
+use std::{fs::File, io::Read, path::PathBuf};
+
+pub const RLN_STATE_PATH: &str = "RLN_STATE_PATH";
+
+#[derive(Default, Serialize, Deserialize)]
+pub(crate) struct Config {
+    pub inner: Option<InnerConfig>,
+}
+
+#[derive(Default, Serialize, Deserialize)]
+pub(crate) struct InnerConfig {
+    pub file: PathBuf,
+    pub tree_height: usize,
+}
+
+impl Config {
+    pub(crate) fn load_config() -> Result<Config> {
+        let path = PathBuf::from(std::env::var(RLN_STATE_PATH)?);
+
+        let mut file = File::open(path)?;
+
+        let mut contents = String::new();
+        file.read_to_string(&mut contents)?;
+        let state: Config = serde_json::from_str(&contents)?;
+        Ok(state)
+    }
+}

rln-cli/src/main.rs

@@ -0,0 +1,157 @@
+use std::{fs::File, io::Read, path::Path};
+
+use clap::Parser;
+use color_eyre::{Report, Result};
+use commands::Commands;
+use rln::public::RLN;
+use state::State;
+
+mod commands;
+mod config;
+mod state;
+
+#[derive(Parser)]
+#[command(author, version, about, long_about = None)]
+struct Cli {
+    #[command(subcommand)]
+    command: Option<Commands>,
+}
+
+fn main() -> Result<()> {
+    let cli = Cli::parse();
+
+    let mut state = State::load_state()?;
+
+    match &cli.command {
+        Some(Commands::New {
+            tree_height,
+            config,
+        }) => {
+            let resources = File::open(&config)?;
+            state.rln = Some(RLN::new(*tree_height, resources)?);
+            Ok(())
+        }
+        Some(Commands::NewWithParams {
+            tree_height,
+            config,
+            tree_config_input,
+        }) => {
+            let mut resources: Vec<Vec<u8>> = Vec::new();
+            for filename in ["rln.wasm", "rln_final.zkey", "verification_key.json"] {
+                let fullpath = config.join(Path::new(filename));
+                let mut file = File::open(&fullpath)?;
+                let metadata = std::fs::metadata(&fullpath)?;
+                let mut buffer = vec![0; metadata.len() as usize];
+                file.read_exact(&mut buffer)?;
+                resources.push(buffer);
+            }
+            let tree_config_input_file = File::open(&tree_config_input)?;
+            state.rln = Some(RLN::new_with_params(
+                *tree_height,
+                resources[0].clone(),
+                resources[1].clone(),
+                resources[2].clone(),
+                tree_config_input_file,
+            )?);
+            Ok(())
+        }
+        Some(Commands::SetTree { tree_height }) => {
+            state
+                .rln
+                .ok_or(Report::msg("no RLN instance initialized"))?
+                .set_tree(*tree_height)?;
+            Ok(())
+        }
+        Some(Commands::SetLeaf { index, file }) => {
+            let input_data = File::open(&file)?;
+            state
+                .rln
+                .ok_or(Report::msg("no RLN instance initialized"))?
+                .set_leaf(*index, input_data)?;
+            Ok(())
+        }
+        Some(Commands::SetMultipleLeaves { index, file }) => {
+            let input_data = File::open(&file)?;
+            state
+                .rln
+                .ok_or(Report::msg("no RLN instance initialized"))?
+                .set_leaves_from(*index, input_data)?;
+            Ok(())
+        }
+        Some(Commands::ResetMultipleLeaves { file }) => {
+            let input_data = File::open(&file)?;
+            state
+                .rln
+                .ok_or(Report::msg("no RLN instance initialized"))?
+                .init_tree_with_leaves(input_data)?;
+            Ok(())
+        }
+        Some(Commands::SetNextLeaf { file }) => {
+            let input_data = File::open(&file)?;
+            state
+                .rln
+                .ok_or(Report::msg("no RLN instance initialized"))?
+                .set_next_leaf(input_data)?;
+            Ok(())
+        }
+        Some(Commands::DeleteLeaf { index }) => {
+            state
+                .rln
+                .ok_or(Report::msg("no RLN instance initialized"))?
+                .delete_leaf(*index)?;
+            Ok(())
+        }
+        Some(Commands::GetRoot) => {
+            let writer = std::io::stdout();
+            state
+                .rln
+                .ok_or(Report::msg("no RLN instance initialized"))?
+                .get_root(writer)?;
+            Ok(())
+        }
+        Some(Commands::GetProof { index }) => {
+            let writer = std::io::stdout();
+            state
+                .rln
+                .ok_or(Report::msg("no RLN instance initialized"))?
+                .get_proof(*index, writer)?;
+            Ok(())
+        }
+        Some(Commands::Prove { input }) => {
+            let input_data = File::open(&input)?;
+            let writer = std::io::stdout();
+            state
+                .rln
+                .ok_or(Report::msg("no RLN instance initialized"))?
+                .prove(input_data, writer)?;
+            Ok(())
+        }
+        Some(Commands::Verify { file }) => {
+            let input_data = File::open(&file)?;
+            state
+                .rln
+                .ok_or(Report::msg("no RLN instance initialized"))?
+                .verify(input_data)?;
+            Ok(())
+        }
+        Some(Commands::GenerateProof { input }) => {
+            let input_data = File::open(&input)?;
+            let writer = std::io::stdout();
+            state
+                .rln
+                .ok_or(Report::msg("no RLN instance initialized"))?
+                .generate_rln_proof(input_data, writer)?;
+            Ok(())
+        }
+        Some(Commands::VerifyWithRoots { input, roots }) => {
+            let input_data = File::open(&input)?;
+            let roots_data = File::open(&roots)?;
+            state
+                .rln
+                .ok_or(Report::msg("no RLN instance initialized"))?
+                .verify_with_roots(input_data, roots_data)?;
+            Ok(())
+        }
+        None => Ok(()),
+    }
+}

rln-cli/src/state.rs

@@ -0,0 +1,23 @@
+use color_eyre::Result;
+use rln::public::RLN;
+use std::fs::File;
+
+use crate::config::{Config, InnerConfig};
+
+#[derive(Default)]
+pub(crate) struct State<'a> {
+    pub rln: Option<RLN<'a>>,
+}
+
+impl<'a> State<'a> {
+    pub(crate) fn load_state() -> Result<State<'a>> {
+        let config = Config::load_config()?;
+        let rln = if let Some(InnerConfig { file, tree_height }) = config.inner {
+            let resources = File::open(&file)?;
+            Some(RLN::new(tree_height, resources)?)
+        } else {
+            None
+        };
+        Ok(State { rln })
+    }
+}

rln-wasm/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "rln-wasm"
-version = "0.0.7"
+version = "0.0.9"
 edition = "2021"
 license = "MIT or Apache2"
 

rln-wasm/Makefile.toml

@@ -9,9 +9,14 @@ script = "sed -i.bak 's/rln-wasm/zerokit-rln-wasm/g' pkg/package.json && rm pkg/
 clear = true
 dependencies = [
     "pack-build",
-    "pack-rename"
+    "pack-rename",
+    "post-build"
 ]
 
+[tasks.post-build]
+command = "wasm-strip"
+args = ["./pkg/rln_wasm_bg.wasm"]
+
 [tasks.test]
 command = "wasm-pack"
 args = ["test", "--release", "--node"]

rln-wasm/README.md

@@ -21,6 +21,11 @@ make installdeps
 cd rln-wasm
 cargo make build
 ```
+4. Compile a slimmer version of zerokit for `wasm32-unknown-unknown`:
+```
+cd rln-wasm
+cargo make post-build
+```
 
 ## Running tests
 ```

rln-wasm/src/lib.rs

@@ -3,9 +3,11 @@
 extern crate wasm_bindgen;
 extern crate web_sys;
 
+use std::vec::Vec;
+
 use js_sys::{BigInt as JsBigInt, Object, Uint8Array};
 use num_bigint::BigInt;
-use rln::public::RLN;
+use rln::public::{hash, poseidon_hash, RLN};
 use wasm_bindgen::prelude::*;
 
 #[wasm_bindgen]
@@ -20,6 +22,163 @@ pub struct RLNWrapper {
     instance: RLN<'static>,
 }
 
+// Macro to call methods with arbitrary amount of arguments,
+// which have the last argument is output buffer pointer
+// First argument to the macro is context,
+// second is the actual method on `RLN`
+// third is the aforementioned output buffer argument
+// rest are all other arguments to the method
+macro_rules! call_with_output_and_error_msg {
+    // this variant is needed for the case when
+    // there are zero other arguments
+    ($instance:expr, $method:ident, $error_msg:expr) => {
+        {
+            let mut output_data: Vec<u8> = Vec::new();
+            let new_instance = $instance.process();
+            if let Err(err) = new_instance.instance.$method(&mut output_data) {
+                std::mem::forget(output_data);
+                Err(format!("Msg: {:#?}, Error: {:#?}", $error_msg, err))
+            } else {
+                let result = Uint8Array::from(&output_data[..]);
+                std::mem::forget(output_data);
+                Ok(result)
+            }
+        }
+    };
+    ($instance:expr, $method:ident, $error_msg:expr, $( $arg:expr ),* ) => {
+        {
+            let mut output_data: Vec<u8> = Vec::new();
+            let new_instance = $instance.process();
+            if let Err(err) = new_instance.instance.$method($($arg.process()),*, &mut output_data) {
+                std::mem::forget(output_data);
+                Err(format!("Msg: {:#?}, Error: {:#?}", $error_msg, err))
+            } else {
+                let result = Uint8Array::from(&output_data[..]);
+                std::mem::forget(output_data);
+                Ok(result)
+            }
+        }
+    };
+}
+
+// Macro to call_with_error_msg methods with arbitrary amount of arguments,
+// First argument to the macro is context,
+// second is the actual method on `RLNWrapper`
+// rest are all other arguments to the method
+macro_rules! call_with_error_msg {
+    ($instance:expr, $method:ident, $error_msg:expr $(, $arg:expr)*) => {
+        {
+            let new_instance: &mut RLNWrapper = $instance.process();
+            if let Err(err) = new_instance.instance.$method($($arg.process()),*) {
+                Err(format!("Msg: {:#?}, Error: {:#?}", $error_msg, err))
+            } else {
+                Ok(())
+            }
+
+        }
+    }
+}
+
+macro_rules! call {
+    ($instance:expr, $method:ident $(, $arg:expr)*) => {
+        {
+            let new_instance: &mut RLNWrapper = $instance.process();
+            new_instance.instance.$method($($arg.process()),*)
+        }
+    }
+}
+
+macro_rules! call_bool_method_with_error_msg {
+    ($instance:expr, $method:ident, $error_msg:expr $(, $arg:expr)*) => {
+        {
+            let new_instance: &RLNWrapper = $instance.process();
+            new_instance.instance.$method($($arg.process()),*).map_err(|err| format!("Msg: {:#?}, Error: {:#?}", $error_msg, err))
+        }
+    }
+}
+
+// Macro to execute a function with arbitrary amount of arguments,
+// First argument is the function to execute
+// Rest are all other arguments to the method
+macro_rules! fn_call_with_output_and_error_msg {
+    // this variant is needed for the case when
+    // there are zero other arguments
+    ($func:ident, $error_msg:expr) => {
+        {
+            let mut output_data: Vec<u8> = Vec::new();
+            if let Err(err) = $func(&mut output_data) {
+                std::mem::forget(output_data);
+                Err(format!("Msg: {:#?}, Error: {:#?}", $error_msg, err))
+            } else {
+                let result = Uint8Array::from(&output_data[..]);
+                std::mem::forget(output_data);
+                Ok(result)
+            }
+        }
+    };
+    ($func:ident, $error_msg:expr, $( $arg:expr ),* ) => {
+        {
+            let mut output_data: Vec<u8> = Vec::new();
+            if let Err(err) = $func($($arg.process()),*, &mut output_data) {
+                std::mem::forget(output_data);
+                Err(format!("Msg: {:#?}, Error: {:#?}", $error_msg, err))
+            } else {
+                let result = Uint8Array::from(&output_data[..]);
+                std::mem::forget(output_data);
+                Ok(result)
+            }
+        }
+    };
+}
+
+trait ProcessArg {
+    type ReturnType;
+    fn process(self) -> Self::ReturnType;
+}
+
+impl ProcessArg for usize {
+    type ReturnType = usize;
+    fn process(self) -> Self::ReturnType {
+        self
+    }
+}
+
+impl<T> ProcessArg for Vec<T> {
+    type ReturnType = Vec<T>;
+    fn process(self) -> Self::ReturnType {
+        self
+    }
+}
+
+impl<'a> ProcessArg for *const RLN<'a> {
+    type ReturnType = &'a RLN<'a>;
+    fn process(self) -> Self::ReturnType {
+        unsafe { &*self }
+    }
+}
+
+impl ProcessArg for *const RLNWrapper {
+    type ReturnType = &'static RLNWrapper;
+    fn process(self) -> Self::ReturnType {
+        unsafe { &*self }
+    }
+}
+
+impl ProcessArg for *mut RLNWrapper {
+    type ReturnType = &'static mut RLNWrapper;
+    fn process(self) -> Self::ReturnType {
+        unsafe { &mut *self }
+    }
+}
+
+impl<'a> ProcessArg for &'a [u8] {
+    type ReturnType = &'a [u8];
+
+    fn process(self) -> Self::ReturnType {
+        self
+    }
+}
+
 #[allow(clippy::not_unsafe_ptr_arg_deref)]
 #[wasm_bindgen(js_name = newRLN)]
 pub fn wasm_new(
@@ -39,24 +198,20 @@ pub fn wasm_get_serialized_rln_witness(
     ctx: *mut RLNWrapper,
     input: Uint8Array,
 ) -> Result<Uint8Array, String> {
-    let wrapper = unsafe { &mut *ctx };
-    let rln_witness = wrapper
-        .instance
-        .get_serialized_rln_witness(&input.to_vec()[..])
+    let rln_witness = call!(ctx, get_serialized_rln_witness, &input.to_vec()[..])
         .map_err(|err| format!("{:#?}", err))?;
-
     Ok(Uint8Array::from(&rln_witness[..]))
 }
 
 #[allow(clippy::not_unsafe_ptr_arg_deref)]
 #[wasm_bindgen(js_name = insertMember)]
 pub fn wasm_set_next_leaf(ctx: *mut RLNWrapper, input: Uint8Array) -> Result<(), String> {
-    let wrapper = unsafe { &mut *ctx };
-    if wrapper.instance.set_next_leaf(&input.to_vec()[..]).is_ok() {
-        Ok(())
-    } else {
-        Err("could not insert member into merkle tree".into())
-    }
+    call_with_error_msg!(
+        ctx,
+        set_next_leaf,
+        "could not insert member into merkle tree".to_string(),
+        &input.to_vec()[..]
+    )
 }
 
 #[allow(clippy::not_unsafe_ptr_arg_deref)]
@@ -66,31 +221,47 @@ pub fn wasm_set_leaves_from(
     index: usize,
     input: Uint8Array,
 ) -> Result<(), String> {
-    let wrapper = unsafe { &mut *ctx };
-    if wrapper
-        .instance
-        .set_leaves_from(index as usize, &input.to_vec()[..])
-        .is_ok()
-    {
-        Ok(())
-    } else {
-        Err("could not set multiple leaves".into())
-    }
+    call_with_error_msg!(
+        ctx,
+        set_leaves_from,
+        "could not set multiple leaves".to_string(),
+        index,
+        &*input.to_vec()
+    )
+}
+
+#[allow(clippy::not_unsafe_ptr_arg_deref)]
+#[wasm_bindgen(js_name = deleteLeaf)]
+pub fn wasm_delete_leaf(ctx: *mut RLNWrapper, index: usize) -> Result<(), String> {
+    call_with_error_msg!(ctx, delete_leaf, "could not delete leaf".to_string(), index)
+}
+
+#[allow(clippy::not_unsafe_ptr_arg_deref)]
+#[wasm_bindgen(js_name = setMetadata)]
+pub fn wasm_set_metadata(ctx: *mut RLNWrapper, input: Uint8Array) -> Result<(), String> {
+    call_with_error_msg!(
+        ctx,
+        set_metadata,
+        "could not set metadata".to_string(),
+        &*input.to_vec()
+    )
+}
+
+#[allow(clippy::not_unsafe_ptr_arg_deref)]
+#[wasm_bindgen(js_name = getMetadata)]
+pub fn wasm_get_metadata(ctx: *mut RLNWrapper) -> Result<Uint8Array, String> {
+    call_with_output_and_error_msg!(ctx, get_metadata, "could not get metadata".to_string())
 }
 
 #[allow(clippy::not_unsafe_ptr_arg_deref)]
 #[wasm_bindgen(js_name = initTreeWithLeaves)]
 pub fn wasm_init_tree_with_leaves(ctx: *mut RLNWrapper, input: Uint8Array) -> Result<(), String> {
-    let wrapper = unsafe { &mut *ctx };
-    if wrapper
-        .instance
-        .init_tree_with_leaves(&input.to_vec()[..])
-        .is_ok()
-    {
-        Ok(())
-    } else {
-        Err("could not init merkle tree".into())
-    }
+    call_with_error_msg!(
+        ctx,
+        init_tree_with_leaves,
+        "could not init merkle tree".to_string(),
+        &*input.to_vec()
+    )
 }
 
 #[allow(clippy::not_unsafe_ptr_arg_deref)]
@@ -99,12 +270,8 @@ pub fn rln_witness_to_json(
     ctx: *mut RLNWrapper,
     serialized_witness: Uint8Array,
 ) -> Result<Object, String> {
-    let wrapper = unsafe { &mut *ctx };
-    let inputs = wrapper
-        .instance
-        .get_rln_witness_json(&serialized_witness.to_vec()[..])
+    let inputs = call!(ctx, get_rln_witness_json, &serialized_witness.to_vec()[..])
         .map_err(|err| err.to_string())?;
-
     let js_value = serde_wasm_bindgen::to_value(&inputs).map_err(|err| err.to_string())?;
     Object::from_entries(&js_value).map_err(|err| format!("{:#?}", err))
 }
@@ -116,8 +283,6 @@ pub fn generate_rln_proof_with_witness(
     calculated_witness: Vec<JsBigInt>,
     serialized_witness: Uint8Array,
 ) -> Result<Uint8Array, String> {
-    let wrapper = unsafe { &mut *ctx };
-
     let mut witness_vec: Vec<BigInt> = vec![];
 
     for v in calculated_witness {
@@ -131,69 +296,36 @@ pub fn generate_rln_proof_with_witness(
         );
     }
 
-    let mut output_data: Vec<u8> = Vec::new();
-
-    if wrapper
-        .instance
-        .generate_rln_proof_with_witness(witness_vec, serialized_witness.to_vec(), &mut output_data)
-        .is_ok()
-    {
-        let result = Uint8Array::from(&output_data[..]);
-        std::mem::forget(output_data);
-        Ok(result)
-    } else {
-        std::mem::forget(output_data);
-        Err("could not generate proof".into())
-    }
+    call_with_output_and_error_msg!(
+        ctx,
+        generate_rln_proof_with_witness,
+        "could not generate proof",
+        witness_vec,
+        serialized_witness.to_vec()
+    )
 }
 
 #[allow(clippy::not_unsafe_ptr_arg_deref)]
 #[wasm_bindgen(js_name = generateMembershipKey)]
 pub fn wasm_key_gen(ctx: *const RLNWrapper) -> Result<Uint8Array, String> {
-    let wrapper = unsafe { &*ctx };
-    let mut output_data: Vec<u8> = Vec::new();
-    if wrapper.instance.key_gen(&mut output_data).is_ok() {
-        let result = Uint8Array::from(&output_data[..]);
-        std::mem::forget(output_data);
-        Ok(result)
-    } else {
-        std::mem::forget(output_data);
-        Err("could not generate membership keys".into())
-    }
+    call_with_output_and_error_msg!(ctx, key_gen, "could not generate membership keys")
 }
 
 #[allow(clippy::not_unsafe_ptr_arg_deref)]
 #[wasm_bindgen(js_name = generateExtendedMembershipKey)]
 pub fn wasm_extended_key_gen(ctx: *const RLNWrapper) -> Result<Uint8Array, String> {
-    let wrapper = unsafe { &*ctx };
-    let mut output_data: Vec<u8> = Vec::new();
-    if wrapper.instance.extended_key_gen(&mut output_data).is_ok() {
-        let result = Uint8Array::from(&output_data[..]);
-        std::mem::forget(output_data);
-        Ok(result)
-    } else {
-        std::mem::forget(output_data);
-        Err("could not generate membership keys".into())
-    }
+    call_with_output_and_error_msg!(ctx, extended_key_gen, "could not generate membership keys")
 }
 
 #[allow(clippy::not_unsafe_ptr_arg_deref)]
 #[wasm_bindgen(js_name = generateSeededMembershipKey)]
 pub fn wasm_seeded_key_gen(ctx: *const RLNWrapper, seed: Uint8Array) -> Result<Uint8Array, String> {
-    let wrapper = unsafe { &*ctx };
-    let mut output_data: Vec<u8> = Vec::new();
-    if wrapper
-        .instance
-        .seeded_key_gen(&seed.to_vec()[..], &mut output_data)
-        .is_ok()
-    {
-        let result = Uint8Array::from(&output_data[..]);
-        std::mem::forget(output_data);
-        Ok(result)
-    } else {
-        std::mem::forget(output_data);
-        Err("could not generate membership key".into())
-    }
+    call_with_output_and_error_msg!(
+        ctx,
+        seeded_key_gen,
+        "could not generate membership key",
+        &seed.to_vec()[..]
+    )
 }
 
 #[allow(clippy::not_unsafe_ptr_arg_deref)]
@@ -202,20 +334,12 @@ pub fn wasm_seeded_extended_key_gen(
     ctx: *const RLNWrapper,
     seed: Uint8Array,
 ) -> Result<Uint8Array, String> {
-    let wrapper = unsafe { &*ctx };
-    let mut output_data: Vec<u8> = Vec::new();
-    if wrapper
-        .instance
-        .seeded_extended_key_gen(&seed.to_vec()[..], &mut output_data)
-        .is_ok()
-    {
-        let result = Uint8Array::from(&output_data[..]);
-        std::mem::forget(output_data);
-        Ok(result)
-    } else {
-        std::mem::forget(output_data);
-        Err("could not generate membership key".into())
-    }
+    call_with_output_and_error_msg!(
+        ctx,
+        seeded_extended_key_gen,
+        "could not generate membership key",
+        &seed.to_vec()[..]
+    )
 }
 
 #[allow(clippy::not_unsafe_ptr_arg_deref)]
@@ -225,38 +349,24 @@ pub fn wasm_recover_id_secret(
     input_proof_data_1: Uint8Array,
     input_proof_data_2: Uint8Array,
 ) -> Result<Uint8Array, String> {
-    let wrapper = unsafe { &*ctx };
-    let mut output_data: Vec<u8> = Vec::new();
-    if wrapper
-        .instance
-        .recover_id_secret(
-            &input_proof_data_1.to_vec()[..],
-            &input_proof_data_2.to_vec()[..],
-            &mut output_data,
-        )
-        .is_ok()
-    {
-        let result = Uint8Array::from(&output_data[..]);
-        std::mem::forget(output_data);
-        Ok(result)
-    } else {
-        std::mem::forget(output_data);
-        Err("could not recover id secret".into())
-    }
+    call_with_output_and_error_msg!(
+        ctx,
+        recover_id_secret,
+        "could not recover id secret",
+        &input_proof_data_1.to_vec()[..],
+        &input_proof_data_2.to_vec()[..]
+    )
 }
 
 #[allow(clippy::not_unsafe_ptr_arg_deref)]
 #[wasm_bindgen(js_name = verifyRLNProof)]
 pub fn wasm_verify_rln_proof(ctx: *const RLNWrapper, proof: Uint8Array) -> Result<bool, String> {
-    let wrapper = unsafe { &*ctx };
-    if match wrapper.instance.verify_rln_proof(&proof.to_vec()[..]) {
-        Ok(verified) => verified,
-        Err(_) => return Err("error while verifying rln proof".into()),
-    } {
-        return Ok(true);
-    }
-
-    Ok(false)
+    call_bool_method_with_error_msg!(
+        ctx,
+        verify_rln_proof,
+        "error while verifying rln proof".to_string(),
+        &proof.to_vec()[..]
+    )
 }
 
 #[allow(clippy::not_unsafe_ptr_arg_deref)]
@@ -266,31 +376,31 @@ pub fn wasm_verify_with_roots(
     proof: Uint8Array,
     roots: Uint8Array,
 ) -> Result<bool, String> {
-    let wrapper = unsafe { &*ctx };
-    if match wrapper
-        .instance
-        .verify_with_roots(&proof.to_vec()[..], &roots.to_vec()[..])
-    {
-        Ok(verified) => verified,
-        Err(_) => return Err("error while verifying proof with roots".into()),
-    } {
-        return Ok(true);
-    }
-
-    Ok(false)
+    call_bool_method_with_error_msg!(
+        ctx,
+        verify_with_roots,
+        "error while verifying proof with roots".to_string(),
+        &proof.to_vec()[..],
+        &roots.to_vec()[..]
+    )
 }
 
 #[allow(clippy::not_unsafe_ptr_arg_deref)]
 #[wasm_bindgen(js_name = getRoot)]
 pub fn wasm_get_root(ctx: *const RLNWrapper) -> Result<Uint8Array, String> {
-    let wrapper = unsafe { &*ctx };
-    let mut output_data: Vec<u8> = Vec::new();
-    if wrapper.instance.get_root(&mut output_data).is_ok() {
-        let result = Uint8Array::from(&output_data[..]);
-        std::mem::forget(output_data);
-        Ok(result)
-    } else {
-        std::mem::forget(output_data);
-        Err("could not obtain root".into())
-    }
+    call_with_output_and_error_msg!(ctx, get_root, "could not obtain root")
+}
+
+#[wasm_bindgen(js_name = hash)]
+pub fn wasm_hash(input: Uint8Array) -> Result<Uint8Array, String> {
+    fn_call_with_output_and_error_msg!(hash, "could not generate hash", &input.to_vec()[..])
+}
+
+#[wasm_bindgen(js_name = poseidonHash)]
+pub fn wasm_poseidon_hash(input: Uint8Array) -> Result<Uint8Array, String> {
+    fn_call_with_output_and_error_msg!(
+        poseidon_hash,
+        "could not generate poseidon hash",
+        &input.to_vec()[..]
+    )
 }

rln-wasm/tests/rln-wasm.rs

@@ -6,8 +6,7 @@ mod tests {
     use rln::circuit::TEST_TREE_HEIGHT;
     use rln::utils::normalize_usize;
     use rln_wasm::*;
-    use wasm_bindgen::prelude::*;
-    use wasm_bindgen::JsValue;
+    use wasm_bindgen::{prelude::*, JsValue};
     use wasm_bindgen_test::wasm_bindgen_test;
 
     #[wasm_bindgen(module = "src/utils.js")]
@@ -30,7 +29,7 @@ mod tests {
         let vk = read_file(&vk_path).unwrap();
 
         // Creating an instance of RLN
-        let rln_instance = wasm_new(tree_height, zkey, vk);
+        let rln_instance = wasm_new(tree_height, zkey, vk).unwrap();
 
         // Creating membership key
         let mem_keys = wasm_key_gen(rln_instance).unwrap();
@@ -59,10 +58,11 @@ mod tests {
         let serialized_message = Uint8Array::from(&serialized_vec[..]);
 
         let serialized_rln_witness =
-            wasm_get_serialized_rln_witness(rln_instance, serialized_message);
+            wasm_get_serialized_rln_witness(rln_instance, serialized_message).unwrap();
 
         // Obtaining inputs that should be sent to circom witness calculator
-        let json_inputs = rln_witness_to_json(rln_instance, serialized_rln_witness.clone());
+        let json_inputs =
+            rln_witness_to_json(rln_instance, serialized_rln_witness.clone()).unwrap();
 
         // Calculating witness with JS
         // (Using a JSON since wasm_bindgen does not like Result<Vec<JsBigInt>,JsValue>)
@@ -108,4 +108,28 @@ mod tests {
         let is_proof_valid = wasm_verify_with_roots(rln_instance, proof_with_signal, roots);
         assert!(is_proof_valid.unwrap(), "verifying proof with roots failed");
     }
+
+    #[wasm_bindgen_test]
+    fn test_metadata() {
+        let tree_height = TEST_TREE_HEIGHT;
+        let circom_path = format!("../rln/resources/tree_height_{TEST_TREE_HEIGHT}/rln.wasm");
+        let zkey_path = format!("../rln/resources/tree_height_{TEST_TREE_HEIGHT}/rln_final.zkey");
+        let vk_path =
+            format!("../rln/resources/tree_height_{TEST_TREE_HEIGHT}/verification_key.json");
+        let zkey = read_file(&zkey_path).unwrap();
+        let vk = read_file(&vk_path).unwrap();
+
+        // Creating an instance of RLN
+        let rln_instance = wasm_new(tree_height, zkey, vk).unwrap();
+
+
+        let test_metadata = Uint8Array::new(&JsValue::from_str("test"));
+        // Inserting random metadata
+        wasm_set_metadata(rln_instance, test_metadata.clone()).unwrap();
+
+        // Getting metadata
+        let metadata = wasm_get_metadata(rln_instance).unwrap();
+
+        assert_eq!(metadata.to_vec(), test_metadata.to_vec());
+    }
 }

rln/Cargo.toml

@@ -1,11 +1,16 @@
 [package]
 name = "rln"
-version = "0.1.0"
+version = "0.3.2"
 edition = "2021"
 license = "MIT OR Apache-2.0"
+description = "APIs to manage, compute and verify zkSNARK proofs and RLN primitives"
+documentation = "https://github.com/vacp2p/zerokit"
+homepage = "https://vac.dev"
+repository = "https://github.com/vacp2p/zerokit"
 
 [lib]
 crate-type = ["rlib", "staticlib"]
+bench = false
 
 # This flag disable cargo doctests, i.e. testing example code-snippets in documentation
 doctest = false
@@ -13,46 +18,51 @@ doctest = false
 
 [dependencies]
 # ZKP Generation
-ark-ec = { version = "=0.3.0", default-features = false }
-ark-ff = { version = "=0.3.0", default-features = false, features = [ "asm"] }
-ark-std = { version = "=0.3.0", default-features = false }
-ark-bn254 = { version = "=0.3.0" }
-ark-groth16 = { git = "https://github.com/arkworks-rs/groth16", rev = "765817f", default-features = false }
-ark-relations = { version = "=0.3.0", default-features = false, features = [ "std" ] }
-ark-serialize = { version = "=0.3.0", default-features = false }
-ark-circom = { git = "https://github.com/vacp2p/ark-circom", rev = "0e587145cb05e08b2d1a01509eb578670088eb2f", default-features = false, features = ["circom-2"] }
+ark-ec = { version = "=0.4.1", default-features = false }
+ark-ff = { version = "=0.4.1", default-features = false, features = [ "asm"] }
+ark-std = { version = "=0.4.0", default-features = false }
+ark-bn254 = { version = "=0.4.0" }
+ark-groth16 = { version = "=0.4.0", features = ["parallel"], default-features = false }
+ark-relations = { version = "=0.4.0", default-features = false, features = [ "std" ] }
+ark-serialize = { version = "=0.4.1", default-features = false }
+ark-circom = { version = "=0.1.0", default-features = false, features = ["circom-2"] }
 
 # WASM
-wasmer = { version = "2.3.0", default-features = false }
+wasmer = { version = "=2.3.0", default-features = false }
 
 # error handling
 color-eyre = "=0.6.2"
-thiserror = "=1.0.38"
+thiserror = "=1.0.39"
 
 # utilities
 cfg-if = "=1.0"
 num-bigint = { version = "=0.4.3", default-features = false, features = ["rand"] }
-num-traits = "0.2.11"
-once_cell = "1.14.0"
+num-traits = "=0.2.15"
+once_cell = "=1.17.1"
 rand = "=0.8.5"
 rand_chacha = "=0.3.1"
 tiny-keccak = { version = "=2.0.2", features = ["keccak"] }
-utils = { path = "../utils/", default-features = false }
-pmtree = { git = "https://github.com/Rate-Limiting-Nullifier/pmtree", rev = "f6d1a1fecad72cd39e6808e78085091d541dc882", optional = true}
+utils = { package = "zerokit_utils", version = "=0.3.2", path = "../utils/", default-features = false }
 
 # serialization
-serde_json = "1.0.48"
+serde_json = "=1.0.96"
+serde = { version = "=1.0.163", features = ["derive"] }
 
 include_dir = "=0.7.3"
 
 [dev-dependencies]
 sled = "=0.34.7"
+criterion = { version = "=0.4.0", features = ["html_reports"] }
 
 [features]
-default = ["parallel", "wasmer/sys-default"]
+default = ["parallel", "wasmer/sys-default", "pmtree-ft"]
 parallel = ["ark-ec/parallel", "ark-ff/parallel", "ark-std/parallel", "ark-groth16/parallel", "utils/parallel"]
 wasm = ["wasmer/js", "wasmer/std"]
 fullmerkletree = ["default"]
 
 # Note: pmtree feature is still experimental
-pmtree-ft = ["default", "pmtree"]
+pmtree-ft = ["utils/pmtree-ft"]
+
+[[bench]]
+name = "pmtree_benchmark"
+harness = false

rln/Makefile.toml

@@ -5,3 +5,7 @@ args = ["build", "--release"]
 [tasks.test]
 command = "cargo"
 args = ["test", "--release"]
+
+[tasks.bench]
+command = "cargo"
+args = ["bench"]

rln/benches/pmtree_benchmark.rs

@@ -0,0 +1,44 @@
+use criterion::{criterion_group, criterion_main, Criterion};
+use utils::ZerokitMerkleTree;
+
+use rln::{circuit::Fr, pm_tree_adapter::PmTree};
+
+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();
+
+    c.bench_function("Pmtree::set", |b| {
+        b.iter(|| {
+            tree.set(0, leaves[0]).unwrap();
+        })
+    });
+
+    c.bench_function("Pmtree:delete", |b| {
+        b.iter(|| {
+            tree.delete(0).unwrap();
+        })
+    });
+
+    c.bench_function("Pmtree::override_range", |b| {
+        b.iter(|| {
+            tree.override_range(0, leaves.clone(), [0, 1, 2, 3])
+                .unwrap();
+        })
+    });
+
+    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();
+        })
+    });
+}
+
+criterion_group!(benches, pmtree_benchmark);
+criterion_main!(benches);

rln/src/circuit.rs

@@ -12,7 +12,6 @@ use color_eyre::{Report, Result};
 use num_bigint::BigUint;
 use serde_json::Value;
 use std::io::Cursor;
-use std::path::Path;
 use std::str::FromStr;
 
 cfg_if! {
@@ -22,6 +21,7 @@ cfg_if! {
         use std::sync::Mutex;
         use wasmer::{Module, Store};
         use include_dir::{include_dir, Dir};
+        use std::path::Path;
     }
 }
 

rln/src/ffi.rs

@@ -12,7 +12,15 @@ macro_rules! call {
     ($instance:expr, $method:ident $(, $arg:expr)*) => {
         {
             let new_instance: &mut RLN = $instance.process();
-            new_instance.$method($($arg.process()),*).is_ok()
+            match new_instance.$method($($arg.process()),*) {
+                Ok(()) => {
+                    true
+                }
+                Err(err) => {
+                    eprintln!("execution error: {err}");
+                    false
+                }
+            }
         }
     }
 }
@@ -30,13 +38,17 @@ macro_rules! call_with_output_arg {
         {
             let mut output_data: Vec<u8> = Vec::new();
             let new_instance = $instance.process();
-            if new_instance.$method(&mut output_data).is_ok() {
-                unsafe { *$output_arg = Buffer::from(&output_data[..]) };
-                std::mem::forget(output_data);
-                true
-            } else {
-                std::mem::forget(output_data);
-                false
+            match new_instance.$method(&mut output_data) {
+                Ok(()) => {
+                    unsafe { *$output_arg = Buffer::from(&output_data[..]) };
+                    std::mem::forget(output_data);
+                    true
+                }
+                Err(err) => {
+                    std::mem::forget(output_data);
+                    eprintln!("execution error: {err}");
+                    false
+                }
             }
         }
     };
@@ -44,13 +56,17 @@ macro_rules! call_with_output_arg {
         {
             let mut output_data: Vec<u8> = Vec::new();
             let new_instance = $instance.process();
-            if new_instance.$method($($arg.process()),*, &mut output_data).is_ok() {
-                unsafe { *$output_arg = Buffer::from(&output_data[..]) };
-                std::mem::forget(output_data);
-                true
-            } else {
-                std::mem::forget(output_data);
-                false
+            match new_instance.$method($($arg.process()),*, &mut output_data) {
+                Ok(()) => {
+                    unsafe { *$output_arg = Buffer::from(&output_data[..]) };
+                    std::mem::forget(output_data);
+                    true
+                }
+                Err(err) => {
+                    std::mem::forget(output_data);
+                    eprintln!("execution error: {err}");
+                    false
+                }
             }
         }
     };
@@ -66,13 +82,17 @@ macro_rules! no_ctx_call_with_output_arg {
     ($method:ident, $output_arg:expr, $( $arg:expr ),* ) => {
         {
             let mut output_data: Vec<u8> = Vec::new();
-            if $method($($arg.process()),*, &mut output_data).is_ok() {
-                unsafe { *$output_arg = Buffer::from(&output_data[..]) };
-                std::mem::forget(output_data);
-                true
-            } else {
-                std::mem::forget(output_data);
-                false
+            match $method($($arg.process()),*, &mut output_data) {
+                Ok(()) => {
+                    unsafe { *$output_arg = Buffer::from(&output_data[..]) };
+                    std::mem::forget(output_data);
+                    true
+                }
+                Err(err) => {
+                    std::mem::forget(output_data);
+                    eprintln!("execution error: {err}");
+                    false
+                }
             }
         }
     }
@@ -89,8 +109,11 @@ macro_rules! call_with_bool_arg {
         {
             let new_instance = $instance.process();
             if match new_instance.$method($($arg.process()),*,) {
-                Ok(verified) => verified,
-                Err(_) => return false,
+                Ok(result) => result,
+                Err(err) => {
+                    eprintln!("execution error: {err}");
+                    return false
+                },
             } {
                 unsafe { *$bool_arg = true };
             } else {
@@ -171,11 +194,15 @@ impl<'a> From<&Buffer> for &'a [u8] {
 #[allow(clippy::not_unsafe_ptr_arg_deref)]
 #[no_mangle]
 pub extern "C" fn new(tree_height: usize, input_buffer: *const Buffer, ctx: *mut *mut RLN) -> bool {
-    if let Ok(rln) = RLN::new(tree_height, input_buffer.process()) {
-        unsafe { *ctx = Box::into_raw(Box::new(rln)) };
-        true
-    } else {
-        false
+    match RLN::new(tree_height, input_buffer.process()) {
+        Ok(rln) => {
+            unsafe { *ctx = Box::into_raw(Box::new(rln)) };
+            true
+        }
+        Err(err) => {
+            eprintln!("could not instantiate rln: {err}");
+            false
+        }
     }
 }
 
@@ -186,18 +213,24 @@ pub extern "C" fn new_with_params(
     circom_buffer: *const Buffer,
     zkey_buffer: *const Buffer,
     vk_buffer: *const Buffer,
+    tree_config: *const Buffer,
     ctx: *mut *mut RLN,
 ) -> bool {
-    if let Ok(rln) = RLN::new_with_params(
+    match RLN::new_with_params(
         tree_height,
         circom_buffer.process().to_vec(),
         zkey_buffer.process().to_vec(),
         vk_buffer.process().to_vec(),
+        tree_config.process(),
     ) {
-        unsafe { *ctx = Box::into_raw(Box::new(rln)) };
-        true
-    } else {
-        false
+        Ok(rln) => {
+            unsafe { *ctx = Box::into_raw(Box::new(rln)) };
+            true
+        }
+        Err(err) => {
+            eprintln!("could not instantiate rln: {err}");
+            false
+        }
     }
 }
 
@@ -222,6 +255,12 @@ pub extern "C" fn set_leaf(ctx: *mut RLN, index: usize, input_buffer: *const Buf
     call!(ctx, set_leaf, index, input_buffer)
 }
 
+#[allow(clippy::not_unsafe_ptr_arg_deref)]
+#[no_mangle]
+pub extern "C" fn get_leaf(ctx: *mut RLN, index: usize, output_buffer: *mut Buffer) -> bool {
+    call_with_output_arg!(ctx, get_leaf, output_buffer, index)
+}
+
 #[allow(clippy::not_unsafe_ptr_arg_deref)]
 #[no_mangle]
 pub extern "C" fn set_next_leaf(ctx: *mut RLN, input_buffer: *const Buffer) -> bool {
@@ -244,6 +283,33 @@ pub extern "C" fn init_tree_with_leaves(ctx: *mut RLN, input_buffer: *const Buff
     call!(ctx, init_tree_with_leaves, input_buffer)
 }
 
+#[allow(clippy::not_unsafe_ptr_arg_deref)]
+#[no_mangle]
+pub extern "C" fn atomic_operation(
+    ctx: *mut RLN,
+    index: usize,
+    leaves_buffer: *const Buffer,
+    indices_buffer: *const Buffer,
+) -> bool {
+    call!(ctx, atomic_operation, index, leaves_buffer, indices_buffer)
+}
+
+#[allow(clippy::not_unsafe_ptr_arg_deref)]
+#[no_mangle]
+pub extern "C" fn seq_atomic_operation(
+    ctx: *mut RLN,
+    leaves_buffer: *const Buffer,
+    indices_buffer: *const Buffer,
+) -> bool {
+    call!(
+        ctx,
+        atomic_operation,
+        ctx.process().leaves_set(),
+        leaves_buffer,
+        indices_buffer
+    )
+}
+
 #[allow(clippy::not_unsafe_ptr_arg_deref)]
 #[no_mangle]
 pub extern "C" fn get_root(ctx: *const RLN, output_buffer: *mut Buffer) -> bool {
@@ -368,6 +434,28 @@ pub extern "C" fn recover_id_secret(
     )
 }
 
+////////////////////////////////////////////////////////
+// Persistent metadata APIs
+////////////////////////////////////////////////////////
+
+#[allow(clippy::not_unsafe_ptr_arg_deref)]
+#[no_mangle]
+pub extern "C" fn set_metadata(ctx: *mut RLN, input_buffer: *const Buffer) -> bool {
+    call!(ctx, set_metadata, input_buffer)
+}
+
+#[allow(clippy::not_unsafe_ptr_arg_deref)]
+#[no_mangle]
+pub extern "C" fn get_metadata(ctx: *const RLN, output_buffer: *mut Buffer) -> bool {
+    call_with_output_arg!(ctx, get_metadata, output_buffer)
+}
+
+#[allow(clippy::not_unsafe_ptr_arg_deref)]
+#[no_mangle]
+pub extern "C" fn flush(ctx: *mut RLN) -> bool {
+    call!(ctx, flush)
+}
+
 #[allow(clippy::not_unsafe_ptr_arg_deref)]
 #[no_mangle]
 pub extern "C" fn hash(input_buffer: *const Buffer, output_buffer: *mut Buffer) -> bool {

rln/src/hashers.rs

@@ -0,0 +1,58 @@
+/// This crate instantiates the Poseidon hash algorithm.
+use crate::{circuit::Fr, utils::bytes_le_to_fr};
+use once_cell::sync::Lazy;
+use tiny_keccak::{Hasher, Keccak};
+use utils::poseidon::Poseidon;
+
+/// These indexed constants hardcode the supported round parameters tuples (t, RF, RN, SKIP_MATRICES) for the Bn254 scalar field.
+/// SKIP_MATRICES is the index of the randomly generated secure MDS matrix.
+/// TODO: generate these parameters
+pub const ROUND_PARAMS: [(usize, usize, usize, usize); 8] = [
+    (2, 8, 56, 0),
+    (3, 8, 57, 0),
+    (4, 8, 56, 0),
+    (5, 8, 60, 0),
+    (6, 8, 60, 0),
+    (7, 8, 63, 0),
+    (8, 8, 64, 0),
+    (9, 8, 63, 0),
+];
+
+/// Poseidon Hash wrapper over above implementation.
+static POSEIDON: Lazy<Poseidon<Fr>> = Lazy::new(|| Poseidon::<Fr>::from(&ROUND_PARAMS));
+
+pub fn poseidon_hash(input: &[Fr]) -> Fr {
+    POSEIDON
+        .hash(input.to_vec())
+        .expect("hash with fixed input size can't fail")
+}
+
+/// The zerokit RLN Merkle tree Hasher.
+#[derive(Clone, Copy, PartialEq, Eq)]
+pub struct PoseidonHash;
+
+/// The default Hasher trait used by Merkle tree implementation in utils.
+impl utils::merkle_tree::Hasher for PoseidonHash {
+    type Fr = Fr;
+
+    fn default_leaf() -> Self::Fr {
+        Self::Fr::from(0)
+    }
+
+    fn hash(inputs: &[Self::Fr]) -> Self::Fr {
+        poseidon_hash(inputs)
+    }
+}
+
+/// Hashes arbitrary signal to the underlying prime field.
+pub fn hash_to_field(signal: &[u8]) -> Fr {
+    // We hash the input signal using Keccak256
+    let mut hash = [0; 32];
+    let mut hasher = Keccak::v256();
+    hasher.update(signal);
+    hasher.finalize(&mut hash);
+
+    // We export the hash as a field element
+    let (el, _) = bytes_le_to_fr(hash.as_ref());
+    el
+}

rln/src/lib.rs

@@ -1,7 +1,9 @@
 #![allow(dead_code)]
 
 pub mod circuit;
-pub mod poseidon_hash;
+pub mod hashers;
+#[cfg(feature = "pmtree-ft")]
+pub mod pm_tree_adapter;
 pub mod poseidon_tree;
 pub mod protocol;
 pub mod public;

rln/src/pm_tree_adapter.rs

@@ -0,0 +1,339 @@
+use std::fmt::Debug;
+use std::path::PathBuf;
+use std::str::FromStr;
+
+use color_eyre::{Report, Result};
+use serde_json::Value;
+
+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};
+
+const METADATA_KEY: [u8; 8] = *b"metadata";
+
+pub struct PmTree {
+    tree: pmtree::MerkleTree<SledDB, PoseidonHash>,
+    // metadata that an application may use to store additional information
+    metadata: Vec<u8>,
+}
+
+pub struct PmTreeProof {
+    proof: pmtree::tree::MerkleProof<PoseidonHash>,
+}
+
+pub type FrOf<H> = <H as Hasher>::Fr;
+
+// The pmtree Hasher trait used by pmtree Merkle tree
+impl Hasher for PoseidonHash {
+    type Fr = Fr;
+
+    fn serialize(value: Self::Fr) -> pmtree::Value {
+        fr_to_bytes_le(&value)
+    }
+
+    fn deserialize(value: pmtree::Value) -> Self::Fr {
+        let (fr, _) = bytes_le_to_fr(&value);
+        fr
+    }
+
+    fn default_leaf() -> Self::Fr {
+        Fr::from(0)
+    }
+
+    fn hash(inputs: &[Self::Fr]) -> Self::Fr {
+        poseidon_hash(inputs)
+    }
+}
+
+fn get_tmp_path() -> PathBuf {
+    std::env::temp_dir().join(format!("pmtree-{}", rand::random::<u64>()))
+}
+
+fn get_tmp() -> bool {
+    true
+}
+
+pub struct PmtreeConfig(Config);
+
+impl FromStr for PmtreeConfig {
+    type Err = Report;
+
+    fn from_str(s: &str) -> Result<Self> {
+        let config: Value = serde_json::from_str(s)?;
+
+        let path = config["path"].as_str();
+        let path = path.map(PathBuf::from);
+        let temporary = config["temporary"].as_bool();
+        let cache_capacity = config["cache_capacity"].as_u64();
+        let flush_every_ms = config["flush_every_ms"].as_u64();
+        let mode = match config["mode"].as_str() {
+            Some("HighThroughput") => Mode::HighThroughput,
+            Some("LowSpace") => Mode::LowSpace,
+            _ => Mode::HighThroughput,
+        };
+        let use_compression = config["use_compression"].as_bool();
+
+        if temporary.is_some()
+            && path.is_some()
+            && temporary.unwrap()
+            && path.as_ref().unwrap().exists()
+        {
+            return Err(Report::msg(format!(
+                "Path {:?} already exists, cannot use temporary",
+                path.unwrap()
+            )));
+        }
+
+        let config = Config::new()
+            .temporary(temporary.unwrap_or(get_tmp()))
+            .path(path.unwrap_or(get_tmp_path()))
+            .cache_capacity(cache_capacity.unwrap_or(1024 * 1024 * 1024))
+            .flush_every_ms(flush_every_ms)
+            .mode(mode)
+            .use_compression(use_compression.unwrap_or(false));
+        Ok(PmtreeConfig(config))
+    }
+}
+
+impl Default for PmtreeConfig {
+    fn default() -> Self {
+        let tmp_path = get_tmp_path();
+        PmtreeConfig(
+            Config::new()
+                .temporary(true)
+                .path(tmp_path)
+                .cache_capacity(150_000)
+                .mode(Mode::HighThroughput)
+                .use_compression(false)
+                .flush_every_ms(Some(12_000)),
+        )
+    }
+}
+impl Debug for PmtreeConfig {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        self.0.fmt(f)
+    }
+}
+
+impl Clone for PmtreeConfig {
+    fn clone(&self) -> Self {
+        PmtreeConfig(self.0.clone())
+    }
+}
+
+impl ZerokitMerkleTree for PmTree {
+    type Proof = PmTreeProof;
+    type Hasher = PoseidonHash;
+    type Config = PmtreeConfig;
+
+    fn default(depth: usize) -> Result<Self> {
+        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> {
+        let tree_loaded = pmtree::MerkleTree::load(config.clone().0);
+        let tree = match tree_loaded {
+            Ok(tree) => tree,
+            Err(_) => pmtree::MerkleTree::new(depth, config.0)?,
+        };
+
+        Ok(PmTree {
+            tree,
+            metadata: Vec::new(),
+        })
+    }
+
+    fn depth(&self) -> usize {
+        self.tree.depth()
+    }
+
+    fn capacity(&self) -> usize {
+        self.tree.capacity()
+    }
+
+    fn leaves_set(&mut self) -> usize {
+        self.tree.leaves_set()
+    }
+
+    fn root(&self) -> FrOf<Self::Hasher> {
+        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_range<I: IntoIterator<Item = FrOf<Self::Hasher>>>(
+        &mut self,
+        start: usize,
+        values: I,
+    ) -> Result<()> {
+        self.tree
+            .set_range(start, values)
+            .map_err(|e| Report::msg(e.to_string()))
+    }
+
+    fn get(&self, index: usize) -> Result<FrOf<Self::Hasher>> {
+        self.tree.get(index).map_err(|e| Report::msg(e.to_string()))
+    }
+
+    fn override_range<I: IntoIterator<Item = FrOf<Self::Hasher>>, J: IntoIterator<Item = usize>>(
+        &mut self,
+        start: usize,
+        leaves: I,
+        indices: J,
+    ) -> Result<()> {
+        let leaves = leaves.into_iter().collect::<Vec<_>>();
+        let mut indices = indices.into_iter().collect::<Vec<_>>();
+        indices.sort();
+
+        match (leaves.is_empty(), indices.is_empty()) {
+            (true, true) => Err(Report::msg("no leaves or indices to be removed")),
+            (false, true) => self.set_range_with_leaves(start, leaves),
+            (true, false) => self.remove_indices(indices),
+            (false, false) => self.remove_indices_and_set_leaves(start, leaves, indices),
+        }
+    }
+
+    fn update_next(&mut self, leaf: FrOf<Self::Hasher>) -> Result<()> {
+        self.tree
+            .update_next(leaf)
+            .map_err(|e| Report::msg(e.to_string()))
+    }
+
+    fn delete(&mut self, index: usize) -> Result<()> {
+        self.tree
+            .delete(index)
+            .map_err(|e| Report::msg(e.to_string()))
+    }
+
+    fn proof(&self, index: usize) -> Result<Self::Proof> {
+        let proof = self.tree.proof(index)?;
+        Ok(PmTreeProof { proof })
+    }
+
+    fn verify(&self, leaf: &FrOf<Self::Hasher>, witness: &Self::Proof) -> Result<bool> {
+        if self.tree.verify(leaf, &witness.proof) {
+            Ok(true)
+        } else {
+            Err(Report::msg("verify failed"))
+        }
+    }
+
+    fn set_metadata(&mut self, metadata: &[u8]) -> Result<()> {
+        self.tree.db.put(METADATA_KEY, metadata.to_vec())?;
+        self.metadata = metadata.to_vec();
+        Ok(())
+    }
+
+    fn metadata(&self) -> Result<Vec<u8>> {
+        if !self.metadata.is_empty() {
+            return Ok(self.metadata.clone());
+        }
+        // if empty, try searching the db
+        let data = self.tree.db.get(METADATA_KEY)?;
+
+        if data.is_none() {
+            return Err(Report::msg("metadata does not exist"));
+        }
+        Ok(data.unwrap())
+    }
+
+    fn close_db_connection(&mut self) -> Result<()> {
+        self.tree.db.close().map_err(|e| Report::msg(e.to_string()))
+    }
+}
+
+type PmTreeHasher = <PmTree as ZerokitMerkleTree>::Hasher;
+type FrOfPmTreeHasher = FrOf<PmTreeHasher>;
+
+impl PmTree {
+    fn set_range_with_leaves(&mut self, start: usize, leaves: Vec<FrOfPmTreeHasher>) -> Result<()> {
+        self.tree
+            .set_range(start, leaves)
+            .map_err(|e| Report::msg(e.to_string()))
+    }
+
+    fn remove_indices(&mut self, indices: Vec<usize>) -> Result<()> {
+        let start = indices[0];
+        let end = indices.last().unwrap() + 1;
+
+        let mut new_leaves: Vec<_> = (start..end)
+            .map(|i| self.tree.get(i))
+            .collect::<Result<_, _>>()?;
+
+        new_leaves
+            .iter_mut()
+            .take(indices.len())
+            .for_each(|leaf| *leaf = PmTreeHasher::default_leaf());
+
+        self.tree
+            .set_range(start, new_leaves)
+            .map_err(|e| Report::msg(e.to_string()))
+    }
+
+    fn remove_indices_and_set_leaves(
+        &mut self,
+        start: usize,
+        leaves: Vec<FrOfPmTreeHasher>,
+        indices: Vec<usize>,
+    ) -> Result<()> {
+        let min_index = *indices.first().unwrap();
+        let max_index = start + leaves.len();
+
+        // Generated a placeholder with the exact size needed,
+        // Initiated with default values to be overridden throughout the method
+        let mut set_values = vec![PmTreeHasher::default_leaf(); max_index - min_index];
+
+        // If the index is not in indices list, keep the original value
+        for i in min_index..start {
+            if !indices.contains(&i) {
+                let value = self.tree.get(i)?;
+                set_values[i - min_index] = value;
+            }
+        }
+
+        // Insert new leaves after 'start' position
+        for (i, &leaf) in leaves.iter().enumerate() {
+            set_values[start - min_index + i] = leaf;
+        }
+
+        self.tree
+            .set_range(min_index, set_values)
+            .map_err(|e| Report::msg(e.to_string()))
+    }
+}
+
+impl ZerokitMerkleProof for PmTreeProof {
+    type Index = u8;
+    type Hasher = PoseidonHash;
+
+    fn length(&self) -> usize {
+        self.proof.length()
+    }
+
+    fn leaf_index(&self) -> usize {
+        self.proof.leaf_index()
+    }
+
+    fn get_path_elements(&self) -> Vec<FrOf<Self::Hasher>> {
+        self.proof.get_path_elements()
+    }
+
+    fn get_path_index(&self) -> Vec<Self::Index> {
+        self.proof.get_path_index()
+    }
+    fn compute_root_from(&self, leaf: &FrOf<Self::Hasher>) -> FrOf<Self::Hasher> {
+        self.proof.compute_root_from(leaf)
+    }
+}

rln/src/poseidon_hash.rs

@@ -1,28 +0,0 @@
-// This crate instantiate the Poseidon hash algorithm
-
-use crate::circuit::Fr;
-use once_cell::sync::Lazy;
-use utils::poseidon::Poseidon;
-
-// These indexed constants hardcodes the supported round parameters tuples (t, RF, RN, SKIP_MATRICES) for the Bn254 scalar field
-// SKIP_MATRICES is the index of the randomly generated secure MDS matrix. See security note in the zerokit_utils::poseidon::poseidon_constants crate on this.
-// TODO: generate these parameters
-pub const ROUND_PARAMS: [(usize, usize, usize, usize); 8] = [
-    (2, 8, 56, 0),
-    (3, 8, 57, 0),
-    (4, 8, 56, 0),
-    (5, 8, 60, 0),
-    (6, 8, 60, 0),
-    (7, 8, 63, 0),
-    (8, 8, 64, 0),
-    (9, 8, 63, 0),
-];
-
-// Poseidon Hash wrapper over above implementation. Adapted from semaphore-rs poseidon hash wrapper.
-static POSEIDON: Lazy<Poseidon<Fr>> = Lazy::new(|| Poseidon::<Fr>::from(&ROUND_PARAMS));
-
-pub fn poseidon_hash(input: &[Fr]) -> Fr {
-    POSEIDON
-        .hash(input.to_vec())
-        .expect("hash with fixed input size can't fail")
-}

rln/src/poseidon_tree.rs

@@ -2,15 +2,16 @@
 
 // Implementation inspired by https://github.com/worldcoin/semaphore-rs/blob/d462a4372f1fd9c27610f2acfe4841fab1d396aa/src/poseidon_tree.rs (no differences)
 
-use crate::circuit::Fr;
-use crate::poseidon_hash::poseidon_hash;
 use cfg_if::cfg_if;
-use utils::merkle_tree::*;
 
-#[cfg(feature = "pmtree-ft")]
-use crate::utils::{bytes_le_to_fr, fr_to_bytes_le};
-#[cfg(feature = "pmtree-ft")]
-use pmtree::*;
+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
@@ -19,48 +20,11 @@ cfg_if! {
     if #[cfg(feature = "fullmerkletree")] {
         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 {
         pub type PoseidonTree = OptimalMerkleTree<PoseidonHash>;
         pub type MerkleProof = OptimalMerkleProof<PoseidonHash>;
     }
 }
-
-// The zerokit RLN Merkle tree Hasher
-#[derive(Clone, Copy, PartialEq, Eq)]
-pub struct PoseidonHash;
-
-// The default Hasher trait used by Merkle tree implementation in utils
-impl utils::merkle_tree::Hasher for PoseidonHash {
-    type Fr = Fr;
-
-    fn default_leaf() -> Self::Fr {
-        Self::Fr::from(0)
-    }
-
-    fn hash(inputs: &[Self::Fr]) -> Self::Fr {
-        poseidon_hash(inputs)
-    }
-}
-
-#[cfg(feature = "pmtree-ft")]
-// The pmtree Hasher trait used by pmtree Merkle tree
-impl pmtree::Hasher for PoseidonHash {
-    type Fr = Fr;
-
-    fn default_leaf() -> Self::Fr {
-        Fr::from(0)
-    }
-
-    fn serialize(value: Self::Fr) -> Value {
-        fr_to_bytes_le(&value)
-    }
-
-    fn deserialize(value: Value) -> Self::Fr {
-        let (fr, _) = bytes_le_to_fr(&value);
-        fr
-    }
-
-    fn hash(inputs: &[Self::Fr]) -> Self::Fr {
-        poseidon_hash(inputs)
-    }
-}

rln/src/protocol.rs

@@ -1,10 +1,7 @@
 // This crate collects all the underlying primitives used to implement RLN
 
 use ark_circom::{CircomReduction, WitnessCalculator};
-use ark_groth16::{
-    create_proof_with_reduction_and_matrices, prepare_verifying_key,
-    verify_proof as ark_verify_proof, Proof as ArkProof, ProvingKey, VerifyingKey,
-};
+use ark_groth16::{prepare_verifying_key, Groth16, Proof as ArkProof, ProvingKey, VerifyingKey};
 use ark_relations::r1cs::ConstraintMatrices;
 use ark_relations::r1cs::SynthesisError;
 use ark_std::{rand::thread_rng, UniformRand};
@@ -20,11 +17,13 @@ use thiserror::Error;
 use tiny_keccak::{Hasher as _, Keccak};
 
 use crate::circuit::{Curve, Fr};
-use crate::poseidon_hash::poseidon_hash;
+use crate::hashers::hash_to_field;
+use crate::hashers::poseidon_hash;
 use crate::poseidon_tree::*;
 use crate::public::RLN_IDENTIFIER;
 use crate::utils::*;
 use cfg_if::cfg_if;
+use utils::{ZerokitMerkleProof, ZerokitMerkleTree};
 
 ///////////////////////////////////////////////////////
 // RLN Witness data structure and utility functions
@@ -156,13 +155,17 @@ pub fn proof_inputs_to_rln_witness(
     let (identity_secret, read) = bytes_le_to_fr(&serialized[all_read..]);
     all_read += read;
 
-    let id_index = usize::from_le_bytes(serialized[all_read..all_read + 8].try_into()?);
+    let id_index = usize::try_from(u64::from_le_bytes(
+        serialized[all_read..all_read + 8].try_into()?,
+    ))?;
     all_read += 8;
 
     let (epoch, read) = bytes_le_to_fr(&serialized[all_read..]);
     all_read += read;
 
-    let signal_len = usize::from_le_bytes(serialized[all_read..all_read + 8].try_into()?);
+    let signal_len = usize::try_from(u64::from_le_bytes(
+        serialized[all_read..all_read + 8].try_into()?,
+    ))?;
     all_read += 8;
 
     let signal: Vec<u8> = serialized[all_read..all_read + signal_len].to_vec();
@@ -485,20 +488,6 @@ pub fn extended_seeded_keygen(signal: &[u8]) -> (Fr, Fr, Fr, Fr) {
     )
 }
 
-// Hashes arbitrary signal to the underlying prime field
-pub fn hash_to_field(signal: &[u8]) -> Fr {
-    // We hash the input signal using Keccak256
-    // (note that a bigger curve order might require a bigger hash blocksize)
-    let mut hash = [0; 32];
-    let mut hasher = Keccak::v256();
-    hasher.update(signal);
-    hasher.finalize(&mut hash);
-
-    // We export the hash as a field element
-    let (el, _) = bytes_le_to_fr(hash.as_ref());
-    el
-}
-
 pub fn compute_id_secret(
     share1: (Fr, Fr),
     share2: (Fr, Fr),
@@ -541,9 +530,11 @@ pub enum ProofError {
     SynthesisError(#[from] SynthesisError),
 }
 
-fn calculate_witness_element<E: ark_ec::PairingEngine>(witness: Vec<BigInt>) -> Result<Vec<E::Fr>> {
-    use ark_ff::{FpParameters, PrimeField};
-    let modulus = <<E::Fr as PrimeField>::Params as FpParameters>::MODULUS;
+fn calculate_witness_element<E: ark_ec::pairing::Pairing>(
+    witness: Vec<BigInt>,
+) -> Result<Vec<E::ScalarField>> {
+    use ark_ff::PrimeField;
+    let modulus = <E::ScalarField as PrimeField>::MODULUS;
 
     // convert it to field elements
     use num_traits::Signed;
@@ -558,7 +549,7 @@ fn calculate_witness_element<E: ark_ec::PairingEngine>(witness: Vec<BigInt>) ->
         } else {
             w.to_biguint().ok_or(Report::msg("not a biguint value"))?
         };
-        witness_vec.push(E::Fr::from(w))
+        witness_vec.push(E::ScalarField::from(w))
     }
 
     Ok(witness_vec)
@@ -587,7 +578,7 @@ pub fn generate_proof_with_witness(
     #[cfg(debug_assertions)]
     let now = Instant::now();
 
-    let proof = create_proof_with_reduction_and_matrices::<_, CircomReduction>(
+    let proof = Groth16::<_, CircomReduction>::create_proof_with_reduction_and_matrices(
         &proving_key.0,
         r,
         s,
@@ -681,7 +672,7 @@ pub fn generate_proof(
     #[cfg(debug_assertions)]
     let now = Instant::now();
 
-    let proof = create_proof_with_reduction_and_matrices::<_, CircomReduction>(
+    let proof = Groth16::<_, CircomReduction>::create_proof_with_reduction_and_matrices(
         &proving_key.0,
         r,
         s,
@@ -726,7 +717,7 @@ pub fn verify_proof(
     #[cfg(debug_assertions)]
     let now = Instant::now();
 
-    let verified = ark_verify_proof(&pvk, proof, &inputs)?;
+    let verified = Groth16::<_, CircomReduction>::verify_proof(&pvk, proof, &inputs)?;
 
     #[cfg(debug_assertions)]
     println!("verify took: {:.2?}", now.elapsed());

rln/src/public.rs

@@ -1,5 +1,5 @@
 use crate::circuit::{vk_from_raw, zkey_from_raw, Curve, Fr};
-use crate::poseidon_hash::poseidon_hash as utils_poseidon_hash;
+use crate::hashers::{hash_to_field, poseidon_hash as utils_poseidon_hash};
 use crate::poseidon_tree::PoseidonTree;
 use crate::protocol::*;
 use crate::utils::*;
@@ -10,9 +10,10 @@ use ark_groth16::{ProvingKey, VerifyingKey};
 use ark_relations::r1cs::ConstraintMatrices;
 use ark_serialize::{CanonicalDeserialize, CanonicalSerialize, Read, Write};
 use cfg_if::cfg_if;
-use color_eyre::Result;
+use color_eyre::{Report, Result};
 use num_bigint::BigInt;
 use std::io::Cursor;
+use utils::{ZerokitMerkleProof, ZerokitMerkleTree};
 
 cfg_if! {
     if #[cfg(not(target_arch = "wasm32"))] {
@@ -20,6 +21,9 @@ cfg_if! {
         use std::sync::Mutex;
         use crate::circuit::{circom_from_folder, vk_from_folder, circom_from_raw, zkey_from_folder, TEST_RESOURCES_FOLDER, TEST_TREE_HEIGHT};
         use ark_circom::WitnessCalculator;
+        use serde_json::{json, Value};
+        use utils::{Hasher};
+        use std::str::FromStr;
     } else {
         use std::marker::*;
     }
@@ -61,7 +65,7 @@ impl RLN<'_> {
     /// use std::io::Cursor;
     ///
     /// let tree_height = 20;
-    /// let resources = Cursor::new("tree_height_20");
+    /// let resources = Cursor::new(json!({"resources_folder": "tree_height_20"});
     ///
     /// // We create a new RLN instance
     /// let mut rln = RLN::new(tree_height, resources);
@@ -72,15 +76,29 @@ impl RLN<'_> {
         let mut input: Vec<u8> = Vec::new();
         input_data.read_to_end(&mut input)?;
 
-        let resources_folder = String::from_utf8(input)?;
+        let rln_config: Value = serde_json::from_str(&String::from_utf8(input)?)?;
+        let resources_folder = rln_config["resources_folder"]
+            .as_str()
+            .unwrap_or(TEST_RESOURCES_FOLDER);
+        let tree_config = rln_config["tree_config"].to_string();
 
-        let witness_calculator = circom_from_folder(&resources_folder)?;
+        let witness_calculator = circom_from_folder(resources_folder)?;
 
-        let proving_key = zkey_from_folder(&resources_folder)?;
-        let verification_key = vk_from_folder(&resources_folder)?;
+        let proving_key = zkey_from_folder(resources_folder)?;
+        let verification_key = vk_from_folder(resources_folder)?;
+
+        let tree_config: <PoseidonTree as ZerokitMerkleTree>::Config = if tree_config.is_empty() {
+            <PoseidonTree as ZerokitMerkleTree>::Config::default()
+        } else {
+            <PoseidonTree as ZerokitMerkleTree>::Config::from_str(&tree_config)?
+        };
 
         // We compute a default empty tree
-        let tree = PoseidonTree::default(tree_height);
+        let tree = PoseidonTree::new(
+            tree_height,
+            <PoseidonTree as ZerokitMerkleTree>::Hasher::default_leaf(),
+            tree_config,
+        )?;
 
         Ok(RLN {
             witness_calculator,
@@ -99,6 +117,7 @@ impl RLN<'_> {
     /// - `circom_vec`: a byte vector containing the ZK circuit (`rln.wasm`) as binary file
     /// - `zkey_vec`: a byte vector containing to the proving key (`rln_final.zkey`) as binary file
     /// - `vk_vec`: a byte vector containing to the verification key (`verification_key.json`) as binary file
+    /// - `tree_config`: a reader for a string containing a json with the merkle tree configuration
     ///
     /// Example:
     /// ```
@@ -116,6 +135,8 @@ impl RLN<'_> {
     ///     let mut buffer = vec![0; metadata.len() as usize];
     ///     file.read_exact(&mut buffer).expect("buffer overflow");
     ///     resources.push(buffer);
+    ///     let tree_config = "{}".to_string();
+    ///     let tree_config_buffer = &Buffer::from(tree_config.as_bytes());
     /// }
     ///
     /// let mut rln = RLN::new_with_params(
@@ -123,11 +144,51 @@ impl RLN<'_> {
     ///     resources[0].clone(),
     ///     resources[1].clone(),
     ///     resources[2].clone(),
+    ///     tree_config_buffer,
     /// );
     /// ```
+    #[cfg(not(target_arch = "wasm32"))]
+    pub fn new_with_params<R: Read>(
+        tree_height: usize,
+        circom_vec: Vec<u8>,
+        zkey_vec: Vec<u8>,
+        vk_vec: Vec<u8>,
+        mut tree_config_input: R,
+    ) -> Result<RLN<'static>> {
+        #[cfg(not(target_arch = "wasm32"))]
+        let witness_calculator = circom_from_raw(circom_vec)?;
+
+        let proving_key = zkey_from_raw(&zkey_vec)?;
+        let verification_key = vk_from_raw(&vk_vec, &zkey_vec)?;
+
+        let mut tree_config_vec: Vec<u8> = Vec::new();
+        tree_config_input.read_to_end(&mut tree_config_vec)?;
+        let tree_config_str = String::from_utf8(tree_config_vec)?;
+        let tree_config: <PoseidonTree as ZerokitMerkleTree>::Config = if tree_config_str.is_empty()
+        {
+            <PoseidonTree as ZerokitMerkleTree>::Config::default()
+        } else {
+            <PoseidonTree as ZerokitMerkleTree>::Config::from_str(&tree_config_str)?
+        };
+
+        // We compute a default empty tree
+        let tree = PoseidonTree::new(
+            tree_height,
+            <PoseidonTree as ZerokitMerkleTree>::Hasher::default_leaf(),
+            tree_config,
+        )?;
+
+        Ok(RLN {
+            witness_calculator,
+            proving_key,
+            verification_key,
+            tree,
+        })
+    }
+
+    #[cfg(target_arch = "wasm32")]
     pub fn new_with_params(
         tree_height: usize,
-        #[cfg(not(target_arch = "wasm32"))] circom_vec: Vec<u8>,
         zkey_vec: Vec<u8>,
         vk_vec: Vec<u8>,
     ) -> Result<RLN<'static>> {
@@ -138,15 +199,12 @@ impl RLN<'_> {
         let verification_key = vk_from_raw(&vk_vec, &zkey_vec)?;
 
         // We compute a default empty tree
-        let tree = PoseidonTree::default(tree_height);
+        let tree = PoseidonTree::default(tree_height)?;
 
         Ok(RLN {
-            #[cfg(not(target_arch = "wasm32"))]
-            witness_calculator,
             proving_key,
             verification_key,
             tree,
-            #[cfg(target_arch = "wasm32")]
             _marker: PhantomData,
         })
     }
@@ -162,7 +220,7 @@ impl RLN<'_> {
     /// - `tree_height`: the height of the Merkle tree.
     pub fn set_tree(&mut self, tree_height: usize) -> Result<()> {
         // We compute a default empty tree of desired height
-        self.tree = PoseidonTree::default(tree_height);
+        self.tree = PoseidonTree::default(tree_height)?;
 
         Ok(())
     }
@@ -199,6 +257,31 @@ impl RLN<'_> {
         Ok(())
     }
 
+    /// Gets a leaf value at position index in the internal Merkle tree.
+    /// The leaf value is written to output_data.
+    /// Input values are:
+    /// - `index`: the index of the leaf
+    ///
+    /// Example:
+    /// ```
+    /// use crate::protocol::*;
+    /// use std::io::Cursor;
+    ///
+    /// let id_index = 10;
+    /// let mut buffer = Cursor::new(Vec::<u8>::new());
+    /// rln.get_leaf(id_index, &mut buffer).unwrap();
+    /// let id_commitment = deserialize_field_element(&buffer.into_inner()).unwrap();
+    pub fn get_leaf<W: Write>(&self, index: usize, mut output_data: W) -> Result<()> {
+        // We get the leaf at input index
+        let leaf = self.tree.get(index)?;
+
+        // We serialize the leaf and write it to output
+        let leaf_byte = fr_to_bytes_le(&leaf);
+        output_data.write_all(&leaf_byte)?;
+
+        Ok(())
+    }
+
     /// Sets multiple leaves starting from position index in the internal Merkle tree.
     ///
     /// If n leaves are passed as input, these will be set at positions `index`, `index+1`, ..., `index+n-1` respectively.
@@ -237,7 +320,10 @@ impl RLN<'_> {
         let (leaves, _) = bytes_le_to_vec_fr(&leaves_byte)?;
 
         // We set the leaves
-        self.tree.set_range(index, leaves)
+        self.tree
+            .override_range(index, leaves, [])
+            .map_err(|_| Report::msg("Could not set leaves"))?;
+        Ok(())
     }
 
     /// Resets the tree state to default and sets multiple leaves starting from index 0.
@@ -254,6 +340,76 @@ impl RLN<'_> {
         self.set_leaves_from(0, input_data)
     }
 
+    /// Sets multiple leaves starting from position index in the internal Merkle tree.
+    /// Also accepts an array of indices to remove from the tree.
+    ///
+    /// If n leaves are passed as input, these will be set at positions `index`, `index+1`, ..., `index+n-1` respectively.
+    /// If m indices are passed as input, these will be removed from the tree.
+    ///
+    /// This function updates the internal Merkle tree `next_index value indicating the next available index corresponding to a never-set leaf as `next_index = max(next_index, index + n)`.
+    ///
+    /// Input values are:
+    /// - `index`: the index of the first leaf to be set
+    /// - `input_leaves`: 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))
+    /// - `input_indices`: a reader for the serialization of multiple indices to remove (serialization done with [`rln::utils::vec_u8_to_bytes_le`](crate::utils::vec_u8_to_bytes_le))
+    ///
+    /// Example:
+    /// ```
+    /// use rln::circuit::Fr;
+    /// use rln::utils::*;
+    ///
+    /// let start_index = 10;
+    /// let no_of_leaves = 256;
+    ///
+    /// // We generate a vector of random leaves
+    /// let mut leaves: Vec<Fr> = Vec::new();
+    /// let mut rng = thread_rng();
+    /// for _ in 0..no_of_leaves {
+    ///     let (_, id_commitment) = keygen();
+    ///     leaves.push(id_commitment);
+    /// }
+    ///
+    /// let mut indices: Vec<u8> = Vec::new();
+    /// for i in 0..no_of_leaves {
+    ///    if i % 2 == 0 {
+    ///       indices.push(i as u8);
+    ///   }
+    /// }
+    ///
+    /// // We atomically add leaves and remove indices from the tree
+    /// let mut leaves_buffer = Cursor::new(vec_fr_to_bytes_le(&leaves));
+    /// let mut indices_buffer = Cursor::new(vec_u8_to_bytes_le(&indices));
+    /// rln.set_leaves_from(index, &mut leaves_buffer, indices_buffer).unwrap();
+    /// ```
+    pub fn atomic_operation<R: Read>(
+        &mut self,
+        index: usize,
+        mut input_leaves: R,
+        mut input_indices: R,
+    ) -> Result<()> {
+        // We read input
+        let mut leaves_byte: Vec<u8> = Vec::new();
+        input_leaves.read_to_end(&mut leaves_byte)?;
+
+        let (leaves, _) = bytes_le_to_vec_fr(&leaves_byte)?;
+
+        let mut indices_byte: Vec<u8> = Vec::new();
+        input_indices.read_to_end(&mut indices_byte)?;
+
+        let (indices, _) = bytes_le_to_vec_u8(&indices_byte)?;
+        let indices: Vec<usize> = indices.iter().map(|x| *x as usize).collect();
+
+        // We set the leaves
+        self.tree
+            .override_range(index, leaves, indices)
+            .map_err(|e| Report::msg(format!("Could not perform the batch operation: {e}")))?;
+        Ok(())
+    }
+
+    pub fn leaves_set(&mut self) -> usize {
+        self.tree.leaves_set()
+    }
+
     /// Sets a leaf value at the next available never-set leaf index.
     ///
     /// This function updates the internal Merkle tree `next_index` value indicating the next available index corresponding to a never-set leaf as `next_index = next_index + 1`.
@@ -325,6 +481,43 @@ impl RLN<'_> {
         Ok(())
     }
 
+    /// Sets some metadata that a consuming application may want to store in the RLN object.
+    /// This metadata is not used by the RLN module.
+    ///
+    /// Input values are:
+    /// - `metadata`: a byte vector containing the metadata
+    ///
+    /// Example
+    ///
+    /// ```
+    /// let metadata = b"some metadata";
+    /// rln.set_metadata(metadata).unwrap();
+    /// ```
+    pub fn set_metadata(&mut self, metadata: &[u8]) -> Result<()> {
+        self.tree.set_metadata(metadata)?;
+        Ok(())
+    }
+
+    /// Returns the metadata stored in the RLN object.
+    ///
+    /// Output values are:
+    /// - `output_data`: a writer receiving the serialization of the metadata
+    ///
+    /// Example
+    ///
+    /// ```
+    /// use std::io::Cursor;
+    ///
+    /// let mut buffer = Cursor::new(Vec::<u8>::new());
+    /// rln.get_metadata(&mut buffer).unwrap();
+    /// let metadata = buffer.into_inner();
+    /// ```
+    pub fn get_metadata<W: Write>(&self, mut output_data: W) -> Result<()> {
+        let metadata = self.tree.metadata()?;
+        output_data.write_all(&metadata)?;
+        Ok(())
+    }
+
     /// Returns the Merkle tree root
     ///
     /// Output values are:
@@ -407,7 +600,7 @@ impl RLN<'_> {
         mut input_data: R,
         mut output_data: W,
     ) -> Result<()> {
-        // We read input RLN witness and we deserialize it
+        // We read input RLN witness and we serialize_compressed it
         let mut serialized: Vec<u8> = Vec::new();
         input_data.read_to_end(&mut serialized)?;
         let (rln_witness, _) = deserialize_witness(&serialized)?;
@@ -421,7 +614,7 @@ impl RLN<'_> {
         let proof = generate_proof(self.witness_calculator, &self.proving_key, &rln_witness)?;
 
         // Note: we export a serialization of ark-groth16::Proof not semaphore::Proof
-        proof.serialize(&mut output_data)?;
+        proof.serialize_compressed(&mut output_data)?;
 
         Ok(())
     }
@@ -467,7 +660,7 @@ impl RLN<'_> {
         // [ proof<128> | root<32> | epoch<32> | share_x<32> | share_y<32> | nullifier<32> | rln_identifier<32> ]
         let mut input_byte: Vec<u8> = Vec::new();
         input_data.read_to_end(&mut input_byte)?;
-        let proof = ArkProof::deserialize(&mut Cursor::new(&input_byte[..128]))?;
+        let proof = ArkProof::deserialize_compressed(&mut Cursor::new(&input_byte[..128]))?;
 
         let (proof_values, _) = deserialize_proof_values(&input_byte[128..]);
 
@@ -526,7 +719,7 @@ impl RLN<'_> {
         mut input_data: R,
         mut output_data: W,
     ) -> Result<()> {
-        // We read input RLN witness and we deserialize it
+        // 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)?;
@@ -536,7 +729,7 @@ impl RLN<'_> {
 
         // Note: we export a serialization of ark-groth16::Proof not semaphore::Proof
         // This proof is compressed, i.e. 128 bytes long
-        proof.serialize(&mut output_data)?;
+        proof.serialize_compressed(&mut output_data)?;
         output_data.write_all(&serialize_proof_values(&proof_values))?;
 
         Ok(())
@@ -561,7 +754,7 @@ impl RLN<'_> {
 
         // Note: we export a serialization of ark-groth16::Proof not semaphore::Proof
         // This proof is compressed, i.e. 128 bytes long
-        proof.serialize(&mut output_data)?;
+        proof.serialize_compressed(&mut output_data)?;
         output_data.write_all(&serialize_proof_values(&proof_values))?;
         Ok(())
     }
@@ -597,12 +790,15 @@ impl RLN<'_> {
         let mut serialized: Vec<u8> = Vec::new();
         input_data.read_to_end(&mut serialized)?;
         let mut all_read = 0;
-        let proof = ArkProof::deserialize(&mut Cursor::new(&serialized[..128].to_vec()))?;
+        let proof =
+            ArkProof::deserialize_compressed(&mut Cursor::new(&serialized[..128].to_vec()))?;
         all_read += 128;
         let (proof_values, read) = deserialize_proof_values(&serialized[all_read..]);
         all_read += read;
 
-        let signal_len = usize::from_le_bytes(serialized[all_read..all_read + 8].try_into()?);
+        let signal_len = usize::try_from(u64::from_le_bytes(
+            serialized[all_read..all_read + 8].try_into()?,
+        ))?;
         all_read += 8;
 
         let signal: Vec<u8> = serialized[all_read..all_read + signal_len].to_vec();
@@ -672,12 +868,15 @@ impl RLN<'_> {
         let mut serialized: Vec<u8> = Vec::new();
         input_data.read_to_end(&mut serialized)?;
         let mut all_read = 0;
-        let proof = ArkProof::deserialize(&mut Cursor::new(&serialized[..128].to_vec()))?;
+        let proof =
+            ArkProof::deserialize_compressed(&mut Cursor::new(&serialized[..128].to_vec()))?;
         all_read += 128;
         let (proof_values, read) = deserialize_proof_values(&serialized[all_read..]);
         all_read += read;
 
-        let signal_len = usize::from_le_bytes(serialized[all_read..all_read + 8].try_into()?);
+        let signal_len = usize::try_from(u64::from_le_bytes(
+            serialized[all_read..all_read + 8].try_into()?,
+        ))?;
         all_read += 8;
 
         let signal: Vec<u8> = serialized[all_read..all_read + signal_len].to_vec();
@@ -745,7 +944,7 @@ impl RLN<'_> {
     /// let mut buffer = Cursor::new(Vec::<u8>::new());
     /// rln.key_gen(&mut buffer).unwrap();
     ///
-    /// // We deserialize the keygen output
+    /// // 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<()> {
@@ -775,7 +974,7 @@ impl RLN<'_> {
     /// let mut buffer = Cursor::new(Vec::<u8>::new());
     /// rln.extended_key_gen(&mut buffer).unwrap();
     ///
-    /// // We deserialize the keygen output
+    /// // 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<()> {
@@ -810,7 +1009,7 @@ impl RLN<'_> {
     /// rln.seeded_key_gen(&mut input_buffer, &mut output_buffer)
     ///     .unwrap();
     ///
-    /// // We deserialize the keygen output
+    /// // We serialize_compressed the keygen output
     /// let (identity_secret_hash, id_commitment) = deserialize_identity_pair(output_buffer.into_inner());
     /// ```
     pub fn seeded_key_gen<R: Read, W: Write>(
@@ -853,7 +1052,7 @@ impl RLN<'_> {
     /// rln.seeded_key_gen(&mut input_buffer, &mut output_buffer)
     ///     .unwrap();
     ///
-    /// // We deserialize the keygen output
+    /// // We serialize_compressed the keygen output
     /// let (identity_trapdoor, identity_nullifier, identity_secret_hash, id_commitment) = deserialize_identity_tuple(buffer.into_inner());
     /// ```
     pub fn seeded_extended_key_gen<R: Read, W: Write>(
@@ -912,7 +1111,7 @@ impl RLN<'_> {
         mut input_proof_data_2: R,
         mut output_data: W,
     ) -> Result<()> {
-        // We deserialize the two proofs and we get the corresponding RLNProofValues objects
+        // 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)?;
         // We skip deserialization of the zk-proof at the beginning
@@ -956,7 +1155,7 @@ impl RLN<'_> {
     ///
     /// The function returns the corresponding [`RLNWitnessInput`](crate::protocol::RLNWitnessInput) object serialized using [`rln::protocol::serialize_witness`](crate::protocol::serialize_witness)).
     pub fn get_serialized_rln_witness<R: Read>(&mut self, mut input_data: R) -> Result<Vec<u8>> {
-        // We read input RLN witness and we deserialize it
+        // 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)?;
@@ -974,13 +1173,21 @@ impl RLN<'_> {
         let (rln_witness, _) = deserialize_witness(serialized_witness)?;
         get_json_inputs(&rln_witness)
     }
+
+    /// Closes the connection to the Merkle tree database.
+    /// This function should be called before the RLN object is dropped.
+    /// If not called, the connection will be closed when the RLN object is dropped.
+    /// This improves robustness of the tree.
+    pub fn flush(&mut self) -> Result<()> {
+        self.tree.close_db_connection()
+    }
 }
 
 #[cfg(not(target_arch = "wasm32"))]
 impl Default for RLN<'_> {
     fn default() -> Self {
         let tree_height = TEST_TREE_HEIGHT;
-        let buffer = Cursor::new(TEST_RESOURCES_FOLDER);
+        let buffer = Cursor::new(json!({ "resources_folder": TEST_RESOURCES_FOLDER }).to_string());
         Self::new(tree_height, buffer).unwrap()
     }
 }
@@ -1004,7 +1211,7 @@ impl Default for RLN<'_> {
 /// hash(&mut input_buffer, &mut output_buffer)
 ///     .unwrap();
 ///
-/// // We deserialize the keygen output
+/// // 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<()> {
@@ -1037,7 +1244,7 @@ pub fn hash<R: Read, W: Write>(mut input_data: R, mut output_data: W) -> Result<
 /// poseidon_hash(&mut input_buffer, &mut output_buffer)
 ///     .unwrap();
 ///
-/// // We deserialize the hash output
+/// // 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<()> {
@@ -1056,6 +1263,8 @@ mod test {
     use super::*;
     use ark_std::{rand::thread_rng, UniformRand};
     use rand::Rng;
+    use serde_json::json;
+    use utils::ZerokitMerkleTree;
 
     #[test]
     // We test merkle batch Merkle tree additions
@@ -1071,7 +1280,8 @@ mod test {
         }
 
         // We create a new tree
-        let input_buffer = Cursor::new(TEST_RESOURCES_FOLDER);
+        let input_buffer =
+            Cursor::new(json!({ "resources_folder": TEST_RESOURCES_FOLDER }).to_string());
         let mut rln = RLN::new(tree_height, input_buffer).unwrap();
 
         // We first add leaves one by one specifying the index
@@ -1126,7 +1336,7 @@ mod test {
 
         // We now delete all leaves set and check if the root corresponds to the empty tree root
         // delete calls over indexes higher than no_of_leaves are ignored and will not increase self.tree.next_index
-        for i in 0..2 * no_of_leaves {
+        for i in 0..no_of_leaves {
             rln.delete_leaf(i).unwrap();
         }
 
@@ -1166,7 +1376,8 @@ mod test {
         let set_index = rng.gen_range(0..no_of_leaves) as usize;
 
         // We create a new tree
-        let input_buffer = Cursor::new(TEST_RESOURCES_FOLDER);
+        let input_buffer =
+            Cursor::new(json!({ "resources_folder": TEST_RESOURCES_FOLDER }).to_string());
         let mut rln = RLN::new(tree_height, input_buffer).unwrap();
 
         // We add leaves in a batch into the tree
@@ -1219,6 +1430,158 @@ mod test {
         let (root_single_additions, _) = bytes_le_to_fr(&buffer.into_inner());
 
         assert_eq!(root_batch_with_init, root_single_additions);
+
+        rln.flush().unwrap();
+    }
+
+    #[test]
+    // Tests the atomic_operation fn, which set_leaves_from uses internally
+    fn test_atomic_operation() {
+        let tree_height = TEST_TREE_HEIGHT;
+        let no_of_leaves = 256;
+
+        // We generate a vector of random leaves
+        let mut leaves: Vec<Fr> = Vec::new();
+        let mut rng = thread_rng();
+        for _ in 0..no_of_leaves {
+            leaves.push(Fr::rand(&mut rng));
+        }
+
+        // We create a new tree
+        let input_buffer =
+            Cursor::new(json!({ "resources_folder": TEST_RESOURCES_FOLDER }).to_string());
+        let mut rln = RLN::new(tree_height, input_buffer).unwrap();
+
+        // We add leaves in a batch into the tree
+        let mut buffer = Cursor::new(vec_fr_to_bytes_le(&leaves).unwrap());
+        rln.init_tree_with_leaves(&mut buffer).unwrap();
+
+        // We check if number of leaves set is consistent
+        assert_eq!(rln.tree.leaves_set(), no_of_leaves);
+
+        // We get the root of the tree obtained adding leaves in batch
+        let mut buffer = Cursor::new(Vec::<u8>::new());
+        rln.get_root(&mut buffer).unwrap();
+        let (root_after_insertion, _) = bytes_le_to_fr(&buffer.into_inner());
+
+        // We check if number of leaves set is consistent
+        assert_eq!(rln.tree.leaves_set(), no_of_leaves);
+
+        let last_leaf = leaves.last().unwrap();
+        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());
+
+        rln.atomic_operation(last_leaf_index, leaves_buffer, indices_buffer)
+            .unwrap();
+
+        // We get the root of the tree obtained after a no-op
+        let mut buffer = Cursor::new(Vec::<u8>::new());
+        rln.get_root(&mut buffer).unwrap();
+        let (root_after_noop, _) = bytes_le_to_fr(&buffer.into_inner());
+
+        assert_eq!(root_after_insertion, root_after_noop);
+    }
+
+    #[test]
+    fn test_atomic_operation_zero_indexed() {
+        // Test duplicated from https://github.com/waku-org/go-zerokit-rln/pull/12/files
+        let tree_height = TEST_TREE_HEIGHT;
+        let no_of_leaves = 256;
+
+        // We generate a vector of random leaves
+        let mut leaves: Vec<Fr> = Vec::new();
+        let mut rng = thread_rng();
+        for _ in 0..no_of_leaves {
+            leaves.push(Fr::rand(&mut rng));
+        }
+
+        // We create a new tree
+        let input_buffer =
+            Cursor::new(json!({ "resources_folder": TEST_RESOURCES_FOLDER }).to_string());
+        let mut rln = RLN::new(tree_height, input_buffer).unwrap();
+
+        // We add leaves in a batch into the tree
+        let mut buffer = Cursor::new(vec_fr_to_bytes_le(&leaves).unwrap());
+        rln.init_tree_with_leaves(&mut buffer).unwrap();
+
+        // We check if number of leaves set is consistent
+        assert_eq!(rln.tree.leaves_set(), no_of_leaves);
+
+        // We get the root of the tree obtained adding leaves in batch
+        let mut buffer = Cursor::new(Vec::<u8>::new());
+        rln.get_root(&mut buffer).unwrap();
+        let (root_after_insertion, _) = bytes_le_to_fr(&buffer.into_inner());
+
+        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());
+        rln.atomic_operation(0, leaves_buffer, indices_buffer)
+            .unwrap();
+
+        // We get the root of the tree obtained after a deletion
+        let mut buffer = Cursor::new(Vec::<u8>::new());
+        rln.get_root(&mut buffer).unwrap();
+        let (root_after_deletion, _) = bytes_le_to_fr(&buffer.into_inner());
+
+        assert_ne!(root_after_insertion, root_after_deletion);
+    }
+
+    #[test]
+    fn test_atomic_operation_consistency() {
+        // Test duplicated from https://github.com/waku-org/go-zerokit-rln/pull/12/files
+        let tree_height = TEST_TREE_HEIGHT;
+        let no_of_leaves = 256;
+
+        // We generate a vector of random leaves
+        let mut leaves: Vec<Fr> = Vec::new();
+        let mut rng = thread_rng();
+        for _ in 0..no_of_leaves {
+            leaves.push(Fr::rand(&mut rng));
+        }
+
+        // We create a new tree
+        let input_buffer =
+            Cursor::new(json!({ "resources_folder": TEST_RESOURCES_FOLDER }).to_string());
+        let mut rln = RLN::new(tree_height, input_buffer).unwrap();
+
+        // We add leaves in a batch into the tree
+        let mut buffer = Cursor::new(vec_fr_to_bytes_le(&leaves).unwrap());
+        rln.init_tree_with_leaves(&mut buffer).unwrap();
+
+        // We check if number of leaves set is consistent
+        assert_eq!(rln.tree.leaves_set(), no_of_leaves);
+
+        // We get the root of the tree obtained adding leaves in batch
+        let mut buffer = Cursor::new(Vec::<u8>::new());
+        rln.get_root(&mut buffer).unwrap();
+        let (root_after_insertion, _) = bytes_le_to_fr(&buffer.into_inner());
+
+        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());
+        rln.atomic_operation(0, leaves_buffer, indices_buffer)
+            .unwrap();
+
+        // We get the root of the tree obtained after a deletion
+        let mut buffer = Cursor::new(Vec::<u8>::new());
+        rln.get_root(&mut buffer).unwrap();
+        let (root_after_deletion, _) = bytes_le_to_fr(&buffer.into_inner());
+
+        assert_ne!(root_after_insertion, root_after_deletion);
+
+        // We get the leaf
+        let mut output_buffer = Cursor::new(Vec::<u8>::new());
+        rln.get_leaf(set_index, &mut output_buffer).unwrap();
+        let (received_leaf, _) = bytes_le_to_fr(output_buffer.into_inner().as_ref());
+
+        assert_eq!(received_leaf, Fr::from(0));
     }
 
     #[allow(unused_must_use)]
@@ -1237,7 +1600,8 @@ mod test {
         let bad_index = (1 << tree_height) - rng.gen_range(0..no_of_leaves) as usize;
 
         // We create a new tree
-        let input_buffer = Cursor::new(TEST_RESOURCES_FOLDER);
+        let input_buffer =
+            Cursor::new(json!({ "resources_folder": TEST_RESOURCES_FOLDER }).to_string());
         let mut rln = RLN::new(tree_height, input_buffer).unwrap();
 
         // Get root of empty tree
@@ -1266,7 +1630,8 @@ mod test {
     fn test_groth16_proof() {
         let tree_height = TEST_TREE_HEIGHT;
 
-        let input_buffer = Cursor::new(TEST_RESOURCES_FOLDER);
+        let input_buffer =
+            Cursor::new(json!({ "resources_folder": TEST_RESOURCES_FOLDER }).to_string());
         let mut rln = RLN::new(tree_height, input_buffer).unwrap();
 
         // Note: we only test Groth16 proof generation, so we ignore setting the tree in the RLN object
@@ -1280,7 +1645,7 @@ mod test {
         let serialized_proof = output_buffer.into_inner();
 
         // Before checking public verify API, we check that the (deserialized) proof generated by prove is actually valid
-        let proof = ArkProof::deserialize(&mut Cursor::new(&serialized_proof)).unwrap();
+        let proof = ArkProof::deserialize_compressed(&mut Cursor::new(&serialized_proof)).unwrap();
         let verified = verify_proof(&rln.verification_key, &proof, &proof_values);
         assert!(verified.unwrap());
 
@@ -1310,7 +1675,8 @@ mod test {
         }
 
         // We create a new RLN instance
-        let input_buffer = Cursor::new(TEST_RESOURCES_FOLDER);
+        let input_buffer =
+            Cursor::new(json!({ "resources_folder": TEST_RESOURCES_FOLDER }).to_string());
         let mut rln = RLN::new(tree_height, input_buffer).unwrap();
 
         // We add leaves in a batch into the tree
@@ -1374,7 +1740,8 @@ mod test {
         }
 
         // We create a new RLN instance
-        let input_buffer = Cursor::new(TEST_RESOURCES_FOLDER);
+        let input_buffer =
+            Cursor::new(json!({ "resources_folder": TEST_RESOURCES_FOLDER }).to_string());
         let mut rln = RLN::new(tree_height, input_buffer).unwrap();
 
         // We add leaves in a batch into the tree
@@ -1407,7 +1774,7 @@ mod test {
 
         let mut input_buffer = Cursor::new(serialized);
 
-        // We read input RLN witness and we deserialize it
+        // We read input RLN witness and we serialize_compressed it
         let mut witness_byte: Vec<u8> = Vec::new();
         input_buffer.read_to_end(&mut witness_byte).unwrap();
         let (rln_witness, _) = proof_inputs_to_rln_witness(&mut rln.tree, &witness_byte).unwrap();
@@ -1470,7 +1837,8 @@ mod test {
         }
 
         // We create a new RLN instance
-        let input_buffer = Cursor::new(TEST_RESOURCES_FOLDER);
+        let input_buffer =
+            Cursor::new(json!({ "resources_folder": TEST_RESOURCES_FOLDER }).to_string());
         let mut rln = RLN::new(tree_height, input_buffer).unwrap();
 
         // We add leaves in a batch into the tree
@@ -1556,7 +1924,8 @@ mod test {
         let tree_height = TEST_TREE_HEIGHT;
 
         // We create a new RLN instance
-        let input_buffer = Cursor::new(TEST_RESOURCES_FOLDER);
+        let input_buffer =
+            Cursor::new(json!({ "resources_folder": TEST_RESOURCES_FOLDER }).to_string());
         let mut rln = RLN::new(tree_height, input_buffer).unwrap();
 
         // Generate identity pair
@@ -1675,4 +2044,50 @@ mod test {
         // We ensure that an empty value was written to output_buffer, i.e. no secret is recovered
         assert!(serialized_identity_secret_hash.is_empty());
     }
+
+    #[test]
+    fn test_get_leaf() {
+        // We generate a random tree
+        let tree_height = 10;
+        let mut rng = thread_rng();
+        let input_buffer =
+            Cursor::new(json!({ "resources_folder": TEST_RESOURCES_FOLDER }).to_string());
+        let mut rln = RLN::new(tree_height, input_buffer).unwrap();
+
+        // We generate a random leaf
+        let leaf = Fr::rand(&mut rng);
+
+        // We generate a random index
+        let index = rng.gen_range(0..rln.tree.capacity());
+
+        // We add the leaf to the tree
+        let mut buffer = Cursor::new(fr_to_bytes_le(&leaf));
+        rln.set_leaf(index, &mut buffer).unwrap();
+
+        // We get the leaf
+        let mut output_buffer = Cursor::new(Vec::<u8>::new());
+        rln.get_leaf(index, &mut output_buffer).unwrap();
+
+        // We ensure that the leaf is the same as the one we added
+        let (received_leaf, _) = bytes_le_to_fr(output_buffer.into_inner().as_ref());
+        assert_eq!(received_leaf, leaf);
+    }
+
+    #[test]
+    fn test_metadata() {
+        let tree_height = TEST_TREE_HEIGHT;
+
+        let input_buffer =
+            Cursor::new(json!({ "resources_folder": TEST_RESOURCES_FOLDER }).to_string());
+        let mut rln = RLN::new(tree_height, input_buffer).unwrap();
+
+        let arbitrary_metadata: &[u8] = b"block_number:200000";
+        rln.set_metadata(arbitrary_metadata).unwrap();
+
+        let mut buffer = Cursor::new(Vec::<u8>::new());
+        rln.get_metadata(&mut buffer).unwrap();
+        let received_metadata = buffer.into_inner();
+
+        assert_eq!(arbitrary_metadata, received_metadata);
+    }
 }

rln/src/utils.rs

@@ -13,8 +13,8 @@ pub fn to_bigint(el: &Fr) -> Result<BigInt> {
 }
 
 pub fn fr_byte_size() -> usize {
-    let mbs = <Fr as PrimeField>::size_in_bits();
-    (mbs + 64 - (mbs % 64)) / 8
+    let mbs = <Fr as PrimeField>::MODULUS_BIT_SIZE;
+    ((mbs + 64 - (mbs % 64)) / 8) as usize
 }
 
 pub fn str_to_fr(input: &str, radix: u32) -> Result<Fr> {
@@ -75,7 +75,7 @@ pub fn fr_to_bytes_be(input: &Fr) -> Vec<u8> {
 pub fn vec_fr_to_bytes_le(input: &[Fr]) -> Result<Vec<u8>> {
     let mut bytes: Vec<u8> = Vec::new();
     //We store the vector length
-    bytes.extend(input.len().to_le_bytes().to_vec());
+    bytes.extend(u64::try_from(input.len())?.to_le_bytes().to_vec());
 
     // We store each element
     input.iter().for_each(|el| bytes.extend(fr_to_bytes_le(el)));
@@ -86,7 +86,7 @@ pub fn vec_fr_to_bytes_le(input: &[Fr]) -> Result<Vec<u8>> {
 pub fn vec_fr_to_bytes_be(input: &[Fr]) -> Result<Vec<u8>> {
     let mut bytes: Vec<u8> = Vec::new();
     //We store the vector length
-    bytes.extend(input.len().to_be_bytes().to_vec());
+    bytes.extend(u64::try_from(input.len())?.to_be_bytes().to_vec());
 
     // We store each element
     input.iter().for_each(|el| bytes.extend(fr_to_bytes_be(el)));
@@ -97,7 +97,7 @@ pub fn vec_fr_to_bytes_be(input: &[Fr]) -> Result<Vec<u8>> {
 pub fn vec_u8_to_bytes_le(input: &[u8]) -> Result<Vec<u8>> {
     let mut bytes: Vec<u8> = Vec::new();
     //We store the vector length
-    bytes.extend(input.len().to_le_bytes().to_vec());
+    bytes.extend(u64::try_from(input.len())?.to_le_bytes().to_vec());
 
     bytes.extend(input);
 
@@ -107,7 +107,7 @@ pub fn vec_u8_to_bytes_le(input: &[u8]) -> Result<Vec<u8>> {
 pub fn vec_u8_to_bytes_be(input: Vec<u8>) -> Result<Vec<u8>> {
     let mut bytes: Vec<u8> = Vec::new();
     //We store the vector length
-    bytes.extend(input.len().to_be_bytes().to_vec());
+    bytes.extend(u64::try_from(input.len())?.to_be_bytes().to_vec());
 
     bytes.extend(input);
 
@@ -117,7 +117,7 @@ pub fn vec_u8_to_bytes_be(input: Vec<u8>) -> Result<Vec<u8>> {
 pub fn bytes_le_to_vec_u8(input: &[u8]) -> Result<(Vec<u8>, usize)> {
     let mut read: usize = 0;
 
-    let len = usize::from_le_bytes(input[0..8].try_into()?);
+    let len = usize::try_from(u64::from_le_bytes(input[0..8].try_into()?))?;
     read += 8;
 
     let res = input[8..8 + len].to_vec();
@@ -129,7 +129,7 @@ pub fn bytes_le_to_vec_u8(input: &[u8]) -> Result<(Vec<u8>, usize)> {
 pub fn bytes_be_to_vec_u8(input: &[u8]) -> Result<(Vec<u8>, usize)> {
     let mut read: usize = 0;
 
-    let len = usize::from_be_bytes(input[0..8].try_into()?);
+    let len = usize::try_from(u64::from_be_bytes(input[0..8].try_into()?))?;
     read += 8;
 
     let res = input[8..8 + len].to_vec();
@@ -143,7 +143,7 @@ pub fn bytes_le_to_vec_fr(input: &[u8]) -> Result<(Vec<Fr>, usize)> {
     let mut read: usize = 0;
     let mut res: Vec<Fr> = Vec::new();
 
-    let len = usize::from_le_bytes(input[0..8].try_into()?);
+    let len = usize::try_from(u64::from_le_bytes(input[0..8].try_into()?))?;
     read += 8;
 
     let el_size = fr_byte_size();
@@ -160,7 +160,7 @@ pub fn bytes_be_to_vec_fr(input: &[u8]) -> Result<(Vec<Fr>, usize)> {
     let mut read: usize = 0;
     let mut res: Vec<Fr> = Vec::new();
 
-    let len = usize::from_be_bytes(input[0..8].try_into()?);
+    let len = usize::try_from(u64::from_be_bytes(input[0..8].try_into()?))?;
     read += 8;
 
     let el_size = fr_byte_size();

rln/tests/ffi.rs

@@ -4,10 +4,11 @@ mod test {
     use rand::Rng;
     use rln::circuit::*;
     use rln::ffi::{hash as ffi_hash, poseidon_hash as ffi_poseidon_hash, *};
-    use rln::poseidon_hash::{poseidon_hash as utils_poseidon_hash, ROUND_PARAMS};
+    use rln::hashers::{hash_to_field, poseidon_hash as utils_poseidon_hash, ROUND_PARAMS};
     use rln::protocol::*;
     use rln::public::RLN;
     use rln::utils::*;
+    use serde_json::json;
     use std::fs::File;
     use std::io::Read;
     use std::mem::MaybeUninit;
@@ -28,7 +29,8 @@ mod test {
 
         // We create a RLN instance
         let mut rln_pointer = MaybeUninit::<*mut RLN>::uninit();
-        let input_buffer = &Buffer::from(TEST_RESOURCES_FOLDER.as_bytes());
+        let input_config = json!({ "resource_folder": TEST_RESOURCES_FOLDER }).to_string();
+        let input_buffer = &Buffer::from(input_config.as_bytes());
         let success = new(tree_height, input_buffer, rln_pointer.as_mut_ptr());
         assert!(success, "RLN object creation failed");
         let rln_pointer = unsafe { &mut *rln_pointer.assume_init() };
@@ -96,8 +98,7 @@ mod test {
 
         // We now delete all leaves set and check if the root corresponds to the empty tree root
         // delete calls over indexes higher than no_of_leaves are ignored and will not increase self.tree.next_index
-        let delete_range = 2 * no_of_leaves;
-        for i in 0..delete_range {
+        for i in 0..no_of_leaves {
             let success = delete_leaf(rln_pointer, i);
             assert!(success, "delete leaf call failed");
         }
@@ -136,7 +137,8 @@ mod test {
 
         // We create a RLN instance
         let mut rln_pointer = MaybeUninit::<*mut RLN>::uninit();
-        let input_buffer = &Buffer::from(TEST_RESOURCES_FOLDER.as_bytes());
+        let input_config = json!({ "resources_folder": TEST_RESOURCES_FOLDER }).to_string();
+        let input_buffer = &Buffer::from(input_config.as_bytes());
         let success = new(tree_height, input_buffer, rln_pointer.as_mut_ptr());
         assert!(success, "RLN object creation failed");
         let rln_pointer = unsafe { &mut *rln_pointer.assume_init() };
@@ -216,6 +218,71 @@ mod test {
         assert_eq!(root_batch_with_init, root_single_additions);
     }
 
+    #[test]
+    // This test is similar to the one in public.rs but it uses the RLN object as a pointer
+    fn test_atomic_operation_ffi() {
+        let tree_height = TEST_TREE_HEIGHT;
+        let no_of_leaves = 256;
+
+        // We create a RLN instance
+        let mut rln_pointer = MaybeUninit::<*mut RLN>::uninit();
+        let input_config = json!({ "resources_folder": TEST_RESOURCES_FOLDER }).to_string();
+        let input_buffer = &Buffer::from(input_config.as_bytes());
+        let success = new(tree_height, input_buffer, rln_pointer.as_mut_ptr());
+        assert!(success, "RLN object creation failed");
+        let rln_pointer = unsafe { &mut *rln_pointer.assume_init() };
+
+        // We generate a vector of random leaves
+        let mut leaves: Vec<Fr> = Vec::new();
+        let mut rng = thread_rng();
+        for _ in 0..no_of_leaves {
+            leaves.push(Fr::rand(&mut rng));
+        }
+
+        // We add leaves in a batch into the tree
+        let leaves_ser = vec_fr_to_bytes_le(&leaves).unwrap();
+        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");
+
+        // We get the root of the tree obtained adding leaves in batch
+        let mut output_buffer = MaybeUninit::<Buffer>::uninit();
+        let success = get_root(rln_pointer, output_buffer.as_mut_ptr());
+        assert!(success, "get root call failed");
+
+        let output_buffer = unsafe { output_buffer.assume_init() };
+        let result_data = <&[u8]>::from(&output_buffer).to_vec();
+        let (root_after_insertion, _) = bytes_le_to_fr(&result_data);
+
+        let last_leaf = leaves.last().unwrap();
+        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_buffer = &Buffer::from(indices.as_ref());
+        let leaves = vec_fr_to_bytes_le(&last_leaf).unwrap();
+        let leaves_buffer = &Buffer::from(leaves.as_ref());
+
+        let success = atomic_operation(
+            rln_pointer,
+            last_leaf_index as usize,
+            leaves_buffer,
+            indices_buffer,
+        );
+        assert!(success, "atomic operation call failed");
+
+        // We get the root of the tree obtained after a no-op
+        let mut output_buffer = MaybeUninit::<Buffer>::uninit();
+        let success = get_root(rln_pointer, output_buffer.as_mut_ptr());
+        assert!(success, "get root call failed");
+
+        let output_buffer = unsafe { output_buffer.assume_init() };
+        let result_data = <&[u8]>::from(&output_buffer).to_vec();
+        let (root_after_noop, _) = bytes_le_to_fr(&result_data);
+
+        assert_eq!(root_after_insertion, root_after_noop);
+    }
+
     #[test]
     // This test is similar to the one in public.rs but it uses the RLN object as a pointer
     fn test_set_leaves_bad_index_ffi() {
@@ -233,7 +300,8 @@ mod test {
 
         // We create a RLN instance
         let mut rln_pointer = MaybeUninit::<*mut RLN>::uninit();
-        let input_buffer = &Buffer::from(TEST_RESOURCES_FOLDER.as_bytes());
+        let input_config = json!({ "resources_folder": TEST_RESOURCES_FOLDER }).to_string();
+        let input_buffer = &Buffer::from(input_config.as_bytes());
         let success = new(tree_height, input_buffer, rln_pointer.as_mut_ptr());
         assert!(success, "RLN object creation failed");
         let rln_pointer = unsafe { &mut *rln_pointer.assume_init() };
@@ -273,7 +341,8 @@ mod test {
 
         // We create a RLN instance
         let mut rln_pointer = MaybeUninit::<*mut RLN>::uninit();
-        let input_buffer = &Buffer::from(TEST_RESOURCES_FOLDER.as_bytes());
+        let input_config = json!({ "resources_folder": TEST_RESOURCES_FOLDER }).to_string();
+        let input_buffer = &Buffer::from(input_config.as_bytes());
         let success = new(tree_height, input_buffer, rln_pointer.as_mut_ptr());
         assert!(success, "RLN object creation failed");
         let rln_pointer = unsafe { &mut *rln_pointer.assume_init() };
@@ -441,7 +510,8 @@ mod test {
 
         // We create a RLN instance
         let mut rln_pointer = MaybeUninit::<*mut RLN>::uninit();
-        let input_buffer = &Buffer::from(TEST_RESOURCES_FOLDER.as_bytes());
+        let input_config = json!({ "resources_folder": TEST_RESOURCES_FOLDER }).to_string();
+        let input_buffer = &Buffer::from(input_config.as_bytes());
         let success = new(tree_height, input_buffer, rln_pointer.as_mut_ptr());
         assert!(success, "RLN object creation failed");
         let rln_pointer = unsafe { &mut *rln_pointer.assume_init() };
@@ -503,7 +573,8 @@ mod test {
 
         // We create a RLN instance using a resource folder path
         let mut rln_pointer = MaybeUninit::<*mut RLN>::uninit();
-        let input_buffer = &Buffer::from(TEST_RESOURCES_FOLDER.as_bytes());
+        let input_config = json!({ "resources_folder": TEST_RESOURCES_FOLDER }).to_string();
+        let input_buffer = &Buffer::from(input_config.as_bytes());
         let success = new(tree_height, input_buffer, rln_pointer.as_mut_ptr());
         assert!(success, "RLN object creation failed");
         let rln_pointer = unsafe { &mut *rln_pointer.assume_init() };
@@ -546,11 +617,14 @@ mod test {
 
         // Creating a RLN instance passing the raw data
         let mut rln_pointer_raw_bytes = MaybeUninit::<*mut RLN>::uninit();
+        let tree_config = "".to_string();
+        let tree_config_buffer = &Buffer::from(tree_config.as_bytes());
         let success = new_with_params(
             tree_height,
             circom_data,
             zkey_data,
             vk_data,
+            tree_config_buffer,
             rln_pointer_raw_bytes.as_mut_ptr(),
         );
         assert!(success, "RLN object creation failed");
@@ -583,7 +657,8 @@ mod test {
 
         // We create a RLN instance
         let mut rln_pointer = MaybeUninit::<*mut RLN>::uninit();
-        let input_buffer = &Buffer::from(TEST_RESOURCES_FOLDER.as_bytes());
+        let input_config = json!({ "resources_folder": TEST_RESOURCES_FOLDER }).to_string();
+        let input_buffer = &Buffer::from(input_config.as_bytes());
         let success = new(tree_height, input_buffer, rln_pointer.as_mut_ptr());
         assert!(success, "RLN object creation failed");
         let rln_pointer = unsafe { &mut *rln_pointer.assume_init() };
@@ -667,7 +742,8 @@ mod test {
 
         // We create a RLN instance
         let mut rln_pointer = MaybeUninit::<*mut RLN>::uninit();
-        let input_buffer = &Buffer::from(TEST_RESOURCES_FOLDER.as_bytes());
+        let input_config = json!({ "resources_folder": TEST_RESOURCES_FOLDER }).to_string();
+        let input_buffer = &Buffer::from(input_config.as_bytes());
         let success = new(tree_height, input_buffer, rln_pointer.as_mut_ptr());
         assert!(success, "RLN object creation failed");
         let rln_pointer = unsafe { &mut *rln_pointer.assume_init() };
@@ -785,7 +861,8 @@ mod test {
 
         // We create a RLN instance
         let mut rln_pointer = MaybeUninit::<*mut RLN>::uninit();
-        let input_buffer = &Buffer::from(TEST_RESOURCES_FOLDER.as_bytes());
+        let input_config = json!({ "resources_folder": TEST_RESOURCES_FOLDER }).to_string();
+        let input_buffer = &Buffer::from(input_config.as_bytes());
         let success = new(tree_height, input_buffer, rln_pointer.as_mut_ptr());
         assert!(success, "RLN object creation failed");
         let rln_pointer = unsafe { &mut *rln_pointer.assume_init() };
@@ -944,7 +1021,8 @@ mod test {
 
         // We create a RLN instance
         let mut rln_pointer = MaybeUninit::<*mut RLN>::uninit();
-        let input_buffer = &Buffer::from(TEST_RESOURCES_FOLDER.as_bytes());
+        let input_config = json!({ "resources_folder": TEST_RESOURCES_FOLDER }).to_string();
+        let input_buffer = &Buffer::from(input_config.as_bytes());
         let success = new(tree_height, input_buffer, rln_pointer.as_mut_ptr());
         assert!(success, "RLN object creation failed");
         let rln_pointer = unsafe { &mut *rln_pointer.assume_init() };
@@ -983,7 +1061,8 @@ mod test {
         let tree_height = TEST_TREE_HEIGHT;
         // We create a RLN instance
         let mut rln_pointer = MaybeUninit::<*mut RLN>::uninit();
-        let input_buffer = &Buffer::from(TEST_RESOURCES_FOLDER.as_bytes());
+        let input_config = json!({ "resources_folder": TEST_RESOURCES_FOLDER }).to_string();
+        let input_buffer = &Buffer::from(input_config.as_bytes());
         let success = new(tree_height, input_buffer, rln_pointer.as_mut_ptr());
         assert!(success, "RLN object creation failed");
         let rln_pointer = unsafe { &mut *rln_pointer.assume_init() };
@@ -1080,4 +1159,77 @@ mod test {
 
         assert_eq!(received_hash, expected_hash);
     }
+
+    #[test]
+    fn test_get_leaf() {
+        // We create a RLN instance
+        let tree_height = TEST_TREE_HEIGHT;
+        let no_of_leaves = 1 << TEST_TREE_HEIGHT;
+
+        let mut rln_pointer = MaybeUninit::<*mut RLN>::uninit();
+        let input_config = json!({ "resources_folder": TEST_RESOURCES_FOLDER }).to_string();
+        let input_buffer = &Buffer::from(input_config.as_bytes());
+        let success = new(tree_height, input_buffer, rln_pointer.as_mut_ptr());
+        assert!(success, "RLN object creation failed");
+        let rln_pointer = unsafe { &mut *rln_pointer.assume_init() };
+
+        // We generate a new identity tuple from an input seed
+        let seed_bytes: &[u8] = &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9];
+        let input_buffer = &Buffer::from(seed_bytes);
+        let mut output_buffer = MaybeUninit::<Buffer>::uninit();
+        let success =
+            seeded_extended_key_gen(rln_pointer, input_buffer, output_buffer.as_mut_ptr());
+        assert!(success, "seeded key gen call failed");
+
+        let output_buffer = unsafe { output_buffer.assume_init() };
+        let result_data = <&[u8]>::from(&output_buffer).to_vec();
+        let (_, _, _, id_commitment) = deserialize_identity_tuple(result_data);
+
+        // We insert the id_commitment into the tree at a random index
+        let mut rng = thread_rng();
+        let index = rng.gen_range(0..no_of_leaves) as usize;
+        let leaf = fr_to_bytes_le(&id_commitment);
+        let input_buffer = &Buffer::from(leaf.as_ref());
+        let success = set_leaf(rln_pointer, index, input_buffer);
+        assert!(success, "set leaf call failed");
+
+        // We get the leaf at the same index
+        let mut output_buffer = MaybeUninit::<Buffer>::uninit();
+        let success = get_leaf(rln_pointer, index, output_buffer.as_mut_ptr());
+        assert!(success, "get leaf call failed");
+        let output_buffer = unsafe { output_buffer.assume_init() };
+        let result_data = <&[u8]>::from(&output_buffer).to_vec();
+        let (received_id_commitment, _) = bytes_le_to_fr(&result_data);
+
+        // We check that the received id_commitment is the same as the one we inserted
+        assert_eq!(received_id_commitment, id_commitment);
+    }
+
+    #[test]
+    fn test_metadata() {
+        // We create a RLN instance
+        let tree_height = TEST_TREE_HEIGHT;
+
+        let mut rln_pointer = MaybeUninit::<*mut RLN>::uninit();
+        let input_config = json!({ "resources_folder": TEST_RESOURCES_FOLDER }).to_string();
+        let input_buffer = &Buffer::from(input_config.as_bytes());
+        let success = new(tree_height, input_buffer, rln_pointer.as_mut_ptr());
+        assert!(success, "RLN object creation failed");
+        let rln_pointer = unsafe { &mut *rln_pointer.assume_init() };
+
+        let seed_bytes: &[u8] = &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9];
+        let input_buffer = &Buffer::from(seed_bytes);
+
+        let success = set_metadata(rln_pointer, input_buffer);
+        assert!(success, "set_metadata call failed");
+
+        let mut output_buffer = MaybeUninit::<Buffer>::uninit();
+        let success = get_metadata(rln_pointer, output_buffer.as_mut_ptr());
+        assert!(success, "get_metadata call failed");
+
+        let output_buffer = unsafe { output_buffer.assume_init() };
+        let result_data = <&[u8]>::from(&output_buffer).to_vec();
+
+        assert_eq!(result_data, seed_bytes.to_vec());
+    }
 }

rln/tests/poseidon_tree.rs

@@ -5,8 +5,8 @@
 #[cfg(test)]
 mod test {
     use rln::circuit::*;
-    use rln::poseidon_tree::*;
-    use utils::{FullMerkleTree, OptimalMerkleTree};
+    use rln::hashers::PoseidonHash;
+    use utils::{FullMerkleTree, OptimalMerkleTree, ZerokitMerkleProof, ZerokitMerkleTree};
 
     #[test]
     /// A basic performance comparison between the two supported Merkle Tree implementations
@@ -25,16 +25,16 @@ mod test {
 
         for _ in 0..sample_size.try_into().unwrap() {
             let now = Instant::now();
-            FullMerkleTree::<PoseidonHash>::default(tree_height);
+            FullMerkleTree::<PoseidonHash>::default(tree_height).unwrap();
             gen_time_full += now.elapsed().as_nanos();
 
             let now = Instant::now();
-            OptimalMerkleTree::<PoseidonHash>::default(tree_height);
+            OptimalMerkleTree::<PoseidonHash>::default(tree_height).unwrap();
             gen_time_opt += now.elapsed().as_nanos();
         }
 
-        let mut tree_full = FullMerkleTree::<PoseidonHash>::default(tree_height);
-        let mut tree_opt = OptimalMerkleTree::<PoseidonHash>::default(tree_height);
+        let mut tree_full = FullMerkleTree::<PoseidonHash>::default(tree_height).unwrap();
+        let mut tree_opt = OptimalMerkleTree::<PoseidonHash>::default(tree_height).unwrap();
 
         for i in 0..sample_size.try_into().unwrap() {
             let now = Instant::now();
@@ -78,540 +78,3 @@ mod test {
         );
     }
 }
-
-// Test module for testing pmtree integration and features in zerokit
-// enabled only if the pmtree feature is enabled
-
-#[cfg(feature = "pmtree-ft")]
-#[cfg(test)]
-mod pmtree_test {
-
-    use pmtree::*;
-    use rln::circuit::Fr;
-    use rln::poseidon_hash::poseidon_hash;
-    use rln::poseidon_tree::PoseidonHash;
-    use rln::protocol::hash_to_field;
-    use rln::utils::str_to_fr;
-    use sled::Db as Sled;
-    use std::fs;
-    use std::{collections::HashMap, path::PathBuf};
-    use utils::{FullMerkleTree, OptimalMerkleTree};
-
-    // pmtree supports in-memory and on-disk databases (Database trait) for storing the Merkle tree state
-
-    // We implement Database for hashmaps, an in-memory database
-    struct MemoryDB(HashMap<DBKey, Value>);
-
-    #[derive(Default)]
-    struct MemoryDBConfig {}
-
-    impl Database for MemoryDB {
-        type Config = MemoryDBConfig;
-
-        fn new(_config: Self::Config) -> Result<Self> {
-            Ok(MemoryDB(HashMap::new()))
-        }
-
-        fn load(_config: Self::Config) -> Result<Self> {
-            Err(Box::new(Error("Cannot load in-memory DB".to_string())))
-        }
-
-        fn get(&self, key: DBKey) -> Result<Option<Value>> {
-            Ok(self.0.get(&key).cloned())
-        }
-
-        fn put(&mut self, key: DBKey, value: Value) -> Result<()> {
-            self.0.insert(key, value);
-
-            Ok(())
-        }
-
-        fn put_batch(&mut self, subtree: HashMap<DBKey, Value>) -> Result<()> {
-            self.0.extend(subtree);
-            Ok(())
-        }
-    }
-
-    // We implement Database for sled DB, an on-disk database
-    struct SledDB(Sled);
-
-    impl Database for SledDB {
-        type Config = sled::Config;
-
-        fn new(config: Self::Config) -> Result<Self> {
-            let dbpath = config.path;
-            if config.dbpath.exists() {
-                match fs::remove_dir_all(&config.dbpath) {
-                    Ok(x) => x,
-                    Err(e) => return Err(Box::new(Error(e.to_string()))),
-                }
-            }
-
-            let db: Sled = match config.open() {
-                Ok(db) => db,
-                Err(e) => return Err(Box::new(Error(e.to_string()))),
-            };
-
-            Ok(SledDB(db))
-        }
-
-        fn load(config: Self::Config) -> Result<Self> {
-            let db: Sled = match sled::open(config.dbpath) {
-                Ok(db) => db,
-                Err(e) => return Err(Box::new(Error(e.to_string()))),
-            };
-
-            if !db.was_recovered() {
-                return Err(Box::new(Error(
-                    "Trying to load non-existing database!".to_string(),
-                )));
-            }
-
-            Ok(SledDB(db))
-        }
-
-        fn get(&self, key: DBKey) -> Result<Option<Value>> {
-            match self.0.get(key) {
-                Ok(value) => Ok(value.map(|val| val.to_vec())),
-                Err(e) => Err(Box::new(Error(e.to_string()))),
-            }
-        }
-
-        fn put(&mut self, key: DBKey, value: Value) -> Result<()> {
-            match self.0.insert(key, value) {
-                Ok(_) => Ok(()),
-                Err(e) => Err(Box::new(Error(e.to_string()))),
-            }
-        }
-
-        fn put_batch(&mut self, subtree: HashMap<DBKey, Value>) -> Result<()> {
-            let mut batch = sled::Batch::default();
-
-            for (key, value) in subtree {
-                batch.insert(&key, value);
-            }
-
-            self.0.apply_batch(batch)?;
-            Ok(())
-        }
-    }
-
-    #[test]
-    /// A basic performance comparison between the two supported Merkle Tree implementations and in-memory/on-disk pmtree implementations
-    fn test_zerokit_and_pmtree_merkle_implementations_performances() {
-        use std::time::{Duration, Instant};
-
-        let tree_height = 20;
-        let sample_size = 100;
-
-        let leaves: Vec<Fr> = (0..sample_size).map(|s| Fr::from(s)).collect();
-
-        let mut gen_time_full: u128 = 0;
-        let mut upd_time_full: u128 = 0;
-        let mut gen_time_opt: u128 = 0;
-        let mut upd_time_opt: u128 = 0;
-        let mut gen_time_pm_memory: u128 = 0;
-        let mut upd_time_pm_memory: u128 = 0;
-        let mut gen_time_pm_sled: u128 = 0;
-        let mut upd_time_pm_sled: u128 = 0;
-
-        for _ in 0..sample_size.try_into().unwrap() {
-            let now = Instant::now();
-            FullMerkleTree::<PoseidonHash>::default(tree_height);
-            gen_time_full += now.elapsed().as_nanos();
-
-            let now = Instant::now();
-            OptimalMerkleTree::<PoseidonHash>::default(tree_height);
-            gen_time_opt += now.elapsed().as_nanos();
-
-            let now = Instant::now();
-            pmtree::MerkleTree::<MemoryDB, PoseidonHash>::default(tree_height).unwrap();
-            gen_time_pm_memory += now.elapsed().as_nanos();
-
-            let now = Instant::now();
-            pmtree::MerkleTree::<SledDB, PoseidonHash>::default(tree_height).unwrap();
-            gen_time_pm_sled += now.elapsed().as_nanos();
-        }
-
-        let mut tree_full = FullMerkleTree::<PoseidonHash>::default(tree_height);
-        let mut tree_opt = OptimalMerkleTree::<PoseidonHash>::default(tree_height);
-        let mut tree_pm_memory =
-            pmtree::MerkleTree::<MemoryDB, PoseidonHash>::default(tree_height).unwrap();
-        let mut tree_pm_sled =
-            pmtree::MerkleTree::<SledDB, PoseidonHash>::default(tree_height).unwrap();
-
-        for i in 0..sample_size.try_into().unwrap() {
-            let now = Instant::now();
-            tree_full.set(i, leaves[i]).unwrap();
-            upd_time_full += now.elapsed().as_nanos();
-            let proof = tree_full.proof(i).expect("index should be set");
-            assert_eq!(proof.leaf_index(), i);
-
-            let now = Instant::now();
-            tree_opt.set(i, leaves[i]).unwrap();
-            upd_time_opt += now.elapsed().as_nanos();
-            let proof = tree_opt.proof(i).expect("index should be set");
-            assert_eq!(proof.leaf_index(), i);
-
-            let now = Instant::now();
-            tree_pm_memory.set(i, leaves[i]).unwrap();
-            upd_time_pm_memory += now.elapsed().as_nanos();
-            let proof = tree_pm_memory.proof(i).expect("index should be set");
-            assert_eq!(proof.leaf_index(), i);
-
-            let now = Instant::now();
-            tree_pm_sled.set(i, leaves[i]).unwrap();
-            upd_time_pm_sled += now.elapsed().as_nanos();
-            let proof = tree_pm_sled.proof(i).expect("index should be set");
-            assert_eq!(proof.leaf_index(), i);
-        }
-
-        // We check all roots are the same
-        let tree_full_root = tree_full.root();
-        let tree_opt_root = tree_opt.root();
-        let tree_pm_memory_root = tree_pm_memory.root();
-        let tree_pm_sled_root = tree_pm_sled.root();
-
-        assert_eq!(tree_full_root, tree_opt_root);
-        assert_eq!(tree_opt_root, tree_pm_memory_root);
-        assert_eq!(tree_pm_memory_root, tree_pm_sled_root);
-
-        println!(" Average tree generation time:");
-        println!(
-            "   - Full Merkle Tree:  {:?}",
-            Duration::from_nanos((gen_time_full / sample_size).try_into().unwrap())
-        );
-        println!(
-            "   - Optimal Merkle Tree: {:?}",
-            Duration::from_nanos((gen_time_opt / sample_size).try_into().unwrap())
-        );
-
-        println!(
-            "   - Pmtree-HashMap Merkle Tree: {:?}",
-            Duration::from_nanos((gen_time_pm_memory / sample_size).try_into().unwrap())
-        );
-
-        println!(
-            "   - Pmtree-Sled Merkle Tree: {:?}",
-            Duration::from_nanos((gen_time_pm_sled / sample_size).try_into().unwrap())
-        );
-
-        println!(" Average update_next execution time:");
-        println!(
-            "   - Full Merkle Tree: {:?}",
-            Duration::from_nanos((upd_time_full / sample_size).try_into().unwrap())
-        );
-
-        println!(
-            "   - Optimal Merkle Tree: {:?}",
-            Duration::from_nanos((upd_time_opt / sample_size).try_into().unwrap())
-        );
-
-        println!(
-            "   - Pmtree-HashMap Merkle Tree: {:?}",
-            Duration::from_nanos((upd_time_pm_memory / sample_size).try_into().unwrap())
-        );
-
-        println!(
-            "   - Pmtree-Sled Merkle Tree: {:?}",
-            Duration::from_nanos((upd_time_pm_sled / sample_size).try_into().unwrap())
-        );
-    }
-
-    // The following two tests contain values that come from public::test_merkle_proof test
-    // We check that pmtree and zerokit Merkle tree implementations match.
-
-    #[test]
-    fn test_pmtree_hashmap() -> Result<()> {
-        let tree_height = 20;
-
-        let mut tree = pmtree::MerkleTree::<MemoryDB, PoseidonHash>::default(tree_height).unwrap();
-
-        let leaf_index = 3;
-
-        let identity_secret = hash_to_field(b"test-merkle-proof");
-        let id_commitment = poseidon_hash(&[identity_secret]);
-
-        // let default_leaf = Fr::from(0);
-        tree.set(leaf_index, id_commitment).unwrap();
-
-        // We check correct computation of the root
-        let root = tree.root();
-
-        assert_eq!(
-            root,
-            str_to_fr(
-                "0x21947ffd0bce0c385f876e7c97d6a42eec5b1fe935aab2f01c1f8a8cbcc356d2",
-                16
-            )
-            .unwrap()
-        );
-
-        let merkle_proof = tree.proof(leaf_index).expect("proof should exist");
-        let path_elements = merkle_proof.get_path_elements();
-        let identity_path_index = merkle_proof.get_path_index();
-
-        // We check correct computation of the path and indexes
-        // These values refers to tree height = 20
-        let expected_path_elements = vec![
-            str_to_fr(
-                "0x0000000000000000000000000000000000000000000000000000000000000000",
-                16,
-            )
-            .unwrap(),
-            str_to_fr(
-                "0x2098f5fb9e239eab3ceac3f27b81e481dc3124d55ffed523a839ee8446b64864",
-                16,
-            )
-            .unwrap(),
-            str_to_fr(
-                "0x1069673dcdb12263df301a6ff584a7ec261a44cb9dc68df067a4774460b1f1e1",
-                16,
-            )
-            .unwrap(),
-            str_to_fr(
-                "0x18f43331537ee2af2e3d758d50f72106467c6eea50371dd528d57eb2b856d238",
-                16,
-            )
-            .unwrap(),
-            str_to_fr(
-                "0x07f9d837cb17b0d36320ffe93ba52345f1b728571a568265caac97559dbc952a",
-                16,
-            )
-            .unwrap(),
-            str_to_fr(
-                "0x2b94cf5e8746b3f5c9631f4c5df32907a699c58c94b2ad4d7b5cec1639183f55",
-                16,
-            )
-            .unwrap(),
-            str_to_fr(
-                "0x2dee93c5a666459646ea7d22cca9e1bcfed71e6951b953611d11dda32ea09d78",
-                16,
-            )
-            .unwrap(),
-            str_to_fr(
-                "0x078295e5a22b84e982cf601eb639597b8b0515a88cb5ac7fa8a4aabe3c87349d",
-                16,
-            )
-            .unwrap(),
-            str_to_fr(
-                "0x2fa5e5f18f6027a6501bec864564472a616b2e274a41211a444cbe3a99f3cc61",
-                16,
-            )
-            .unwrap(),
-            str_to_fr(
-                "0x0e884376d0d8fd21ecb780389e941f66e45e7acce3e228ab3e2156a614fcd747",
-                16,
-            )
-            .unwrap(),
-            str_to_fr(
-                "0x1b7201da72494f1e28717ad1a52eb469f95892f957713533de6175e5da190af2",
-                16,
-            )
-            .unwrap(),
-            str_to_fr(
-                "0x1f8d8822725e36385200c0b201249819a6e6e1e4650808b5bebc6bface7d7636",
-                16,
-            )
-            .unwrap(),
-            str_to_fr(
-                "0x2c5d82f66c914bafb9701589ba8cfcfb6162b0a12acf88a8d0879a0471b5f85a",
-                16,
-            )
-            .unwrap(),
-            str_to_fr(
-                "0x14c54148a0940bb820957f5adf3fa1134ef5c4aaa113f4646458f270e0bfbfd0",
-                16,
-            )
-            .unwrap(),
-            str_to_fr(
-                "0x190d33b12f986f961e10c0ee44d8b9af11be25588cad89d416118e4bf4ebe80c",
-                16,
-            )
-            .unwrap(),
-            str_to_fr(
-                "0x22f98aa9ce704152ac17354914ad73ed1167ae6596af510aa5b3649325e06c92",
-                16,
-            )
-            .unwrap(),
-            str_to_fr(
-                "0x2a7c7c9b6ce5880b9f6f228d72bf6a575a526f29c66ecceef8b753d38bba7323",
-                16,
-            )
-            .unwrap(),
-            str_to_fr(
-                "0x2e8186e558698ec1c67af9c14d463ffc470043c9c2988b954d75dd643f36b992",
-                16,
-            )
-            .unwrap(),
-            str_to_fr(
-                "0x0f57c5571e9a4eab49e2c8cf050dae948aef6ead647392273546249d1c1ff10f",
-                16,
-            )
-            .unwrap(),
-            str_to_fr(
-                "0x1830ee67b5fb554ad5f63d4388800e1cfe78e310697d46e43c9ce36134f72cca",
-                16,
-            )
-            .unwrap(),
-        ];
-
-        let expected_identity_path_index: Vec<u8> =
-            vec![1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
-
-        assert_eq!(path_elements, expected_path_elements);
-        assert_eq!(identity_path_index, expected_identity_path_index);
-
-        // We check correct verification of the proof
-        assert!(tree.verify(&id_commitment, &merkle_proof));
-
-        Ok(())
-    }
-
-    #[test]
-    fn test_pmtree_sled() -> Result<()> {
-        let tree_height = 20;
-
-        let mut tree = pmtree::MerkleTree::<SledDB, PoseidonHash>::default(tree_height).unwrap();
-
-        let leaf_index = 3;
-
-        let identity_secret = hash_to_field(b"test-merkle-proof");
-        let id_commitment = poseidon_hash(&[identity_secret]);
-
-        // let default_leaf = Fr::from(0);
-        tree.set(leaf_index, id_commitment).unwrap();
-
-        // We check correct computation of the root
-        let root = tree.root();
-
-        assert_eq!(
-            root,
-            str_to_fr(
-                "0x21947ffd0bce0c385f876e7c97d6a42eec5b1fe935aab2f01c1f8a8cbcc356d2",
-                16
-            )
-            .unwrap()
-        );
-
-        let merkle_proof = tree.proof(leaf_index).expect("proof should exist");
-        let path_elements = merkle_proof.get_path_elements();
-        let identity_path_index = merkle_proof.get_path_index();
-
-        // We check correct computation of the path and indexes
-        // These values refers to tree height = 20
-        let expected_path_elements = vec![
-            str_to_fr(
-                "0x0000000000000000000000000000000000000000000000000000000000000000",
-                16,
-            )
-            .unwrap(),
-            str_to_fr(
-                "0x2098f5fb9e239eab3ceac3f27b81e481dc3124d55ffed523a839ee8446b64864",
-                16,
-            )
-            .unwrap(),
-            str_to_fr(
-                "0x1069673dcdb12263df301a6ff584a7ec261a44cb9dc68df067a4774460b1f1e1",
-                16,
-            )
-            .unwrap(),
-            str_to_fr(
-                "0x18f43331537ee2af2e3d758d50f72106467c6eea50371dd528d57eb2b856d238",
-                16,
-            )
-            .unwrap(),
-            str_to_fr(
-                "0x07f9d837cb17b0d36320ffe93ba52345f1b728571a568265caac97559dbc952a",
-                16,
-            )
-            .unwrap(),
-            str_to_fr(
-                "0x2b94cf5e8746b3f5c9631f4c5df32907a699c58c94b2ad4d7b5cec1639183f55",
-                16,
-            )
-            .unwrap(),
-            str_to_fr(
-                "0x2dee93c5a666459646ea7d22cca9e1bcfed71e6951b953611d11dda32ea09d78",
-                16,
-            )
-            .unwrap(),
-            str_to_fr(
-                "0x078295e5a22b84e982cf601eb639597b8b0515a88cb5ac7fa8a4aabe3c87349d",
-                16,
-            )
-            .unwrap(),
-            str_to_fr(
-                "0x2fa5e5f18f6027a6501bec864564472a616b2e274a41211a444cbe3a99f3cc61",
-                16,
-            )
-            .unwrap(),
-            str_to_fr(
-                "0x0e884376d0d8fd21ecb780389e941f66e45e7acce3e228ab3e2156a614fcd747",
-                16,
-            )
-            .unwrap(),
-            str_to_fr(
-                "0x1b7201da72494f1e28717ad1a52eb469f95892f957713533de6175e5da190af2",
-                16,
-            )
-            .unwrap(),
-            str_to_fr(
-                "0x1f8d8822725e36385200c0b201249819a6e6e1e4650808b5bebc6bface7d7636",
-                16,
-            )
-            .unwrap(),
-            str_to_fr(
-                "0x2c5d82f66c914bafb9701589ba8cfcfb6162b0a12acf88a8d0879a0471b5f85a",
-                16,
-            )
-            .unwrap(),
-            str_to_fr(
-                "0x14c54148a0940bb820957f5adf3fa1134ef5c4aaa113f4646458f270e0bfbfd0",
-                16,
-            )
-            .unwrap(),
-            str_to_fr(
-                "0x190d33b12f986f961e10c0ee44d8b9af11be25588cad89d416118e4bf4ebe80c",
-                16,
-            )
-            .unwrap(),
-            str_to_fr(
-                "0x22f98aa9ce704152ac17354914ad73ed1167ae6596af510aa5b3649325e06c92",
-                16,
-            )
-            .unwrap(),
-            str_to_fr(
-                "0x2a7c7c9b6ce5880b9f6f228d72bf6a575a526f29c66ecceef8b753d38bba7323",
-                16,
-            )
-            .unwrap(),
-            str_to_fr(
-                "0x2e8186e558698ec1c67af9c14d463ffc470043c9c2988b954d75dd643f36b992",
-                16,
-            )
-            .unwrap(),
-            str_to_fr(
-                "0x0f57c5571e9a4eab49e2c8cf050dae948aef6ead647392273546249d1c1ff10f",
-                16,
-            )
-            .unwrap(),
-            str_to_fr(
-                "0x1830ee67b5fb554ad5f63d4388800e1cfe78e310697d46e43c9ce36134f72cca",
-                16,
-            )
-            .unwrap(),
-        ];
-
-        let expected_identity_path_index: Vec<u8> =
-            vec![1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
-
-        assert_eq!(path_elements, expected_path_elements);
-        assert_eq!(identity_path_index, expected_identity_path_index);
-
-        // We check correct verification of the proof
-        assert!(tree.verify(&id_commitment, &merkle_proof));
-
-        Ok(())
-    }
-}

rln/tests/protocol.rs

@@ -4,10 +4,13 @@ mod test {
         circom_from_folder, vk_from_folder, zkey_from_folder, Fr, TEST_RESOURCES_FOLDER,
         TEST_TREE_HEIGHT,
     };
-    use rln::poseidon_hash::poseidon_hash;
+    use rln::hashers::{hash_to_field, poseidon_hash};
     use rln::poseidon_tree::PoseidonTree;
     use rln::protocol::*;
     use rln::utils::str_to_fr;
+    use utils::{ZerokitMerkleProof, ZerokitMerkleTree};
+
+    type ConfigOf<T> = <T as ZerokitMerkleTree>::Config;
 
     // Input generated with https://github.com/oskarth/zk-kit/commit/b6a872f7160c7c14e10a0ea40acab99cbb23c9a8
     const WITNESS_JSON_15: &str = r#"
@@ -171,7 +174,12 @@ mod test {
 
         // generate merkle tree
         let default_leaf = Fr::from(0);
-        let mut tree = PoseidonTree::new(tree_height, default_leaf);
+        let mut tree = PoseidonTree::new(
+            tree_height,
+            default_leaf,
+            ConfigOf::<PoseidonTree>::default(),
+        )
+        .unwrap();
         tree.set(leaf_index, id_commitment.into()).unwrap();
 
         // We check correct computation of the root
@@ -381,7 +389,12 @@ mod test {
 
         //// generate merkle tree
         let default_leaf = Fr::from(0);
-        let mut tree = PoseidonTree::new(tree_height, default_leaf);
+        let mut tree = PoseidonTree::new(
+            tree_height,
+            default_leaf,
+            ConfigOf::<PoseidonTree>::default(),
+        )
+        .unwrap();
         tree.set(leaf_index, id_commitment.into()).unwrap();
 
         let merkle_proof = tree.proof(leaf_index).expect("proof should exist");

rln/tests/public.rs

@@ -3,10 +3,11 @@ mod test {
     use ark_std::{rand::thread_rng, UniformRand};
     use rand::Rng;
     use rln::circuit::{Fr, TEST_RESOURCES_FOLDER, TEST_TREE_HEIGHT};
-    use rln::poseidon_hash::{poseidon_hash as utils_poseidon_hash, ROUND_PARAMS};
-    use rln::protocol::{compute_tree_root, deserialize_identity_tuple, hash_to_field};
+    use rln::hashers::{hash_to_field, poseidon_hash as utils_poseidon_hash, ROUND_PARAMS};
+    use rln::protocol::{compute_tree_root, deserialize_identity_tuple};
     use rln::public::{hash as public_hash, poseidon_hash as public_poseidon_hash, RLN};
     use rln::utils::*;
+    use serde_json::json;
     use std::io::Cursor;
 
     #[test]
@@ -15,7 +16,8 @@ mod test {
         let tree_height = TEST_TREE_HEIGHT;
         let leaf_index = 3;
 
-        let input_buffer = Cursor::new(TEST_RESOURCES_FOLDER);
+        let input_buffer =
+            Cursor::new(json!({ "resources_folder": TEST_RESOURCES_FOLDER }).to_string());
         let mut rln = RLN::new(tree_height, input_buffer).unwrap();
 
         // generate identity

semaphore/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "semaphore-wrapper"
-version = "0.1.0"
+version = "0.3.0"
 edition = "2021"
 license = "MIT OR Apache-2.0"
 
@@ -21,7 +21,7 @@ color-eyre = "0.6.1"
 once_cell = "1.8"
 rand = "0.8.4"
 semaphore = { git = "https://github.com/worldcoin/semaphore-rs", rev = "ee658c2"}
-ethers-core = { git = "https://github.com/gakonst/ethers-rs", default-features = false }
+ethers-core = { version = "2.0.8", default-features = false }
 ruint = { version = "1.2.0", features = [ "serde", "num-bigint", "ark-ff" ] }
 serde = "1.0"
 thiserror = "1.0.0"
@@ -47,4 +47,4 @@ opt-level = 3
 # Dependencies are optimized, even in a dev build. This improves dev performance
 # while having neglible impact on incremental build times.
 [profile.dev.package."*"]
-opt-level = 3
+opt-level = 3

utils/Cargo.toml

@@ -1,20 +1,36 @@
 [package]
-name = "utils"
-version = "0.1.0"
+name = "zerokit_utils"
+version = "0.3.2"
 edition = "2021"
 license = "MIT OR Apache-2.0"
+description = "Various utilities for Zerokit"
+documentation = "https://github.com/vacp2p/zerokit"
+homepage = "https://vac.dev"
+repository = "https://github.com/vacp2p/zerokit"
+
+[lib]
+bench = false
 
 [dependencies]
-ark-ff = { version = "=0.3.0", default-features = false, features = ["asm"] }
+ark-ff = { version = "=0.4.1", default-features = false, features = ["asm"] }
 num-bigint = { version = "=0.4.3", default-features = false, features = ["rand"] }
 color-eyre = "=0.6.2"
+pmtree = { package = "pmtree", version = "=2.0.0", optional = true}
+sled = "=0.34.7"
+serde = "=1.0.163"
 
 [dev-dependencies]
-ark-bn254 = "=0.3.0"
-num-traits = "0.2.11"
-hex-literal = "0.3.4"
-tiny-keccak = { version = "2.0.2", features = ["keccak"] }
+ark-bn254 = "=0.4.0"
+num-traits = "=0.2.15"
+hex-literal = "=0.3.4"
+tiny-keccak = { version = "=2.0.2", features = ["keccak"] }
+criterion = { version = "=0.4.0", features = ["html_reports"] }
 
 [features]
 default = ["parallel"]
 parallel = ["ark-ff/parallel"]
+pmtree-ft = ["pmtree"]
+
+[[bench]]
+name = "merkle_tree_benchmark"
+harness = false

utils/Makefile.toml

@@ -5,3 +5,7 @@ args = ["build", "--release"]
 [tasks.test]
 command = "cargo"
 args = ["test", "--release"]
+
+[tasks.bench]
+command = "cargo"
+args = ["bench"]

utils/README.md

@@ -0,0 +1,15 @@
+# Utils crate
+
+## Building
+
+1. `cargo build`
+
+## Testing
+
+1. `cargo test`
+
+## Benchmarking
+
+1. `cargo bench`
+
+To view the results of the benchmark, open the `target/criterion/report/index.html` file generated after the bench

utils/benches/merkle_tree_benchmark.rs

@@ -0,0 +1,119 @@
+use criterion::{criterion_group, criterion_main, Criterion};
+use hex_literal::hex;
+use tiny_keccak::{Hasher as _, Keccak};
+use zerokit_utils::{
+    FullMerkleConfig, FullMerkleTree, Hasher, OptimalMerkleConfig, OptimalMerkleTree,
+    ZerokitMerkleTree,
+};
+
+#[derive(Clone, Copy, Eq, PartialEq)]
+struct Keccak256;
+
+impl Hasher for Keccak256 {
+    type Fr = [u8; 32];
+
+    fn default_leaf() -> Self::Fr {
+        [0; 32]
+    }
+
+    fn hash(inputs: &[Self::Fr]) -> Self::Fr {
+        let mut output = [0; 32];
+        let mut hasher = Keccak::v256();
+        for element in inputs {
+            hasher.update(element);
+        }
+        hasher.finalize(&mut output);
+        output
+    }
+}
+
+pub fn optimal_merkle_tree_benchmark(c: &mut Criterion) {
+    let mut tree =
+        OptimalMerkleTree::<Keccak256>::new(2, [0; 32], OptimalMerkleConfig::default()).unwrap();
+
+    let leaves = [
+        hex!("0000000000000000000000000000000000000000000000000000000000000001"),
+        hex!("0000000000000000000000000000000000000000000000000000000000000002"),
+        hex!("0000000000000000000000000000000000000000000000000000000000000003"),
+        hex!("0000000000000000000000000000000000000000000000000000000000000004"),
+    ];
+
+    c.bench_function("OptimalMerkleTree::set", |b| {
+        b.iter(|| {
+            tree.set(0, leaves[0]).unwrap();
+        })
+    });
+
+    c.bench_function("OptimalMerkleTree::delete", |b| {
+        b.iter(|| {
+            tree.delete(0).unwrap();
+        })
+    });
+
+    c.bench_function("OptimalMerkleTree::override_range", |b| {
+        b.iter(|| {
+            tree.override_range(0, leaves, [0, 1, 2, 3]).unwrap();
+        })
+    });
+
+    c.bench_function("OptimalMerkleTree::compute_root", |b| {
+        b.iter(|| {
+            tree.compute_root().unwrap();
+        })
+    });
+
+    c.bench_function("OptimalMerkleTree::get", |b| {
+        b.iter(|| {
+            tree.get(0).unwrap();
+        })
+    });
+}
+
+pub fn full_merkle_tree_benchmark(c: &mut Criterion) {
+    let mut tree =
+        FullMerkleTree::<Keccak256>::new(2, [0; 32], FullMerkleConfig::default()).unwrap();
+
+    let leaves = [
+        hex!("0000000000000000000000000000000000000000000000000000000000000001"),
+        hex!("0000000000000000000000000000000000000000000000000000000000000002"),
+        hex!("0000000000000000000000000000000000000000000000000000000000000003"),
+        hex!("0000000000000000000000000000000000000000000000000000000000000004"),
+    ];
+
+    c.bench_function("FullMerkleTree::set", |b| {
+        b.iter(|| {
+            tree.set(0, leaves[0]).unwrap();
+        })
+    });
+
+    c.bench_function("FullMerkleTree::delete", |b| {
+        b.iter(|| {
+            tree.delete(0).unwrap();
+        })
+    });
+
+    c.bench_function("FullMerkleTree::override_range", |b| {
+        b.iter(|| {
+            tree.override_range(0, leaves, [0, 1, 2, 3]).unwrap();
+        })
+    });
+
+    c.bench_function("FullMerkleTree::compute_root", |b| {
+        b.iter(|| {
+            tree.compute_root().unwrap();
+        })
+    });
+
+    c.bench_function("FullMerkleTree::get", |b| {
+        b.iter(|| {
+            tree.get(0).unwrap();
+        })
+    });
+}
+
+criterion_group!(
+    benches,
+    optimal_merkle_tree_benchmark,
+    full_merkle_tree_benchmark
+);
+criterion_main!(benches);

utils/src/lib.rs

@@ -3,3 +3,8 @@ pub use self::poseidon::*;
 
 pub mod merkle_tree;
 pub use self::merkle_tree::*;
+
+#[cfg(feature = "pmtree-ft")]
+pub mod pm_tree;
+#[cfg(feature = "pmtree-ft")]
+pub use self::pm_tree::*;

utils/src/merkle_tree/full_merkle_tree.rs

@@ -0,0 +1,373 @@
+use crate::merkle_tree::{FrOf, Hasher, ZerokitMerkleProof, ZerokitMerkleTree};
+use color_eyre::{Report, Result};
+use std::{
+    cmp::max,
+    fmt::Debug,
+    iter::{once, repeat, successors},
+    str::FromStr,
+};
+
+////////////////////////////////////////////////////////////
+/// Full Merkle Tree Implementation
+////////////////////////////////////////////////////////////
+
+/// Merkle tree with all leaf and intermediate hashes stored
+#[derive(Clone, PartialEq, Eq, Debug)]
+pub struct FullMerkleTree<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>,
+
+    // 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: Vec<u8>,
+}
+
+/// Element of a Merkle proof
+#[derive(Clone, Copy, PartialEq, Eq)]
+pub enum FullMerkleBranch<H: Hasher> {
+    /// Left branch taken, value is the right sibling hash.
+    Left(H::Fr),
+
+    /// Right branch taken, value is the left sibling hash.
+    Right(H::Fr),
+}
+
+/// Merkle proof path, bottom to top.
+#[derive(Clone, PartialEq, Eq)]
+pub struct FullMerkleProof<H: Hasher>(pub Vec<FullMerkleBranch<H>>);
+
+#[derive(Default)]
+pub struct FullMerkleConfig(());
+
+impl FromStr for FullMerkleConfig {
+    type Err = Report;
+
+    fn from_str(_s: &str) -> Result<Self> {
+        Ok(FullMerkleConfig::default())
+    }
+}
+
+/// Implementations
+impl<H: Hasher> ZerokitMerkleTree for FullMerkleTree<H>
+where
+    H: Hasher,
+{
+    type Proof = FullMerkleProof<H>;
+    type Hasher = H;
+    type Config = FullMerkleConfig;
+
+    fn default(depth: usize) -> Result<Self> {
+        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> {
+        // 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<_>>();
+
+        // Compute node values
+        let nodes = cached_nodes
+            .iter()
+            .rev()
+            .enumerate()
+            .flat_map(|(levels, hash)| repeat(hash).take(1 << levels))
+            .cloned()
+            .collect::<Vec<_>>();
+        debug_assert!(nodes.len() == (1 << (depth + 1)) - 1);
+
+        let next_index = 0;
+
+        Ok(Self {
+            depth,
+            cached_nodes,
+            nodes,
+            next_index,
+            metadata: Vec::new(),
+        })
+    }
+
+    fn close_db_connection(&mut self) -> Result<()> {
+        Ok(())
+    }
+
+    // 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
+    fn capacity(&self) -> usize {
+        1 << self.depth
+    }
+
+    // Returns the total number of leaves set
+    fn leaves_set(&mut self) -> usize {
+        self.next_index
+    }
+
+    #[must_use]
+    // 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<()> {
+        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>> {
+        if leaf >= self.capacity() {
+            return Err(Report::msg("leaf index out of bounds"));
+        }
+        Ok(self.nodes[self.capacity() + leaf - 1])
+    }
+
+    // Sets tree nodes, starting from start index
+    // Function proper of FullMerkleTree implementation
+    fn set_range<I: IntoIterator<Item = FrOf<Self::Hasher>>>(
+        &mut self,
+        start: usize,
+        hashes: I,
+    ) -> Result<()> {
+        let index = self.capacity() + start - 1;
+        let mut count = 0;
+        // first count number of hashes, 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"));
+        }
+        hashes.into_iter().for_each(|hash| {
+            self.nodes[index + count] = hash;
+            count += 1;
+        });
+        if count != 0 {
+            self.update_nodes(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, to_remove_indices: J) -> Result<()>
+    where
+        I: IntoIterator<Item = FrOf<Self::Hasher>>,
+        J: IntoIterator<Item = usize>,
+    {
+        let index = self.capacity() + start - 1;
+        let mut count = 0;
+        let leaves = leaves.into_iter().collect::<Vec<_>>();
+        let to_remove_indices = to_remove_indices.into_iter().collect::<Vec<_>>();
+        // first count number of hashes, and check that they fit in the tree
+        // then insert into the tree
+        if leaves.len() + start - to_remove_indices.len() > self.capacity() {
+            return Err(Report::msg("provided hashes do not fit in the tree"));
+        }
+
+        // remove leaves
+        for i in &to_remove_indices {
+            self.delete(*i)?;
+        }
+
+        // insert new leaves
+        for hash in leaves {
+            self.nodes[index + count] = hash;
+            count += 1;
+        }
+
+        if count != 0 {
+            self.update_nodes(index, index + (count - 1))?;
+            self.next_index = max(self.next_index, start + count - to_remove_indices.len());
+        }
+        Ok(())
+    }
+
+    // Sets a leaf at the next available index
+    fn update_next(&mut self, leaf: FrOf<Self::Hasher>) -> Result<()> {
+        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<()> {
+        // 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())?;
+        }
+        Ok(())
+    }
+
+    // Computes a merkle proof the the leaf at the specified index
+    fn proof(&self, leaf: usize) -> Result<FullMerkleProof<H>> {
+        if leaf >= self.capacity() {
+            return Err(Report::msg("index exceeds set size"));
+        }
+        let mut index = self.capacity() + leaf - 1;
+        let mut path = Vec::with_capacity(self.depth + 1);
+        while let Some(parent) = self.parent(index) {
+            // Add proof for node at index to parent
+            path.push(match index & 1 {
+                1 => FullMerkleBranch::Left(self.nodes[index + 1]),
+                0 => FullMerkleBranch::Right(self.nodes[index - 1]),
+                _ => unreachable!(),
+            });
+            index = parent;
+        }
+        Ok(FullMerkleProof(path))
+    }
+
+    // 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> {
+        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<()> {
+        self.metadata = metadata.to_vec();
+        Ok(())
+    }
+
+    fn metadata(&self) -> Result<Vec<u8>> {
+        Ok(self.metadata.to_vec())
+    }
+}
+
+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> {
+        if index == 0 {
+            None
+        } else {
+            Some(((index + 1) >> 1) - 1)
+        }
+    }
+
+    /// For a given node index, return index of the first (left) child.
+    fn first_child(&self, index: usize) -> usize {
+        (index << 1) + 1
+    }
+
+    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)"));
+        }
+        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]]);
+            }
+            self.update_nodes(start, end)?;
+        }
+        Ok(())
+    }
+}
+
+impl<H: Hasher> ZerokitMerkleProof for FullMerkleProof<H> {
+    type Index = u8;
+    type Hasher = H;
+
+    #[must_use]
+    // Returns the length of a Merkle proof
+    fn length(&self) -> usize {
+        self.0.len()
+    }
+
+    /// Computes the leaf index corresponding to a Merkle proof
+    #[must_use]
+    fn leaf_index(&self) -> usize {
+        self.0.iter().rev().fold(0, |index, branch| match branch {
+            FullMerkleBranch::Left(_) => index << 1,
+            FullMerkleBranch::Right(_) => (index << 1) + 1,
+        })
+    }
+
+    #[must_use]
+    /// Returns the path elements forming a Merkle proof
+    fn get_path_elements(&self) -> Vec<FrOf<Self::Hasher>> {
+        self.0
+            .iter()
+            .map(|x| match x {
+                FullMerkleBranch::Left(value) | FullMerkleBranch::Right(value) => *value,
+            })
+            .collect()
+    }
+
+    /// Returns the path indexes forming a Merkle proof
+    #[must_use]
+    fn get_path_index(&self) -> Vec<Self::Index> {
+        self.0
+            .iter()
+            .map(|branch| match branch {
+                FullMerkleBranch::Left(_) => 0,
+                FullMerkleBranch::Right(_) => 1,
+            })
+            .collect()
+    }
+
+    /// Computes the Merkle root corresponding by iteratively hashing a Merkle proof with a given input leaf
+    #[must_use]
+    fn compute_root_from(&self, hash: &FrOf<Self::Hasher>) -> FrOf<Self::Hasher> {
+        self.0.iter().fold(*hash, |hash, branch| match branch {
+            FullMerkleBranch::Left(sibling) => H::hash(&[hash, *sibling]),
+            FullMerkleBranch::Right(sibling) => H::hash(&[*sibling, hash]),
+        })
+    }
+}
+
+// Debug formatting for printing a (Full) Merkle Proof Branch
+impl<H> Debug for FullMerkleBranch<H>
+where
+    H: Hasher,
+    H::Fr: Debug,
+{
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Self::Left(arg0) => f.debug_tuple("Left").field(arg0).finish(),
+            Self::Right(arg0) => f.debug_tuple("Right").field(arg0).finish(),
+        }
+    }
+}
+
+// Debug formatting for printing a (Full) Merkle Proof
+impl<H> Debug for FullMerkleProof<H>
+where
+    H: Hasher,
+    H::Fr: Debug,
+{
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_tuple("Proof").field(&self.0).finish()
+    }
+}

utils/src/merkle_tree/merkle_tree.rs

@@ -13,22 +13,15 @@
 //! * Disk based storage backend (using mmaped files should be easy)
 //! * Implement serialization for tree and Merkle proof
 
-#![allow(dead_code)]
+use std::str::FromStr;
 
-use std::collections::HashMap;
-use std::{
-    cmp::max,
-    fmt::Debug,
-    iter::{once, repeat, successors},
-};
-
-use color_eyre::{Report, Result};
+use color_eyre::Result;
 
 /// 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: Copy + Clone + Eq;
+    type Fr: Clone + Copy + Eq;
 
     /// Returns the default tree leaf
     fn default_leaf() -> Self::Fr;
@@ -37,528 +30,51 @@ pub trait Hasher {
     fn hash(input: &[Self::Fr]) -> Self::Fr;
 }
 
-////////////////////////////////////////////////////////////
-/// Optimal Merkle Tree Implementation
-////////////////////////////////////////////////////////////
+pub type FrOf<H> = <H as Hasher>::Fr;
 
-/// The Merkle tree structure
-#[derive(Clone, PartialEq, Eq, Debug)]
-pub struct OptimalMerkleTree<H>
-where
-    H: Hasher,
-{
-    /// The depth of the tree, i.e. the number of levels from leaf to root
-    depth: usize,
+/// In the ZerokitMerkleTree trait we define the methods that are required to be implemented by a Merkle tree
+/// Including, OptimalMerkleTree, FullMerkleTree
+pub trait ZerokitMerkleTree {
+    type Proof: ZerokitMerkleProof;
+    type Hasher: Hasher;
+    type Config: Default + FromStr;
 
-    /// 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: HashMap<(usize, usize), H::Fr>,
-
-    // 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,
+    fn default(depth: usize) -> Result<Self>
+    where
+        Self: Sized;
+    fn new(depth: usize, default_leaf: FrOf<Self::Hasher>, config: Self::Config) -> Result<Self>
+    where
+        Self: Sized;
+    fn depth(&self) -> usize;
+    fn capacity(&self) -> usize;
+    fn leaves_set(&mut self) -> usize;
+    fn root(&self) -> FrOf<Self::Hasher>;
+    fn compute_root(&mut self) -> 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<()>
+    where
+        I: IntoIterator<Item = FrOf<Self::Hasher>>;
+    fn get(&self, index: usize) -> Result<FrOf<Self::Hasher>>;
+    fn override_range<I, J>(&mut self, start: usize, leaves: I, to_remove_indices: J) -> Result<()>
+    where
+        I: IntoIterator<Item = FrOf<Self::Hasher>>,
+        J: IntoIterator<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<()>;
 }
 
-/// The Merkle proof
-/// Contains a vector of (node, branch_index) that defines the proof path elements and branch direction (1 or 0)
-#[derive(Clone, PartialEq, Eq)]
-pub struct OptimalMerkleProof<H: Hasher>(pub Vec<(H::Fr, u8)>);
+pub trait ZerokitMerkleProof {
+    type Index;
+    type Hasher: Hasher;
 
-/// Implementations
-
-impl<H: Hasher> OptimalMerkleTree<H> {
-    pub fn default(depth: usize) -> Self {
-        OptimalMerkleTree::<H>::new(depth, H::default_leaf())
-    }
-
-    /// 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
-    pub fn new(depth: usize, default_leaf: H::Fr) -> Self {
-        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();
-        OptimalMerkleTree {
-            cached_nodes: cached_nodes.clone(),
-            depth,
-            nodes: HashMap::new(),
-            next_index: 0,
-        }
-    }
-
-    // Returns the depth of the tree
-    pub fn depth(&self) -> usize {
-        self.depth
-    }
-
-    // Returns the capacity of the tree, i.e. the maximum number of accumulatable leaves
-    pub fn capacity(&self) -> usize {
-        1 << self.depth
-    }
-
-    // Returns the total number of leaves set
-    pub fn leaves_set(&mut self) -> usize {
-        self.next_index
-    }
-
-    #[must_use]
-    // Returns the root of the tree
-    pub fn root(&self) -> H::Fr {
-        self.get_node(0, 0)
-    }
-
-    // Sets a leaf at the specified tree index
-    pub 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.recalculate_from(index)?;
-        self.next_index = max(self.next_index, index + 1);
-        Ok(())
-    }
-
-    // Sets multiple leaves from the specified tree index
-    pub fn set_range<I: IntoIterator<Item = H::Fr>>(
-        &mut self,
-        start: usize,
-        leaves: I,
-    ) -> Result<()> {
-        let leaves = leaves.into_iter().collect::<Vec<_>>();
-        // check if the range is valid
-        if start + leaves.len() > self.capacity() {
-            return Err(Report::msg("provided range exceeds set size"));
-        }
-        for (i, leaf) in leaves.iter().enumerate() {
-            self.nodes.insert((self.depth, start + i), *leaf);
-            self.recalculate_from(start + i)?;
-        }
-        self.next_index = max(self.next_index, start + leaves.len());
-        Ok(())
-    }
-
-    // Sets a leaf at the next available index
-    pub fn update_next(&mut self, leaf: H::Fr) -> Result<()> {
-        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)
-    pub fn delete(&mut self, index: usize) -> Result<()> {
-        // 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())?;
-        }
-        Ok(())
-    }
-
-    // Computes a merkle proof the the leaf at the specified index
-    pub fn proof(&self, index: usize) -> Result<OptimalMerkleProof<H>> {
-        if index >= self.capacity() {
-            return Err(Report::msg("index exceeds set size"));
-        }
-        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()));
-            i >>= 1;
-            depth -= 1;
-            if depth == 0 {
-                break;
-            }
-        }
-        if i != 0 {
-            Err(Report::msg("i != 0"))
-        } else {
-            Ok(OptimalMerkleProof(witness))
-        }
-    }
-
-    // Verifies a Merkle proof with respect to the input leaf and the tree root
-    pub fn verify(&self, leaf: &H::Fr, witness: &OptimalMerkleProof<H>) -> Result<bool> {
-        if witness.length() != self.depth {
-            return Err(Report::msg("witness length doesn't match tree depth"));
-        }
-        let expected_root = witness.compute_root_from(leaf);
-        Ok(expected_root.eq(&self.root()))
-    }
-
-    // Utilities for updating the tree nodes
-
-    fn get_node(&self, depth: usize, index: usize) -> H::Fr {
-        let node = *self
-            .nodes
-            .get(&(depth, index))
-            .unwrap_or_else(|| &self.cached_nodes[depth]);
-        node
-    }
-
-    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 {
-        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);
-            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(())
-    }
-}
-
-impl<H: Hasher> OptimalMerkleProof<H> {
-    #[must_use]
-    // Returns the length of a Merkle proof
-    pub fn length(&self) -> usize {
-        self.0.len()
-    }
-
-    /// Computes the leaf index corresponding to a Merkle proof
-    #[must_use]
-    pub fn leaf_index(&self) -> usize {
-        // In current implementation the path indexes in a proof correspond to the binary representation of the leaf index
-        let mut binary_repr = self.get_path_index();
-        binary_repr.reverse();
-        binary_repr
-            .into_iter()
-            .fold(0, |acc, digit| (acc << 1) + usize::from(digit))
-    }
-
-    #[must_use]
-    /// Returns the path elements forming a Merkle proof
-    pub fn get_path_elements(&self) -> Vec<H::Fr> {
-        self.0.iter().map(|x| x.0).collect()
-    }
-
-    /// Returns the path indexes forming a Merkle proof
-    #[must_use]
-    pub fn get_path_index(&self) -> Vec<u8> {
-        self.0.iter().map(|x| x.1).collect()
-    }
-
-    #[must_use]
-    /// Computes the Merkle root corresponding by iteratively hashing a Merkle proof with a given input leaf
-    pub fn compute_root_from(&self, leaf: &H::Fr) -> H::Fr {
-        let mut acc: H::Fr = *leaf;
-        for w in self.0.iter() {
-            if w.1 == 0 {
-                acc = H::hash(&[acc, w.0]);
-            } else {
-                acc = H::hash(&[w.0, acc]);
-            }
-        }
-        acc
-    }
-}
-
-// Debug formatting for printing a (Optimal) Merkle Proof
-impl<H> Debug for OptimalMerkleProof<H>
-where
-    H: Hasher,
-    H::Fr: Debug,
-{
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        f.debug_tuple("Proof").field(&self.0).finish()
-    }
-}
-
-////////////////////////////////////////////////////////////
-/// Full Merkle Tree Implementation
-////////////////////////////////////////////////////////////
-
-/// Merkle tree with all leaf and intermediate hashes stored
-#[derive(Clone, PartialEq, Eq, Debug)]
-pub struct FullMerkleTree<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>,
-
-    // 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,
-}
-
-/// Element of a Merkle proof
-#[derive(Clone, Copy, PartialEq, Eq)]
-pub enum FullMerkleBranch<H: Hasher> {
-    /// Left branch taken, value is the right sibling hash.
-    Left(H::Fr),
-
-    /// Right branch taken, value is the left sibling hash.
-    Right(H::Fr),
-}
-
-/// Merkle proof path, bottom to top.
-#[derive(Clone, PartialEq, Eq)]
-pub struct FullMerkleProof<H: Hasher>(pub Vec<FullMerkleBranch<H>>);
-
-/// Implementations
-
-impl<H: Hasher> FullMerkleTree<H> {
-    pub fn default(depth: usize) -> Self {
-        FullMerkleTree::<H>::new(depth, H::default_leaf())
-    }
-
-    /// 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
-    pub fn new(depth: usize, initial_leaf: H::Fr) -> Self {
-        // 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<_>>();
-
-        // Compute node values
-        let nodes = cached_nodes
-            .iter()
-            .rev()
-            .enumerate()
-            .flat_map(|(levels, hash)| repeat(hash).take(1 << levels))
-            .cloned()
-            .collect::<Vec<_>>();
-        debug_assert!(nodes.len() == (1 << (depth + 1)) - 1);
-
-        let next_index = 0;
-
-        Self {
-            depth,
-            cached_nodes,
-            nodes,
-            next_index,
-        }
-    }
-
-    // Returns the depth of the tree
-    pub fn depth(&self) -> usize {
-        self.depth
-    }
-
-    // Returns the capacity of the tree, i.e. the maximum number of accumulatable leaves
-    pub fn capacity(&self) -> usize {
-        1 << self.depth
-    }
-
-    // Returns the total number of leaves set
-    pub fn leaves_set(&mut self) -> usize {
-        self.next_index
-    }
-
-    #[must_use]
-    // Returns the root of the tree
-    pub fn root(&self) -> H::Fr {
-        self.nodes[0]
-    }
-
-    // Sets a leaf at the specified tree index
-    pub fn set(&mut self, leaf: usize, hash: H::Fr) -> Result<()> {
-        self.set_range(leaf, once(hash))?;
-        self.next_index = max(self.next_index, leaf + 1);
-        Ok(())
-    }
-
-    // Sets tree nodes, starting from start index
-    // Function proper of FullMerkleTree implementation
-    fn set_range<I: IntoIterator<Item = H::Fr>>(&mut self, start: usize, hashes: I) -> Result<()> {
-        let index = self.capacity() + start - 1;
-        let mut count = 0;
-        // first count number of hashes, 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"));
-        }
-        hashes.into_iter().for_each(|hash| {
-            self.nodes[index + count] = hash;
-            count += 1;
-        });
-        if count != 0 {
-            self.update_nodes(index, index + (count - 1))?;
-            self.next_index = max(self.next_index, start + count);
-        }
-        Ok(())
-    }
-
-    // Sets a leaf at the next available index
-    pub fn update_next(&mut self, leaf: H::Fr) -> Result<()> {
-        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)
-    pub fn delete(&mut self, index: usize) -> Result<()> {
-        // 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())?;
-        }
-        Ok(())
-    }
-
-    // Computes a merkle proof the the leaf at the specified index
-    pub fn proof(&self, leaf: usize) -> Result<FullMerkleProof<H>> {
-        if leaf >= self.capacity() {
-            return Err(Report::msg("index exceeds set size"));
-        }
-        let mut index = self.capacity() + leaf - 1;
-        let mut path = Vec::with_capacity(self.depth + 1);
-        while let Some(parent) = self.parent(index) {
-            // Add proof for node at index to parent
-            path.push(match index & 1 {
-                1 => FullMerkleBranch::Left(self.nodes[index + 1]),
-                0 => FullMerkleBranch::Right(self.nodes[index - 1]),
-                _ => unreachable!(),
-            });
-            index = parent;
-        }
-        Ok(FullMerkleProof(path))
-    }
-
-    // Verifies a Merkle proof with respect to the input leaf and the tree root
-    pub fn verify(&self, hash: &H::Fr, proof: &FullMerkleProof<H>) -> Result<bool> {
-        Ok(proof.compute_root_from(hash) == self.root())
-    }
-
-    // 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> {
-        if index == 0 {
-            None
-        } else {
-            Some(((index + 1) >> 1) - 1)
-        }
-    }
-
-    /// For a given node index, return index of the first (left) child.
-    fn first_child(&self, index: usize) -> usize {
-        (index << 1) + 1
-    }
-
-    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)"));
-        }
-        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]]);
-            }
-            self.update_nodes(start, end)?;
-        }
-        Ok(())
-    }
-}
-
-impl<H: Hasher> FullMerkleProof<H> {
-    #[must_use]
-    // Returns the length of a Merkle proof
-    pub fn length(&self) -> usize {
-        self.0.len()
-    }
-
-    /// Computes the leaf index corresponding to a Merkle proof
-    #[must_use]
-    pub fn leaf_index(&self) -> usize {
-        self.0.iter().rev().fold(0, |index, branch| match branch {
-            FullMerkleBranch::Left(_) => index << 1,
-            FullMerkleBranch::Right(_) => (index << 1) + 1,
-        })
-    }
-
-    #[must_use]
-    /// Returns the path elements forming a Merkle proof
-    pub fn get_path_elements(&self) -> Vec<H::Fr> {
-        self.0
-            .iter()
-            .map(|x| match x {
-                FullMerkleBranch::Left(value) | FullMerkleBranch::Right(value) => *value,
-            })
-            .collect()
-    }
-
-    /// Returns the path indexes forming a Merkle proof
-    #[must_use]
-    pub fn get_path_index(&self) -> Vec<u8> {
-        self.0
-            .iter()
-            .map(|branch| match branch {
-                FullMerkleBranch::Left(_) => 0,
-                FullMerkleBranch::Right(_) => 1,
-            })
-            .collect()
-    }
-
-    /// Computes the Merkle root corresponding by iteratively hashing a Merkle proof with a given input leaf
-    #[must_use]
-    pub fn compute_root_from(&self, hash: &H::Fr) -> H::Fr {
-        self.0.iter().fold(*hash, |hash, branch| match branch {
-            FullMerkleBranch::Left(sibling) => H::hash(&[hash, *sibling]),
-            FullMerkleBranch::Right(sibling) => H::hash(&[*sibling, hash]),
-        })
-    }
-}
-
-// Debug formatting for printing a (Full) Merkle Proof Branch
-impl<H> Debug for FullMerkleBranch<H>
-where
-    H: Hasher,
-    H::Fr: Debug,
-{
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        match self {
-            Self::Left(arg0) => f.debug_tuple("Left").field(arg0).finish(),
-            Self::Right(arg0) => f.debug_tuple("Right").field(arg0).finish(),
-        }
-    }
-}
-
-// Debug formatting for printing a (Full) Merkle Proof
-impl<H> Debug for FullMerkleProof<H>
-where
-    H: Hasher,
-    H::Fr: Debug,
-{
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        f.debug_tuple("Proof").field(&self.0).finish()
-    }
+    fn length(&self) -> usize;
+    fn leaf_index(&self) -> usize;
+    fn get_path_elements(&self) -> Vec<FrOf<Self::Hasher>>;
+    fn get_path_index(&self) -> Vec<Self::Index>;
+    fn compute_root_from(&self, leaf: &FrOf<Self::Hasher>) -> FrOf<Self::Hasher>;
 }

utils/src/merkle_tree/mod.rs

@@ -1,3 +1,7 @@
+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::*;

utils/src/merkle_tree/optimal_merkle_tree.rs

@@ -0,0 +1,343 @@
+use crate::merkle_tree::{Hasher, ZerokitMerkleProof, ZerokitMerkleTree};
+use crate::FrOf;
+use color_eyre::{Report, Result};
+use std::collections::HashMap;
+use std::str::FromStr;
+use std::{cmp::max, fmt::Debug};
+
+////////////////////////////////////////////////////////////
+/// Optimal Merkle Tree Implementation
+////////////////////////////////////////////////////////////
+
+/// The Merkle tree structure
+#[derive(Clone, PartialEq, Eq, Debug)]
+pub struct OptimalMerkleTree<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: HashMap<(usize, usize), H::Fr>,
+
+    // 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: Vec<u8>,
+}
+
+/// The Merkle proof
+/// Contains a vector of (node, branch_index) that defines the proof path elements and branch direction (1 or 0)
+#[derive(Clone, PartialEq, Eq)]
+pub struct OptimalMerkleProof<H: Hasher>(pub Vec<(H::Fr, u8)>);
+
+#[derive(Default)]
+pub struct OptimalMerkleConfig(());
+
+impl FromStr for OptimalMerkleConfig {
+    type Err = Report;
+
+    fn from_str(_s: &str) -> Result<Self> {
+        Ok(OptimalMerkleConfig::default())
+    }
+}
+
+/// Implementations
+
+impl<H: Hasher> ZerokitMerkleTree for OptimalMerkleTree<H>
+where
+    H: Hasher,
+{
+    type Proof = OptimalMerkleProof<H>;
+    type Hasher = H;
+    type Config = OptimalMerkleConfig;
+
+    fn default(depth: usize) -> Result<Self> {
+        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> {
+        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: cached_nodes.clone(),
+            depth,
+            nodes: HashMap::new(),
+            next_index: 0,
+            metadata: Vec::new(),
+        })
+    }
+
+    fn close_db_connection(&mut self) -> Result<()> {
+        Ok(())
+    }
+
+    // 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
+    fn capacity(&self) -> usize {
+        1 << self.depth
+    }
+
+    // Returns the total number of leaves set
+    fn leaves_set(&mut self) -> usize {
+        self.next_index
+    }
+
+    #[must_use]
+    // Returns the root of the tree
+    fn root(&self) -> H::Fr {
+        self.get_node(0, 0)
+    }
+
+    // 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.recalculate_from(index)?;
+        self.next_index = max(self.next_index, 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))
+    }
+
+    // Sets multiple leaves from the specified tree index
+    fn set_range<I: IntoIterator<Item = H::Fr>>(&mut self, start: usize, leaves: I) -> Result<()> {
+        let leaves = leaves.into_iter().collect::<Vec<_>>();
+        // check if the range is valid
+        if start + leaves.len() > self.capacity() {
+            return Err(Report::msg("provided range exceeds set size"));
+        }
+        for (i, leaf) in leaves.iter().enumerate() {
+            self.nodes.insert((self.depth, start + i), *leaf);
+            self.recalculate_from(start + i)?;
+        }
+        self.next_index = max(self.next_index, start + leaves.len());
+        Ok(())
+    }
+
+    fn override_range<I, J>(&mut self, start: usize, leaves: I, to_remove_indices: J) -> Result<()>
+    where
+        I: IntoIterator<Item = FrOf<Self::Hasher>>,
+        J: IntoIterator<Item = usize>,
+    {
+        let leaves = leaves.into_iter().collect::<Vec<_>>();
+        let to_remove_indices = to_remove_indices.into_iter().collect::<Vec<_>>();
+        // check if the range is valid
+        if leaves.len() + start - to_remove_indices.len() > self.capacity() {
+            return Err(Report::msg("provided range exceeds set size"));
+        }
+
+        // remove leaves
+        for i in &to_remove_indices {
+            self.delete(*i)?;
+        }
+
+        // add leaves
+        for (i, leaf) in leaves.iter().enumerate() {
+            self.nodes.insert((self.depth, start + i), *leaf);
+            self.recalculate_from(start + i)?;
+        }
+
+        self.next_index = max(
+            self.next_index,
+            start + leaves.len() - to_remove_indices.len(),
+        );
+        Ok(())
+    }
+
+    // Sets a leaf at the next available index
+    fn update_next(&mut self, leaf: H::Fr) -> Result<()> {
+        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<()> {
+        // 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())?;
+        }
+        Ok(())
+    }
+
+    // Computes a merkle proof the the leaf at the specified index
+    fn proof(&self, index: usize) -> Result<Self::Proof> {
+        if index >= self.capacity() {
+            return Err(Report::msg("index exceeds set size"));
+        }
+        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()));
+            i >>= 1;
+            depth -= 1;
+            if depth == 0 {
+                break;
+            }
+        }
+        if i != 0 {
+            Err(Report::msg("i != 0"))
+        } 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> {
+        if witness.length() != self.depth {
+            return Err(Report::msg("witness length doesn't match tree depth"));
+        }
+        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<()> {
+        self.metadata = metadata.to_vec();
+        Ok(())
+    }
+
+    fn metadata(&self) -> Result<Vec<u8>> {
+        Ok(self.metadata.to_vec())
+    }
+}
+
+impl<H: Hasher> OptimalMerkleTree<H>
+where
+    H: Hasher,
+{
+    // Utilities for updating the tree nodes
+
+    fn get_node(&self, depth: usize, index: usize) -> H::Fr {
+        let node = *self
+            .nodes
+            .get(&(depth, index))
+            .unwrap_or_else(|| &self.cached_nodes[depth]);
+        node
+    }
+
+    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 {
+        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);
+            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(())
+    }
+}
+
+impl<H: Hasher> ZerokitMerkleProof for OptimalMerkleProof<H>
+where
+    H: Hasher,
+{
+    type Index = u8;
+    type Hasher = H;
+
+    #[must_use]
+    // Returns the length of a Merkle proof
+    fn length(&self) -> usize {
+        self.0.len()
+    }
+
+    /// Computes the leaf index corresponding to a Merkle proof
+    #[must_use]
+    fn leaf_index(&self) -> usize {
+        // In current implementation the path indexes in a proof correspond to the binary representation of the leaf index
+        let mut binary_repr = self.get_path_index();
+        binary_repr.reverse();
+        binary_repr
+            .into_iter()
+            .fold(0, |acc, digit| (acc << 1) + usize::from(digit))
+    }
+
+    #[must_use]
+    /// Returns the path elements forming a Merkle proof
+    fn get_path_elements(&self) -> Vec<H::Fr> {
+        self.0.iter().map(|x| x.0).collect()
+    }
+
+    /// Returns the path indexes forming a Merkle proof
+    #[must_use]
+    fn get_path_index(&self) -> Vec<u8> {
+        self.0.iter().map(|x| x.1).collect()
+    }
+
+    #[must_use]
+    /// Computes the Merkle root corresponding by iteratively hashing a Merkle proof with a given input leaf
+    fn compute_root_from(&self, leaf: &H::Fr) -> H::Fr {
+        let mut acc: H::Fr = *leaf;
+        for w in self.0.iter() {
+            if w.1 == 0 {
+                acc = H::hash(&[acc, w.0]);
+            } else {
+                acc = H::hash(&[w.0, acc]);
+            }
+        }
+        acc
+    }
+}
+
+// Debug formatting for printing a (Optimal) Merkle Proof
+impl<H> Debug for OptimalMerkleProof<H>
+where
+    H: Hasher,
+    H::Fr: Debug,
+{
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_tuple("Proof").field(&self.0).finish()
+    }
+}

utils/src/pm_tree/mod.rs

@@ -0,0 +1,4 @@
+pub mod sled_adapter;
+pub use self::sled_adapter::*;
+pub use pmtree;
+pub use sled::*;

utils/src/pm_tree/sled_adapter.rs

@@ -0,0 +1,105 @@
+use pmtree::*;
+
+use sled::Db as Sled;
+use std::collections::HashMap;
+use std::thread;
+use std::time::Duration;
+
+pub struct SledDB(Sled);
+
+impl SledDB {
+    fn new_with_tries(config: <SledDB as Database>::Config, tries: u32) -> PmtreeResult<Self> {
+        // If we've tried more than 10 times, we give up and return an error.
+        if tries >= 10 {
+            return Err(PmtreeErrorKind::DatabaseError(
+                DatabaseErrorKind::CustomError(format!(
+                    "Cannot create database: exceeded maximum retry attempts. {config:#?}"
+                )),
+            ));
+        }
+        match config.open() {
+            Ok(db) => Ok(SledDB(db)),
+            Err(e) if e.to_string().contains("WouldBlock") => {
+                // try till the fd is freed
+                // sleep for 10^tries milliseconds, then recursively try again
+                thread::sleep(Duration::from_millis(10u64.pow(tries)));
+                Self::new_with_tries(config, tries + 1)
+            }
+            Err(e) => {
+                // On any other error, we return immediately.
+                Err(PmtreeErrorKind::DatabaseError(
+                    DatabaseErrorKind::CustomError(format!(
+                        "Cannot create database: {e} {config:#?}"
+                    )),
+                ))
+            }
+        }
+    }
+}
+
+impl Database for SledDB {
+    type Config = sled::Config;
+
+    fn new(config: Self::Config) -> PmtreeResult<Self> {
+        let db = Self::new_with_tries(config, 0)?;
+        Ok(db)
+    }
+
+    fn load(config: Self::Config) -> PmtreeResult<Self> {
+        let db = match config.open() {
+            Ok(db) => db,
+            Err(e) => {
+                return Err(PmtreeErrorKind::DatabaseError(
+                    DatabaseErrorKind::CustomError(format!("Cannot load database: {e}")),
+                ))
+            }
+        };
+
+        if !db.was_recovered() {
+            return Err(PmtreeErrorKind::DatabaseError(
+                DatabaseErrorKind::CustomError(format!(
+                    "Database was not recovered: {}",
+                    config.path.display()
+                )),
+            ));
+        }
+
+        Ok(SledDB(db))
+    }
+
+    fn close(&mut self) -> PmtreeResult<()> {
+        let _ = self.0.flush().map_err(|_| {
+            PmtreeErrorKind::DatabaseError(DatabaseErrorKind::CustomError(
+                "Cannot flush database".to_string(),
+            ))
+        })?;
+        Ok(())
+    }
+
+    fn get(&self, key: DBKey) -> PmtreeResult<Option<Value>> {
+        match self.0.get(key) {
+            Ok(value) => Ok(value.map(|val| val.to_vec())),
+            Err(_e) => Err(PmtreeErrorKind::TreeError(TreeErrorKind::InvalidKey)),
+        }
+    }
+
+    fn put(&mut self, key: DBKey, value: Value) -> PmtreeResult<()> {
+        match self.0.insert(key, value) {
+            Ok(_) => Ok(()),
+            Err(_e) => Err(PmtreeErrorKind::TreeError(TreeErrorKind::InvalidKey)),
+        }
+    }
+
+    fn put_batch(&mut self, subtree: HashMap<DBKey, Value>) -> PmtreeResult<()> {
+        let mut batch = sled::Batch::default();
+
+        for (key, value) in subtree {
+            batch.insert(&key, value);
+        }
+
+        self.0
+            .apply_batch(batch)
+            .map_err(|_| PmtreeErrorKind::TreeError(TreeErrorKind::InvalidKey))?;
+        Ok(())
+    }
+}

utils/src/poseidon/poseidon_constants.rs

@@ -11,7 +11,7 @@
 
 #![allow(dead_code)]
 
-use ark_ff::{FpParameters, PrimeField};
+use ark_ff::PrimeField;
 use num_bigint::BigUint;
 
 pub struct PoseidonGrainLFSR {
@@ -36,20 +36,12 @@ impl PoseidonGrainLFSR {
         assert!(is_field == 1);
 
         // b0, b1 describes the field
-        if is_field == 1 {
-            state[1] = true;
-        } else {
-            state[1] = false;
-        }
+        state[1] = is_field == 1;
 
         assert!(is_sbox_an_inverse == 0 || is_sbox_an_inverse == 1);
 
         // b2, ..., b5 describes the S-BOX
-        if is_sbox_an_inverse == 1 {
-            state[5] = true;
-        } else {
-            state[5] = false;
-        }
+        state[5] = is_sbox_an_inverse == 1;
 
         // b6, ..., b17 are the binary representation of n (prime_num_bits)
         {
@@ -129,8 +121,8 @@ impl PoseidonGrainLFSR {
         &mut self,
         num_elems: usize,
     ) -> Vec<F> {
-        assert_eq!(F::Params::MODULUS_BITS as u64, self.prime_num_bits);
-        let modulus: BigUint = F::Params::MODULUS.into();
+        assert_eq!(F::MODULUS_BIT_SIZE as u64, self.prime_num_bits);
+        let modulus: BigUint = F::MODULUS.into();
 
         let mut res = Vec::new();
         for _ in 0..num_elems {
@@ -163,7 +155,7 @@ impl PoseidonGrainLFSR {
     }
 
     pub fn get_field_elements_mod_p<F: PrimeField>(&mut self, num_elems: usize) -> Vec<F> {
-        assert_eq!(F::Params::MODULUS_BITS as u64, self.prime_num_bits);
+        assert_eq!(F::MODULUS_BIT_SIZE as u64, self.prime_num_bits);
 
         let mut res = Vec::new();
         for _ in 0..num_elems {

utils/src/poseidon/poseidon_hash.rs

@@ -4,7 +4,7 @@
 // and adapted to work over arkworks field traits and custom data structures
 
 use crate::poseidon_constants::find_poseidon_ark_and_mds;
-use ark_ff::{FpParameters, PrimeField};
+use ark_ff::PrimeField;
 
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct RoundParamenters<F: PrimeField> {
@@ -32,7 +32,7 @@ impl<F: PrimeField> Poseidon<F> {
             let (ark, mds) = find_poseidon_ark_and_mds::<F>(
                 1, // is_field = 1
                 0, // is_sbox_inverse = 0
-                F::Params::MODULUS_BITS as u64,
+                F::MODULUS_BIT_SIZE as u64,
                 t,
                 n_rounds_f as u64,
                 n_rounds_p as u64,

utils/tests/merkle_tree.rs

@@ -3,8 +3,11 @@
 mod test {
     use hex_literal::hex;
     use tiny_keccak::{Hasher as _, Keccak};
-    use utils::{FullMerkleTree, Hasher, OptimalMerkleTree};
-
+    use zerokit_utils::{
+        FullMerkleConfig, FullMerkleTree, Hasher, OptimalMerkleConfig, OptimalMerkleTree,
+        ZerokitMerkleProof, ZerokitMerkleTree,
+    };
+    #[derive(Clone, Copy, Eq, PartialEq)]
     struct Keccak256;
 
     impl Hasher for Keccak256 {
@@ -44,14 +47,17 @@ mod test {
             hex!("a9bb8c3f1f12e9aa903a50c47f314b57610a3ab32f2d463293f58836def38d36"),
         ];
 
-        let mut tree = FullMerkleTree::<Keccak256>::new(2, [0; 32]);
+        let mut tree =
+            FullMerkleTree::<Keccak256>::new(2, [0; 32], FullMerkleConfig::default()).unwrap();
         assert_eq!(tree.root(), default_tree_root);
         for i in 0..leaves.len() {
             tree.set(i, leaves[i]).unwrap();
             assert_eq!(tree.root(), roots[i]);
         }
 
-        let mut tree = OptimalMerkleTree::<Keccak256>::new(2, [0; 32]);
+        let mut tree =
+            OptimalMerkleTree::<Keccak256>::new(2, [0; 32], OptimalMerkleConfig::default())
+                .unwrap();
         assert_eq!(tree.root(), default_tree_root);
         for i in 0..leaves.len() {
             tree.set(i, leaves[i]).unwrap();
@@ -69,7 +75,8 @@ mod test {
         ];
 
         // We thest the FullMerkleTree implementation
-        let mut tree = FullMerkleTree::<Keccak256>::new(2, [0; 32]);
+        let mut tree =
+            FullMerkleTree::<Keccak256>::new(2, [0; 32], FullMerkleConfig::default()).unwrap();
         for i in 0..leaves.len() {
             // We set the leaves
             tree.set(i, leaves[i]).unwrap();
@@ -93,7 +100,9 @@ mod test {
         }
 
         // We test the OptimalMerkleTree implementation
-        let mut tree = OptimalMerkleTree::<Keccak256>::new(2, [0; 32]);
+        let mut tree =
+            OptimalMerkleTree::<Keccak256>::new(2, [0; 32], OptimalMerkleConfig::default())
+                .unwrap();
         for i in 0..leaves.len() {
             // We set the leaves
             tree.set(i, leaves[i]).unwrap();
@@ -116,4 +125,41 @@ mod test {
                 .unwrap());
         }
     }
+
+    #[test]
+    fn test_override_range() {
+        let initial_leaves = [
+            hex!("0000000000000000000000000000000000000000000000000000000000000001"),
+            hex!("0000000000000000000000000000000000000000000000000000000000000002"),
+            hex!("0000000000000000000000000000000000000000000000000000000000000003"),
+            hex!("0000000000000000000000000000000000000000000000000000000000000004"),
+        ];
+
+        let mut tree =
+            OptimalMerkleTree::<Keccak256>::new(2, [0; 32], OptimalMerkleConfig::default())
+                .unwrap();
+
+        // We set the leaves
+        tree.set_range(0, initial_leaves.iter().cloned()).unwrap();
+
+        let new_leaves = [
+            hex!("0000000000000000000000000000000000000000000000000000000000000005"),
+            hex!("0000000000000000000000000000000000000000000000000000000000000006"),
+        ];
+
+        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();
+
+        // ensure that the leaves are set correctly
+        for i in 0..new_leaves.len() {
+            assert_eq!(tree.get_leaf(i), new_leaves[i]);
+        }
+    }
 }

utils/tests/poseidon_constants.rs

@@ -3,7 +3,7 @@ mod test {
     use ark_bn254::Fr;
     use num_bigint::BigUint;
     use num_traits::Num;
-    use utils::poseidon_hash::Poseidon;
+    use zerokit_utils::poseidon_hash::Poseidon;
 
     const ROUND_PARAMS: [(usize, usize, usize, usize); 8] = [
         (2, 8, 56, 0),