diff --git a/src/BaseTree.ts b/src/BaseTree.ts new file mode 100644 index 0000000..b6d78fc --- /dev/null +++ b/src/BaseTree.ts @@ -0,0 +1,172 @@ +import { Element, HashFunction, ProofPath } from './index' + +export class BaseTree { + levels: number + protected _hashFn: HashFunction + protected zeroElement: Element + protected _zeros: Element[] + protected _layers: Array + + get capacity() { + return 2 ** this.levels + } + + get layers(): Array { + return this._layers.slice() + } + + get zeros(): Element[] { + return this._zeros.slice() + } + + get elements(): Element[] { + return this._layers[0].slice() + } + + get root(): Element { + return this._layers[this.levels][0] ?? this._zeros[this.levels] + } + + /** + * Find an element in the tree + * @param element An element to find + * @param comparator A function that checks leaf value equality + * @returns {number} Index if element is found, otherwise -1 + */ + indexOf(element: Element, comparator?: (arg0: T, arg1: T) => boolean): number { + if (comparator) { + return this._layers[0].findIndex((el) => comparator(element, el)) + } else { + return this._layers[0].indexOf(element) + } + } + + /** + * 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 + const left = this._layers[level - 1][index * 2] + const right = this._layers[level - 1][index * 2 + 1] + this._layers[level][index] = this._hashFn(left, right) + } + } + 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 + this._processUpdate(index) + } + + proof(element: Element): ProofPath { + const index = this.indexOf(element) + return this.path(index) + } + + /** + * Get merkle path to a leaf + * @param {number} index Leaf index to generate path for + * @returns {{pathElements: Object[], pathIndex: number[]}} An object containing adjacent elements and left-right index + */ + path(index: number): 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, + pathRoot: this.root, + } + } + + protected _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]) + } + } + + protected _processNodes(nodes: Element[], layerIndex: number) { + const length = nodes.length + let currentLength = Math.ceil(length / 2) + const currentLayer = new Array(currentLength) + currentLength-- + const starFrom = length - ((length % 2) ^ 1) + let j = 0 + for (let i = starFrom; i >= 0; i -= 2) { + if (nodes[i - 1] === undefined) break + const left = nodes[i - 1] + const right = (i === starFrom && length % 2 === 1) ? this._zeros[layerIndex - 1] : nodes[i] + currentLayer[currentLength - j] = this._hashFn(left, right) + j++ + } + return currentLayer + } + + protected _processUpdate(index: number) { + 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] + this._layers[level][index] = this._hashFn(left, right) + } + } + +} diff --git a/src/FixedMerkleTree.ts b/src/FixedMerkleTree.ts index 79fc007..4177531 100644 --- a/src/FixedMerkleTree.ts +++ b/src/FixedMerkleTree.ts @@ -1,96 +1,34 @@ -import { - Element, - HashFunction, - MerkleTreeOptions, - ProofPath, - SerializedTreeState, - simpleHash, - TreeEdge, - TreeSlice, -} from './' +import { Element, HashFunction, MerkleTreeOptions, SerializedTreeState, TreeEdge, TreeSlice } from './' +import defaultHash from './simpleHash' +import { BaseTree } from './BaseTree' -const defaultHash = (left: Element, right: Element): string => simpleHash([left, right]) - -export default class MerkleTree { - levels: number - private _hashFn: HashFunction - private zeroElement: Element - private _zeros: Element[] - private _layers: Array +export default class MerkleTree extends BaseTree { constructor(levels: number, elements: Element[] = [], { hashFunction = defaultHash, zeroElement = 0, }: MerkleTreeOptions = {}) { + super() this.levels = levels if (elements.length > this.capacity) { throw new Error('Tree is full') } this._hashFn = hashFunction this.zeroElement = zeroElement - this._layers = [] const leaves = elements.slice() this._layers = [leaves] this._buildZeros() this._buildHashes() - // this._buildHashes2(leaves) } - get capacity() { - return 2 ** this.levels - } - - get layers(): Array { - 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++) { - this._zeros[i] = this._hashFn(this._zeros[i - 1], this._zeros[i - 1]) + private _buildHashes() { + for (let layerIndex = 1; layerIndex <= this.levels; layerIndex++) { + const nodes = this._layers[layerIndex - 1] + this._layers[layerIndex] = this._processNodes(nodes, layerIndex) } } - _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++) { - 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], - ) - } - } - } - - /** - * Get tree root - */ - get root(): Element { - return this._layers[this.levels][0] ?? this._zeros[this.levels] - } - - /** - * 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. @@ -123,78 +61,6 @@ export default class MerkleTree { 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] - this._layers[level][index] = this._hashFn(left, right) - } - } - - /** - * Get merkle path to a leaf - * @param {number} index Leaf index to generate path for - * @returns {{pathElements: Object[], pathIndex: number[]}} An object containing adjacent elements and left-right index - */ - path(index: number): 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, - pathRoot: this.root, - } - } - - /** - * Find an element in the tree - * @param element An element to find - * @param comparator A function that checks leaf value equality - * @returns {number} Index if element is found, otherwise -1 - */ - indexOf(element: Element, comparator?: (arg0: T, arg1: T) => boolean): number { - if (comparator) { - return this._layers[0].findIndex((el) => comparator(element, el)) - } else { - return this._layers[0].indexOf(element) - } - } - - proof(element: Element): ProofPath { - const index = this.indexOf(element) - return this.path(index) - } - getTreeEdge(edgeIndex: number): TreeEdge { const edgeElement = this._layers[0][edgeIndex] if (edgeElement === undefined) { diff --git a/src/PartialMerkleTree.ts b/src/PartialMerkleTree.ts index cc2bfe4..2f4addd 100644 --- a/src/PartialMerkleTree.ts +++ b/src/PartialMerkleTree.ts @@ -5,26 +5,16 @@ import { MerkleTreeOptions, ProofPath, SerializedPartialTreeState, - simpleHash, TreeEdge, } from './' +import defaultHash from './simpleHash' +import { BaseTree } from './BaseTree' -export const defaultHash = (left: Element, right: Element): string => simpleHash([left, right]) - -export class PartialMerkleTree { - get edgeLeafProof(): ProofPath { - return this._edgeLeafProof - } - - levels: number - private zeroElement: Element - private _zeros: Element[] - private _layers: Array +export class PartialMerkleTree extends BaseTree { private _leaves: Element[] private _leavesAfterEdge: Element[] private _edgeLeaf: LeafWithIndex private _initialRoot: Element - private _hashFn: HashFunction private _edgeLeafProof: ProofPath private _proofMap: Map @@ -34,6 +24,7 @@ export class PartialMerkleTree { edgeIndex, edgeElementsCount, }: TreeEdge, leaves: Element[], { hashFunction, zeroElement }: MerkleTreeOptions = {}) { + super() if (edgeIndex + leaves.length !== edgeElementsCount) throw new Error('Invalid number of elements') this._edgeLeafProof = edgePath this._initialRoot = edgePath.pathRoot @@ -46,26 +37,6 @@ export class PartialMerkleTree { this._buildTree() } - get capacity() { - return 2 ** this.levels - } - - get layers(): Array { - return this._layers.slice() - } - - get zeros(): Element[] { - return this._zeros.slice() - } - - get elements(): Element[] { - return this._layers[0].slice() - } - - get root(): Element { - return this._layers[this.levels][0] ?? this._zeros[this.levels] - } - get edgeIndex(): number { return this._edgeLeaf.index } @@ -74,6 +45,10 @@ export class PartialMerkleTree { return this._edgeLeaf.data } + get edgeLeafProof(): ProofPath { + return this._edgeLeafProof + } + private _createProofMap() { this._proofMap = this.edgeLeafProof.pathPositions.reduce((p, c, i) => { p.set(i, [c, this.edgeLeafProof.pathElements[i]]) @@ -94,29 +69,10 @@ export class PartialMerkleTree { 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() { + private _buildHashes() { for (let layerIndex = 1; layerIndex <= this.levels; layerIndex++) { const nodes = this._layers[layerIndex - 1] - const length = nodes.length - let currentLength = Math.ceil(length / 2) - const currentLayer = new Array(currentLength) - currentLength-- - const starFrom = length - ((length % 2) ^ 1) - let j = 0 - for (let i = starFrom; i >= 0; i -= 2) { - if (nodes[i - 1] === undefined) break - const left = nodes[i - 1] - const right = (i === starFrom && length % 2 === 1) ? this._zeros[layerIndex - 1] : nodes[i] - currentLayer[currentLength - j] = this._hashFn(left, right) - j++ - } + const currentLayer = this._processNodes(nodes, layerIndex) if (this._proofMap.has(layerIndex)) { const [proofPos, proofEl] = this._proofMap.get(layerIndex) if (!currentLayer[proofPos]) currentLayer[proofPos] = proofEl @@ -125,46 +81,6 @@ export class PartialMerkleTree { } } - /** - * 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 - const left = this._layers[level - 1][index * 2] - const right = this._layers[level - 1][index * 2 + 1] - this._layers[level][index] = this._hashFn(left, right) - } - } - this.insert(elements[elements.length - 1]) - } /** * Change an element in the tree @@ -179,14 +95,7 @@ export class PartialMerkleTree { throw new Error(`Index ${index} is below the edge: ${this._edgeLeaf.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] - this._layers[level][index] = this._hashFn(left, right) - } + this._processUpdate(index) } path(index: number): ProofPath { @@ -221,18 +130,6 @@ export class PartialMerkleTree { } } - indexOf(element: Element, comparator?: (arg0: T, arg1: T) => boolean): number { - if (comparator) { - return this._layers[0].findIndex((el) => comparator(element, el)) - } else { - return this._layers[0].indexOf(element) - } - } - - proof(element: Element): ProofPath { - const index = this.indexOf(element) - return this.path(index) - } /** * Shifts edge of tree to left diff --git a/src/simpleHash.ts b/src/simpleHash.ts index 9c56848..fd3e56a 100644 --- a/src/simpleHash.ts +++ b/src/simpleHash.ts @@ -1,3 +1,5 @@ +import { Element } from './index' + /*** * This is insecure hash function, just for example only * @param data @@ -17,3 +19,4 @@ export function simpleHash(data: T[], seed?: number, hashLength = 40): string return BigInt('0x' + hash.padEnd(hashLength - (hash.length - 1), '0')).toString(10) } +export default (left: Element, right: Element): string => simpleHash([left, right])