Mst efficiency improvement

zk_prover/src/merkle_sum_tree/mst.rs

@@ -2,6 +2,7 @@ use crate::merkle_sum_tree::utils::{
     build_leaves_from_entries, build_merkle_tree_from_leaves, parse_csv_to_entries,
 };
 use crate::merkle_sum_tree::{Entry, Node, Tree};
+use halo2_proofs::halo2curves::bn256::Fr as Fp;
 use num_bigint::BigUint;
 
 /// Merkle Sum Tree Data Structure.
@@ -154,7 +155,15 @@ impl<const N_ASSETS: usize, const N_BYTES: usize> MerkleSumTree<N_ASSETS, N_BYTE
             let parent_index = current_index / 2;
             let left_child = &self.nodes[depth - 1][2 * parent_index];
             let right_child = &self.nodes[depth - 1][2 * parent_index + 1];
-            self.nodes[depth][parent_index] = Node::<N_ASSETS>::middle(left_child, right_child);
+
+            let mut hash_preimage = [Fp::zero(); N_ASSETS + 2];
+            for (i, balance) in hash_preimage.iter_mut().enumerate().take(N_ASSETS) {
+                *balance = left_child.balances[i] + right_child.balances[i];
+            }
+            hash_preimage[N_ASSETS] = left_child.hash;
+            hash_preimage[N_ASSETS + 1] = right_child.hash;
+
+            self.nodes[depth][parent_index] = Node::middle_node_from_preimage(&hash_preimage);
             current_index = parent_index;
         }
 

zk_prover/src/merkle_sum_tree/node.rs

@@ -8,24 +8,6 @@ pub struct Node<const N_ASSETS: usize> {
     pub balances: [Fp; N_ASSETS],
 }
 impl<const N_ASSETS: usize> Node<N_ASSETS> {
-    /// Builds a "middle" (non-leaf-level) node of the MST
-    /// The middle node hash is equal to `H(LeftChild.balance[0] + RightChild.balance[0], LeftChild.balance[1] + RightChild.balance[1], ..., LeftChild.balance[N_ASSETS - 1] + RightChild.balance[N_ASSETS - 1], LeftChild.hash, RightChild.hash)`
-    /// The balances are equal to `LeftChild.balance[0] + RightChild.balance[0], LeftChild.balance[1] + RightChild.balance[1], ..., LeftChild.balance[N_ASSETS - 1] + RightChild.balance[N_ASSETS - 1]`
-    pub fn middle(child_l: &Node<N_ASSETS>, child_r: &Node<N_ASSETS>) -> Node<N_ASSETS>
-    where
-        [(); N_ASSETS + 2]: Sized,
-    {
-        let mut balances_sum = [Fp::zero(); N_ASSETS];
-        for (i, balance) in balances_sum.iter_mut().enumerate() {
-            *balance = child_l.balances[i] + child_r.balances[i];
-        }
-
-        Node {
-            hash: Self::poseidon_hash_middle(balances_sum, child_l.hash, child_r.hash),
-            balances: balances_sum,
-        }
-    }
-
     pub fn init_empty() -> Node<N_ASSETS>
     where
         [usize; N_ASSETS + 1]: Sized,
@@ -49,7 +31,10 @@ impl<const N_ASSETS: usize> Node<N_ASSETS> {
         }
     }
 
-    pub fn middle_node_from_preimage(preimage: [Fp; N_ASSETS + 2]) -> Node<N_ASSETS>
+    /// Builds a middle-level node of the MST
+    /// The hash preimage must be equal to `LeftChild.balance[0] + RightChild.balance[0], LeftChild.balance[1] + RightChild.balance[1], ..., LeftChild.balance[N_ASSETS - 1] + RightChild.balance[N_ASSETS - 1], LeftChild.hash, RightChild.hash`
+    /// The balances are equal to `LeftChild.balance[0] + RightChild.balance[0], LeftChild.balance[1] + RightChild.balance[1], ..., LeftChild.balance[N_ASSETS - 1] + RightChild.balance[N_ASSETS - 1]`
+    pub fn middle_node_from_preimage(preimage: &[Fp; N_ASSETS + 2]) -> Node<N_ASSETS>
     where
         [usize; N_ASSETS + 2]: Sized,
     {

zk_prover/src/merkle_sum_tree/tests.rs

@@ -219,7 +219,7 @@ mod test {
             .get_middle_node_hash_preimage(level, index)
             .unwrap();
 
-        let computed_middle_node = Node::<N_ASSETS>::middle_node_from_preimage(hash_preimage);
+        let computed_middle_node = Node::<N_ASSETS>::middle_node_from_preimage(&hash_preimage);
 
         // The hash of the middle node should match the hash computed from the hash preimage
         assert_eq!(middle_node.hash, computed_middle_node.hash);

zk_prover/src/merkle_sum_tree/tree.rs

@@ -1,7 +1,7 @@
 use crate::merkle_sum_tree::utils::big_uint_to_fp;
+use crate::merkle_sum_tree::Cryptocurrency;
 use crate::merkle_sum_tree::{Entry, MerkleProof, Node};
 use halo2_proofs::halo2curves::bn256::Fr as Fp;
-use crate::merkle_sum_tree::Cryptocurrency;
 
 /// A trait representing the basic operations for a Merkle-Sum-like Tree.
 pub trait Tree<const N_ASSETS: usize, const N_BYTES: usize> {
@@ -152,21 +152,39 @@ pub trait Tree<const N_ASSETS: usize, const N_BYTES: usize> {
         let sibling_leaf_node =
             Node::<N_ASSETS>::leaf_node_from_preimage(proof.sibling_leaf_node_hash_preimage);
 
+        let mut hash_preimage = [Fp::zero(); N_ASSETS + 2];
+        for (i, balance) in hash_preimage.iter_mut().enumerate().take(N_ASSETS) {
+            *balance = node.balances[i] + sibling_leaf_node.balances[i];
+        }
+
         if proof.path_indices[0] == 0.into() {
-            node = Node::middle(&node, &sibling_leaf_node);
+            hash_preimage[N_ASSETS] = node.hash;
+            hash_preimage[N_ASSETS + 1] = sibling_leaf_node.hash;
+            node = Node::middle_node_from_preimage(&hash_preimage);
         } else {
-            node = Node::middle(&sibling_leaf_node, &node);
+            hash_preimage[N_ASSETS] = sibling_leaf_node.hash;
+            hash_preimage[N_ASSETS + 1] = node.hash;
+            node = Node::middle_node_from_preimage(&hash_preimage);
         }
 
         for i in 1..proof.path_indices.len() {
             let sibling_node = Node::<N_ASSETS>::middle_node_from_preimage(
-                proof.sibling_middle_node_hash_preimages[i - 1],
+                &proof.sibling_middle_node_hash_preimages[i - 1],
             );
 
+            let mut hash_preimage = [Fp::zero(); N_ASSETS + 2];
+            for (i, balance) in hash_preimage.iter_mut().enumerate().take(N_ASSETS) {
+                *balance = node.balances[i] + sibling_node.balances[i];
+            }
+
             if proof.path_indices[i] == 0.into() {
-                node = Node::middle(&node, &sibling_node);
+                hash_preimage[N_ASSETS] = node.hash;
+                hash_preimage[N_ASSETS + 1] = sibling_node.hash;
+                node = Node::middle_node_from_preimage(&hash_preimage);
             } else {
-                node = Node::middle(&sibling_node, &node);
+                hash_preimage[N_ASSETS] = sibling_node.hash;
+                hash_preimage[N_ASSETS + 1] = node.hash;
+                node = Node::middle_node_from_preimage(&hash_preimage);
             }
         }
 

zk_prover/src/merkle_sum_tree/utils/build_tree.rs

@@ -70,7 +70,18 @@ where
     let results: Vec<Node<N_ASSETS>> = (0..tree[level - 1].len())
         .into_par_iter()
         .step_by(2)
-        .map(|index| Node::middle(&tree[level - 1][index], &tree[level - 1][index + 1]))
+        .map(|index| {
+            let mut hash_preimage = [Fp::zero(); N_ASSETS + 2];
+
+            for (i, balance) in hash_preimage.iter_mut().enumerate().take(N_ASSETS) {
+                *balance =
+                    tree[level - 1][index].balances[i] + tree[level - 1][index + 1].balances[i];
+            }
+
+            hash_preimage[N_ASSETS] = tree[level - 1][index].hash;
+            hash_preimage[N_ASSETS + 1] = tree[level - 1][index + 1].hash;
+            Node::middle_node_from_preimage(&hash_preimage)
+        })
         .collect();
 
     for (index, new_node) in results.into_iter().enumerate() {