implement the specified sparse merkletree (#82)

* wip

* prototype custom predicates 1b

* feat: implement custom pred recursion

* files reorg, add github CI for rustfmt checks

* start sparsemerkletree. impl add_leaf method, initial Leaf & Intermediate types with methods

* mt: add hash computation of all the nodes in the tree, add method to print the tree to visualize it as a graphviz

* mt: add  (till the leaf) method which is used by get,contains,prove methods

* mt: add verify (of inclusion) method

* mt: update 'down' method to reuse siblings, update get,contains,prove methods (the three use 'down' under the hood)

* Add nonexistence proofs and iterator

* Add iterator test

* migrate usage of old merkletree to the new merkletree impl in POD2 code

---------

Co-authored-by: Eduard S. <eduardsanou@posteo.net>
Co-authored-by: Ahmad <root@ahmadafuni.com>

src/backends/mock_main/mod.rs

@@ -485,7 +485,7 @@ pub mod tests {
             pk: "ZooDeel".into(),
         };
         let pay_stub_pod = pay_stub_builder.sign(&mut signer)?;
-        let kyc_builder = zu_kyc_pod_builder(&params, &gov_id_pod, &pay_stub_pod);
+        let kyc_builder = zu_kyc_pod_builder(&params, &gov_id_pod, &pay_stub_pod)?;
 
         let mut prover = MockProver {};
         let kyc_pod = kyc_builder.prove(&mut prover)?;
@@ -501,7 +501,7 @@ pub mod tests {
 
     #[test]
     fn test_mock_main_great_boy() -> Result<()> {
-        let great_boy_builder = great_boy_pod_full_flow();
+        let great_boy_builder = great_boy_pod_full_flow()?;
 
         let mut prover = MockProver {};
         let great_boy_pod = great_boy_builder.prove(&mut prover)?;
@@ -520,7 +520,7 @@ pub mod tests {
 
     #[test]
     fn test_mock_main_tickets() -> Result<()> {
-        let tickets_builder = tickets_pod_full_flow();
+        let tickets_builder = tickets_pod_full_flow()?;
         let mut prover = MockProver {};
         let proof_pod = tickets_builder.prove(&mut prover)?;
         let pod = proof_pod.pod.into_any().downcast::<MockMainPod>().unwrap();

src/backends/mock_signed.rs

@@ -2,6 +2,7 @@ use anyhow::Result;
 use std::any::Any;
 use std::collections::HashMap;
 
+use crate::constants::MAX_DEPTH;
 use crate::middleware::{
     containers::Dictionary, hash_str, AnchoredKey, Hash, Params, Pod, PodId, PodSigner, PodType,
     Statement, Value, KEY_SIGNER, KEY_TYPE,
@@ -19,7 +20,7 @@ impl PodSigner for MockSigner {
         kvs.insert(hash_str(&KEY_SIGNER), Value(pk_hash.0));
         kvs.insert(hash_str(&KEY_TYPE), Value::from(PodType::MockSigned));
 
-        let dict = Dictionary::new(&kvs);
+        let dict = Dictionary::new(&kvs)?;
         let id = PodId(dict.commitment());
         let signature = format!("{}_signed_by_{}", id, pk_hash);
         Ok(Box::new(MockSignedPod {
@@ -49,13 +50,17 @@ impl Pod for MockSignedPod {
         }
 
         // Verify id
-        let mt = MerkleTree::new(
+        let mt = match MerkleTree::new(
+            MAX_DEPTH,
             &self
                 .dict
                 .iter()
                 .map(|(&k, &v)| (k, v))
                 .collect::<HashMap<Value, Value>>(),
-        );
+        ) {
+            Ok(mt) => mt,
+            Err(_) => return false,
+        };
         let id = PodId(mt.root());
         if id != self.id {
             return false;
@@ -93,14 +98,16 @@ impl Pod for MockSignedPod {
 
 #[cfg(test)]
 pub mod tests {
-    use super::*;
-    use crate::frontend;
-    use crate::middleware::{self, F, NULL};
     use plonky2::field::types::Field;
     use std::iter;
 
+    use super::*;
+    use crate::constants::MAX_DEPTH;
+    use crate::frontend;
+    use crate::middleware::{self, F, NULL};
+
     #[test]
-    fn test_mock_signed_0() {
+    fn test_mock_signed_0() -> Result<()> {
         let params = middleware::Params::default();
         let mut pod = frontend::SignedPodBuilder::new(&params);
         pod.insert("idNumber", "4242424242");
@@ -131,7 +138,7 @@ pub mod tests {
             .map(|(AnchoredKey(_, k), v)| (Value(k.0), v))
             .chain(iter::once(bad_kv))
             .collect::<HashMap<Value, Value>>();
-        let bad_mt = MerkleTree::new(&bad_kvs_mt);
+        let bad_mt = MerkleTree::new(MAX_DEPTH, &bad_kvs_mt)?;
         bad_pod.dict.mt = bad_mt;
         assert_eq!(bad_pod.verify(), false);
 
@@ -143,8 +150,10 @@ pub mod tests {
             .map(|(AnchoredKey(_, k), v)| (Value(k.0), v))
             .chain(iter::once(bad_kv))
             .collect::<HashMap<Value, Value>>();
-        let bad_mt = MerkleTree::new(&bad_kvs_mt);
+        let bad_mt = MerkleTree::new(MAX_DEPTH, &bad_kvs_mt)?;
         bad_pod.dict.mt = bad_mt;
         assert_eq!(bad_pod.verify(), false);
+
+        Ok(())
     }
 }

src/constants.rs

@@ -0,0 +1 @@
+pub const MAX_DEPTH: usize = 32;

src/examples.rs

@@ -1,3 +1,4 @@
+use anyhow::Result;
 use std::collections::HashMap;
 
 use crate::backends::mock_signed::MockSigner;
@@ -24,8 +25,8 @@ pub fn zu_kyc_pod_builder(
     params: &Params,
     gov_id: &SignedPod,
     pay_stub: &SignedPod,
-) -> MainPodBuilder {
+) -> Result<MainPodBuilder> {
-    let sanction_list = Value::Dictionary(Dictionary::new(&HashMap::new())); // empty dictionary
+    let sanction_list = Value::Dictionary(Dictionary::new(&HashMap::new())?); // empty dictionary
     let now_minus_18y: i64 = 1169909388;
     let now_minus_1y: i64 = 1706367566;
 
@@ -41,7 +42,7 @@ pub fn zu_kyc_pod_builder(
     ));
     kyc.pub_op(op!(eq, (pay_stub, "startDate"), now_minus_1y));
 
-    kyc
+    Ok(kyc)
 }
 
 // GreatBoy
@@ -130,7 +131,7 @@ pub fn great_boy_pod_builder(
     great_boy
 }
 
-pub fn great_boy_pod_full_flow() -> MainPodBuilder {
+pub fn great_boy_pod_full_flow() -> Result<MainPodBuilder> {
     let params = Params {
         max_input_signed_pods: 6,
         max_statements: 100,
@@ -179,8 +180,8 @@ pub fn great_boy_pod_full_flow() -> MainPodBuilder {
     alice_friend_pods.push(friend.sign(&mut bob_signer).unwrap());
     alice_friend_pods.push(friend.sign(&mut charlie_signer).unwrap());
 
-    let good_boy_issuers_dict = Value::Dictionary(Dictionary::new(&HashMap::new())); // empty
+    let good_boy_issuers_dict = Value::Dictionary(Dictionary::new(&HashMap::new())?); // empty
-    great_boy_pod_builder(
+    Ok(great_boy_pod_builder(
         &params,
         [
             &bob_good_boys[0],
@@ -191,7 +192,7 @@ pub fn great_boy_pod_full_flow() -> MainPodBuilder {
         [&alice_friend_pods[0], &alice_friend_pods[1]],
         &good_boy_issuers_dict,
         alice,
-    )
+    ))
 }
 
 // Tickets
@@ -229,15 +230,15 @@ pub fn tickets_pod_builder(
     builder
 }
 
-pub fn tickets_pod_full_flow() -> MainPodBuilder {
+pub fn tickets_pod_full_flow() -> Result<MainPodBuilder> {
     let params = Params::default();
     let builder = tickets_sign_pod_builder(&params);
     let signed_pod = builder.sign(&mut MockSigner { pk: "test".into() }).unwrap();
-    tickets_pod_builder(
+    Ok(tickets_pod_builder(
         &params,
         &signed_pod,
         123,
         true,
-        &Value::Dictionary(Dictionary::new(&HashMap::new())),
+        &Value::Dictionary(Dictionary::new(&HashMap::new())?),
-    )
+    ))
 }

src/frontend/mod.rs

@@ -510,7 +510,7 @@ pub mod tests {
         let pay_stub = pay_stub.sign(&mut signer).unwrap();
         println!("{}", pay_stub);
 
-        let kyc = zu_kyc_pod_builder(&params, &gov_id, &pay_stub);
+        let kyc = zu_kyc_pod_builder(&params, &gov_id, &pay_stub)?;
         println!("{}", kyc);
 
         // TODO: prove kyc with MockProver and print it
@@ -520,7 +520,7 @@ pub mod tests {
 
     #[test]
     fn test_front_great_boy() -> Result<()> {
-        let great_boy = great_boy_pod_full_flow();
+        let great_boy = great_boy_pod_full_flow()?;
         println!("{}", great_boy);
 
         // TODO: prove kyc with MockProver and print it
@@ -530,7 +530,7 @@ pub mod tests {
 
     #[test]
     fn test_front_tickets() -> Result<()> {
-        let builder = tickets_pod_full_flow();
+        let builder = tickets_pod_full_flow()?;
         println!("{}", builder);
 
         Ok(())

src/lib.rs

@@ -1,4 +1,5 @@
 pub mod backends;
+pub mod constants;
 pub mod frontend;
 pub mod middleware;
 pub mod primitives;

src/middleware/containers.rs

@@ -6,6 +6,7 @@ use plonky2::plonk::config::Hasher;
 use std::collections::HashMap;
 
 use super::{Hash, Value, EMPTY};
+use crate::constants::MAX_DEPTH;
 use crate::primitives::merkletree::{MerkleProof, MerkleTree};
 
 /// Dictionary: the user original keys and values are hashed to be used in the leaf.
@@ -18,11 +19,11 @@ pub struct Dictionary {
 }
 
 impl Dictionary {
-    pub fn new(kvs: &HashMap<Hash, Value>) -> Self {
+    pub fn new(kvs: &HashMap<Hash, Value>) -> Result<Self> {
         let kvs: HashMap<Value, Value> = kvs.into_iter().map(|(&k, &v)| (Value(k.0), v)).collect();
-        Self {
+        Ok(Self {
-            mt: MerkleTree::new(&kvs),
+            mt: MerkleTree::new(MAX_DEPTH, &kvs)?,
-        }
+        })
     }
     pub fn commitment(&self) -> Hash {
         self.mt.root()
@@ -30,25 +31,25 @@ impl Dictionary {
     pub fn get(&self, key: &Value) -> Result<Value> {
         self.mt.get(key)
     }
-    pub fn prove(&self, key: &Value) -> Result<MerkleProof> {
+    pub fn prove(&self, key: &Value) -> Result<(Value, MerkleProof)> {
         self.mt.prove(key)
     }
     pub fn prove_nonexistence(&self, key: &Value) -> Result<MerkleProof> {
         self.mt.prove_nonexistence(key)
     }
     pub fn verify(root: Hash, proof: &MerkleProof, key: &Value, value: &Value) -> Result<()> {
-        MerkleTree::verify(root, proof, key, value)
+        MerkleTree::verify(MAX_DEPTH, root, proof, key, value)
     }
     pub fn verify_nonexistence(root: Hash, proof: &MerkleProof, key: &Value) -> Result<()> {
-        MerkleTree::verify_nonexistence(root, proof, key)
+        MerkleTree::verify_nonexistence(MAX_DEPTH, root, proof, key)
     }
-    pub fn iter(&self) -> std::collections::hash_map::Iter<Value, Value> {
+    pub fn iter(&self) -> crate::primitives::merkletree::Iter {
         self.mt.iter()
     }
 }
 impl<'a> IntoIterator for &'a Dictionary {
     type Item = (&'a Value, &'a Value);
-    type IntoIter = std::collections::hash_map::Iter<'a, Value, Value>;
+    type IntoIter = crate::primitives::merkletree::Iter<'a>;
 
     fn into_iter(self) -> Self::IntoIter {
         self.mt.iter()
@@ -71,7 +72,7 @@ pub struct Set {
 }
 
 impl Set {
-    pub fn new(set: &Vec<Value>) -> Self {
+    pub fn new(set: &Vec<Value>) -> Result<Self> {
         let kvs: HashMap<Value, Value> = set
             .into_iter()
             .map(|e| {
@@ -79,29 +80,30 @@ impl Set {
                 (Value(h), EMPTY)
             })
             .collect();
-        Self {
+        Ok(Self {
-            mt: MerkleTree::new(&kvs),
+            mt: MerkleTree::new(MAX_DEPTH, &kvs)?,
-        }
+        })
     }
     pub fn commitment(&self) -> Hash {
         self.mt.root()
     }
-    pub fn contains(&self, value: &Value) -> bool {
+    pub fn contains(&self, value: &Value) -> Result<bool> {
         self.mt.contains(value)
     }
     pub fn prove(&self, value: &Value) -> Result<MerkleProof> {
-        self.mt.prove(value)
+        let (_, proof) = self.mt.prove(value)?;
+        Ok(proof)
     }
     pub fn prove_nonexistence(&self, value: &Value) -> Result<MerkleProof> {
         self.mt.prove_nonexistence(value)
     }
     pub fn verify(root: Hash, proof: &MerkleProof, value: &Value) -> Result<()> {
-        MerkleTree::verify(root, proof, value, &EMPTY)
+        MerkleTree::verify(MAX_DEPTH, root, proof, value, &EMPTY)
     }
     pub fn verify_nonexistence(root: Hash, proof: &MerkleProof, value: &Value) -> Result<()> {
-        MerkleTree::verify_nonexistence(root, proof, value)
+        MerkleTree::verify_nonexistence(MAX_DEPTH, root, proof, value)
     }
-    pub fn iter(&self) -> std::collections::hash_map::Iter<Value, Value> {
+    pub fn iter(&self) -> crate::primitives::merkletree::Iter {
         self.mt.iter()
     }
 }
@@ -123,16 +125,16 @@ pub struct Array {
 }
 
 impl Array {
-    pub fn new(array: &Vec<Value>) -> Self {
+    pub fn new(array: &Vec<Value>) -> Result<Self> {
         let kvs: HashMap<Value, Value> = array
             .into_iter()
             .enumerate()
             .map(|(i, &e)| (Value::from(i as i64), e))
             .collect();
 
-        Self {
+        Ok(Self {
-            mt: MerkleTree::new(&kvs),
+            mt: MerkleTree::new(MAX_DEPTH, &kvs)?,
-        }
+        })
     }
     pub fn commitment(&self) -> Hash {
         self.mt.root()
@@ -140,13 +142,13 @@ impl Array {
     pub fn get(&self, i: usize) -> Result<Value> {
         self.mt.get(&Value::from(i as i64))
     }
-    pub fn prove(&self, i: usize) -> Result<MerkleProof> {
+    pub fn prove(&self, i: usize) -> Result<(Value, MerkleProof)> {
         self.mt.prove(&Value::from(i as i64))
     }
     pub fn verify(root: Hash, proof: &MerkleProof, i: usize, value: &Value) -> Result<()> {
-        MerkleTree::verify(root, proof, &Value::from(i as i64), value)
+        MerkleTree::verify(MAX_DEPTH, root, proof, &Value::from(i as i64), value)
     }
-    pub fn iter(&self) -> std::collections::hash_map::Iter<Value, Value> {
+    pub fn iter(&self) -> crate::primitives::merkletree::Iter {
         self.mt.iter()
     }
 }

src/middleware/mod.rs

@@ -48,6 +48,15 @@ pub type Entry = (String, Value);
 #[derive(Clone, Copy, Debug, Default, Hash, PartialEq, Eq)]
 pub struct Value(pub [F; 4]);
 
+impl Value {
+    pub fn to_bytes(self) -> Vec<u8> {
+        self.0
+            .iter()
+            .flat_map(|e| e.to_canonical_u64().to_le_bytes())
+            .collect()
+    }
+}
+
 impl Ord for Value {
     fn cmp(&self, other: &Self) -> Ordering {
         for (lhs, rhs) in self.0.iter().zip(other.0.iter()).rev() {

src/primitives/merkletree.rs

@@ -1,214 +1,642 @@
-/// MerkleTree implementation for POD2.
+//! Module that implements the MerkleTree specified at
-///
+//! https://0xparc.github.io/pod2/merkletree.html .
-/// Current implementation is a wrapper on top of Plonky2's MerkleTree, but the future iteration
-/// will replace it by the MerkleTree specified at https://0xparc.github.io/pod2/merkletree.html .
 use anyhow::{anyhow, Result};
-use itertools::Itertools;
+use plonky2::field::goldilocks_field::GoldilocksField;
-use plonky2::field::types::Field;
+use plonky2::hash::poseidon::PoseidonHash;
-use plonky2::hash::{
-    hash_types::HashOut,
-    merkle_proofs::{verify_merkle_proof, MerkleProof as PlonkyMerkleProof},
-    merkle_tree::MerkleTree as PlonkyMerkleTree,
-    poseidon::PoseidonHash,
-};
-use plonky2::plonk::config::GenericConfig;
 use plonky2::plonk::config::Hasher;
 use std::collections::HashMap;
+use std::fmt;
 use std::iter::IntoIterator;
 
-use crate::middleware::{Hash, Value, C, D, F};
+use crate::middleware::{Hash, Value, F, NULL};
 
-const CAP_HEIGHT: usize = 0;
+/// Implements the MerkleTree specified at
-
+/// https://0xparc.github.io/pod2/merkletree.html
-/// MerkleTree currently is a wrapper on top of Plonky2's MerkleTree. A future iteration will
-/// replace it by the MerkleTree specified at https://0xparc.github.io/pod2/merkletree.html .
 #[derive(Clone, Debug)]
 pub struct MerkleTree {
-    tree: PlonkyMerkleTree<F, <C as GenericConfig<D>>::Hasher>,
+    max_depth: usize,
-    // keyindex: key -> index mapping. This is just for the current plonky-tree wrapper
+    root: Node,
-    keyindex: HashMap<Value, usize>,
-    // kvs are a field in the MerkleTree in order to be able to iterate over the keyvalues. This is
-    // specific of the current implementation (Plonky2's tree wrapper), in the next iteration this
-    // will not be needed since the tree implementation itself will offer the hashmap
-    // functionality.
-    pub kvs: HashMap<Value, Value>,
-    // leaves_map is a map between the leaf (leaf=Hash(key,value)) and the actual (key, value). It
-    // is used to get the actual value from a leaf for a given key (through the method
-    // `MerkleTree.get`.
-    leaves_map: HashMap<Hash, (Value, Value)>,
-}
-
-pub struct MerkleProof {
-    existence: bool,
-    index: usize,
-    proof: PlonkyMerkleProof<F, <C as GenericConfig<D>>::Hasher>,
 }
 
 impl MerkleTree {
     /// builds a new `MerkleTree` where the leaves contain the given key-values
-    pub fn new(kvs: &HashMap<Value, Value>) -> Self {
+    pub fn new(max_depth: usize, kvs: &HashMap<Value, Value>) -> Result<Self> {
-        let mut keyindex: HashMap<Value, usize> = HashMap::new();
+        let mut root = Node::Intermediate(Intermediate::empty());
-        let mut leaves: Vec<Vec<F>> = Vec::new();
+
-        let mut leaves_map: HashMap<Hash, (Value, Value)> = HashMap::new();
+        for (k, v) in kvs.iter() {
-        // Note: current version iterates sorting by keys of the kvs, but the merkletree defined at
+            let leaf = Leaf::new(max_depth, *k, *v)?;
-        // https://0xparc.github.io/pod2/merkletree.html will not need it since it will be
+            root.add_leaf(0, max_depth, leaf)?;
-        // deterministic based on the keys values not on the order of the keys when added into the
-        // tree.
-        for (i, (k, v)) in kvs.iter().sorted_by_key(|kv| kv.0).enumerate() {
-            let input: Vec<F> = [k.0, v.0].concat();
-            let leaf = PoseidonHash::hash_no_pad(&input).elements;
-            leaves.push(leaf.into());
-            keyindex.insert(*k, i);
-            leaves_map.insert(Hash(leaf), (*k, *v));
         }
 
-        // pad to a power of two if needed
+        let _ = root.compute_hash();
-        let leaf_empty: Vec<F> = vec![F::ZERO, F::ZERO, F::ZERO, F::ZERO];
+        Ok(Self { max_depth, root })
-        for _ in leaves.len()..leaves.len().next_power_of_two() {
-            leaves.push(leaf_empty.clone());
     }
 
-        let tree = PlonkyMerkleTree::<F, <C as GenericConfig<D>>::Hasher>::new(leaves, CAP_HEIGHT);
-        Self {
-            tree,
-            keyindex,
-            kvs: kvs.clone(),
-            leaves_map,
-        }
-    }
-}
-
-impl MerkleTree {
     /// returns the root of the tree
     pub fn root(&self) -> Hash {
-        if self.tree.cap.is_empty() {
+        self.root.hash()
-            return crate::middleware::NULL;
     }
-        Hash(self.tree.cap.0[0].elements)
+
+    pub fn max_depth(&self) -> usize {
+        self.max_depth
     }
 
     /// returns the value at the given key
     pub fn get(&self, key: &Value) -> Result<Value> {
-        let i = self.keyindex.get(&key).ok_or(anyhow!("key not in tree"))?;
+        let path = keypath(self.max_depth, *key)?;
-        let leaf_hash_raw = self.tree.get(*i);
+        let key_resolution = self.root.down(0, self.max_depth, path, None)?;
-        let leaf_hash_f: [F; 4] = leaf_hash_raw
+        match key_resolution {
-            .try_into()
+            Some((k, v)) if &k == key => Ok(v),
-            .map_err(|_| anyhow!("unexpected length (len!=4)"))?;
+            _ => Err(anyhow!("key not found")),
-        let leaf_hash: Hash = Hash(leaf_hash_f);
+        }
-        let (_, value) = self.leaves_map.get(&leaf_hash).unwrap();
-        Ok(*value)
     }
 
     /// returns a boolean indicating whether the key exists in the tree
-    pub fn contains(&self, key: &Value) -> bool {
+    pub fn contains(&self, key: &Value) -> Result<bool> {
-        self.keyindex.get(&key).is_some()
+        let path = keypath(self.max_depth, *key)?;
+        match self.root.down(0, self.max_depth, path, None) {
+            Ok(Some((k, _))) => {
+                if &k == key {
+                    Ok(true)
+                } else {
+                    Ok(false)
+                }
+            }
+            _ => Ok(false),
+        }
     }
 
     /// returns a proof of existence, which proves that the given key exists in
-    /// the tree. It returns the `MerkleProof`.
+    /// the tree. It returns the `value` of the leaf at the given `key`, and the
-    pub fn prove(&self, key: &Value) -> Result<MerkleProof> {
+    /// `MerkleProof`.
-        let i = self.keyindex.get(&key).ok_or(anyhow!("key not in tree"))?;
+    pub fn prove(&self, key: &Value) -> Result<(Value, MerkleProof)> {
-        let proof = self.tree.prove(*i);
+        let path = keypath(self.max_depth, *key)?;
-        Ok(MerkleProof {
+
+        let mut siblings: Vec<Hash> = Vec::new();
+
+        match self
+            .root
+            .down(0, self.max_depth, path, Some(&mut siblings))?
+        {
+            Some((k, v)) if &k == key => Ok((
+                v,
+                MerkleProof {
                     existence: true,
-            index: *i,
+                    siblings,
-            proof,
+                    other_leaf: None,
-        })
+                },
+            )),
+            _ => Err(anyhow!("key not found")),
+        }
     }
 
-    /// returns a proof of non-existence, which proves that the given `key`
+    /// returns a proof of non-existence, which proves that the given
-    /// does not exist in the tree
+    /// `key` does not exist in the tree. The return value specifies
-    pub fn prove_nonexistence(&self, _key: &Value) -> Result<MerkleProof> {
+    /// the key-value pair in the leaf reached as a result of
-        // mock method
+    /// resolving `key` as well as a `MerkleProof`.
-        println!("WARNING: MerkleTree::verify_nonexistence is currently a mock");
+    pub fn prove_nonexistence(&self, key: &Value) -> Result<MerkleProof> {
-        Ok(MerkleProof {
+        let path = keypath(self.max_depth, *key)?;
+
+        let mut siblings: Vec<Hash> = Vec::new();
+
+        // note: non-existence of a key can be in 2 cases:
+        match self
+            .root
+            .down(0, self.max_depth, path, Some(&mut siblings))?
+        {
+            // case i) the expected leaf does not exist
+            None => Ok(MerkleProof {
                 existence: false,
-            index: 0,
+                siblings,
-            proof: PlonkyMerkleProof { siblings: vec![] },
+                other_leaf: None,
-        })
+            }),
+            // case ii) the expected leaf does exist in the tree, but it has a different `key`
+            Some((k, v)) if &k != key => Ok(MerkleProof {
+                existence: false,
+                siblings,
+                other_leaf: Some((k, v)),
+            }),
+            _ => Err(anyhow!("key found")),
+        }
+        // both cases prove that the given key don't exist in the tree. ∎
     }
 
     /// verifies an inclusion proof for the given `key` and `value`
-    pub fn verify(root: Hash, proof: &MerkleProof, key: &Value, value: &Value) -> Result<()> {
+    pub fn verify(
-        if !proof.existence {
+        max_depth: usize,
-            return Err(anyhow!(
+        root: Hash,
-                "expected proof of existence, found proof of non-existence"
+        proof: &MerkleProof,
-            ));
+        key: &Value,
+        value: &Value,
+    ) -> Result<()> {
+        let h = proof.compute_root_from_leaf(max_depth, key, Some(*value))?;
+
+        if h != root {
+            return Err(anyhow!("proof of inclusion does not verify"));
+        } else {
+            Ok(())
         }
-        let leaf = PoseidonHash::hash_no_pad(&[key.0, value.0].concat()).elements;
-        let root = HashOut::from_vec(root.0.to_vec());
-        verify_merkle_proof(leaf.into(), proof.index, root, &proof.proof)
     }
 
     /// verifies a non-inclusion proof for the given `key`, that is, the given
     /// `key` does not exist in the tree
-    pub fn verify_nonexistence(_root: Hash, proof: &MerkleProof, _key: &Value) -> Result<()> {
+    pub fn verify_nonexistence(
-        // mock method
+        max_depth: usize,
-        if proof.existence {
+        root: Hash,
-            return Err(anyhow!(
+        proof: &MerkleProof,
-                "expected proof of non-existence, found proof of existence"
+        key: &Value,
-            ));
+    ) -> Result<()> {
-        }
+        match proof.other_leaf {
-        println!("WARNING: MerkleTree::verify_nonexistence is currently a mock");
+            Some((k, _v)) if &k == key => Err(anyhow!("Invalid non-existence proof.")),
+            _ => {
+                let k = proof.other_leaf.map(|(k, _)| k).unwrap_or(*key);
+                let v: Option<Value> = proof.other_leaf.map(|(_, v)| v);
+                let h = proof.compute_root_from_leaf(max_depth, &k, v)?;
+
+                if h != root {
+                    return Err(anyhow!("proof of exclusion does not verify"));
+                } else {
                     Ok(())
                 }
+            }
+        }
+    }
 
     /// returns an iterator over the leaves of the tree
-    pub fn iter(&self) -> std::collections::hash_map::Iter<Value, Value> {
+    pub fn iter(&self) -> Iter {
-        self.kvs.iter()
+        Iter {
+            state: vec![&self.root],
+        }
     }
 }
 
 impl<'a> IntoIterator for &'a MerkleTree {
     type Item = (&'a Value, &'a Value);
-    type IntoIter = std::collections::hash_map::Iter<'a, Value, Value>;
+    type IntoIter = Iter<'a>;
 
     fn into_iter(self) -> Self::IntoIter {
-        self.kvs.iter()
+        self.iter()
+    }
+}
+
+impl fmt::Display for MerkleTree {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        writeln!(
+            f,
+            "\nPaste in GraphViz (https://dreampuf.github.io/GraphvizOnline/):\n-----"
+        )?;
+        writeln!(f, "digraph hierarchy {{")?;
+        writeln!(f, "node [fontname=Monospace,fontsize=10,shape=box]")?;
+        write!(f, "{}", self.root)?;
+        writeln!(f, "\n}}\n-----")
+    }
+}
+
+#[derive(Clone, Debug)]
+pub struct MerkleProof {
+    // note: currently we don't use the `_existence` field, we would use if we merge the methods
+    // `verify` and `verify_nonexistence` into a single one
+    #[allow(unused)]
+    existence: bool,
+    siblings: Vec<Hash>,
+    // other_leaf is used for non-existence proofs
+    other_leaf: Option<(Value, Value)>,
+}
+
+impl fmt::Display for MerkleProof {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        for (i, s) in self.siblings.iter().enumerate() {
+            if i > 0 {
+                write!(f, ", ")?;
+            }
+            write!(f, "{}", s)?;
+        }
+        Ok(())
+    }
+}
+
+impl MerkleProof {
+    /// Computes the root of the Merkle tree suggested by a Merkle proof given a
+    /// key & value. If a value is not provided, the terminal node is assumed to
+    /// be empty.
+    fn compute_root_from_leaf(
+        &self,
+        max_depth: usize,
+        key: &Value,
+        value: Option<Value>,
+    ) -> Result<Hash> {
+        if self.siblings.len() >= max_depth {
+            return Err(anyhow!("max depth reached"));
+        }
+
+        let path = keypath(max_depth, *key)?;
+        let mut h = value
+            .map(|v| Hash(PoseidonHash::hash_no_pad(&[key.0, v.0].concat()).elements))
+            .unwrap_or(Hash([GoldilocksField(0); 4]));
+        for (i, sibling) in self.siblings.iter().enumerate().rev() {
+            let input: Vec<F> = if path[i] {
+                [sibling.0, h.0].concat()
+            } else {
+                [h.0, sibling.0].concat()
+            };
+            h = Hash(PoseidonHash::hash_no_pad(&input).elements);
+        }
+        Ok(h)
+    }
+}
+
+#[derive(Clone, Debug)]
+enum Node {
+    None,
+    Leaf(Leaf),
+    Intermediate(Intermediate),
+}
+
+impl fmt::Display for Node {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            Self::Intermediate(n) => {
+                writeln!(
+                    f,
+                    "\"{}\" -> {{ \"{}\" \"{}\" }}",
+                    n.hash(),
+                    n.left.hash(),
+                    n.right.hash()
+                )?;
+                write!(f, "{}", n.left)?;
+                write!(f, "{}", n.right)
+            }
+            Self::Leaf(l) => {
+                writeln!(f, "\"{}\" [style=filled]", l.hash())?;
+                writeln!(f, "\"k:{}\\nv:{}\" [style=dashed]", l.key, l.value)?;
+                writeln!(
+                    f,
+                    "\"{}\" -> {{ \"k:{}\\nv:{}\" }}",
+                    l.hash(),
+                    l.key,
+                    l.value,
+                )
+            }
+            Self::None => Ok(()),
+        }
+    }
+}
+
+impl Node {
+    fn is_empty(&self) -> bool {
+        match self {
+            Self::None => true,
+            Self::Leaf(_l) => false,
+            Self::Intermediate(_n) => false,
+        }
+    }
+    fn compute_hash(&mut self) -> Hash {
+        match self {
+            Self::None => NULL,
+            Self::Leaf(l) => l.compute_hash(),
+            Self::Intermediate(n) => n.compute_hash(),
+        }
+    }
+    fn hash(&self) -> Hash {
+        match self {
+            Self::None => NULL,
+            Self::Leaf(l) => l.hash(),
+            Self::Intermediate(n) => n.hash(),
+        }
+    }
+
+    /// Goes down from the current node until it encounters a terminal node,
+    /// viz. a leaf or empty node, or until it reaches the maximum depth. The
+    /// `siblings` parameter is used to store the siblings while going down to
+    /// the leaf, if the given parameter is set to `None`, then no siblings are
+    /// stored. In this way, the same method `down` can be used by MerkleTree
+    /// methods `get`, `contains`, `prove` and `prove_nonexistence`.
+    ///
+    /// Be aware that this method will return the found leaf at the given path,
+    /// which may contain a different key and value than the expected one.
+    fn down(
+        &self,
+        lvl: usize,
+        max_depth: usize,
+        path: Vec<bool>,
+        mut siblings: Option<&mut Vec<Hash>>,
+    ) -> Result<Option<(Value, Value)>> {
+        if lvl >= max_depth {
+            return Err(anyhow!("max depth reached"));
+        }
+
+        match self {
+            Self::Intermediate(n) => {
+                if path[lvl] {
+                    if let Some(s) = siblings.as_mut() {
+                        s.push(n.left.hash());
+                    }
+                    return n.right.down(lvl + 1, max_depth, path, siblings);
+                } else {
+                    if let Some(s) = siblings.as_mut() {
+                        s.push(n.right.hash());
+                    }
+                    return n.left.down(lvl + 1, max_depth, path, siblings);
+                }
+            }
+            Self::Leaf(Leaf {
+                key,
+                value,
+                path: _p,
+                hash: _h,
+            }) => Ok(Some((key.clone(), value.clone()))),
+            _ => Ok(None),
+        }
+    }
+
+    // adds the leaf at the tree from the current node (self), without computing any hash
+    fn add_leaf(&mut self, lvl: usize, max_depth: usize, leaf: Leaf) -> Result<()> {
+        if lvl >= max_depth {
+            return Err(anyhow!("max depth reached"));
+        }
+
+        match self {
+            Self::Intermediate(n) => {
+                if leaf.path[lvl] {
+                    if n.right.is_empty() {
+                        // empty sub-node, add the leaf here
+                        n.right = Box::new(Node::Leaf(leaf));
+                        return Ok(());
+                    }
+                    n.right.add_leaf(lvl + 1, max_depth, leaf)?;
+                } else {
+                    if n.left.is_empty() {
+                        // empty sub-node, add the leaf here
+                        n.left = Box::new(Node::Leaf(leaf));
+                        return Ok(());
+                    }
+                    n.left.add_leaf(lvl + 1, max_depth, leaf)?;
+                }
+            }
+            Self::Leaf(l) => {
+                // in this case, it means that we found a leaf in the new-leaf
+                // path, thus we need to push both leaves (old-leaf and
+                // new-leaf) down the path till their paths diverge.
+
+                // first check that keys of both leaves are different
+                // (l=old-leaf, leaf=new-leaf)
+                if l.key == leaf.key {
+                    // Note: current approach returns an error when trying to
+                    // add to a leaf where the key already exists. We could also
+                    // ignore it if needed.
+                    return Err(anyhow!("key already exists"));
+                }
+                let old_leaf = l.clone();
+                // set self as an intermediate node
+                *self = Node::Intermediate(Intermediate::empty());
+                return self.down_till_divergence(lvl, max_depth, old_leaf, leaf);
+            }
+            Self::None => {
+                return Err(anyhow!("reached empty node, should not have entered"));
+            }
+        }
+        Ok(())
+    }
+
+    /// goes down through a 'virtual' path till finding a divergence. This
+    /// method is used for when adding a new leaf another already existing leaf
+    /// is found, so that both leaves (new and old) are pushed down the path
+    /// till their keys diverge.
+    fn down_till_divergence(
+        &mut self,
+        lvl: usize,
+        max_depth: usize,
+        old_leaf: Leaf,
+        new_leaf: Leaf,
+    ) -> Result<()> {
+        if lvl >= max_depth {
+            return Err(anyhow!("max depth reached"));
+        }
+
+        if let Node::Intermediate(ref mut n) = self {
+            if old_leaf.path[lvl] != new_leaf.path[lvl] {
+                // reached divergence in next level, set the leaves as children
+                // at the current node
+                if new_leaf.path[lvl] {
+                    n.left = Box::new(Node::Leaf(old_leaf));
+                    n.right = Box::new(Node::Leaf(new_leaf));
+                } else {
+                    n.left = Box::new(Node::Leaf(new_leaf));
+                    n.right = Box::new(Node::Leaf(old_leaf));
+                }
+                return Ok(());
+            }
+
+            // no divergence yet, continue going down
+            if new_leaf.path[lvl] {
+                n.right = Box::new(Node::Intermediate(Intermediate::empty()));
+                return n
+                    .right
+                    .down_till_divergence(lvl + 1, max_depth, old_leaf, new_leaf);
+            } else {
+                n.left = Box::new(Node::Intermediate(Intermediate::empty()));
+                return n
+                    .left
+                    .down_till_divergence(lvl + 1, max_depth, old_leaf, new_leaf);
+            }
+        }
+        Ok(())
+    }
+}
+
+#[derive(Clone, Debug)]
+struct Intermediate {
+    hash: Option<Hash>,
+    left: Box<Node>,
+    right: Box<Node>,
+}
+impl Intermediate {
+    fn empty() -> Self {
+        Self {
+            hash: None,
+            left: Box::new(Node::None),
+            right: Box::new(Node::None),
+        }
+    }
+    fn compute_hash(&mut self) -> Hash {
+        if self.left.clone().is_empty() && self.right.clone().is_empty() {
+            self.hash = Some(NULL);
+            return NULL;
+        }
+        let l_hash = self.left.compute_hash();
+        let r_hash = self.right.compute_hash();
+        let input: Vec<F> = [l_hash.0, r_hash.0].concat();
+        let h = Hash(PoseidonHash::hash_no_pad(&input).elements);
+        self.hash = Some(h);
+        h
+    }
+    fn hash(&self) -> Hash {
+        self.hash.unwrap()
+    }
+}
+
+#[derive(Clone, Debug)]
+struct Leaf {
+    hash: Option<Hash>,
+    path: Vec<bool>,
+    key: Value,
+    value: Value,
+}
+impl Leaf {
+    fn new(max_depth: usize, key: Value, value: Value) -> Result<Self> {
+        Ok(Self {
+            hash: None,
+            path: keypath(max_depth, key)?,
+            key,
+            value,
+        })
+    }
+    fn compute_hash(&mut self) -> Hash {
+        let input: Vec<F> = [self.key.0, self.value.0].concat();
+        let h = Hash(PoseidonHash::hash_no_pad(&input).elements);
+        self.hash = Some(h);
+        h
+    }
+    fn hash(&self) -> Hash {
+        self.hash.unwrap()
+    }
+}
+
+// NOTE 1: think if maybe the length of the returned vector can be <256
+// (8*bytes.len()), so that we can do fewer iterations. For example, if the
+// tree.max_depth is set to 20, we just need 20 iterations of the loop, not 256.
+// NOTE 2: which approach do we take with keys that are longer than the
+// max-depth? ie, what happens when two keys share the same path for more bits
+// than the max_depth?
+/// returns the path of the given key
+fn keypath(max_depth: usize, k: Value) -> Result<Vec<bool>> {
+    let bytes = k.to_bytes();
+    if max_depth > 8 * bytes.len() {
+        // note that our current keys are of Value type, which are 4 Goldilocks
+        // field elements, ie ~256 bits, therefore the max_depth can not be
+        // bigger than 256.
+        return Err(anyhow!(
+            "key to short (key length: {}) for the max_depth: {}",
+            8 * bytes.len(),
+            max_depth
+        ));
+    }
+    Ok((0..max_depth)
+        .map(|n| bytes[n / 8] & (1 << (n % 8)) != 0)
+        .collect())
+}
+
+pub struct Iter<'a> {
+    state: Vec<&'a Node>,
+}
+
+impl<'a> Iterator for Iter<'a> {
+    type Item = (&'a Value, &'a Value);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let node = self.state.pop();
+        match node {
+            Some(Node::None) => self.next(),
+            Some(Node::Leaf(Leaf {
+                hash: _,
+                path: _,
+                key,
+                value,
+            })) => Some((key, value)),
+            Some(Node::Intermediate(Intermediate {
+                hash: _,
+                left,
+                right,
+            })) => {
+                self.state.push(&right);
+                self.state.push(&left);
+                self.next()
+            }
+            _ => None,
+        }
     }
 }
 
 #[cfg(test)]
 pub mod tests {
-    use super::*;
+    use itertools::Itertools;
+    use std::cmp::Ordering;
 
-    use crate::middleware::hash_str;
+    use super::*;
 
     #[test]
     fn test_merkletree() -> Result<()> {
-        let (k0, v0) = (
-            Value(hash_str("key_0".into()).0),
-            Value(hash_str("value_0".into()).0),
-        );
-        let (k1, v1) = (
-            Value(hash_str("key_1".into()).0),
-            Value(hash_str("value_1".into()).0),
-        );
-        let (k2, v2) = (
-            Value(hash_str("key_2".into()).0),
-            Value(hash_str("value_2".into()).0),
-        );
-
         let mut kvs = HashMap::new();
-        kvs.insert(k0, v0);
+        for i in 0..8 {
-        kvs.insert(k1, v1);
+            if i == 1 {
-        kvs.insert(k2, v2);
+                continue;
+            }
+            kvs.insert(Value::from(i), Value::from(1000 + i));
+        }
+        let key = Value::from(13);
+        let value = Value::from(1013);
+        kvs.insert(key, value);
 
-        let tree = MerkleTree::new(&kvs);
+        let tree = MerkleTree::new(32, &kvs)?;
+        // when printing the tree, it should print the same tree as in
+        // https://0xparc.github.io/pod2/merkletree.html#example-2
+        println!("{}", tree);
 
-        let proof = tree.prove(&k2)?;
+        // Inclusion checks
-        MerkleTree::verify(tree.root(), &proof, &k2, &v2)?;
+        let (v, proof) = tree.prove(&Value::from(13))?;
+        assert_eq!(v, Value::from(1013));
+        println!("{}", proof);
 
-        // expect verification to fail with different key / value
+        MerkleTree::verify(32, tree.root(), &proof, &key, &value)?;
-        assert!(MerkleTree::verify(tree.root(), &proof, &k2, &v0).is_err());
-        assert!(MerkleTree::verify(tree.root(), &proof, &k0, &v2).is_err());
 
-        // non-existence proofs
+        // Exclusion checks
-        let proof_ne = tree.prove_nonexistence(&k2)?;
+        let key = Value::from(12);
-        let _ = MerkleTree::verify_nonexistence(tree.root(), &proof_ne, &k2)?;
+        let proof = tree.prove_nonexistence(&key)?;
+        assert_eq!(
+            proof.other_leaf.unwrap(),
+            (Value::from(4), Value::from(1004))
+        );
+        println!("{}", proof);
 
-        // expect verification of existence fail for nonexistence proof
+        MerkleTree::verify_nonexistence(32, tree.root(), &proof, &key)?;
-        let _ = MerkleTree::verify(tree.root(), &proof_ne, &k2, &v2).is_err();
+
+        let key = Value::from(1);
+        let proof = tree.prove_nonexistence(&Value::from(1))?;
+        assert_eq!(proof.other_leaf, None);
+        println!("{}", proof);
+
+        MerkleTree::verify_nonexistence(32, tree.root(), &proof, &key)?;
+
+        // Check iterator
+        let collected_kvs: Vec<_> = tree.into_iter().collect::<Vec<_>>();
+
+        // Expected key ordering
+        let cmp = |max_depth: usize| {
+            move |k1, k2| {
+                let path1 = keypath(max_depth, k1).unwrap();
+                let path2 = keypath(max_depth, k2).unwrap();
+
+                let first_unequal_bits = std::iter::zip(path1, path2).find(|(b1, b2)| b1 != b2);
+
+                match first_unequal_bits {
+                    Some((b1, b2)) => {
+                        if b1 < b2 {
+                            Ordering::Less
+                        } else {
+                            Ordering::Greater
+                        }
+                    }
+                    _ => Ordering::Equal,
+                }
+            }
+        };
+
+        let sorted_kvs = kvs
+            .iter()
+            .sorted_by(|(k1, _), (k2, _)| cmp(32)(**k1, **k2))
+            .collect::<Vec<_>>();
+
+        assert_eq!(collected_kvs, sorted_kvs);
 
         Ok(())
     }