diff --git a/src/MerkleTree.ts b/src/MerkleTree.ts index 36ceadb..ae5beda 100644 --- a/src/MerkleTree.ts +++ b/src/MerkleTree.ts @@ -29,6 +29,8 @@ export interface Options { sort?: boolean /** If defined, the resulting hash of this function will be used to fill in odd numbered layers. */ fillDefaultHash?: TFillDefaultHash | Buffer | string + /** If set to `true`, the resulting tree will be a complete tree. Recommended for use of multiProofs. */ + complete?: boolean; } /** @@ -46,6 +48,7 @@ export class MerkleTree extends Base { private sortPairs: boolean = false private sort: boolean = false private fillDefaultHash: TFillDefaultHash | null = null + private complete: boolean = false /** * @desc Constructs a Merkle Tree. @@ -71,10 +74,21 @@ export class MerkleTree extends Base { */ constructor (leaves: any[], hashFn = SHA256, options: Options = {}) { super() + + if (options.complete) { + if (options.isBitcoinTree) { + throw new Error('option "complete" is incompatible with "isBitcoinTree"') + } + if (options.duplicateOdd) { + throw new Error('option "complete" is incompatible with "duplicateOdd"') + } + } + this.isBitcoinTree = !!options.isBitcoinTree this.hashLeaves = !!options.hashLeaves this.sortLeaves = !!options.sortLeaves this.sortPairs = !!options.sortPairs + this.complete = !!options.complete if (options.fillDefaultHash) { if (typeof options.fillDefaultHash === 'function') { @@ -126,8 +140,15 @@ export class MerkleTree extends Base { this.layers.push([]) + const layerLimit = this.complete && layerIndex === 1 && !Number.isInteger(Math.log2(nodes.length)) + ? (2 * nodes.length) - (2 ** Math.ceil(Math.log2(nodes.length))) + : nodes.length + for (let i = 0; i < nodes.length; i += 2) { - if (i + 1 === nodes.length) { + if (i >= layerLimit) { + this.layers[layerIndex].push(...nodes.slice(layerLimit)) + break + } else if (i + 1 === nodes.length) { if (nodes.length % 2 === 1) { let data = nodes[nodes.length - 1] let hash = data @@ -744,6 +765,10 @@ export class MerkleTree extends Base { *``` */ getMultiProof (tree?: any[], indices?: any[]):Buffer[] { + if (!this.complete) { + console.warn('for correct use of multiProof it\'s strongly recommended to set complete: true') + } + if (!indices) { indices = tree tree = this.getLayersFlat() diff --git a/test/MerkleTree.test.js b/test/MerkleTree.test.js index c7f932c..082e02a 100644 --- a/test/MerkleTree.test.js +++ b/test/MerkleTree.test.js @@ -1158,3 +1158,28 @@ test('ethereum-cryptography/keccak256', t => { const tree2 = new MerkleTree(leaves, ethCryptoKeccak256, { hashLeaves: true }) t.equal(tree2.getHexRoot(), expectedRoot) }) + +test('keccak256 with complete option', t => { + t.plan(1) + + const leaves = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i'].map(v => keccak256(Buffer.from(v))) + const tree = new MerkleTree(leaves, keccak256, { complete: true }) + + const root = '581ddb9a48c5ec60fd5a0023d4673e2b33c256f8886342bef35a8eebeda51b44' + + t.equal(tree.getRoot().toString('hex'), root) +}) + +test('complete option with incompatible options', t => { + t.plan(2) + + const leaves = ['a', 'b', 'c'].map(v => keccak256(Buffer.from(v))) + t.throws( + () => new MerkleTree(leaves, keccak256, { complete: true, isBitcoinTree: true }), + { message: 'option "complete" is incompatible with "isBitcoinTree"' }, + ) + t.throws( + () => new MerkleTree(leaves, keccak256, { complete: true, duplicateOdd: true }), + { message: 'option "complete" is incompatible with "duplicateOdd"' }, + ) +})