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