From 65fd0202c5534702f9231cfd62647c7161a9429e Mon Sep 17 00:00:00 2001
From: Alexey <peretsmobil@gmail.com>
Date: Thu, 11 Jul 2019 13:38:22 +0300
Subject: [PATCH] merkle tree lib

---
 contracts/test/MerkleTreeWithHistory.test.js |  74 +++++-
 lib/MerkleTree.js                            | 262 +++++++++++++++++++
 lib/MiMC.js                                  |  13 +
 lib/Storage.js                               |  36 +++
 4 files changed, 377 insertions(+), 8 deletions(-)
 create mode 100644 lib/MerkleTree.js
 create mode 100644 lib/MiMC.js
 create mode 100644 lib/Storage.js

diff --git a/contracts/test/MerkleTreeWithHistory.test.js b/contracts/test/MerkleTreeWithHistory.test.js
index 800436d..cf48a46 100644
--- a/contracts/test/MerkleTreeWithHistory.test.js
+++ b/contracts/test/MerkleTreeWithHistory.test.js
@@ -9,6 +9,10 @@ const { takeSnapshot, revertSnapshot, increaseTime } = require('../scripts/ganac
 const MerkleTreeWithHistory = artifacts.require('./MerkleTreeWithHistoryMock.sol')
 const MiMC = artifacts.require('./MiMC.sol')
 
+const JsStorage = require('../../lib/Storage')
+const MerkleTree = require('../../lib/MerkleTree')
+const MimcHacher = require('../../lib/MiMC')
+
 function BNArrayToStringArray(array) {
   const arrayToPrint = []
   array.forEach(item => {
@@ -25,10 +29,22 @@ contract('MerkleTreeWithHistory', async accounts => {
   const levels = 5
   const zeroValue = 1337
   let snapshotId
+  let prefix = 'test'
+  let tree
+  let hasher
 
   before(async () => {
+    const storage = new JsStorage()
+    hasher = new MimcHacher()
+    tree = new MerkleTree(
+      prefix,
+      storage,
+      hasher,
+      levels,
+      zeroValue,
+    )
     miMC = MiMC.deployed()
-    await MerkleTreeWithHistory.link(MiMC, miMC.address);
+    await MerkleTreeWithHistory.link(MiMC, miMC.address)
     merkleTreeWithHistory = await MerkleTreeWithHistory.new(levels, zeroValue)
     snapshotId = await takeSnapshot()
   })
@@ -36,9 +52,9 @@ contract('MerkleTreeWithHistory', async accounts => {
   describe('#constuctor', async () => {
     it('should initialize', async () => {
       const filled_subtrees = await merkleTreeWithHistory.filled_subtrees()
-      console.log('filled_subtrees', BNArrayToStringArray(filled_subtrees))
+      // console.log('filled_subtrees', BNArrayToStringArray(filled_subtrees))
       const root = await merkleTreeWithHistory.getLastRoot()
-      console.log('root', root.toString())
+      // console.log('root', root.toString())
       filled_subtrees[0].should.be.eq.BN(zeroValue)
       const zeros = await merkleTreeWithHistory.zeros()
       // console.log('zeros', BNArrayToStringArray(zeros))
@@ -48,24 +64,66 @@ contract('MerkleTreeWithHistory', async accounts => {
     })
   })
 
+  describe('merkleTreeLib', async () => {
+    it('index_to_key', async () => {
+      assert.equal(
+        MerkleTree.index_to_key('test', 5, 20),
+        "test_tree_5_20",
+      )
+    })
+
+    it('tests insert', async () => {
+      const storage = new JsStorage()
+      hasher = new MimcHacher()
+      tree = new MerkleTree(
+        prefix,
+        storage,
+        hasher,
+        2,
+        zeroValue,
+      )
+      await tree.update(0, '5')
+      let {root, path_elements, path_index} = await tree.path(0)
+      const calculated_root = hasher.hash(null,
+        hasher.hash(null, '5', path_elements[0]),
+        path_elements[1]
+      )
+      // console.log(root)
+      assert.equal(root, calculated_root)
+    })
+  })
+
   describe('#insert', async () => {
     it('should insert', async () => {
       let filled_subtrees
-      let root
+      let rootFromContract
 
       for (i = 1; i < 11; i++) {
         await merkleTreeWithHistory.insert(i)
+        await tree.update(i - 1, i)
         filled_subtrees = await merkleTreeWithHistory.filled_subtrees()
-        console.log('filled_subtrees', BNArrayToStringArray(filled_subtrees))
-        root = await merkleTreeWithHistory.getLastRoot()
-        console.log('root', root.toString())
+        let {root, path_elements, path_index} = await tree.path(i - 1)
+        // console.log('path_elements  ', path_elements)
+        // console.log('filled_subtrees', BNArrayToStringArray(filled_subtrees))
+        // console.log('rootFromLib', root)
+        rootFromContract = await merkleTreeWithHistory.getLastRoot()
+        root.should.be.equal(rootFromContract.toString())
+        // console.log('rootFromCon', root.toString())
       }
-
     })
   })
 
   afterEach(async () => {
     await revertSnapshot(snapshotId.result)
     snapshotId = await takeSnapshot()
+    const storage = new JsStorage()
+    hasher = new MimcHacher()
+    tree = new MerkleTree(
+      prefix,
+      storage,
+      hasher,
+      levels,
+      zeroValue,
+    )
   })
 })
diff --git a/lib/MerkleTree.js b/lib/MerkleTree.js
new file mode 100644
index 0000000..25fc471
--- /dev/null
+++ b/lib/MerkleTree.js
@@ -0,0 +1,262 @@
+
+// const AwaitLock = require('await-lock');
+
+class MerkleTree {
+
+  constructor(prefix, storage, hasher, n_levels, zero_value) {
+    this.prefix = prefix;
+    this.storage = storage;
+    this.hasher = hasher;
+    this.n_levels = n_levels;
+    this.zero_values = [];
+
+    let current_zero_value = zero_value;
+    this.zero_values.push(current_zero_value);
+    for (let i = 0; i < n_levels; i++) {
+      current_zero_value = this.hasher.hash(i, current_zero_value, current_zero_value);
+      this.zero_values.push(
+        current_zero_value.toString(),
+      );
+    }
+    // this.lock = new AwaitLock();
+  }
+
+  static index_to_key(prefix, level, index) {
+    const key = `${prefix}_tree_${level}_${index}`;
+    return key;
+  }
+
+  static element_to_key(prefix, element) {
+    const key = `${prefix}_element_${element}`;
+    return key;
+  }
+
+
+
+  static update_log_to_key(prefix) {
+    return `${prefix}_update_log_index`;
+  }
+
+  static update_log_element_to_key(prefix, update_log_index) {
+    return `${prefix}_update_log_element_${update_log_index}`;
+  }
+
+  async update_log(index, old_element, new_element, update_log_index, should_put_element_update) {
+    let ops = [];
+
+    const update_log_key = MerkleTree.update_log_to_key(this.prefix);
+    ops.push({
+      type: 'put',
+      key: update_log_key,
+      value: update_log_index.toString(),
+    });
+
+    if (should_put_element_update) {
+      const update_log_element_key = MerkleTree.update_log_element_to_key(this.prefix, update_log_index);
+      ops.push({
+        type: 'put',
+        key: update_log_element_key,
+        value: JSON.stringify({
+          index,
+          old_element,
+          new_element,
+        })
+      });
+    }
+    await this.storage.put_batch(ops);
+  }
+
+  async root() {
+    let root = await this.storage.get_or_element(
+      MerkleTree.index_to_key(this.prefix, this.n_levels, 0),
+      this.zero_values[this.n_levels],
+    );
+
+    return root;
+  }
+
+  async element_index(element) {
+    const element_key = MerkleTree.element_to_key(this.prefix, element);
+    const index = await this.storage.get_or_element(element_key, -1);
+    return index;
+  }
+
+  async path(index) {
+    class PathTraverser {
+      constructor(prefix, storage, zero_values) {
+        this.prefix = prefix;
+        this.storage = storage;
+        this.zero_values = zero_values;
+        this.path_elements = [];
+        this.path_index = [];
+      }
+
+      async handle_index(level, element_index, sibling_index) {
+        const sibling = await this.storage.get_or_element(
+          MerkleTree.index_to_key(this.prefix, level, sibling_index),
+          this.zero_values[level],
+        );
+        this.path_elements.push(sibling);
+        this.path_index.push(element_index % 2);
+      }
+    }
+    let traverser = new PathTraverser(this.prefix, this.storage, this.zero_values);
+    const root = await this.storage.get_or_element(
+      MerkleTree.index_to_key(this.prefix, this.n_levels, 0),
+      this.zero_values[this.n_levels],
+    );
+
+    const element = await this.storage.get_or_element(
+      MerkleTree.index_to_key(this.prefix, 0, index),
+      this.zero_values[0],
+    );
+
+    await this.traverse(index, traverser);
+    return {
+      root,
+      path_elements: traverser.path_elements,
+      path_index: traverser.path_index,
+      element
+    };
+  }
+
+  async update(index, element, update_log_index) {
+    // await this.lock.acquireAsync();
+    try {
+      //console.log(`updating ${index}, ${element}`);
+      class UpdateTraverser {
+        constructor(prefix, storage, hasher, element, zero_values) {
+          this.prefix = prefix;
+          this.current_element = element;
+          this.zero_values = zero_values;
+          this.storage = storage;
+          this.hasher = hasher;
+          this.key_values_to_put = [];
+        }
+
+        async handle_index(level, element_index, sibling_index) {
+          if (level == 0) {
+            this.original_element = await this.storage.get_or_element(
+              MerkleTree.index_to_key(this.prefix, level, element_index),
+              this.zero_values[level],
+            );
+            this.key_values_to_put.push({
+              key: MerkleTree.element_to_key(this.prefix, element),
+              value: index.toString(),
+            });
+
+          }
+          const sibling = await this.storage.get_or_element(
+            MerkleTree.index_to_key(this.prefix, level, sibling_index),
+            this.zero_values[level],
+          );
+          let left, right;
+          if (element_index % 2 == 0) {
+            left = this.current_element;
+            right = sibling;
+          } else {
+            left = sibling;
+            right = this.current_element;
+          }
+
+          this.key_values_to_put.push({
+            key: MerkleTree.index_to_key(this.prefix, level, element_index),
+            value: this.current_element,
+          });
+          //console.log(`left: ${left}, right: ${right}`);
+          this.current_element = this.hasher.hash(level, left, right);
+          //console.log(`current_element: ${this.current_element}`);
+        }
+      }
+      let traverser = new UpdateTraverser(
+        this.prefix,
+        this.storage,
+        this.hasher,
+        element,
+        this.zero_values
+      );
+
+      await this.traverse(index, traverser);
+      //console.log(`traverser.current_element: ${traverser.current_element}`);
+      traverser.key_values_to_put.push({
+        key: MerkleTree.index_to_key(this.prefix, this.n_levels, 0),
+        value: traverser.current_element,
+      });
+
+      if (update_log_index == undefined) {
+        const update_log_key = MerkleTree.update_log_to_key(this.prefix);
+        let update_log_index_from_db = await this.storage.get_or_element(update_log_key, -1);
+        update_log_index = parseInt(update_log_index_from_db) + 1;
+        await this.update_log(index, traverser.original_element, element, update_log_index, true);
+      } else {
+        await this.update_log(index, traverser.original_element, element, update_log_index, false);
+      }
+
+      await this.storage.del(MerkleTree.element_to_key(this.prefix, traverser.original_element));
+      //traverser.key_values_to_put.forEach((e) => console.log(`key_values: ${JSON.stringify(e)}`));
+      await this.storage.put_batch(traverser.key_values_to_put);
+
+      const root = await this.root();
+      //console.log(`updated root ${root}`);
+    } finally {
+    //   this.lock.release();
+    }
+  }
+
+  async traverse(index, handler) {
+    let current_index = index;
+    for (let i = 0; i < this.n_levels; i++) {
+        let sibling_index = current_index;
+        if (current_index % 2 == 0) {
+          sibling_index += 1;
+        } else {
+          sibling_index -= 1;
+        }
+        await handler.handle_index(i, current_index, sibling_index);
+        current_index = Math.floor(current_index / 2);
+      }
+  }
+
+  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);
+      for (let i = 0; i < updates; i++) {
+        const update_log_element_key = MerkleTree.update_log_element_to_key(this.prefix, update_log_index - i);
+        const update_element_log = JSON.parse(await this.storage.get(update_log_element_key));
+
+        await this.update(update_element_log.index, update_element_log.old_element, update_log_index - i - 1);
+      }
+    } finally {
+    //   this.lock.release();
+    }
+  }
+
+  async rollback_to_root(root) {
+    // await this.lock.acquireAsync();
+    try {
+      const update_log_key = MerkleTree.update_log_to_key(this.prefix);
+      let update_log_index = await this.storage.get(update_log_key);
+      while (update_log_index >= 0) {
+        update_log_index -= 1;
+        const update_log_element_key = MerkleTree.update_log_element_to_key(this.prefix, update_log_index - i);
+        const update_element_log = JSON.parse(await this.storage.get(update_log_element_key));
+
+        await this.update(update_element_log.index, update_element_log.old_element, update_log_index);
+        const current_root = await this.root();
+        if (current_root == root) {
+          break;
+        }
+      }
+      if (await this.root() != root) {
+        throw new Error(`could not rollback to root ${root}`);
+      }
+    } finally {
+    //   this.lock.release();
+    }
+
+  }
+}
+
+module.exports = MerkleTree;
\ No newline at end of file
diff --git a/lib/MiMC.js b/lib/MiMC.js
new file mode 100644
index 0000000..843ec91
--- /dev/null
+++ b/lib/MiMC.js
@@ -0,0 +1,13 @@
+const circomlib = require('circomlib');
+const mimcsponge = circomlib.mimcsponge;
+const snarkjs = require('snarkjs');
+
+const bigInt = snarkjs.bigInt;
+
+class MimcSpongeHasher {
+    hash(level, left, right) {
+        return mimcsponge.multiHash([bigInt(left), bigInt(right)]).toString();
+    }
+}
+
+module.exports = MimcSpongeHasher;
\ No newline at end of file
diff --git a/lib/Storage.js b/lib/Storage.js
new file mode 100644
index 0000000..9d45112
--- /dev/null
+++ b/lib/Storage.js
@@ -0,0 +1,36 @@
+
+
+class JsStorage {
+  constructor() {
+    this.db = {};
+  }
+
+  get(key) {
+    return this.db[key];
+  }
+
+  get_or_element(key, defaultElement) {
+    const element = this.db[key];
+    if (element === undefined) {
+        return defaultElement;
+    } else {
+        return element
+    }
+  }
+
+  put(key, value) {
+    this.db[key] = value;
+  }
+
+  del(key) {
+    delete this.db[key];
+  }
+
+  put_batch(key_values) {
+    key_values.forEach(element => {
+        this.db[element.key] = element.value;
+    });
+  }
+}
+
+module.exports = JsStorage;
\ No newline at end of file