Integrate pmtree into rln (#147)

* feat: integrate pmtree with zerokit

* fix: tests

Cargo.lock

@@ -2110,7 +2110,7 @@ dependencies = [
 [[package]]
 name = "pmtree"
 version = "1.0.0"
-source = "git+https://github.com/Rate-Limiting-Nullifier/pmtree?rev=f6d1a1fecad72cd39e6808e78085091d541dc882#f6d1a1fecad72cd39e6808e78085091d541dc882"
+source = "git+https://github.com/Rate-Limiting-Nullifier/pmtree?rev=b3a02216cece3e9c24e1754ea381bf784fd1df48#b3a02216cece3e9c24e1754ea381bf784fd1df48"
 dependencies = [
  "rayon",
 ]
@@ -2393,7 +2393,6 @@ dependencies = [
  "num-bigint",
  "num-traits",
  "once_cell",
- "pmtree",
  "rand",
  "rand_chacha",
  "serde_json",
@@ -3117,6 +3116,8 @@ dependencies = [
  "hex-literal",
  "num-bigint",
  "num-traits",
+ "pmtree",
+ "sled",
  "tiny-keccak",
 ]
 

rln/Cargo.toml

@@ -38,7 +38,6 @@ 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}
 
 # serialization
 serde_json = "1.0.48"
@@ -55,4 +54,4 @@ wasm = ["wasmer/js", "wasmer/std"]
 fullmerkletree = ["default"]
 
 # Note: pmtree feature is still experimental
-pmtree-ft = ["default", "pmtree"]
+pmtree-ft = ["default", "utils/pmtree-ft"]

rln/src/hashers.rs

@@ -28,6 +28,23 @@ pub fn poseidon_hash(input: &[Fr]) -> Fr {
         .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

rln/src/lib.rs

@@ -2,6 +2,8 @@
 
 pub mod circuit;
 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,159 @@
+use crate::circuit::Fr;
+use crate::hashers::{poseidon_hash, PoseidonHash};
+use crate::utils::{bytes_le_to_fr, fr_to_bytes_le};
+use color_eyre::{Report, Result};
+use std::fmt::Debug;
+use utils::*;
+
+pub struct PmTree {
+    tree: pmtree::MerkleTree<SledDB, PoseidonHash>,
+}
+
+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 pmtree::Hasher for PoseidonHash {
+    type Fr = Fr;
+
+    fn default_leaf() -> Self::Fr {
+        Fr::from(0)
+    }
+
+    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 hash(inputs: &[Self::Fr]) -> Self::Fr {
+        poseidon_hash(inputs)
+    }
+}
+
+pub struct PmtreeConfig(pm_tree::Config);
+impl Default for PmtreeConfig {
+    fn default() -> Self {
+        let tmp_path = std::env::temp_dir().join(format!("pmtree-{}", rand::random::<u64>()));
+        PmtreeConfig(pm_tree::Config::new().temporary(true).path(tmp_path))
+    }
+}
+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 })
+    }
+
+    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 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 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"))
+        }
+    }
+}
+
+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_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::hashers::poseidon_hash;
 use cfg_if::cfg_if;
-use utils::merkle_tree::*;
 
-#[cfg(feature = "pmtree-ft")]
+cfg_if! {
-use crate::utils::{bytes_le_to_fr, fr_to_bytes_le};
+    if #[cfg(feature = "pmtree-ft")] {
-#[cfg(feature = "pmtree-ft")]
+        use crate::pm_tree_adapter::*;
-use pmtree::*;
+    } 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/public.rs

@@ -10,11 +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};
-// use rkyv::Deserialize;
 
 cfg_if! {
     if #[cfg(not(target_arch = "wasm32"))] {
@@ -82,7 +81,7 @@ impl RLN<'_> {
         let verification_key = vk_from_folder(&resources_folder)?;
 
         // We compute a default empty tree
-        let tree = PoseidonTree::default(tree_height);
+        let tree = PoseidonTree::default(tree_height)?;
 
         Ok(RLN {
             witness_calculator,
@@ -140,7 +139,7 @@ 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"))]
@@ -164,7 +163,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(())
     }
@@ -239,7 +238,10 @@ impl RLN<'_> {
         let (leaves, _) = bytes_le_to_vec_fr(&leaves_byte)?;
 
         // We set the leaves
-        self.tree.set_range(index, leaves)
+        self.tree
+            .set_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.
@@ -1061,7 +1063,6 @@ mod test {
     use ark_std::{rand::thread_rng, UniformRand};
     use rand::Rng;
     use utils::ZerokitMerkleTree;
-    // use rkyv::Deserialize;
 
     #[test]
     // We test merkle batch Merkle tree additions
@@ -1132,7 +1133,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();
         }
 

rln/tests/ffi.rs

@@ -96,8 +96,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..no_of_leaves {
-        for i in 0..delete_range {
             let success = delete_leaf(rln_pointer, i);
             assert!(success, "delete leaf call failed");
         }

rln/tests/poseidon_tree.rs

@@ -5,7 +5,7 @@
 #[cfg(test)]
 mod test {
     use rln::circuit::*;
-    use rln::poseidon_tree::*;
+    use rln::hashers::PoseidonHash;
     use utils::{FullMerkleTree, OptimalMerkleTree, ZerokitMerkleProof, ZerokitMerkleTree};
 
     #[test]
@@ -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_full = FullMerkleTree::<PoseidonHash>::default(tree_height).unwrap();
-        let mut tree_opt = OptimalMerkleTree::<PoseidonHash>::default(tree_height);
+        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::hashers::{hash_to_field, 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

@@ -10,6 +10,8 @@ mod test {
     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#"
             {
@@ -172,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
@@ -382,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");

utils/Cargo.toml

@@ -8,6 +8,8 @@ license = "MIT OR Apache-2.0"
 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 = { git = "https://github.com/Rate-Limiting-Nullifier/pmtree", rev = "b3a02216cece3e9c24e1754ea381bf784fd1df48", optional = true}
+sled = "=0.34.7"
 
 [dev-dependencies]
 ark-bn254 = "=0.4.0"
@@ -18,3 +20,4 @@ tiny-keccak = { version = "2.0.2", features = ["keccak"] }
 [features]
 default = ["parallel"]
 parallel = ["ark-ff/parallel"]
+pmtree-ft = ["pmtree"]

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

@@ -1,4 +1,4 @@
-use crate::merkle_tree::{Hasher, ZerokitMerkleProof, ZerokitMerkleTree};
+use crate::merkle_tree::{FrOf, Hasher, ZerokitMerkleProof, ZerokitMerkleTree};
 use color_eyre::{Report, Result};
 use std::{
     cmp::max,
@@ -45,16 +45,21 @@ pub enum FullMerkleBranch<H: Hasher> {
 pub struct FullMerkleProof<H: Hasher>(pub Vec<FullMerkleBranch<H>>);
 
 /// Implementations
-
+impl<H: Hasher> ZerokitMerkleTree for FullMerkleTree<H>
-impl<H: Hasher> ZerokitMerkleTree<H> for FullMerkleTree<H> {
+where
+    H: Hasher,
+{
     type Proof = FullMerkleProof<H>;
-    fn default(depth: usize) -> Self {
+    type Hasher = H;
-        FullMerkleTree::<H>::new(depth, H::default_leaf())
+    type Config = ();
+
+    fn default(depth: usize) -> Result<Self> {
+        FullMerkleTree::<H>::new(depth, Self::Hasher::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
-    fn new(depth: usize, initial_leaf: H::Fr) -> Self {
+    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)
@@ -72,12 +77,12 @@ impl<H: Hasher> ZerokitMerkleTree<H> for FullMerkleTree<H> {
 
         let next_index = 0;
 
-        Self {
+        Ok(Self {
             depth,
             cached_nodes,
             nodes,
             next_index,
-        }
+        })
     }
 
     // Returns the depth of the tree
@@ -97,12 +102,12 @@ impl<H: Hasher> ZerokitMerkleTree<H> for FullMerkleTree<H> {
 
     #[must_use]
     // Returns the root of the tree
-    fn root(&self) -> H::Fr {
+    fn root(&self) -> FrOf<Self::Hasher> {
         self.nodes[0]
     }
 
     // Sets a leaf at the specified tree index
-    fn set(&mut self, leaf: usize, hash: H::Fr) -> Result<()> {
+    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(())
@@ -110,7 +115,11 @@ impl<H: Hasher> ZerokitMerkleTree<H> for FullMerkleTree<H> {
 
     // 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<()> {
+    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
@@ -131,7 +140,7 @@ impl<H: Hasher> ZerokitMerkleTree<H> for FullMerkleTree<H> {
     }
 
     // Sets a leaf at the next available index
-    fn update_next(&mut self, leaf: H::Fr) -> Result<()> {
+    fn update_next(&mut self, leaf: FrOf<Self::Hasher>) -> Result<()> {
         self.set(self.next_index, leaf)?;
         Ok(())
     }
@@ -165,7 +174,7 @@ impl<H: Hasher> ZerokitMerkleTree<H> for FullMerkleTree<H> {
     }
 
     // Verifies a Merkle proof with respect to the input leaf and the tree root
-    fn verify(&self, hash: &H::Fr, proof: &FullMerkleProof<H>) -> Result<bool> {
+    fn verify(&self, hash: &FrOf<Self::Hasher>, proof: &FullMerkleProof<H>) -> Result<bool> {
         Ok(proof.compute_root_from(hash) == self.root())
     }
 }
@@ -212,8 +221,9 @@ where
     }
 }
 
-impl<H: Hasher> ZerokitMerkleProof<H> for FullMerkleProof<H> {
+impl<H: Hasher> ZerokitMerkleProof for FullMerkleProof<H> {
     type Index = u8;
+    type Hasher = H;
 
     #[must_use]
     // Returns the length of a Merkle proof
@@ -232,7 +242,7 @@ impl<H: Hasher> ZerokitMerkleProof<H> for FullMerkleProof<H> {
 
     #[must_use]
     /// Returns the path elements forming a Merkle proof
-    fn get_path_elements(&self) -> Vec<H::Fr> {
+    fn get_path_elements(&self) -> Vec<FrOf<Self::Hasher>> {
         self.0
             .iter()
             .map(|x| match x {
@@ -255,7 +265,7 @@ impl<H: Hasher> ZerokitMerkleProof<H> for FullMerkleProof<H> {
 
     /// Computes the Merkle root corresponding by iteratively hashing a Merkle proof with a given input leaf
     #[must_use]
-    fn compute_root_from(&self, hash: &H::Fr) -> H::Fr {
+    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]),

utils/src/merkle_tree/merkle_tree.rs

@@ -28,39 +28,42 @@ pub trait Hasher {
     fn hash(input: &[Self::Fr]) -> Self::Fr;
 }
 
-/// In the ZerokitMerkleTree trait we define the methods that are required to be implemented by a Merkle tree
+pub type FrOf<H> = <H as Hasher>::Fr;
-/// Including, OptimalMerkleTree, FullMerkleTree, Pmtree
-pub trait ZerokitMerkleTree<H: Hasher>
-where
-    H: Hasher,
-{
-    type Proof;
 
-    fn default(depth: usize) -> Self;
+/// In the ZerokitMerkleTree trait we define the methods that are required to be implemented by a Merkle tree
-    fn new(depth: usize, default_leaf: H::Fr) -> Self;
+/// Including, OptimalMerkleTree, FullMerkleTree
+pub trait ZerokitMerkleTree {
+    type Proof: ZerokitMerkleProof;
+    type Hasher: Hasher;
+    type Config: Default;
+
+    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) -> H::Fr;
+    fn root(&self) -> FrOf<Self::Hasher>;
-    fn set(&mut self, index: usize, leaf: H::Fr) -> Result<()>;
+    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 = H::Fr>;
+        I: IntoIterator<Item = FrOf<Self::Hasher>>;
-    fn update_next(&mut self, leaf: H::Fr) -> Result<()>;
+    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: &H::Fr, witness: &Self::Proof) -> Result<bool>;
+    fn verify(&self, leaf: &FrOf<Self::Hasher>, witness: &Self::Proof) -> Result<bool>;
 }
 
-pub trait ZerokitMerkleProof<H: Hasher>
+pub trait ZerokitMerkleProof {
-where
-    H: Hasher,
-{
     type Index;
+    type Hasher: Hasher;
 
     fn length(&self) -> usize;
     fn leaf_index(&self) -> usize;
-    fn get_path_elements(&self) -> Vec<H::Fr>;
+    fn get_path_elements(&self) -> Vec<FrOf<Self::Hasher>>;
     fn get_path_index(&self) -> Vec<Self::Index>;
-    fn compute_root_from(&self, leaf: &H::Fr) -> H::Fr;
+    fn compute_root_from(&self, leaf: &FrOf<Self::Hasher>) -> FrOf<Self::Hasher>;
 }

utils/src/merkle_tree/optimal_merkle_tree.rs

@@ -37,31 +37,33 @@ pub struct OptimalMerkleProof<H: Hasher>(pub Vec<(H::Fr, u8)>);
 
 /// Implementations
 
-impl<H: Hasher> ZerokitMerkleTree<H> for OptimalMerkleTree<H>
+impl<H: Hasher> ZerokitMerkleTree for OptimalMerkleTree<H>
 where
     H: Hasher,
 {
     type Proof = OptimalMerkleProof<H>;
+    type Hasher = H;
+    type Config = ();
 
-    fn default(depth: usize) -> Self {
+    fn default(depth: usize) -> Result<Self> {
-        OptimalMerkleTree::<H>::new(depth, H::default_leaf())
+        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
-    fn new(depth: usize, default_leaf: H::Fr) -> Self {
+    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();
-        OptimalMerkleTree {
+        Ok(OptimalMerkleTree {
             cached_nodes: cached_nodes.clone(),
             depth,
             nodes: HashMap::new(),
             next_index: 0,
-        }
+        })
     }
 
     // Returns the depth of the tree
@@ -205,11 +207,13 @@ where
     }
 }
 
-impl<H: Hasher> ZerokitMerkleProof<H> for OptimalMerkleProof<H>
+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 {

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,75 @@
+use pmtree::*;
+
+use sled::Db as Sled;
+use std::collections::HashMap;
+
+pub struct SledDB(Sled);
+
+impl Database for SledDB {
+    type Config = sled::Config;
+
+    fn new(config: Self::Config) -> PmtreeResult<Self> {
+        let db: Sled = match config.open() {
+            Ok(db) => db,
+            Err(e) => {
+                return Err(PmtreeErrorKind::DatabaseError(
+                    DatabaseErrorKind::CustomError(format!(
+                        "Cannot create database: {} {:#?}",
+                        e, config
+                    )),
+                ))
+            }
+        };
+
+        Ok(SledDB(db))
+    }
+
+    fn load(config: Self::Config) -> PmtreeResult<Self> {
+        let db: Sled = match sled::open(&config.path) {
+            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 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/tests/merkle_tree.rs

@@ -4,7 +4,7 @@ mod test {
     use hex_literal::hex;
     use tiny_keccak::{Hasher as _, Keccak};
     use utils::{FullMerkleTree, Hasher, OptimalMerkleTree, ZerokitMerkleProof, ZerokitMerkleTree};
-
+    #[derive(Clone, Copy, Eq, PartialEq)]
     struct Keccak256;
 
     impl Hasher for Keccak256 {
@@ -44,14 +44,14 @@ mod test {
             hex!("a9bb8c3f1f12e9aa903a50c47f314b57610a3ab32f2d463293f58836def38d36"),
         ];
 
-        let mut tree = FullMerkleTree::<Keccak256>::new(2, [0; 32]);
+        let mut tree = FullMerkleTree::<Keccak256>::new(2, [0; 32], ()).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], ()).unwrap();
         assert_eq!(tree.root(), default_tree_root);
         for i in 0..leaves.len() {
             tree.set(i, leaves[i]).unwrap();
@@ -69,7 +69,7 @@ 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], ()).unwrap();
         for i in 0..leaves.len() {
             // We set the leaves
             tree.set(i, leaves[i]).unwrap();
@@ -93,7 +93,7 @@ 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], ()).unwrap();
         for i in 0..leaves.len() {
             // We set the leaves
             tree.set(i, leaves[i]).unwrap();