Implemented Batched Proof for multiple openings and Path Pruning

src/merkle_tree/mod.rs

@@ -7,6 +7,7 @@ use ark_serialize::{CanonicalDeserialize, CanonicalSerialize};
 use ark_std::borrow::Borrow;
 use ark_std::hash::Hash;
 use ark_std::vec::Vec;
+use std::collections::{BTreeSet, HashMap};
 
 #[cfg(test)]
 mod tests;
@@ -106,7 +107,7 @@ pub type LeafParam<P> = <<P as Config>::LeafHash as CRHScheme>::Parameters;
 ///    [I] J
 /// ```
 ///  Suppose we want to prove I, then `leaf_sibling_hash` is J, `auth_path` is `[C,D]`
-#[derive(Derivative, CanonicalSerialize, CanonicalDeserialize)]
+#[derive(PartialEq, Derivative, CanonicalSerialize, CanonicalDeserialize)]
 #[derivative(
     Clone(bound = "P: Config"),
     Debug(bound = "P: Config"),
@@ -181,6 +182,212 @@ impl<P: Config> Path<P> {
     }
 }
 
+/// Optimized data structure to store multiple nodes proofs.
+/// For example:
+/// ```tree_diagram
+///         [A]
+///        /   \
+///      [B]    C
+///     / \   /  \
+///    D [E] F    H
+///   / \ / \ ....
+///  [I]J L  M
+/// ```
+///  Suppose we want to prove I and J, then:
+///     `leaf_indexes` is: [2,3] (indexes in Merkle Tree leaves vector)
+///     `leaf_siblings_hashes`: [J,I]
+///     `auth_paths_prefix_lenghts`: [0,2]
+///     `auth_paths_suffixes`: [ [C,D], []]
+///  We can reconstruct the paths incrementally:
+///  First, we reconstruct the first path. The prefix length is 0, hence we do not have any prefix encoding.
+///  The path is thus [C,D].
+///  Once the first path is verified, we can reconstruct the second path.
+///  The prefix length of 2 means that the path prefix will be `previous_path[:2] -> [C,D]`.
+///  Since the Merkle Tree branch is the same, the authentication path is the same (which means in this case that there is no suffix).
+///  The second path is hence `[C,D] + []` (i.e., plus the empty suffix). We can verify the second path as the first one.
+
+#[derive(Derivative, CanonicalSerialize, CanonicalDeserialize)]
+#[derivative(
+    Clone(bound = "P: Config"),
+    Debug(bound = "P: Config"),
+    Default(bound = "P: Config")
+)]
+pub struct MultiPath<P: Config> {
+    /// For node i, stores the hash of node i's sibling
+    pub leaf_siblings_hashes: Vec<P::LeafDigest>,
+    /// For node i path, stores at index i the prefix length of the path, for Incremental encoding
+    pub auth_paths_prefix_lenghts: Vec<usize>,
+    /// For node i path, stores at index i the suffix of the path for Incremental Encoding (as vector of symbols to be resolved with self.lut). Order is from higher layer to lower layer (does not include root node).
+    pub auth_paths_suffixes: Vec<Vec<P::InnerDigest>>,
+    /// stores the leaf indexes of the nodes to prove
+    pub leaf_indexes: Vec<usize>,
+}
+
+impl<P: Config> MultiPath<P> {
+    /// Returns a compressed MultiPath containing multiple encoded authentication paths for `indexes`
+    fn compress(
+        indexes: impl IntoIterator<Item = usize>,
+        leaf_siblings_hashes: Vec<P::LeafDigest>,
+        auth_paths: Vec<Path<P>>,
+    ) -> Result<Self, crate::Error> {
+        let indexes = Vec::from_iter(indexes);
+        // use multipath for more than 1 leaf
+        assert!(
+            indexes.len() > 1,
+            "Expected more than one leaf to verify for MultiPath, got {}",
+            indexes.len()
+        );
+
+        let mut auth_paths_prefix_lenghts: Vec<usize> = Vec::with_capacity(indexes.len());
+        let mut auth_paths_suffixes: Vec<Vec<P::InnerDigest>> = Vec::with_capacity(indexes.len());
+
+        auth_paths_prefix_lenghts.push(0);
+        auth_paths_suffixes.push(auth_paths[0].auth_path.clone());
+
+        let mut prev_path = auth_paths[0].clone();
+
+        // Encode paths with Incremental Encoding (Front compression)
+        for path in auth_paths[1..].to_vec() {
+            let mut prefix_len = 0;
+            if prev_path.auth_path.len() != 0 && path.auth_path.len() != 0 {
+                while prev_path.auth_path[prefix_len] == path.auth_path[prefix_len] {
+                    prefix_len += 1;
+                    if prefix_len == prev_path.auth_path.len() || prefix_len == path.auth_path.len()
+                    {
+                        break;
+                    }
+                }
+            }
+            auth_paths_prefix_lenghts.push(prefix_len);
+            if prefix_len == prev_path.auth_path.len() {
+                auth_paths_suffixes.push(vec![]);
+            } else {
+                auth_paths_suffixes.push(path.auth_path[prefix_len..].to_vec().clone());
+            }
+            prev_path = path;
+        }
+
+        Ok(MultiPath {
+            leaf_indexes: indexes,
+            auth_paths_prefix_lenghts: auth_paths_prefix_lenghts,
+            auth_paths_suffixes: auth_paths_suffixes,
+            leaf_siblings_hashes: leaf_siblings_hashes,
+        })
+    }
+
+    /// Returns the decompressed authentication paths for every leaf index
+    fn decompress(
+        &'_ self,
+    ) -> Result<impl '_ + Iterator<Item = Vec<P::InnerDigest>>, crate::Error> {
+        // Incrementally reconstruct all the paths
+        let mut curr_path = self.auth_paths_suffixes[0].clone();
+        let mut auth_paths = (0..self.leaf_indexes.len())
+            .map(|_| Vec::new())
+            .collect::<Vec<Vec<P::InnerDigest>>>();
+        auth_paths[0] = self.auth_paths_suffixes[0].clone();
+
+        for i in 1..auth_paths.len() {
+            auth_paths[i].extend_from_slice(&curr_path[0..self.auth_paths_prefix_lenghts[i]]);
+            auth_paths[i].extend(self.auth_paths_suffixes[i].clone());
+            curr_path = auth_paths[i].clone();
+        }
+        Ok(auth_paths.into_iter())
+    }
+
+    /// Verify that leaves are at `self.leaf_indexes` of the merkle tree.
+    /// Note that the order of the leaves hashes should match the leaves respective indexes
+    /// * `leaf_size`: leaf size in number of bytes
+    ///
+    /// `verify` infers the tree height by setting `tree_height = self.auth_paths_suffixes[0].len() + 2`
+    pub fn verify<L: Borrow<P::Leaf> + Clone>(
+        &self,
+        leaf_hash_params: &LeafParam<P>,
+        two_to_one_params: &TwoToOneParam<P>,
+        root_hash: &P::InnerDigest,
+        leaves: impl IntoIterator<Item = L>,
+    ) -> Result<bool, crate::Error> {
+        // array of auth paths as arrays of InnerDigests
+
+        let leaves: Vec<L> = leaves.into_iter().collect();
+
+        let auth_paths: Vec<Vec<P::InnerDigest>> = self.decompress()?.collect();
+
+        let tree_height = auth_paths[0].len() + 2;
+
+        // LookUp table to speedup computation avoid redundant hash computations
+        let mut hash_lut: HashMap<usize, P::InnerDigest> = HashMap::new();
+
+        for i in 0..self.leaf_indexes.len() {
+            let leaf_index = self.leaf_indexes[i];
+            let leaf = &leaves[i];
+            let leaf_sibling_hash = &self.leaf_siblings_hashes[i];
+            let auth_path = &auth_paths[i];
+
+            let claimed_leaf_hash = P::LeafHash::evaluate(&leaf_hash_params, leaf.clone())?;
+            let (left_child, right_child) =
+                select_left_right_child(leaf_index, &claimed_leaf_hash, &leaf_sibling_hash)?;
+            // check hash along the path from bottom to root
+
+            // leaf layer to inner layer conversion
+            let left_child = P::LeafInnerDigestConverter::convert(left_child)?;
+            let right_child = P::LeafInnerDigestConverter::convert(right_child)?;
+
+            // we will use `index` variable to track the position of path
+            let mut index = leaf_index;
+            let mut index_in_tree = convert_index_to_last_level(leaf_index, tree_height);
+            index >>= 1;
+            index_in_tree = parent(index_in_tree).unwrap();
+
+            let mut curr_path_node =
+                hash_lut
+                    .entry(index_in_tree)
+                    .or_insert(P::TwoToOneHash::evaluate(
+                        &two_to_one_params,
+                        left_child,
+                        right_child,
+                    )?);
+
+            // Check levels between leaf level and root
+            for level in (0..auth_path.len()).rev() {
+                // check if path node at this level is left or right
+                let (left, right) =
+                    select_left_right_child(index, curr_path_node, &auth_path[level])?;
+                // update curr_path_node
+                index >>= 1;
+                index_in_tree = parent(index_in_tree).unwrap();
+                curr_path_node = hash_lut
+                    .entry(index_in_tree)
+                    .or_insert(P::TwoToOneHash::compress(&two_to_one_params, left, right)?);
+            }
+
+            // check if final hash is root
+            if curr_path_node != root_hash {
+                return Ok(false);
+            }
+        }
+        Ok(true)
+    }
+
+    /// The position of on_path node in `leaf_and_sibling_hash` and `non_leaf_and_sibling_hash_path`.
+    /// `position[i]` is 0 (false) iff `i`th on-path node from top to bottom is on the left.
+    ///
+    /// This function simply converts every index in `self.leaf_indexes` to boolean array in big endian form.
+    #[allow(unused)] // this function is actually used when r1cs feature is on
+    fn position_list(&'_ self) -> impl '_ + Iterator<Item = Vec<bool>> {
+        let path_len = self.auth_paths_suffixes[0].len();
+
+        cfg_into_iter!(self.leaf_indexes.clone())
+            .map(move |i| {
+                (0..path_len + 1)
+                    .map(move |j| ((i >> j) & 1) != 0)
+                    .rev()
+                    .collect()
+            })
+            .collect::<Vec<_>>()
+            .into_iter()
+    }
+}
+
 /// `index` is the first `path.len()` bits of
 /// the position of tree.
 ///
@@ -293,8 +500,8 @@ impl<P: Config> MerkleTree<P> {
                     // leaf in the whole tree (represented as a list in level order). We need to shift it
                     // by `-upper_bound` to get the index in `leaf_nodes` list.
 
-                    //similarly, we need to rescale i by start_index
-                    //to get the index outside the slice and in the level-ordered list of nodes
+                    // similarly, we need to rescale i by start_index
+                    // to get the index outside the slice and in the level-ordered list of nodes
 
                     let current_index = i + start_index;
                     let left_leaf_index = left_child(current_index) - upper_bound;
@@ -329,8 +536,8 @@ impl<P: Config> MerkleTree<P> {
                     // leaf in the whole tree (represented as a list in level order). We need to shift it
                     // by `-upper_bound` to get the index in `leaf_nodes` list.
 
-                    //similarly, we need to rescale i by start_index
-                    //to get the index outside the slice and in the level-ordered list of nodes
+                    // similarly, we need to rescale i by start_index
+                    // to get the index outside the slice and in the level-ordered list of nodes
                     let current_index = i + start_index;
                     let left_leaf_index = left_child(current_index) - upper_bound;
                     let right_leaf_index = right_child(current_index) - upper_bound;
@@ -400,6 +607,37 @@ impl<P: Config> MerkleTree<P> {
         })
     }
 
+    /// Returns a MultiPath (multiple authentication paths in compressed form, with Front Incremental Encoding),
+    /// from every leaf to root.
+    /// Note that for compression efficiency, the indexes are internally sorted.
+    /// When verifying the proof, leaves hashes should be supplied in order, that is:
+    /// let ordered_leaves: Vec<_> = self.leaf_indexes.into_iter().map(|i| leaves[i]).collect();
+    pub fn generate_multi_proof(
+        &self,
+        indexes: impl IntoIterator<Item = usize>,
+    ) -> Result<MultiPath<P>, crate::Error> {
+        // pruned and sorted for encoding efficiency
+        let indexes: BTreeSet<usize> = indexes.into_iter().collect();
+
+        let auth_paths: Vec<Path<P>> = cfg_into_iter!(indexes.clone())
+            .map(|i| self.generate_proof(i))
+            .collect::<Result<Vec<Path<P>>, crate::Error>>()?;
+
+        let leaf_siblings_hashes = cfg_into_iter!(indexes.clone())
+            .map(|i| {
+                if i & 1 == 0 {
+                    // leaf is left child
+                    self.leaf_nodes[i + 1].clone()
+                } else {
+                    // leaf is right child
+                    self.leaf_nodes[i - 1].clone()
+                }
+            })
+            .collect();
+
+        MultiPath::compress(indexes, leaf_siblings_hashes, auth_paths)
+    }
+
     /// Given the index and new leaf, return the hash of leaf and an updated path in order from root to bottom non-leaf level.
     /// This does not mutate the underlying tree.
     fn updated_path<T: Borrow<P::Leaf>>(