From cbe4737439857ca321213c6cda06289950d22b1a Mon Sep 17 00:00:00 2001
From: Hanh <hanh425@gmail.com>
Date: Thu, 27 Oct 2022 20:59:04 +0800
Subject: [PATCH] Batch pedersen hash

---
 src/sapling/hash.rs | 30 ++++++++++++++++++++++++------
 src/sync/tree.rs    | 40 ++++++++++++++++++++++++++++++++++------
 2 files changed, 58 insertions(+), 12 deletions(-)

diff --git a/src/sapling/hash.rs b/src/sapling/hash.rs
index d65cad8..2722421 100644
--- a/src/sapling/hash.rs
+++ b/src/sapling/hash.rs
@@ -1,6 +1,6 @@
 use ff::PrimeField;
 use group::Curve;
-use jubjub::{ExtendedPoint, Fr};
+use jubjub::{AffinePoint, ExtendedPoint, Fr};
 use zcash_primitives::constants::PEDERSEN_HASH_CHUNKS_PER_GENERATOR;
 use crate::sync::{Hasher, Node};
 use super::GENERATORS_EXP;
@@ -35,6 +35,14 @@ fn accumulate_generator(acc: &Fr, idx_generator: u32) -> ExtendedPoint {
 }
 
 pub fn hash_combine(depth: u8, left: &[u8; 32], right: &[u8; 32]) -> [u8; 32] {
+    let p = hash_combine_inner(depth, left, right);
+    p
+        .to_affine()
+        .get_u()
+        .to_repr()
+}
+
+pub fn hash_combine_inner(depth: u8, left: &[u8; 32], right: &[u8; 32]) -> ExtendedPoint {
     let mut hash = ExtendedPoint::identity();
     let mut acc = Fr::zero();
     let mut cur = Fr::one();
@@ -95,10 +103,6 @@ pub fn hash_combine(depth: u8, left: &[u8; 32], right: &[u8; 32]) -> [u8; 32] {
     }
     hash += accumulate_generator(&acc, idx_generator);
 
-    let hash = hash
-        .to_affine()
-        .get_u()
-        .to_repr();
     hash
 }
 
@@ -106,6 +110,8 @@ pub fn hash_combine(depth: u8, left: &[u8; 32], right: &[u8; 32]) -> [u8; 32] {
 pub struct SaplingHasher {}
 
 impl Hasher for SaplingHasher {
+    type Extended = ExtendedPoint;
+
     fn uncommited_node() -> Node {
         [0u8; 32]
     }
@@ -113,8 +119,20 @@ impl Hasher for SaplingHasher {
     fn node_combine(&self, depth: u8, left: &Node, right: &Node) -> Node {
         hash_combine(depth, left, right)
     }
-}
 
+    fn node_combine_extended(&self, depth: u8, left: &Node, right: &Node) -> Self::Extended {
+        hash_combine_inner(depth, left, right)
+    }
+
+    fn normalize(&self, extended: &[Self::Extended]) -> Vec<Node> {
+        let mut hash_affine: Vec<AffinePoint> = vec![AffinePoint::identity(); extended.len()];
+        ExtendedPoint::batch_normalize(extended, &mut hash_affine);
+        hash_affine
+            .iter()
+            .map(|p| p.get_u().to_repr())
+            .collect()
+    }
+}
 
 #[cfg(test)]
 mod tests {
diff --git a/src/sync/tree.rs b/src/sync/tree.rs
index 3ae825d..03c4a61 100644
--- a/src/sync/tree.rs
+++ b/src/sync/tree.rs
@@ -2,13 +2,19 @@ use rayon::prelude::*;
 use std::io::{Read, Write};
 use std::marker::PhantomData;
 use byteorder::WriteBytesExt;
+use group::Curve;
 use zcash_encoding::{Optional, Vector};
 
 pub type Node = [u8; 32];
 
 pub trait Hasher: Clone + Sync {
+    type Extended: Curve + Clone + Send;
+
     fn uncommited_node() -> Node;
     fn node_combine(&self, depth: u8, left: &Node, right: &Node) -> Node;
+
+    fn node_combine_extended(&self, depth: u8, left: &Node, right: &Node) -> Self::Extended;
+    fn normalize(&self, extended: &[Self::Extended]) -> Vec<Node>;
 }
 
 #[derive(Clone)]
@@ -388,9 +394,34 @@ fn combine_level<H: Hasher>(
 ) -> usize {
     assert_eq!(n % 2, 0);
 
-    // TODO: Support batch hash combine
     let nn = n / 2;
-    let next_level: Vec<_> = (0..nn)
+    let next_level = if nn > 100 {
+        batch_level_combine(commitments, offset, nn, depth, hasher)
+    }
+    else {
+        single_level_combine(commitments, offset, nn, depth, hasher)
+    };
+
+    commitments[0..nn].copy_from_slice(&next_level);
+    nn
+}
+
+fn batch_level_combine<H: Hasher>(commitments: &mut [Node], offset: Option<Node>, nn: usize, depth: u8, hasher: &H) -> Vec<Node> {
+    let hash_extended: Vec<_> = (0..nn)
+        .into_par_iter()
+        .map(|i| {
+            hasher.node_combine_extended(
+                depth,
+                CTreeBuilder::<H>::get(commitments, 2 * i, &offset),
+                CTreeBuilder::<H>::get(commitments, 2 * i + 1, &offset),
+            )
+        })
+        .collect();
+    hasher.normalize(&hash_extended)
+}
+
+fn single_level_combine<H: Hasher>(commitments: &mut [Node], offset: Option<Node>, nn: usize, depth: u8, hasher: &H) -> Vec<Node> {
+    (0..nn)
         .into_par_iter()
         .map(|i| {
             hasher.node_combine(
@@ -399,10 +430,7 @@ fn combine_level<H: Hasher>(
                 CTreeBuilder::<H>::get(commitments, 2 * i + 1, &offset),
             )
         })
-        .collect();
-
-    commitments[0..nn].copy_from_slice(&next_level);
-    nn
+        .collect()
 }
 
 struct WitnessBuilder<H: Hasher> {