Merge branch 'frangio-complete'

src/MerkleTree.ts

@@ -21,7 +21,7 @@ export interface Options {
   hashLeaves?: boolean
   /** If set to `true`, constructs the Merkle Tree using the [Bitcoin Merkle Tree implementation](http://www.righto.com/2014/02/bitcoin-mining-hard-way-algorithms.html). Enable it when you need to replicate Bitcoin constructed Merkle Trees. In Bitcoin Merkle Trees, single nodes are combined with themselves, and each output hash is hashed again. */
   isBitcoinTree?: boolean
-  /** If set to `true`, the leaves will be sorted. */
+  /** If set to `true`, the leaves will be sorted. Recommended for use of multiProofs. */
   sortLeaves?: boolean
   /** If set to `true`, the hashing pairs will be sorted. */
   sortPairs?: boolean
@@ -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('Warning: For correct multiProofs it\'s strongly recommended to set complete: true')
+    }
+
     if (!indices) {
       indices = tree
       tree = this.getLayersFlat()
@@ -883,7 +908,7 @@ export class MerkleTree extends Base {
 
     let ids : number[]
     if (leaves.every(Number.isInteger)) {
-      ids = leaves.sort((a, b) => a === b ? 0 : a > b ? 1 : -1) // Indices where passed
+      ids = [...leaves].sort((a, b) => a === b ? 0 : a > b ? 1 : -1) // Indices where passed
     } else {
       ids = leaves.map((el) => this._bufferIndexOf(this.leaves, el)).sort((a, b) => a === b ? 0 : a > b ? 1 : -1)
     }

test/MerkleTree.test.js

@@ -920,7 +920,7 @@ test('sha256 getMultiProof', t => {
 })
 
 test('sha256 getProofFlag with indices', t => {
-  t.plan(2)
+  t.plan(3)
 
   const leaves = ['a', 'b', 'c', 'd'].map(sha256)
   const tree = new MerkleTree(leaves, sha256, { sortPairs: true })
@@ -928,15 +928,17 @@ test('sha256 getProofFlag with indices', t => {
   t.equal(root, '0x4c6aae040ffada3d02598207b8485fcbe161c03f4cb3f660e4d341e7496ff3b2')
 
   const treeFlat = tree.getLayersFlat()
-  const proofIndices = [1, 2]
-  const proof = tree.getMultiProof(treeFlat, proofIndices)
-
+  const proofIndices = [2, 1]
+  const proof = tree.getMultiProof(proofIndices)
   const proofFlags = tree.getProofFlags(proofIndices, proof)
   t.deepEqual(proofFlags, [
     false,
     false,
     true
   ])
+  t.ok(
+    tree.verifyMultiProofWithFlags(root, proofIndices.map(i => leaves[i]), proof, proofFlags)
+  )
 })
 
 test('sha256 getMultiProof - statusim', t => {
@@ -1158,3 +1160,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 }),
+    /option "complete" is incompatible with "isBitcoinTree"/,
+  )
+  t.throws(
+    () => new MerkleTree(leaves, keccak256, { complete: true, duplicateOdd: true }),
+    /option "complete" is incompatible with "duplicateOdd"/,
+  )
+})