Merkle tree API improvements (#162)

* Expose index and leaf element of Merkle proof.

Since MerkleNode is a private type, users cannot pattern match to
extract a MerkleNode::Leaf, so this is the only way to get the leaf
elemenet of a Merkle proof. This is also, therefore, the only way
to verify that a Merkle proof proves inclusion of the desired element,
since `verify` doesn't check this. (Technically you can also use
`remember`, but that includes side-effects that may be undesirable.)

* Make serde usable for MerkleTree types

The Merkle trees are ususally instantiated with cryptographic types,
e.g. field elements, which do not implement serde traits. This commit
changes the bound on the derived serde impls for Merkle tree types to
require ark_serialize impls on the element types and uses #[serde(with)]
to serialize these elements using ark_serialize.

We also add an instantiation of the Rescue hash scheme for field elements
as the _index_ type, since the BigUint instantiation just coverts to field
elements anyways, but BigUint does not implement serde _or_ ark_serialize,
while field elements do.

* Rename 'field_elem' serializer to 'canonical', since it's more general

primitives/Cargo.toml

@@ -50,6 +50,7 @@ ark-bw6-761 = { git = "https://github.com/arkworks-rs/curves", rev = "677b4ae751
 ark-ed-on-bls12-377 = { git = "https://github.com/arkworks-rs/curves", rev = "677b4ae751a274037880ede86e9b6f30f62635af" }
 ark-ed-on-bls12-381-bandersnatch = { git = "https://github.com/arkworks-rs/curves", rev = "677b4ae751a274037880ede86e9b6f30f62635af" }
 ark-ed-on-bn254 = "0.3.0"
+bincode = "1.3"
 criterion = "0.4.0"
 hashbrown = "0.13.1"
 rand_core = { version = "^0.6.0", features = ["getrandom"] }

primitives/src/merkle_tree/append_only.rs

@@ -210,4 +210,30 @@ mod mt_tests {
         assert!(mt.remember(0, elem, &proof).is_ok());
         assert!(mt.lookup(0).expect_ok().is_ok());
     }
+
+    #[test]
+    fn test_mt_serde() {
+        test_mt_serde_helper::<Fq254>();
+        test_mt_serde_helper::<Fq377>();
+        test_mt_serde_helper::<Fq381>();
+    }
+
+    fn test_mt_serde_helper<F: RescueParameter>() {
+        let mt = RescueMerkleTree::<F>::from_elems(2, &[F::from(3u64), F::from(1u64)]).unwrap();
+        let proof = mt.lookup(0).expect_ok().unwrap().1;
+        let node = &proof.proof[0];
+
+        assert_eq!(
+            mt,
+            bincode::deserialize(&bincode::serialize(&mt).unwrap()).unwrap()
+        );
+        assert_eq!(
+            proof,
+            bincode::deserialize(&bincode::serialize(&proof).unwrap()).unwrap()
+        );
+        assert_eq!(
+            *node,
+            bincode::deserialize(&bincode::serialize(node).unwrap()).unwrap()
+        );
+    }
 }

primitives/src/merkle_tree/internal.rs

@@ -16,24 +16,32 @@ use ark_std::{
     vec::Vec,
 };
 use itertools::Itertools;
+use jf_utils::canonical;
 use num_bigint::BigUint;
 use serde::{Deserialize, Serialize};
 use tagged_base64::tagged;
 use typenum::Unsigned;
 
 #[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
+#[serde(bound = "E: CanonicalSerialize + CanonicalDeserialize,
+                 I: CanonicalSerialize + CanonicalDeserialize,")]
 pub enum MerkleNode<E: Element, I: Index, T: NodeValue> {
     Empty,
     Branch {
+        #[serde(with = "canonical")]
         value: T,
         children: Vec<Box<MerkleNode<E, I, T>>>,
     },
     Leaf {
+        #[serde(with = "canonical")]
         value: T,
+        #[serde(with = "canonical")]
         pos: I,
+        #[serde(with = "canonical")]
         elem: E,
     },
     ForgettenSubtree {
+        #[serde(with = "canonical")]
         value: T,
     },
 }
@@ -100,6 +108,8 @@ impl<T: NodeValue> MerkleCommitment<T> for MerkleTreeCommitment<T> {
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
+#[serde(bound = "E: CanonicalSerialize + CanonicalDeserialize,
+             I: CanonicalSerialize + CanonicalDeserialize,")]
 pub struct MerkleProof<E, I, T, Arity>
 where
     E: Element,
@@ -108,6 +118,7 @@ where
     Arity: Unsigned,
 {
     /// Proof of inclusion for element at index `pos`
+    #[serde(with = "canonical")]
     pub pos: I,
     /// Nodes of proof path, from root to leaf
     pub proof: Vec<MerkleNode<E, I, T>>,
@@ -134,6 +145,17 @@ where
             _phantom_arity: PhantomData,
         }
     }
+
+    pub fn index(&self) -> &I {
+        &self.pos
+    }
+
+    pub fn elem(&self) -> Option<&E> {
+        match self.proof.last() {
+            Some(MerkleNode::Leaf { elem, .. }) => Some(elem),
+            _ => None,
+        }
+    }
 }
 
 #[allow(clippy::type_complexity)]

primitives/src/merkle_tree/macros.rs

@@ -12,6 +12,10 @@ macro_rules! impl_merkle_tree_scheme {
     ($name: ident) => {
         /// A standard append only Merkle tree implementation
         #[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
+        #[serde(
+            bound = "E: ark_serialize::CanonicalSerialize + ark_serialize::CanonicalDeserialize,
+                     I: ark_serialize::CanonicalSerialize + ark_serialize::CanonicalDeserialize,"
+        )]
         pub struct $name<E, H, I, Arity, T>
         where
             E: Element,

primitives/src/merkle_tree/mod.rs

@@ -133,6 +133,9 @@ impl_to_traversal_path_primitives!(u32);
 impl_to_traversal_path_primitives!(u64);
 impl_to_traversal_path_biguint!(u128);
 impl_to_traversal_path_biguint!(BigUint);
+impl_to_traversal_path_biguint!(ark_ed_on_bn254::Fq);
+impl_to_traversal_path_biguint!(ark_ed_on_bls12_377::Fq);
+impl_to_traversal_path_biguint!(ark_ed_on_bls12_381::Fq);
 
 /// Trait for a succint merkle tree commitment
 pub trait MerkleCommitment<T: NodeValue>:

primitives/src/merkle_tree/prelude.rs

@@ -22,6 +22,7 @@ use num_bigint::BigUint;
 use typenum::U3;
 
 /// Wrapper for rescue hash function
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
 pub struct RescueHash<F: RescueParameter> {
     phantom_f: PhantomData<F>,
 }
@@ -51,5 +52,16 @@ impl<F: RescueParameter> DigestAlgorithm<F, BigUint, F> for RescueHash<F> {
     }
 }
 
-/// Example instantiation of a SparseMerkleTree indexed by BigUInt
-pub type RescueSparseMerkleTree<E, F> = UniversalMerkleTree<E, RescueHash<F>, BigUint, U3, F>;
+impl<F: RescueParameter> DigestAlgorithm<F, F, F> for RescueHash<F> {
+    fn digest(data: &[F]) -> F {
+        RescueCRHF::<F>::sponge_no_padding(data, 1).unwrap()[0]
+    }
+
+    fn digest_leaf(pos: &F, elem: &F) -> F {
+        let data = [F::zero(), *pos, *elem];
+        RescueCRHF::<F>::sponge_no_padding(&data, 1).unwrap()[0]
+    }
+}
+
+/// Example instantiation of a SparseMerkleTree indexed by I
+pub type RescueSparseMerkleTree<I, E, F> = UniversalMerkleTree<E, RescueHash<F>, I, U3, F>;

primitives/src/merkle_tree/universal_merkle_tree.rs

@@ -105,7 +105,8 @@ where
 mod mt_tests {
     use crate::{
         merkle_tree::{
-            prelude::RescueSparseMerkleTree, MerkleTreeScheme, UniversalMerkleTreeScheme,
+            prelude::RescueSparseMerkleTree, LookupResult, MerkleTreeScheme, ToTraversalPath,
+            UniversalMerkleTreeScheme,
         },
         rescue::RescueParameter,
     };
@@ -114,6 +115,7 @@ mod mt_tests {
     use ark_ed_on_bn254::Fq as Fq254;
     use hashbrown::HashMap;
     use num_bigint::BigUint;
+    use typenum::U3;
 
     #[test]
     fn test_universal_mt_builder() {
@@ -123,16 +125,18 @@ mod mt_tests {
     }
 
     fn test_universal_mt_builder_helper<F: RescueParameter>() {
-        let mt =
-            RescueSparseMerkleTree::<F, F>::from_kv_set(1, &[(BigUint::from(1u64), F::from(1u64))])
-                .unwrap();
+        let mt = RescueSparseMerkleTree::<BigUint, F, F>::from_kv_set(
+            1,
+            &[(BigUint::from(1u64), F::from(1u64))],
+        )
+        .unwrap();
         assert_eq!(mt.num_leaves(), 1);
 
         let mut hashmap = HashMap::new();
         hashmap.insert(BigUint::from(1u64), F::from(2u64));
         hashmap.insert(BigUint::from(2u64), F::from(2u64));
         hashmap.insert(BigUint::from(1u64), F::from(3u64));
-        let mt = RescueSparseMerkleTree::<F, F>::from_kv_set(10, &hashmap).unwrap();
+        let mt = RescueSparseMerkleTree::<BigUint, F, F>::from_kv_set(10, &hashmap).unwrap();
         assert_eq!(mt.num_leaves(), hashmap.len() as u64);
     }
 
@@ -148,7 +152,7 @@ mod mt_tests {
         hashmap.insert(BigUint::from(1u64), F::from(2u64));
         hashmap.insert(BigUint::from(2u64), F::from(2u64));
         hashmap.insert(BigUint::from(1u64), F::from(3u64));
-        let mt = RescueSparseMerkleTree::<F, F>::from_kv_set(10, &hashmap).unwrap();
+        let mt = RescueSparseMerkleTree::<BigUint, F, F>::from_kv_set(10, &hashmap).unwrap();
         assert_eq!(mt.num_leaves(), hashmap.len() as u64);
 
         let mut proof = mt
@@ -165,4 +169,41 @@ mod mt_tests {
         let verify_result = mt.non_membership_verify(BigUint::from(4u64), proof);
         assert!(verify_result.is_err());
     }
+
+    #[test]
+    fn test_universal_mt_serde() {
+        test_universal_mt_serde_helper::<Fq254>();
+        test_universal_mt_serde_helper::<Fq377>();
+        test_universal_mt_serde_helper::<Fq381>();
+    }
+
+    fn test_universal_mt_serde_helper<F: RescueParameter + ToTraversalPath<U3>>() {
+        let mut hashmap = HashMap::new();
+        hashmap.insert(F::from(1u64), F::from(2u64));
+        hashmap.insert(F::from(10u64), F::from(3u64));
+        let mt = RescueSparseMerkleTree::<F, F, F>::from_kv_set(3, &hashmap).unwrap();
+        let mem_proof = mt.lookup(F::from(10u64)).expect_ok().unwrap().1;
+        let node = &mem_proof.proof[0];
+        let non_mem_proof = match mt.lookup(F::from(9u64)) {
+            LookupResult::NotFound(proof) => proof,
+            res => panic!("expected NotFound, got {:?}", res),
+        };
+
+        assert_eq!(
+            mt,
+            bincode::deserialize(&bincode::serialize(&mt).unwrap()).unwrap()
+        );
+        assert_eq!(
+            mem_proof,
+            bincode::deserialize(&bincode::serialize(&mem_proof).unwrap()).unwrap()
+        );
+        assert_eq!(
+            non_mem_proof,
+            bincode::deserialize(&bincode::serialize(&non_mem_proof).unwrap()).unwrap()
+        );
+        assert_eq!(
+            *node,
+            bincode::deserialize(&bincode::serialize(node).unwrap()).unwrap()
+        );
+    }
 }

utilities/src/serialize.rs

@@ -41,15 +41,16 @@ macro_rules! deserialize_canonical_bytes {
     };
 }
 
-/// Serializers for finite field elements.
+/// Serializers for elements that are Ark-Works serializable but not serde
+/// serializable.
 ///
-/// Field elements are typically foreign types that we cannot apply the
-/// [tagged] macro to. Instead, use `#[serde(with = "field_elem")]`
-/// at the point where the field element is used inside a struct or enum
-/// definition.
+/// Many cryptographic objects (e.g. finite field elements) are foreign types
+/// that we cannot apply [tagged] or `#[derive(Deserialize, Serialize)]` to.
+/// Instead, use `#[serde(with = "canonical")]` at the point where the object is
+/// used inside a struct or enum definition.
 ///
 /// [tagged]: tagged_base64::tagged
-pub mod field_elem {
+pub mod canonical {
     use super::*;
     use ark_serialize::{CanonicalDeserialize, CanonicalSerialize};
     use ark_std::format;