wip on updating root

zwaves_primitives/src/hasher.rs

@@ -4,14 +4,18 @@ extern crate bellman;
 extern crate pairing;
 extern crate sapling_crypto;
 
+use pairing::PrimeField;
 use pairing::bls12_381::{Bls12, Fr};
-use pairing::{PrimeField};
+
 use sapling_crypto::jubjub::{JubjubBls12, JubjubEngine};
 use sapling_crypto::pedersen_hash::{pedersen_hash, Personalization};
 
 use crate::bit_iterator::BitIteratorLe;
+use self::pairing::{Field, Engine};
+use std::ptr::hash;
+
 pub struct PedersenHasher<E: JubjubEngine> {
-params: E::Params,
+    params: E::Params,
 }
 
 impl<E: JubjubEngine> PedersenHasher<E> {
@@ -68,6 +72,41 @@ impl<E: JubjubEngine> PedersenHasher<E> {
             })
         }
     }
+
+    pub fn update_root(&self, path: &[&E::Fr], index: usize, elements: &[&E::Fr], merkle_defaults: &[E::Fr]) -> E::Fr {
+        let s = elements.len();
+        let height = path.len() + 1;
+        assert!((index + s) as u32 <= u32::pow(2, (height - 1) as u32), "too many elements");
+
+        let mut offset = index & 0x1;
+        let mut memframesz = s + offset;
+        let zero = <E::Fr as Field>::zero();
+        let mut memframe = vec![zero; (memframesz + 1) as usize];
+
+        (0..s).for_each(|i| memframe[i + offset] = *elements[i]);
+
+        if offset > 0 {
+            memframe[0] = *path[0];
+        }
+
+         (1..height).for_each(|i| {
+            offset = (index >> i) & 0x1;
+            (0..((memframesz + 1) >> 1)).for_each(|j| {
+                memframe[j + offset] = self.compress(&memframe[j * 2], &memframe[j * 2 + 1], Personalization::MerkleTree(i-1));
+            });
+
+            memframesz = offset + ((memframesz + 1) >> 1);
+            if memframesz & 0x1 == 1 {
+                memframe[memframesz] = merkle_defaults[i];
+            }
+
+            if (offset > 0) {
+                memframe[0] = *path[i]
+            }
+        });
+
+        return memframe[0];
+    }
 }
 
 
@@ -86,7 +125,6 @@ impl Default for PedersenHasherBls12 {
 #[test]
 fn test_pedersen_hash() {
     let hasher = PedersenHasherBls12::default();
-    let message = vec![false, false, false, false, false, false, false, false];
     let mut hash = hasher.hash(Fr::from_str("6").unwrap());
 
     println!("testing....");
@@ -96,7 +134,7 @@ fn test_pedersen_hash() {
     }
     println!("Empty root hash: {:?}", hash);
 
-    assert_eq!(hash.to_string(), "Fr(0x5c34420d29ac16e496bfe0b7f9476a2f62633b83470f5999437549e8f20baedd)");
+    assert_eq!(hash.to_string(), "Fr(0x01c2bcb36b2d8126d5812ad211bf90706db31f50bf27f77225d558047571e1aa)");
 }
 
 fn str_to_bin(i: u32) -> Vec<bool> {
@@ -148,3 +186,35 @@ fn test_root_2() {
     assert_eq!(tree[14].to_string(), "Fr(0x4ae608379b1f4b34616934667566fbd43088b5e36ec4e5330b943ba78c273d39)");
 }
 
+//           Merkle tree:
+//            h14 (root)
+//        !h12           h13
+//   h8      h9     h10    !h11
+// h0 h1  h2 h3  !h4 >h5< [h6 h7]
+#[test]
+fn test_update_root() {
+    let hasher = PedersenHasherBls12::default();
+
+    let mut tree: Vec<_> = (1..=15).map(|i| hasher.hash_bits(str_to_bin(i))).collect();
+
+    tree[6] = <Bls12 as Engine>::Fr::zero();
+    tree[7] = <Bls12 as Engine>::Fr::zero();
+
+    tree[8] = hasher.compress(&tree[0], &tree[1], Personalization::MerkleTree(0));
+    tree[9] = hasher.compress(&tree[2], &tree[3], Personalization::MerkleTree(0));
+    tree[10] = hasher.compress(&tree[4], &tree[5], Personalization::MerkleTree(0));
+    tree[11] = hasher.compress(&tree[6], &tree[7], Personalization::MerkleTree(0));
+
+    tree[12] = hasher.compress(&tree[8], &tree[9], Personalization::MerkleTree(1));
+    tree[13] = hasher.compress(&tree[10], &tree[11], Personalization::MerkleTree(1));
+
+    tree[14] = hasher.compress(&tree[12], &tree[13], Personalization::MerkleTree(2));
+
+    let merkle_defaults: Vec<_> = (0..256).scan(&<Bls12 as Engine>::Fr::zero(), |res, _| {
+        Some(hasher.compress(&res, &res, Personalization::MerkleTree(0)))
+    }).collect();
+
+    let res = hasher.update_root(&[&tree[4], &tree[11], &tree[12]], 6, &[&tree[6], &tree[7]], merkle_defaults.as_slice());
+
+    assert_eq!(res.to_string(), "Fr(0x4ae608379b1f4b34616934667566fbd43088b5e36ec4e5330b943ba78c273d39)");
+}