diff --git a/src/backends/mock_main/mod.rs b/src/backends/mock_main/mod.rs index cbb0c7dc..4db4359b 100644 --- a/src/backends/mock_main/mod.rs +++ b/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(¶ms, &gov_id_pod, &pay_stub_pod); + let kyc_builder = zu_kyc_pod_builder(¶ms, &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::().unwrap(); diff --git a/src/backends/mock_signed.rs b/src/backends/mock_signed.rs index afbaf760..51ce54e6 100644 --- a/src/backends/mock_signed.rs +++ b/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::>(), - ); + ) { + 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 plonky2::field::types::Field; + use std::iter; + use super::*; + use crate::constants::MAX_DEPTH; use crate::frontend; use crate::middleware::{self, F, NULL}; - use plonky2::field::types::Field; - use std::iter; #[test] - fn test_mock_signed_0() { + fn test_mock_signed_0() -> Result<()> { let params = middleware::Params::default(); let mut pod = frontend::SignedPodBuilder::new(¶ms); 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::>(); - 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::>(); - 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(()) } } diff --git a/src/constants.rs b/src/constants.rs new file mode 100644 index 00000000..1b5be372 --- /dev/null +++ b/src/constants.rs @@ -0,0 +1 @@ +pub const MAX_DEPTH: usize = 32; diff --git a/src/examples.rs b/src/examples.rs index b617aaea..818807a0 100644 --- a/src/examples.rs +++ b/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 { - let sanction_list = Value::Dictionary(Dictionary::new(&HashMap::new())); // empty dictionary +) -> Result { + 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 { 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 - great_boy_pod_builder( + let good_boy_issuers_dict = Value::Dictionary(Dictionary::new(&HashMap::new())?); // empty + Ok(great_boy_pod_builder( ¶ms, [ &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 { let params = Params::default(); let builder = tickets_sign_pod_builder(¶ms); let signed_pod = builder.sign(&mut MockSigner { pk: "test".into() }).unwrap(); - tickets_pod_builder( + Ok(tickets_pod_builder( ¶ms, &signed_pod, 123, true, - &Value::Dictionary(Dictionary::new(&HashMap::new())), - ) + &Value::Dictionary(Dictionary::new(&HashMap::new())?), + )) } diff --git a/src/frontend/mod.rs b/src/frontend/mod.rs index 085b093c..f1335cbe 100644 --- a/src/frontend/mod.rs +++ b/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(¶ms, &gov_id, &pay_stub); + let kyc = zu_kyc_pod_builder(¶ms, &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(()) diff --git a/src/lib.rs b/src/lib.rs index 95245d9f..963de0b4 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,4 +1,5 @@ pub mod backends; +pub mod constants; pub mod frontend; pub mod middleware; pub mod primitives; diff --git a/src/middleware/containers.rs b/src/middleware/containers.rs index 43204b43..743402db 100644 --- a/src/middleware/containers.rs +++ b/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) -> Self { + pub fn new(kvs: &HashMap) -> Result { let kvs: HashMap = kvs.into_iter().map(|(&k, &v)| (Value(k.0), v)).collect(); - Self { - mt: MerkleTree::new(&kvs), - } + Ok(Self { + 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 { self.mt.get(key) } - pub fn prove(&self, key: &Value) -> Result { + pub fn prove(&self, key: &Value) -> Result<(Value, MerkleProof)> { self.mt.prove(key) } pub fn prove_nonexistence(&self, key: &Value) -> Result { 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 { + 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) -> Self { + pub fn new(set: &Vec) -> Result { let kvs: HashMap = set .into_iter() .map(|e| { @@ -79,29 +80,30 @@ impl Set { (Value(h), EMPTY) }) .collect(); - Self { - mt: MerkleTree::new(&kvs), - } + Ok(Self { + 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 { self.mt.contains(value) } pub fn prove(&self, value: &Value) -> Result { - self.mt.prove(value) + let (_, proof) = self.mt.prove(value)?; + Ok(proof) } pub fn prove_nonexistence(&self, value: &Value) -> Result { 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 { + 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) -> Self { + pub fn new(array: &Vec) -> Result { let kvs: HashMap = array .into_iter() .enumerate() .map(|(i, &e)| (Value::from(i as i64), e)) .collect(); - Self { - mt: MerkleTree::new(&kvs), - } + Ok(Self { + 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 { self.mt.get(&Value::from(i as i64)) } - pub fn prove(&self, i: usize) -> Result { + 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 { + pub fn iter(&self) -> crate::primitives::merkletree::Iter { self.mt.iter() } } diff --git a/src/middleware/mod.rs b/src/middleware/mod.rs index 14cd9f29..34cbbe46 100644 --- a/src/middleware/mod.rs +++ b/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 { + 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() { diff --git a/src/primitives/merkletree.rs b/src/primitives/merkletree.rs index 8a393af6..9a214478 100644 --- a/src/primitives/merkletree.rs +++ b/src/primitives/merkletree.rs @@ -1,214 +1,642 @@ -/// MerkleTree implementation for POD2. -/// -/// 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 . +//! Module that implements the MerkleTree specified at +//! https://0xparc.github.io/pod2/merkletree.html . use anyhow::{anyhow, Result}; -use itertools::Itertools; -use plonky2::field::types::Field; -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::field::goldilocks_field::GoldilocksField; +use plonky2::hash::poseidon::PoseidonHash; 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; - -/// 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 . +/// Implements the MerkleTree specified at +/// https://0xparc.github.io/pod2/merkletree.html #[derive(Clone, Debug)] pub struct MerkleTree { - tree: PlonkyMerkleTree>::Hasher>, - // keyindex: key -> index mapping. This is just for the current plonky-tree wrapper - keyindex: HashMap, - // 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, - // 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, -} - -pub struct MerkleProof { - existence: bool, - index: usize, - proof: PlonkyMerkleProof>::Hasher>, + max_depth: usize, + root: Node, } impl MerkleTree { /// builds a new `MerkleTree` where the leaves contain the given key-values - pub fn new(kvs: &HashMap) -> Self { - let mut keyindex: HashMap = HashMap::new(); - let mut leaves: Vec> = Vec::new(); - let mut leaves_map: HashMap = HashMap::new(); - // Note: current version iterates sorting by keys of the kvs, but the merkletree defined at - // https://0xparc.github.io/pod2/merkletree.html will not need it since it will be - // 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 = [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 leaf_empty: Vec = vec![F::ZERO, F::ZERO, F::ZERO, F::ZERO]; - for _ in leaves.len()..leaves.len().next_power_of_two() { - leaves.push(leaf_empty.clone()); - } - - let tree = PlonkyMerkleTree::>::Hasher>::new(leaves, CAP_HEIGHT); - Self { - tree, - keyindex, - kvs: kvs.clone(), - leaves_map, + pub fn new(max_depth: usize, kvs: &HashMap) -> Result { + let mut root = Node::Intermediate(Intermediate::empty()); + + for (k, v) in kvs.iter() { + let leaf = Leaf::new(max_depth, *k, *v)?; + root.add_leaf(0, max_depth, leaf)?; } + + let _ = root.compute_hash(); + Ok(Self { max_depth, root }) } -} -impl MerkleTree { /// returns the root of the tree pub fn root(&self) -> Hash { - if self.tree.cap.is_empty() { - return crate::middleware::NULL; - } - Hash(self.tree.cap.0[0].elements) + self.root.hash() + } + + pub fn max_depth(&self) -> usize { + self.max_depth } /// returns the value at the given key pub fn get(&self, key: &Value) -> Result { - let i = self.keyindex.get(&key).ok_or(anyhow!("key not in tree"))?; - let leaf_hash_raw = self.tree.get(*i); - let leaf_hash_f: [F; 4] = leaf_hash_raw - .try_into() - .map_err(|_| anyhow!("unexpected length (len!=4)"))?; - let leaf_hash: Hash = Hash(leaf_hash_f); - let (_, value) = self.leaves_map.get(&leaf_hash).unwrap(); - Ok(*value) + let path = keypath(self.max_depth, *key)?; + let key_resolution = self.root.down(0, self.max_depth, path, None)?; + match key_resolution { + Some((k, v)) if &k == key => Ok(v), + _ => Err(anyhow!("key not found")), + } } /// returns a boolean indicating whether the key exists in the tree - pub fn contains(&self, key: &Value) -> bool { - self.keyindex.get(&key).is_some() + pub fn contains(&self, key: &Value) -> Result { + 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`. - pub fn prove(&self, key: &Value) -> Result { - let i = self.keyindex.get(&key).ok_or(anyhow!("key not in tree"))?; - let proof = self.tree.prove(*i); - Ok(MerkleProof { - existence: true, - index: *i, - proof, - }) + /// the tree. It returns the `value` of the leaf at the given `key`, and the + /// `MerkleProof`. + pub fn prove(&self, key: &Value) -> Result<(Value, MerkleProof)> { + let path = keypath(self.max_depth, *key)?; + + let mut siblings: Vec = 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, + siblings, + other_leaf: None, + }, + )), + _ => Err(anyhow!("key not found")), + } } - /// returns a proof of non-existence, which proves that the given `key` - /// does not exist in the tree - pub fn prove_nonexistence(&self, _key: &Value) -> Result { - // mock method - println!("WARNING: MerkleTree::verify_nonexistence is currently a mock"); - Ok(MerkleProof { - existence: false, - index: 0, - proof: PlonkyMerkleProof { siblings: vec![] }, - }) + /// returns a proof of non-existence, which proves that the given + /// `key` does not exist in the tree. The return value specifies + /// the key-value pair in the leaf reached as a result of + /// resolving `key` as well as a `MerkleProof`. + pub fn prove_nonexistence(&self, key: &Value) -> Result { + let path = keypath(self.max_depth, *key)?; + + let mut siblings: Vec = 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, + siblings, + 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<()> { - if !proof.existence { - return Err(anyhow!( - "expected proof of existence, found proof of non-existence" - )); + pub fn verify( + max_depth: usize, + root: Hash, + 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<()> { - // mock method - if proof.existence { - return Err(anyhow!( - "expected proof of non-existence, found proof of existence" - )); - } - println!("WARNING: MerkleTree::verify_nonexistence is currently a mock"); - Ok(()) + pub fn verify_nonexistence( + max_depth: usize, + root: Hash, + proof: &MerkleProof, + key: &Value, + ) -> Result<()> { + match proof.other_leaf { + 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 = 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 { - self.kvs.iter() + pub fn iter(&self) -> 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, + // 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, + ) -> Result { + 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 = 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, + mut siblings: Option<&mut Vec>, + ) -> Result> { + 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, + left: Box, + right: Box, +} +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 = [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, + path: Vec, + key: Value, + value: Value, +} +impl Leaf { + fn new(max_depth: usize, key: Value, value: Value) -> Result { + Ok(Self { + hash: None, + path: keypath(max_depth, key)?, + key, + value, + }) + } + fn compute_hash(&mut self) -> Hash { + let input: Vec = [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> { + 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 { + 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(); + for i in 0..8 { + if i == 1 { + 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(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); + + // Inclusion checks + let (v, proof) = tree.prove(&Value::from(13))?; + assert_eq!(v, Value::from(1013)); + println!("{}", proof); + + MerkleTree::verify(32, tree.root(), &proof, &key, &value)?; + + // Exclusion checks + let key = Value::from(12); + let proof = tree.prove_nonexistence(&key)?; + assert_eq!( + proof.other_leaf.unwrap(), + (Value::from(4), Value::from(1004)) ); + println!("{}", proof); - let mut kvs = HashMap::new(); - kvs.insert(k0, v0); - kvs.insert(k1, v1); - kvs.insert(k2, v2); + MerkleTree::verify_nonexistence(32, tree.root(), &proof, &key)?; + + 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::>(); - let tree = MerkleTree::new(&kvs); + // 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 proof = tree.prove(&k2)?; - MerkleTree::verify(tree.root(), &proof, &k2, &v2)?; + let first_unequal_bits = std::iter::zip(path1, path2).find(|(b1, b2)| b1 != b2); - // expect verification to fail with different key / value - assert!(MerkleTree::verify(tree.root(), &proof, &k2, &v0).is_err()); - assert!(MerkleTree::verify(tree.root(), &proof, &k0, &v2).is_err()); + match first_unequal_bits { + Some((b1, b2)) => { + if b1 < b2 { + Ordering::Less + } else { + Ordering::Greater + } + } + _ => Ordering::Equal, + } + } + }; - // non-existence proofs - let proof_ne = tree.prove_nonexistence(&k2)?; - let _ = MerkleTree::verify_nonexistence(tree.root(), &proof_ne, &k2)?; + let sorted_kvs = kvs + .iter() + .sorted_by(|(k1, _), (k2, _)| cmp(32)(**k1, **k2)) + .collect::>(); - // expect verification of existence fail for nonexistence proof - let _ = MerkleTree::verify(tree.root(), &proof_ne, &k2, &v2).is_err(); + assert_eq!(collected_kvs, sorted_kvs); Ok(()) }