add initial MerkleTree implementation (#13)

Add initial MerkleTree implementation, which is a wrapper on top of
Plonky2's MerkleTree, with the idea that the future iteration will
replace it by the MerkleTree specified at
https://0xparc.github.io/pod2/merkletree.html .

src/lib.rs

@@ -2,13 +2,19 @@ use hex::{FromHex, FromHexError};
 use plonky2::field::goldilocks_field::GoldilocksField;
 use plonky2::field::types::{Field, PrimeField64};
 use plonky2::hash::poseidon::PoseidonHash;
-use plonky2::plonk::config::Hasher;
+use plonky2::plonk::config::{Hasher, PoseidonGoldilocksConfig};
+use std::cmp::Ordering;
 use std::fmt;
 
 pub mod backend;
 pub mod frontend;
+pub mod merkletree;
 
 pub type F = GoldilocksField;
+/// C is the Plonky2 config used in POD2 to work with Plonky2 recursion.
+pub type C = PoseidonGoldilocksConfig;
+/// D defines the extension degree of the field used in the Plonky2 proofs (quadratic extension).
+pub const D: usize = 2;
 
 #[derive(Clone, Copy, Debug, Default, Hash, PartialEq, Eq)]
 pub struct Hash(pub [F; 4]);
@@ -40,6 +46,17 @@ impl FromHex for Hash {
     }
 }
 
+impl Ord for Hash {
+    fn cmp(&self, other: &Self) -> Ordering {
+        self.to_string().cmp(&other.to_string())
+    }
+}
+impl PartialOrd for Hash {
+    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
 #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Default)]
 pub struct PodId(pub Hash);
 pub const SELF: PodId = PodId(Hash([F::ONE, F::ZERO, F::ZERO, F::ZERO]));

src/merkletree.rs

@@ -0,0 +1,163 @@
+/// 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 .
+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::plonk::config::Hasher;
+use std::collections::HashMap;
+
+use crate::backend::Value;
+use crate::{Hash, C, D, F};
+
+const CAP_HEIGHT: usize = 0;
+
+/// MerkleTree currently implements the MerkleTree interface with 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>,
+    // keyindex: key -> index mapping. This is just for the current plonky-tree wrapper
+    keyindex: HashMap<Hash, usize>,
+}
+
+pub struct MerkleProof {
+    existence: bool,
+    index: usize,
+    proof: PlonkyMerkleProof<F, <C as GenericConfig<D>>::Hasher>,
+}
+
+impl MerkleTree {
+    pub fn new(kvs: HashMap<Hash, Value>) -> Self {
+        let mut keyindex: HashMap<Hash, usize> = HashMap::new();
+        let mut leaves: Vec<Vec<F>> = Vec::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.into_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);
+        }
+
+        // pad to a power of two if needed
+        let leaf_empty: Vec<F> = 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::<F, <C as GenericConfig<D>>::Hasher>::new(leaves, CAP_HEIGHT);
+        Self { tree, keyindex }
+    }
+
+    pub fn root(&self) -> Result<Hash> {
+        if self.tree.cap.is_empty() {
+            return Err(anyhow!("empty tree"));
+        }
+        Ok(Hash(self.tree.cap.0[0].elements))
+    }
+
+    pub fn prove(&self, key: &Hash) -> Result<MerkleProof> {
+        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,
+        })
+    }
+
+    pub fn prove_nonexistence(&self, _key: &Hash) -> Result<MerkleProof> {
+        // mock method
+        println!("WARNING: MerkleTree::verify_nonexistence is currently a mock");
+        Ok(MerkleProof {
+            existence: false,
+            index: 0,
+            proof: PlonkyMerkleProof { siblings: vec![] },
+        })
+    }
+
+    pub fn verify(root: Hash, proof: &MerkleProof, key: &Hash, value: &Value) -> Result<()> {
+        if !proof.existence {
+            return Err(anyhow!(
+                "expected proof of existence, found proof of non-existence"
+            ));
+        }
+        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)
+    }
+
+    pub fn verify_nonexistence(
+        _root: Hash,
+        proof: &MerkleProof,
+        _key: &Hash,
+        _value: &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(())
+    }
+}
+
+#[cfg(test)]
+pub mod tests {
+    use super::*;
+
+    use crate::hash_str;
+
+    #[test]
+    fn test_merkletree() -> Result<()> {
+        let (k0, v0) = (
+            hash_str("key_0".into()),
+            Value(hash_str("value_0".into()).0),
+        );
+        let (k1, v1) = (
+            hash_str("key_1".into()),
+            Value(hash_str("value_1".into()).0),
+        );
+        let (k2, v2) = (
+            hash_str("key_2".into()),
+            Value(hash_str("value_2".into()).0),
+        );
+
+        let mut kvs = HashMap::new();
+        kvs.insert(k0, v0);
+        kvs.insert(k1, v1);
+        kvs.insert(k2, v2);
+
+        let tree = MerkleTree::new(kvs);
+
+        let proof = tree.prove(&k2)?;
+        MerkleTree::verify(tree.root()?, &proof, &k2, &v2)?;
+
+        // 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());
+
+        // non-existence proofs
+        let proof_ne = tree.prove_nonexistence(&k2)?;
+        let _ = MerkleTree::verify_nonexistence(tree.root()?, &proof_ne, &k2, &v2)?;
+
+        // expect verification of existence fail for nonexistence proof
+        let _ = MerkleTree::verify(tree.root()?, &proof_ne, &k2, &v2).is_err();
+
+        Ok(())
+    }
+}