tornadocash
/
fixed-merkle-tree
mirror of https://github.com/tornadocash/fixed-merkle-tree.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

implementing features and cover it by tests

This commit is contained in:
Sergei SMART 2022-03-01 15:33:46 +10:00
parent 9536889468
commit 21e673a659
6 changed files with 346 additions and 125 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

50
src/fixedMerkleTree.ts → src/FixedMerkleTree.ts
View File

@ -1,21 +1,8 @@
import {
defaultHash,
Element,
HashFunction,
MerkleTreeOptions,
ProofPath,
SerializedTreeState,
TreeEdge,
} from './'
import { defaultHash, Element, HashFunction, MerkleTreeOptions, ProofPath, SerializedTreeState, TreeEdge } from './'
export default class MerkleTree {
get layers(): Array<Element[]> {
return this._layers.slice()
}
levels: number
capacity: number
private _hashFn: HashFunction<Element>
private zeroElement: Element
private _zeros: Element[]
@ -26,7 +13,6 @@ export default class MerkleTree {
zeroElement = 0,
}: MerkleTreeOptions = {}) {
this.levels = levels
this.capacity = 2 ** levels
if (elements.length > this.capacity) {
throw new Error('Tree is full')
}
@ -39,6 +25,22 @@ export default class MerkleTree {
this._rebuild()
}
get capacity() {
return this.levels ** 2
}
get layers(): Array<Element[]> {
return this._layers.slice()
}
get zeros(): Element[] {
return this._zeros.slice()
}
get elements(): Element[] {
return this._layers[0].slice()
}
private _buildZeros() {
this._zeros = [this.zeroElement]
for (let i = 1; i <= this.levels; i++) {
@ -63,7 +65,7 @@ export default class MerkleTree {
/**
* Get tree root
*/
root(): Element {
get root(): Element {
return this._layers[this.levels][0] ?? this._zeros[this.levels]
}
@ -180,26 +182,12 @@ export default class MerkleTree {
const leaves = this._layers[0]
const edgeIndex = leaves.indexOf(edgeElement)
if (edgeIndex <= -1) {
return null
throw new Error('Element not found')
}
const edgePath = this.path(edgeIndex)
return { edgePath, edgeElement, edgeIndex }
}
/**
* Returns a copy of non-zero tree elements.
*/
get elements() {
return this._layers[0].slice()
}
/**
* Returns a copy of n-th zero elements array
*/
get zeros() {
return this._zeros.slice()
}
/**
* Serialize entire tree state including intermediate layers into a plain object
* Deserializing it back will not require to recompute any hashes

194
src/PartialMerkleTree.ts Normal file
View File

@ -0,0 +1,194 @@
import { defaultHash, Element, HashFunction, MerkleTreeOptions, ProofPath, TreeEdge } from './'
type LeafWithIndex = { index: number, data: Element }
export class PartialMerkleTree {
levels: number
private zeroElement: Element
private _zeros: Element[]
private _layers: Array<Element[]>
private _leaves: Element[]
private _leavesAfterEdge: Element[]
private _edgeLeaf: LeafWithIndex
private _initialRoot: Element
private _hashFn: HashFunction<Element>
private _edgeLeafProof: ProofPath
constructor(levels: number, {
edgePath,
edgeElement,
edgeIndex,
}: TreeEdge, leaves: Element[], root: Element, { hashFunction, zeroElement }: MerkleTreeOptions = {}) {
this._edgeLeafProof = edgePath
this.zeroElement = zeroElement ?? 0
this._edgeLeaf = { data: edgeElement, index: edgeIndex }
this._leavesAfterEdge = leaves
this._initialRoot = root
this.levels = levels
this._hashFn = hashFunction || defaultHash
this._buildTree()
}
get capacity() {
return this.levels ** 2
}
get layers(): Array<Element[]> {
return this._layers.slice()
}
get zeros(): Element[] {
return this._zeros.slice()
}
get elements(): Element[] {
return this._layers[0].slice()
}
private _buildTree(): void {
const edgeLeafIndex = this._edgeLeaf.index
this._leaves = [...Array.from({ length: edgeLeafIndex }, () => null), ...this._leavesAfterEdge]
if (this._edgeLeafProof.pathIndices[0] === 1) {
this._leaves[this._edgeLeafProof.pathPositions[0]] = this._edgeLeafProof.pathElements[0]
}
this._layers = [this._leaves]
this._buildZeros()
this._buildHashes()
}
private _buildZeros() {
this._zeros = [this.zeroElement]
for (let i = 1; i <= this.levels; i++) {
this._zeros[i] = this._hashFn(this._zeros[i - 1], this._zeros[i - 1])
}
}
_buildHashes() {
for (let level = 1; level <= this.levels; level++) {
this._layers[level] = []
for (let i = 0; i < Math.ceil(this._layers[level - 1].length / 2); i++) {
const left = this._layers[level - 1][i * 2]
const right = i * 2 + 1 < this._layers[level - 1].length
? this._layers[level - 1][i * 2 + 1]
: this._zeros[level - 1]
let hash: Element = this._hashFn(left, right)
if (!hash && this._edgeLeafProof.pathPositions[level] === i) hash = this._edgeLeafProof.pathElements[level]
if (level === this.levels) hash = hash || this._initialRoot
this._layers[level][i] = hash
}
}
}
/**
* Insert new element into the tree
* @param element Element to insert
*/
insert(element: Element) {
if (this._layers[0].length >= this.capacity) {
throw new Error('Tree is full')
}
this.update(this._layers[0].length, element)
}
/**
* Insert multiple elements into the tree.
* @param {Array} elements Elements to insert
*/
bulkInsert(elements: Element[]): void {
if (!elements.length) {
return
}
if (this._layers[0].length + elements.length > this.capacity) {
throw new Error('Tree is full')
}
// First we insert all elements except the last one
// updating only full subtree hashes (all layers where inserted element has odd index)
// the last element will update the full path to the root making the tree consistent again
for (let i = 0; i < elements.length - 1; i++) {
this._layers[0].push(elements[i])
let level = 0
let index = this._layers[0].length - 1
while (index % 2 === 1) {
level++
index >>= 1
this._layers[level][index] = this._hashFn(
this._layers[level - 1][index * 2],
this._layers[level - 1][index * 2 + 1],
)
}
}
this.insert(elements[elements.length - 1])
}
/**
* Change an element in the tree
* @param {number} index Index of element to change
* @param element Updated element value
*/
update(index: number, element: Element) {
if (isNaN(Number(index)) || index < 0 || index > this._layers[0].length || index >= this.capacity) {
throw new Error('Insert index out of bounds: ' + index)
}
this._layers[0][index] = element
for (let level = 1; level <= this.levels; level++) {
index >>= 1
const left = this._layers[level - 1][index * 2]
const right = index * 2 + 1 < this._layers[level - 1].length
? this._layers[level - 1][index * 2 + 1]
: this._zeros[level - 1]
let hash: Element = this._hashFn(left, right)
if (!hash && this._edgeLeafProof.pathPositions[level] === index) {
hash = this._edgeLeafProof.pathElements[level]
}
if (level === this.levels) {
hash = hash || this._initialRoot
}
// console.log({ index, level, left, right, hash })
this._layers[level][index] = hash
}
}
path(index: Element): ProofPath {
if (isNaN(Number(index)) || index < 0 || index >= this._layers[0].length) {
throw new Error('Index out of bounds: ' + index)
}
let elIndex = +index
const pathElements: Element[] = []
const pathIndices: number[] = []
const pathPositions: number [] = []
for (let level = 0; level < this.levels; level++) {
pathIndices[level] = elIndex % 2
const leafIndex = elIndex ^ 1
if (leafIndex < this._layers[level].length) {
pathElements[level] = this._layers[level][leafIndex]
pathPositions[level] = leafIndex
} else {
pathElements[level] = this._zeros[level]
pathPositions[level] = 0
}
elIndex >>= 1
}
return {
pathElements,
pathIndices,
pathPositions,
}
}
indexOf(element: Element, comparator?: <T> (arg0: T, arg1: T) => boolean): number {
if (comparator) {
return this._layers[0].findIndex((el) => comparator<Element>(element, el))
} else {
return this._layers[0].indexOf(element)
}
}
/**
* Get tree root
*/
get root(): Element {
return this._layers[this.levels][0] ?? this._zeros[this.levels]
}
}

6
src/index.ts
View File

@ -1,7 +1,7 @@
import { simpleHash } from './simpleHash'
export { default as MerkleTree } from './fixedMerkleTree'
export { PartialMerkleTree } from './partialMerkleTree'
export { default as MerkleTree } from './FixedMerkleTree'
export { PartialMerkleTree } from './PartialMerkleTree'
export { simpleHash } from './simpleHash'
export type HashFunction<T> = {
@ -31,4 +31,4 @@ export type TreeEdge = {
edgePath: ProofPath;
edgeIndex: number
}
export const defaultHash = (left: Element, right: Element): string => simpleHash([left, right])
export const defaultHash = (left: Element, right: Element): string => (left !== null && right !== null) ? simpleHash([left, right]) : null

64
src/partialMerkleTree.ts
View File

@ -1,64 +0,0 @@
import { defaultHash, Element, HashFunction, MerkleTreeOptions, ProofPath, TreeEdge } from './'
type LeafWithIndex = { index: number, data: Element }
export class PartialMerkleTree {
levels: number
private zeroElement: Element
private _zeros: Element[]
private _layers: Array<Element[]>
private _leaves: Element[]
private _leavesAfterEdge: Element[]
private _edgeLeaf: LeafWithIndex
private _root: Element
private _hashFn: HashFunction<Element>
private _edgeLeafProof: ProofPath
constructor({
edgePath,
edgeElement,
edgeIndex,
}: TreeEdge, leaves: Element[], root: Element, { hashFunction, zeroElement }: MerkleTreeOptions = {}) {
this._edgeLeafProof = edgePath
this.zeroElement = zeroElement ?? 0
this._edgeLeaf = { data: edgeElement, index: edgeIndex }
this._leavesAfterEdge = leaves
this._root = root
this._hashFn = hashFunction || defaultHash
this._buildTree()
}
get capacity() {
return this.levels ** 2
}
private _buildTree(): void {
const edgeLeafIndex = this._edgeLeaf.index
this._leaves = [...Array.from({ length: edgeLeafIndex - 1 }, () => null), ...this._leavesAfterEdge]
this._layers = [this._leaves]
this._buildZeros()
this._rebuild()
}
private _buildZeros() {
this._zeros = [this.zeroElement]
for (let i = 1; i <= this.levels; i++) {
this._zeros[i] = this._hashFn(this._zeros[i - 1], this._zeros[i - 1])
}
}
_rebuild() {
for (let level = 1; level <= this.levels; level++) {
this._layers[level] = []
for (let i = 0; i < Math.ceil(this._layers[level - 1].length / 2); i++) {
this._layers[level][i] = this._hashFn(
this._layers[level - 1][i * 2],
i * 2 + 1 < this._layers[level - 1].length
? this._layers[level - 1][i * 2 + 1]
: this._zeros[level - 1],
)
}
}
}
}

49
test/fixedMerkleTree.spec.ts
View File

@ -1,6 +1,7 @@
import { MerkleTree } from '../src'
import { assert, should } from 'chai'
import { it } from 'mocha'
import { TreeEdge } from '../lib'
describe('MerkleTree', () => {
@ -8,22 +9,22 @@ describe('MerkleTree', () => {
it('should have correct zero root', () => {
const tree = new MerkleTree(10, [])
return should().equal(tree.root(), '3060353338620102847451617558650138132480')
return should().equal(tree.root, '3060353338620102847451617558650138132480')
})
it('should have correct 1 element root', () => {
const tree = new MerkleTree(10, [1])
should().equal(tree.root(), '4059654748770657324723044385589999697920')
should().equal(tree.root, '4059654748770657324723044385589999697920')
})
it('should have correct even elements root', () => {
const tree = new MerkleTree(10, [1, 2])
should().equal(tree.root(), '3715471817149864798706576217905179918336')
should().equal(tree.root, '3715471817149864798706576217905179918336')
})
it('should have correct odd elements root', () => {
const tree = new MerkleTree(10, [1, 2, 3])
should().equal(tree.root(), '5199180210167621115778229238102210117632')
should().equal(tree.root, '5199180210167621115778229238102210117632')
})
it('should be able to create a full tree', () => {
@ -40,19 +41,19 @@ describe('MerkleTree', () => {
it('should insert into empty tree', () => {
const tree = new MerkleTree(10)
tree.insert(42)
should().equal(tree.root(), '750572848877730275626358141391262973952')
should().equal(tree.root, '750572848877730275626358141391262973952')
})
it('should insert into odd tree', () => {
const tree = new MerkleTree(10, [1])
tree.insert(42)
should().equal(tree.root(), '5008383558940708447763798816817296703488')
should().equal(tree.root, '5008383558940708447763798816817296703488')
})
it('should insert into even tree', () => {
const tree = new MerkleTree(10, [1, 2])
tree.insert(42)
should().equal(tree.root(), '5005864318873356880627322373636156817408')
should().equal(tree.root, '5005864318873356880627322373636156817408')
})
it('should insert last element', () => {
@ -71,7 +72,7 @@ describe('MerkleTree', () => {
it('should work', () => {
const tree = new MerkleTree(10, [1, 2, 3])
tree.bulkInsert([4, 5, 6])
should().equal(tree.root(), '4066635800770511602067209448381558554624')
should().equal(tree.root, '4066635800770511602067209448381558554624')
})
it('should give the same result as sequential inserts', () => {
@ -95,7 +96,7 @@ describe('MerkleTree', () => {
for (const item of inserted) {
tree2.insert(item)
}
should().equal(tree1.root(), tree2.root())
should().equal(tree1.root, tree2.root)
}
}
}).timeout(10000)
@ -122,31 +123,31 @@ describe('MerkleTree', () => {
it('should update first element', () => {
const tree = new MerkleTree(10, [1, 2, 3, 4, 5])
tree.update(0, 42)
should().equal(tree.root(), '3884161948856565981263417078389340635136')
should().equal(tree.root, '3884161948856565981263417078389340635136')
})
it('should update last element', () => {
const tree = new MerkleTree(10, [1, 2, 3, 4, 5])
tree.update(4, 42)
should().equal(tree.root(), '3564959811529894228734180300843252711424')
should().equal(tree.root, '3564959811529894228734180300843252711424')
})
it('should update odd element', () => {
const tree = new MerkleTree(10, [1, 2, 3, 4, 5])
tree.update(1, 42)
should().equal(tree.root(), '4576704573778433422699674477203122290688')
should().equal(tree.root, '4576704573778433422699674477203122290688')
})
it('should update even element', () => {
const tree = new MerkleTree(10, [1, 2, 3, 4, 5])
tree.update(2, 42)
should().equal(tree.root(), '1807994110952186123819489133812038762496')
should().equal(tree.root, '1807994110952186123819489133812038762496')
})
it('should update extra element', () => {
const tree = new MerkleTree(10, [1, 2, 3, 4])
tree.update(4, 5)
should().equal(tree.root(), '1099080610107164849381389194938128793600')
should().equal(tree.root, '1099080610107164849381389194938128793600')
})
it('should fail to update incorrect index', () => {
@ -248,8 +249,22 @@ describe('MerkleTree', () => {
})
describe('#getTreeEdge', () => {
it('should return correct treeEdge', () => {
const expectedEdge: TreeEdge = {
edgePath: {
pathElements: [
5,
'1390935134112885103361924701261056180224',
'1952916572242076545231119328171167580160',
'938972308169430750202858820582946897920'
],
pathIndices: [ 0, 0, 1, 0 ],
pathPositions: [ 5, 0, 0, 0 ]
},
edgeElement: 4,
edgeIndex: 4
}
const tree = new MerkleTree(4, [0, 1, 2, 3, 4, 5])
console.log(tree.getTreeEdge(4))
assert.deepEqual(tree.getTreeEdge(4), expectedEdge)
})
})
describe('#getters', () => {
@ -319,12 +334,12 @@ describe('MerkleTree', () => {
const src = new MerkleTree(10, [1, 2, 3, 4, 5, 6, 7, 8, 9])
const data = src.serialize()
const dst = MerkleTree.deserialize(data)
should().equal(src.root(), dst.root())
should().equal(src.root, dst.root)
src.insert(10)
dst.insert(10)
should().equal(src.root(), dst.root())
should().equal(src.root, dst.root)
})
})
})

108
test/partialMerkleTree.spec.ts
View File

@ -1,17 +1,105 @@
import { MerkleTree, PartialMerkleTree } from '../src'
import { Element, MerkleTree, PartialMerkleTree } from '../src'
import { it } from 'mocha'
import { should } from 'chai'
describe('PartialMerkleTree', () => {
const getTestTrees = (levels: number, elements: Element[], edgeElement: Element) => {
const fullTree = new MerkleTree(levels, elements)
const edge = fullTree.getTreeEdge(edgeElement)
const leavesAfterEdge = elements.slice(edge.edgeIndex)
const partialTree = new PartialMerkleTree(levels, edge, leavesAfterEdge, fullTree.root)
return { fullTree, partialTree }
}
describe('#constructor', () => {
const leaves = [1, 2, 3, 4, 5]
const fullTree = new MerkleTree(4, leaves)
const root = fullTree.root()
const edge = fullTree.getTreeEdge(3)
const leavesAfterEdge = leaves.splice(edge.edgeIndex)
it('should initialize merkle tree', () => {
const partialTree = new PartialMerkleTree(edge, leavesAfterEdge, root)
return true
const { fullTree, partialTree } = getTestTrees(20, ['0', '1', '2', '3', '4', '5'], '2')
it('should initialize merkle tree with same root', () => {
should().equal(fullTree.root, partialTree.root)
})
it('should initialize merkle tree with same leaves count', () => {
should().equal(fullTree.elements.length, partialTree.elements.length)
})
})
describe('#insert', () => {
it('should have equal root to full tree after insertion ', () => {
const { fullTree, partialTree } = getTestTrees(10, ['0', '1', '2', '3', '4', '5', '6', '7'], '5')
fullTree.insert('9')
partialTree.insert('9')
should().equal(fullTree.root, partialTree.root)
})
it('should fail to insert when tree is full', () => {
const { partialTree } = getTestTrees(3, ['0', '1', '2', '3', '4', '5', '6', '7', '8'], '5')
const call = () => partialTree.insert('9')
should().throw(call, 'Tree is full')
})
})
describe('#bulkInsert', () => {
it('should work like full tree', () => {
const { fullTree, partialTree } = getTestTrees(20, [1, 2, 3, 4, 5], 3)
partialTree.bulkInsert([6, 7, 8])
fullTree.bulkInsert([6, 7, 8])
should().equal(fullTree.root, partialTree.root)
})
it('should give the same result as sequential inserts', () => {
const initialArray = [
[1],
[1, 2],
[1, 2, 3],
[1, 2, 3, 4],
]
const insertedArray = [
[11],
[11, 12],
[11, 12, 13],
[11, 12, 13, 14],
]
for (const initial of initialArray) {
for (const inserted of insertedArray) {
const { partialTree: tree1 } = getTestTrees(10, initial, initial.length > 1 ? initial.length - 1 : initial.length)
const { partialTree: tree2 } = getTestTrees(10, initial, initial.length > 1 ? initial.length - 1 : initial.length)
tree1.bulkInsert(inserted)
for (const item of inserted) {
tree2.insert(item)
}
should().equal(tree1.root, tree2.root)
}
}
}).timeout(10000)
it('should fail to insert too many elements', () => {
const { fullTree, partialTree } = getTestTrees(2, [1, 2, 3, 4], 3)
const call = () => partialTree.bulkInsert([5, 6, 7])
should().throw(call, 'Tree is full')
})
})
describe('#indexOf', () => {
it('should return same result as full tree', () => {
const { fullTree, partialTree } = getTestTrees(10, [1, 2, 3, 4, 5, 6, 7, 8], 4)
should().equal(partialTree.indexOf(5), fullTree.indexOf(5))
})
it('should find index', () => {
const { partialTree } = getTestTrees(10, [1, 2, 3, 4, 5], 3)
should().equal(partialTree.indexOf(3), 2)
})
it('should work with comparator', () => {
const { partialTree } = getTestTrees(10, [1, 2, 3, 4, 5], 3)
should().equal(partialTree.indexOf(4, (arg0, arg1) => arg0 === arg1), 3)
})
it('should return -1 for non existent element', () => {
const { partialTree } = getTestTrees(10, [1, 2, 3, 4, 5], 3)
should().equal(partialTree.indexOf(42), -1)
})
})
})