merkle tree lib

This commit is contained in:
Alexey 2019-07-11 13:38:22 +03:00
parent c0e27ffafc
commit 65fd0202c5
4 changed files with 377 additions and 8 deletions

View File

@ -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,
)
})
})

262
lib/MerkleTree.js Normal file
View File

@ -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;

13
lib/MiMC.js Normal file
View File

@ -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;

36
lib/Storage.js Normal file
View File

@ -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;