mirror of
https://github.com/tornadocash/tornado-core.git
synced 2024-12-23 01:29:53 +01:00
merkle tree tidy
This commit is contained in:
parent
b5b3474165
commit
15b573b5a1
@ -7,6 +7,7 @@ class MerkleTree {
|
||||
this.hasher = hasher;
|
||||
this.n_levels = n_levels;
|
||||
this.zero_values = [];
|
||||
this.totalElements = 0;
|
||||
|
||||
let current_zero_value = zero_value;
|
||||
this.zero_values.push(current_zero_value);
|
||||
@ -18,12 +19,12 @@ class MerkleTree {
|
||||
}
|
||||
if (defaultElements) {
|
||||
let level = 0
|
||||
this.totalElements = defaultElements.length
|
||||
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)
|
||||
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)
|
||||
@ -45,46 +46,6 @@ class MerkleTree {
|
||||
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),
|
||||
@ -94,12 +55,6 @@ class MerkleTree {
|
||||
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) {
|
||||
@ -139,9 +94,13 @@ class MerkleTree {
|
||||
};
|
||||
}
|
||||
|
||||
async update(index, element, update_log_index) {
|
||||
async update(index, element, insert = false) {
|
||||
if (!insert && index >= this.totalElements) {
|
||||
throw Error('Use insert method for new elements.')
|
||||
} else if(insert && index < this.totalElements) {
|
||||
throw Error('Use update method for existing elements.')
|
||||
}
|
||||
try {
|
||||
//console.log(`updating ${index}, ${element}`);
|
||||
class UpdateTraverser {
|
||||
constructor(prefix, storage, hasher, element, zero_values) {
|
||||
this.prefix = prefix;
|
||||
@ -158,11 +117,6 @@ class MerkleTree {
|
||||
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),
|
||||
@ -181,9 +135,7 @@ class MerkleTree {
|
||||
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(
|
||||
@ -195,32 +147,23 @@ class MerkleTree {
|
||||
);
|
||||
|
||||
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}`);
|
||||
} catch(e) {
|
||||
console.error(e)
|
||||
}
|
||||
}
|
||||
|
||||
async insert(element) {
|
||||
const index = this.totalElements
|
||||
await this.update(index, element, true)
|
||||
this.totalElements++
|
||||
}
|
||||
|
||||
async traverse(index, handler) {
|
||||
let current_index = index;
|
||||
for (let i = 0; i < this.n_levels; i++) {
|
||||
@ -234,46 +177,6 @@ class MerkleTree {
|
||||
current_index = Math.floor(current_index / 2);
|
||||
}
|
||||
}
|
||||
|
||||
async rollback(updates) {
|
||||
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);
|
||||
}
|
||||
} catch (e) {
|
||||
console.error(e)
|
||||
}
|
||||
}
|
||||
|
||||
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}`);
|
||||
}
|
||||
} catch (e) {
|
||||
console.log(e)
|
||||
}
|
||||
|
||||
}
|
||||
}
|
||||
|
||||
module.exports = MerkleTree;
|
||||
module.exports = MerkleTree;
|
||||
|
@ -9,9 +9,9 @@ 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')
|
||||
const JsStorage = require('../lib/Storage')
|
||||
const MerkleTree = require('../lib/MerkleTree')
|
||||
const MimcHacher = require('../lib/MiMC')
|
||||
|
||||
function BNArrayToStringArray(array) {
|
||||
const arrayToPrint = []
|
||||
@ -82,7 +82,7 @@ contract('MerkleTreeWithHistory', async accounts => {
|
||||
2,
|
||||
zeroValue,
|
||||
)
|
||||
await tree.update(0, '5')
|
||||
await tree.insert('5')
|
||||
let {root, path_elements, path_index} = await tree.path(0)
|
||||
const calculated_root = hasher.hash(null,
|
||||
hasher.hash(null, '5', path_elements[0]),
|
||||
@ -94,7 +94,7 @@ contract('MerkleTreeWithHistory', async accounts => {
|
||||
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)
|
||||
await tree.insert(el)
|
||||
}
|
||||
|
||||
const storage = new JsStorage()
|
||||
@ -117,7 +117,7 @@ contract('MerkleTreeWithHistory', async accounts => {
|
||||
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)
|
||||
await tree.insert(el)
|
||||
}
|
||||
|
||||
const storage = new JsStorage()
|
||||
@ -136,6 +136,29 @@ contract('MerkleTreeWithHistory', async accounts => {
|
||||
pathViaConstructor.should.be.deep.equal(pathViaUpdate)
|
||||
}
|
||||
})
|
||||
|
||||
it.skip('creation using 1000 elements', async () => {
|
||||
const elements = []
|
||||
for(let i = 2000; i < 2001; i++) {
|
||||
elements.push(i)
|
||||
}
|
||||
const storage = new JsStorage()
|
||||
hasher = new MimcHacher()
|
||||
console.time('MerkleTree');
|
||||
tree = new MerkleTree(
|
||||
prefix,
|
||||
storage,
|
||||
hasher,
|
||||
levels,
|
||||
zeroValue,
|
||||
elements
|
||||
);
|
||||
console.timeEnd('MerkleTree');
|
||||
// 2,7 GHz Intel Core i7
|
||||
// 1000 : 1949.084ms
|
||||
// 10000: 19456.220ms
|
||||
// 30000: 57293.640ms
|
||||
})
|
||||
})
|
||||
|
||||
describe('#insert', async () => {
|
||||
@ -145,7 +168,7 @@ contract('MerkleTreeWithHistory', async accounts => {
|
||||
|
||||
for (i = 1; i < 11; i++) {
|
||||
await merkleTreeWithHistory.insert(i)
|
||||
await tree.update(i - 1, i)
|
||||
await tree.insert(i)
|
||||
filled_subtrees = await merkleTreeWithHistory.filled_subtrees()
|
||||
let {root, path_elements, path_index} = await tree.path(i - 1)
|
||||
// console.log('path_elements ', path_elements)
|
||||
|
Loading…
Reference in New Issue
Block a user