From ea51af1512c586cf17d3c69d1d445cf93f72d49d Mon Sep 17 00:00:00 2001
From: Alexey <peretsmobil@gmail.com>
Date: Thu, 11 Jul 2019 17:04:36 +0300
Subject: [PATCH] batch merkle tree creation

---
 contracts/test/MerkleTreeWithHistory.test.js | 45 ++++++++++++++++++++
 lib/MerkleTree.js                            | 41 ++++++++++++------
 lib/Storage.js                               |  3 ++
 3 files changed, 77 insertions(+), 12 deletions(-)

diff --git a/contracts/test/MerkleTreeWithHistory.test.js b/contracts/test/MerkleTreeWithHistory.test.js
index cf48a46..17903fe 100644
--- a/contracts/test/MerkleTreeWithHistory.test.js
+++ b/contracts/test/MerkleTreeWithHistory.test.js
@@ -91,6 +91,51 @@ contract('MerkleTreeWithHistory', async accounts => {
       // console.log(root)
       assert.equal(root, calculated_root)
     })
+    it('creation odd elements count', async () => {
+      const elements = [12, 13, 14, 15, 16, 17, 18, 19, 20]
+      for(const [i, el] of Object.entries(elements)) {
+        await tree.update(i, el)
+      }
+
+      const storage = new JsStorage()
+      hasher = new MimcHacher()
+      const batchTree = new MerkleTree(
+        prefix,
+        storage,
+        hasher,
+        levels,
+        zeroValue,
+        elements
+      );
+      for(const [i, el] of Object.entries(elements)) {
+        const pathViaConstructor = await batchTree.path(i)
+        const pathViaUpdate = await tree.path(i)
+        pathViaConstructor.should.be.deep.equal(pathViaUpdate)
+      }
+    })
+
+    it('creation even elements count', async () => {
+      const elements = [12, 13, 14, 15, 16, 17]
+      for(const [i, el] of Object.entries(elements)) {
+        await tree.update(i, el)
+      }
+
+      const storage = new JsStorage()
+      hasher = new MimcHacher()
+      const batchTree = new MerkleTree(
+        prefix,
+        storage,
+        hasher,
+        levels,
+        zeroValue,
+        elements
+      );
+      for(const [i, el] of Object.entries(elements)) {
+        const pathViaConstructor = await batchTree.path(i)
+        const pathViaUpdate = await tree.path(i)
+        pathViaConstructor.should.be.deep.equal(pathViaUpdate)
+      }
+    })
   })
 
   describe('#insert', async () => {
diff --git a/lib/MerkleTree.js b/lib/MerkleTree.js
index 25fc471..6c4d0a9 100644
--- a/lib/MerkleTree.js
+++ b/lib/MerkleTree.js
@@ -1,9 +1,7 @@
 
-// const AwaitLock = require('await-lock');
-
 class MerkleTree {
 
-  constructor(prefix, storage, hasher, n_levels, zero_value) {
+  constructor(prefix, storage, hasher, n_levels, zero_value, defaultElements) {
     this.prefix = prefix;
     this.storage = storage;
     this.hasher = hasher;
@@ -18,7 +16,28 @@ class MerkleTree {
         current_zero_value.toString(),
       );
     }
-    // this.lock = new AwaitLock();
+    if (defaultElements) {
+        let level = 0
+        defaultElements.forEach((element, i) => {
+            this.storage.put(MerkleTree.element_to_key(prefix, element), i.toString())
+            this.storage.put(MerkleTree.index_to_key(prefix, level, i), element)
+        })
+        level++
+        let numberOfElementInRow = Math.ceil(defaultElements.length / 2) 
+        for (level; level <= this.n_levels; level++) {
+            for(let i = 0; i < numberOfElementInRow; i++) {
+                const leftKey = MerkleTree.index_to_key(prefix, level - 1, 2 * i)
+                const rightKey = MerkleTree.index_to_key(prefix, level - 1, 2 * i + 1)
+
+                const left = this.storage.get(leftKey)
+                const right = this.storage.get_or_element(rightKey, this.zero_values[level - 1])
+
+                const subRoot = this.hasher.hash(null, left, right);
+                this.storage.put(MerkleTree.index_to_key(prefix, level, i), subRoot)
+            }
+            numberOfElementInRow = Math.max(Math.ceil(numberOfElementInRow / 2), 1)
+        }
+    }
   }
 
   static index_to_key(prefix, level, index) {
@@ -121,7 +140,6 @@ class MerkleTree {
   }
 
   async update(index, element, update_log_index) {
-    // await this.lock.acquireAsync();
     try {
       //console.log(`updating ${index}, ${element}`);
       class UpdateTraverser {
@@ -198,8 +216,8 @@ class MerkleTree {
 
       const root = await this.root();
       //console.log(`updated root ${root}`);
-    } finally {
-    //   this.lock.release();
+    } catch(e) {
+        console.error(e)
     }
   }
 
@@ -218,7 +236,6 @@ class MerkleTree {
   }
 
   async rollback(updates) {
-    // await this.lock.acquireAsync();
     try {
       const update_log_key = MerkleTree.update_log_to_key(this.prefix);
       const update_log_index = await this.storage.get(update_log_key);
@@ -228,8 +245,8 @@ class MerkleTree {
 
         await this.update(update_element_log.index, update_element_log.old_element, update_log_index - i - 1);
       }
-    } finally {
-    //   this.lock.release();
+    } catch (e) {
+        console.error(e)
     }
   }
 
@@ -252,8 +269,8 @@ class MerkleTree {
       if (await this.root() != root) {
         throw new Error(`could not rollback to root ${root}`);
       }
-    } finally {
-    //   this.lock.release();
+    } catch (e) {
+        console.log(e)
     }
 
   }
diff --git a/lib/Storage.js b/lib/Storage.js
index 9d45112..ae86fc0 100644
--- a/lib/Storage.js
+++ b/lib/Storage.js
@@ -19,6 +19,9 @@ class JsStorage {
   }
 
   put(key, value) {
+    if (key === undefined || value === undefined) {
+        throw Error('key or value is undefined')
+    }
     this.db[key] = value;
   }